Добавил выбор из кандидатов, если LLM не уверена в раскладке

This commit is contained in:
2026-06-14 16:43:50 +03:00
parent 4af3ad2dde
commit 7f7f5f69d4
16 changed files with 831 additions and 88 deletions
+179 -5
View File
@@ -2,6 +2,7 @@ package worker
import (
"context"
"database/sql"
"encoding/json"
"io"
"log/slog"
@@ -11,6 +12,7 @@ import (
"time"
"git.vakhrushev.me/av/jellybit/internal/layout"
"git.vakhrushev.me/av/jellybit/internal/metadata"
"git.vakhrushev.me/av/jellybit/internal/qbt"
"git.vakhrushev.me/av/jellybit/internal/recognize"
"git.vakhrushev.me/av/jellybit/internal/store"
@@ -73,11 +75,12 @@ func TestNotifier_FiresOnDone(t *testing.T) {
// memStore — полноценный in-memory store для тестов Ф3.
type memStore struct {
downloads map[int64]*store.Download
recs []*store.Recognition
hints map[int64][]string
overrides map[int64]map[string]string
links []store.FileLink
downloads map[int64]*store.Download
recs []*store.Recognition
hints map[int64][]string
overrides map[int64]map[string]string
links []store.FileLink
candidates []store.MetadataCandidate
}
func newMemStore() *memStore {
@@ -199,6 +202,40 @@ func (m *memStore) DeleteFileLinksByBatch(_ context.Context, batch string) error
return nil
}
func (m *memStore) CreateCandidates(_ context.Context, cands []store.MetadataCandidate) error {
for _, c := range cands {
c.ID = int64(len(m.candidates) + 1)
m.candidates = append(m.candidates, c)
}
return nil
}
func (m *memStore) ListCandidatesByRecognition(_ context.Context, recID int64) ([]store.MetadataCandidate, error) {
var out []store.MetadataCandidate
for _, c := range m.candidates {
if c.RecognitionID == recID {
out = append(out, c)
}
}
return out, nil
}
func (m *memStore) GetCandidate(_ context.Context, id int64) (*store.MetadataCandidate, error) {
for i := range m.candidates {
if m.candidates[i].ID == id {
cp := m.candidates[i]
return &cp, nil
}
}
return nil, nil
}
func (m *memStore) SetCandidateChosen(_ context.Context, recID, id int64) error {
for i := range m.candidates {
if m.candidates[i].RecognitionID == recID {
m.candidates[i].Chosen = m.candidates[i].ID == id
}
}
return nil
}
func jsonMarshal(v any) (string, error) {
b, err := json.Marshal(v)
return string(b), err
@@ -659,6 +696,143 @@ func TestProviderTag(t *testing.T) {
}
}
// reviewWithCandidate готовит memStore: задача в review, одна попытка
// распознавания с одним кандидатом базы.
func reviewWithCandidate(t *testing.T, cand store.MetadataCandidate) (*Worker, *memStore) {
t.Helper()
st := newMemStore()
d := completedDownload(1)
d.State = store.StateReview
st.put(d)
planJSON, _ := json.Marshal(recognize.Plan{Type: recognize.MediaSeries, Title: "Догадка", Year: 2000})
st.recs = append(st.recs, &store.Recognition{
ID: 1, DownloadID: 1, IsCurrent: true, Plan: store.NullString(string(planJSON)),
Provider: store.NullString("none"),
})
cand.RecognitionID = 1
_ = st.CreateCandidates(context.Background(), []store.MetadataCandidate{cand})
w := testWorkerWith(st, &fakeQbt{}, &fakeRecognizer{}, nil)
return w, st
}
func TestRecognizeOne_PersistsCandidates(t *testing.T) {
st := newMemStore()
st.put(completedDownload(1))
qb := &fakeQbt{
torrents: []qbt.Torrent{{Hash: ihTest, Name: "Show", SavePath: "/d"}},
files: []qbt.File{{Name: "e1.mkv", Size: 1}},
}
res := seriesResult()
res.Candidates = []metadata.Candidate{
{Provider: "tvmaze", ID: "1", Title: "Show A", Year: 2006, TagProvider: "tvdb", TagID: "269613"},
{Provider: "tvmaze", ID: "2", Title: "Show B", Year: 2007},
}
w := testWorkerWith(st, qb, &fakeRecognizer{result: res}, nil)
w.recognizeOne(context.Background(), 1)
if len(st.candidates) != 2 {
t.Fatalf("candidates = %d, want 2", len(st.candidates))
}
// Тег-предпочтительный provider/id сохранён (TVMaze → tvdb).
if st.candidates[0].Provider != "tvdb" || st.candidates[0].ProviderID != "269613" {
t.Errorf("candidate[0] = %+v", st.candidates[0])
}
}
func TestChooseCandidate_PinsOverrides(t *testing.T) {
w, st := reviewWithCandidate(t, store.MetadataCandidate{
Provider: "tvdb", ProviderID: "269613",
Title: store.NullString("Fargo"), Year: sql.NullInt64{Int64: 2014, Valid: true},
})
candID := st.candidates[0].ID
if err := w.ChooseCandidate(context.Background(), 1, candID); err != nil {
t.Fatalf("ChooseCandidate: %v", err)
}
ov := st.overrides[1]
if ov[ovrProvider] != "tvdb" || ov[ovrProviderID] != "269613" ||
ov[ovrTitle] != "Fargo" || ov[ovrYear] != "2014" {
t.Errorf("overrides = %v", ov)
}
if !st.candidates[0].Chosen {
t.Error("кандидат не помечен выбранным")
}
// Эффективный план берёт каноническое имя/год и тег [tvdbid-...].
plan, tag, err := w.effectivePlan(context.Background(), 1)
if err != nil {
t.Fatalf("effectivePlan: %v", err)
}
if plan.Title != "Fargo" || plan.Year != 2014 {
t.Errorf("plan = %q (%d)", plan.Title, plan.Year)
}
if tag != "tvdbid-269613" {
t.Errorf("tag = %q", tag)
}
}
func TestChooseCandidate_RejectsForeign(t *testing.T) {
w, _ := reviewWithCandidate(t, store.MetadataCandidate{Provider: "tvdb", ProviderID: "1"})
if err := w.ChooseCandidate(context.Background(), 1, 999); err == nil {
t.Error("чужой кандидат должен отклоняться")
}
}
func TestSetProviderID(t *testing.T) {
w, st := reviewWithCandidate(t, store.MetadataCandidate{Provider: "tvdb", ProviderID: "1"})
if err := w.SetProviderID(context.Background(), 1, "TMDB", " 603 "); err != nil {
t.Fatalf("SetProviderID: %v", err)
}
if st.overrides[1][ovrProvider] != "tmdb" || st.overrides[1][ovrProviderID] != "603" {
t.Errorf("overrides = %v", st.overrides[1])
}
if err := w.SetProviderID(context.Background(), 1, "kinopoisk", "1"); err == nil {
t.Error("недопустимый провайдер должен отклоняться")
}
if err := w.SetProviderID(context.Background(), 1, "tmdb", ""); err == nil {
t.Error("пустой id должен отклоняться")
}
}
func TestClearProvider(t *testing.T) {
w, st := reviewWithCandidate(t, store.MetadataCandidate{Provider: "tvdb", ProviderID: "1"})
_ = st.SetOverride(context.Background(), 1, ovrProvider, "tvdb")
if err := w.ClearProvider(context.Background(), 1); err != nil {
t.Fatalf("ClearProvider: %v", err)
}
if st.overrides[1][ovrProvider] != "none" {
t.Errorf("provider override = %q, want none", st.overrides[1][ovrProvider])
}
// «Без базы» → пустой тег.
_, tag, _ := w.effectivePlan(context.Background(), 1)
if tag != "" {
t.Errorf("tag = %q, want empty", tag)
}
}
func TestReviewData_IncludesCandidates(t *testing.T) {
w, st := reviewWithCandidate(t, store.MetadataCandidate{
Provider: "tvdb", ProviderID: "269613", Title: store.NullString("Fargo"),
})
candID := st.candidates[0].ID
if err := w.ChooseCandidate(context.Background(), 1, candID); err != nil {
t.Fatal(err)
}
rd, err := w.ReviewData(context.Background(), 1)
if err != nil {
t.Fatalf("ReviewData: %v", err)
}
if len(rd.Candidates) != 1 {
t.Fatalf("candidates = %d", len(rd.Candidates))
}
if rd.Provider != "tvdb" || rd.ProviderID != "269613" {
t.Errorf("eff provider = %s/%s", rd.Provider, rd.ProviderID)
}
if rd.Plan.Title != "Fargo" {
t.Errorf("plan title = %q", rd.Plan.Title)
}
}
func TestToLayoutPlan(t *testing.T) {
s, e := 1, 3
plan := recognize.Plan{