Добавил выбор из кандидатов, если 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
+164 -9
View File
@@ -7,9 +7,11 @@ import (
"errors"
"fmt"
"path/filepath"
"strconv"
"strings"
"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"
@@ -19,6 +21,10 @@ import (
const (
ovrMediaType = "media_type"
ovrIgnoredFiles = "ignored_files"
ovrProvider = "provider" // выбранная база ("none" = без базы)
ovrProviderID = "provider_id" // id в выбранной базе
ovrTitle = "title" // запиненное каноническое название
ovrYear = "year" // запиненный год
)
// recognizePending распознаёт завершённые загрузки и перезапускает те, что
@@ -158,10 +164,17 @@ func (w *Worker) finishRecognition(ctx context.Context, id int64, res recognize.
"download_id", id, "state", d.State)
return
}
if _, err := w.store.CreateRecognition(ctx, rec, res.Decision.Reasons); err != nil {
recID, err := w.store.CreateRecognition(ctx, rec, res.Decision.Reasons)
if err != nil {
w.log.Error("recognize: persist", "download_id", id, "err", err)
return
}
// Кандидаты базы — для ручного выбора в review.
if cands := toStoreCandidates(recID, res.Candidates); len(cands) > 0 {
if err := w.store.CreateCandidates(ctx, cands); err != nil {
w.log.Warn("recognize: persist candidates", "download_id", id, "err", err)
}
}
// Авто-раскладка при подтверждённом матче и чистой валидации (Ф4);
// иначе — review. Раскладчик может быть не сконфигурирован.
@@ -413,14 +426,105 @@ func (w *Worker) requireReviewable(ctx context.Context, id int64, op string) (*s
return d, nil
}
// --- Выбор базы метаданных (пиннинг; остаёмся в review, применяет человек) ---
// ChooseCandidate пиннит выбранного кандидата базы как override (провайдер,
// id, каноническое имя/год). Раскладку не запускает — превью обновится, а
// человек подтвердит «Применить».
func (w *Worker) ChooseCandidate(ctx context.Context, id, candidateID int64) error {
w.mu.Lock()
defer w.mu.Unlock()
if _, err := w.requireReviewable(ctx, id, "choose candidate"); err != nil {
return err
}
rec, err := w.store.GetCurrentRecognition(ctx, id)
if err != nil {
return fmt.Errorf("choose candidate: %w", err)
}
cand, err := w.store.GetCandidate(ctx, candidateID)
if err != nil {
return fmt.Errorf("choose candidate: %w", err)
}
if rec == nil || cand == nil || cand.RecognitionID != rec.ID {
return fmt.Errorf("choose candidate: кандидат %d не относится к текущему распознаванию", candidateID)
}
pins := map[string]string{ovrProvider: cand.Provider, ovrProviderID: cand.ProviderID}
if cand.Title.Valid && cand.Title.String != "" {
pins[ovrTitle] = cand.Title.String
}
if cand.Year.Valid {
pins[ovrYear] = strconv.FormatInt(cand.Year.Int64, 10)
}
for field, value := range pins {
if err := w.store.SetOverride(ctx, id, field, value); err != nil {
return fmt.Errorf("choose candidate: %w", err)
}
}
if err := w.store.SetCandidateChosen(ctx, rec.ID, candidateID); err != nil {
return fmt.Errorf("choose candidate: %w", err)
}
w.log.Info("review: candidate chosen",
"download_id", id, "provider", cand.Provider, "provider_id", cand.ProviderID)
return nil
}
// SetProviderID пиннит провайдера и id вручную (без выбора из списка).
func (w *Worker) SetProviderID(ctx context.Context, id int64, provider, providerID string) error {
provider = strings.TrimSpace(strings.ToLower(provider))
providerID = strings.TrimSpace(providerID)
switch provider {
case "tmdb", "tvdb", "imdb":
default:
return fmt.Errorf("set provider: недопустимый провайдер %q (tmdb/tvdb/imdb)", provider)
}
if providerID == "" {
return fmt.Errorf("set provider: пустой id")
}
w.mu.Lock()
defer w.mu.Unlock()
if _, err := w.requireReviewable(ctx, id, "set provider"); err != nil {
return err
}
if err := w.store.SetOverride(ctx, id, ovrProvider, provider); err != nil {
return fmt.Errorf("set provider: %w", err)
}
if err := w.store.SetOverride(ctx, id, ovrProviderID, providerID); err != nil {
return fmt.Errorf("set provider: %w", err)
}
return nil
}
// ClearProvider — «без базы»: снимает матч (тег папки не ставится).
func (w *Worker) ClearProvider(ctx context.Context, id int64) error {
w.mu.Lock()
defer w.mu.Unlock()
if _, err := w.requireReviewable(ctx, id, "clear provider"); err != nil {
return err
}
if err := w.store.SetOverride(ctx, id, ovrProvider, "none"); err != nil {
return fmt.Errorf("clear provider: %w", err)
}
if err := w.store.SetOverride(ctx, id, ovrProviderID, ""); err != nil {
return fmt.Errorf("clear provider: %w", err)
}
return nil
}
// --- Данные для экрана ревью ---
// ReviewData — всё, что нужно транспорту для отрисовки ревью.
type ReviewData struct {
Download store.Download
Recognition *store.Recognition
Plan recognize.Plan // эффективный (с применёнными правками)
Preview []layout.Link // целевые пути (Src — относительный, для показа)
Plan recognize.Plan // эффективный (с применёнными правками)
Preview []layout.Link // целевые пути (Src — относительный, для показа)
Candidates []store.MetadataCandidate // кандидаты базы для ручного выбора
Provider string // эффективный провайдер (с учётом выбора)
ProviderID string // эффективный id в базе
Hints []string
Overrides map[string]string
}
@@ -444,7 +548,16 @@ func (w *Worker) ReviewData(ctx context.Context, id int64) (*ReviewData, error)
return nil, fmt.Errorf("review data: %w", err)
}
rd := &ReviewData{Download: *d, Recognition: rec, Hints: hints, Overrides: overrides}
prov, pid := effectiveProvider(rec, overrides)
rd := &ReviewData{
Download: *d, Recognition: rec, Hints: hints, Overrides: overrides,
Provider: prov, ProviderID: pid,
}
if rec != nil {
if cands, cerr := w.store.ListCandidatesByRecognition(ctx, rec.ID); cerr == nil {
rd.Candidates = cands
}
}
if rec != nil && rec.Plan.Valid {
var plan recognize.Plan
if err := json.Unmarshal([]byte(rec.Plan.String), &plan); err == nil {
@@ -453,7 +566,7 @@ func (w *Worker) ReviewData(ctx context.Context, id int64) (*ReviewData, error)
// Превью строим по относительным путям с provider-тегом; ошибку
// игнорируем — просто покажем причины без превью.
if w.layouter != nil {
tag := providerTag(rec.Provider.String, rec.ProviderID.String)
tag := providerTag(prov, pid)
if links, lerr := w.layouter.BuildLinks(toLayoutPlan(plan, "", tag)); lerr == nil {
rd.Preview = links
}
@@ -481,18 +594,27 @@ func (w *Worker) effectivePlan(ctx context.Context, id int64) (recognize.Plan, s
if err != nil {
return recognize.Plan{}, "", err
}
tag := providerTag(rec.Provider.String, rec.ProviderID.String)
return applyOverrides(plan, overrides), tag, nil
prov, pid := effectiveProvider(rec, overrides)
return applyOverrides(plan, overrides), providerTag(prov, pid), nil
}
// --- Хелперы преобразования ---
// applyOverrides применяет ручные правки к плану: форсит тип и помечает
// игнорируемые файлы ролью ignore (их раскладка пропустит).
// applyOverrides применяет ручные правки к плану: форсит тип, каноническое
// имя/год (из выбранного кандидата базы) и помечает игнорируемые файлы ролью
// ignore (их раскладка пропустит).
func applyOverrides(plan recognize.Plan, overrides map[string]string) recognize.Plan {
if mt := overrides[ovrMediaType]; mt == string(recognize.MediaMovie) || mt == string(recognize.MediaSeries) {
plan.Type = recognize.MediaType(mt)
}
if t := overrides[ovrTitle]; t != "" {
plan.Title = t
}
if y := overrides[ovrYear]; y != "" {
if year, err := strconv.Atoi(y); err == nil {
plan.Year = year
}
}
ignored := parseIgnored(overrides[ovrIgnoredFiles])
if len(ignored) > 0 {
for i := range plan.Files {
@@ -504,6 +626,39 @@ func applyOverrides(plan recognize.Plan, overrides map[string]string) recognize.
return plan
}
// effectiveProvider возвращает провайдера и id для тега папки с учётом
// ручного выбора: запиненный override перекрывает распознанный матч.
// override "none" означает явный отказ от базы.
func effectiveProvider(rec *store.Recognition, overrides map[string]string) (provider, id string) {
if p, ok := overrides[ovrProvider]; ok {
return p, overrides[ovrProviderID]
}
if rec != nil {
return rec.Provider.String, rec.ProviderID.String
}
return "", ""
}
// toStoreCandidates переводит кандидатов распознавания в строки БД,
// подставляя тег-предпочтительный provider/id (внешний из TVMaze и т.п.).
func toStoreCandidates(recognitionID int64, cands []metadata.Candidate) []store.MetadataCandidate {
out := make([]store.MetadataCandidate, 0, len(cands))
for _, c := range cands {
prov, id := recognize.CandidateTag(c)
mc := store.MetadataCandidate{
RecognitionID: recognitionID,
Provider: prov,
ProviderID: id,
Title: store.NullString(c.Title),
}
if c.Year != 0 {
mc.Year = sql.NullInt64{Int64: int64(c.Year), Valid: true}
}
out = append(out, mc)
}
return out
}
// providerTag строит тег папки для Jellyfin из провайдера и id: "tmdbid-…"
// / "tvdbid-…". Пустой id (нет матча) → пустой тег.
func providerTag(provider, id string) string {
+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{
+6
View File
@@ -42,6 +42,12 @@ type Store interface {
LatestBatchID(ctx context.Context, downloadID int64) (string, error)
ListFileLinksByBatch(ctx context.Context, batchID string) ([]store.FileLink, error)
DeleteFileLinksByBatch(ctx context.Context, batchID string) error
// Кандидаты базы метаданных (ручной выбор в review).
CreateCandidates(ctx context.Context, cands []store.MetadataCandidate) error
ListCandidatesByRecognition(ctx context.Context, recognitionID int64) ([]store.MetadataCandidate, error)
GetCandidate(ctx context.Context, id int64) (*store.MetadataCandidate, error)
SetCandidateChosen(ctx context.Context, recognitionID, candidateID int64) error
}
// QBittorrent — нужная worker часть клиента qBittorrent.
+10
View File
@@ -83,6 +83,16 @@ func (f *fakeStore) ListFileLinksByBatch(_ context.Context, _ string) ([]store.F
return nil, nil
}
func (f *fakeStore) DeleteFileLinksByBatch(_ context.Context, _ string) error { return nil }
func (f *fakeStore) CreateCandidates(_ context.Context, _ []store.MetadataCandidate) error {
return nil
}
func (f *fakeStore) ListCandidatesByRecognition(_ context.Context, _ int64) ([]store.MetadataCandidate, error) {
return nil, nil
}
func (f *fakeStore) GetCandidate(_ context.Context, _ int64) (*store.MetadataCandidate, error) {
return nil, nil
}
func (f *fakeStore) SetCandidateChosen(_ context.Context, _, _ int64) error { return nil }
type fakeQbt struct {
torrents []qbt.Torrent