264 lines
9.7 KiB
Go
264 lines
9.7 KiB
Go
// Package recognize по сигналам торрента определяет фильм/сериал, строит
|
|
// план раскладки и оценивает уверенность.
|
|
//
|
|
// Конвейер (см. docs/specs/recognition.md):
|
|
// 1. пред-парс имени релиза (go-ptn) — черновые название/год/сезон/серия;
|
|
// 2. вызов LLM со структурированным выводом → план в нашей схеме;
|
|
// 3. сверка с базами метаданных (TMDB/TVDB, опц.) — единичный сильный матч
|
|
// по названию+году даёт официальный id и каноническое имя;
|
|
// 4. решение «авто или review»: авто только при подтверждённом матче,
|
|
// чистой структурной валидации (для сериала — число серий бьётся с
|
|
// базой), согласованности с пред-парсом и уверенности не ниже порога.
|
|
//
|
|
// Без включённых баз (или без матча) авто-раскладка не делается — задача
|
|
// уходит в review. Выход LLM недоверенный: план принимается только если
|
|
// каждый files[].src совпадает с реальным файлом торрента; итоговая
|
|
// безопасность пути держится на раскладке (layout).
|
|
package recognize
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log/slog"
|
|
|
|
"git.vakhrushev.me/av/jellybit/internal/llm"
|
|
"git.vakhrushev.me/av/jellybit/internal/metadata"
|
|
)
|
|
|
|
// MediaType — вид контента.
|
|
type MediaType string
|
|
|
|
const (
|
|
MediaMovie MediaType = "movie"
|
|
MediaSeries MediaType = "series"
|
|
)
|
|
|
|
// FileRole — роль файла в раздаче.
|
|
type FileRole string
|
|
|
|
const (
|
|
RoleMain FileRole = "main" // основной видеофайл фильма
|
|
RoleEpisode FileRole = "episode" // серия сериала
|
|
RoleSubtitle FileRole = "subtitle" // внешние субтитры
|
|
RoleExtra FileRole = "extra" // допматериалы
|
|
RoleSample FileRole = "sample" // семпл
|
|
RoleIgnore FileRole = "ignore" // мусор/не нужное
|
|
)
|
|
|
|
func (r FileRole) valid() bool {
|
|
switch r {
|
|
case RoleMain, RoleEpisode, RoleSubtitle, RoleExtra, RoleSample, RoleIgnore:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
// File — входной файл торрента (путь относительно save_path и размер).
|
|
type File struct {
|
|
Path string
|
|
Size int64
|
|
}
|
|
|
|
// Input — сигналы для распознавания одной раздачи.
|
|
type Input struct {
|
|
Name string // имя торрента
|
|
Files []File // список файлов с размерами
|
|
Context string // текстовый контекст человека (опц.)
|
|
Hints []string // накопленные подсказки из review (Ф3; в Ф2 обычно пусто)
|
|
}
|
|
|
|
// PlanFile — файл в плане раскладки. Season/Episode заданы на файле, чтобы
|
|
// выражать мультисезонные паки и спецвыпуски (см. recognition.md).
|
|
type PlanFile struct {
|
|
Src string `json:"src"`
|
|
Role FileRole `json:"role"`
|
|
Season *int `json:"season,omitempty"`
|
|
Episode *int `json:"episode,omitempty"`
|
|
}
|
|
|
|
// Plan — структурированный результат распознавания (схема ответа LLM).
|
|
type Plan struct {
|
|
Type MediaType `json:"type"`
|
|
Title string `json:"title"`
|
|
OriginalTitle string `json:"original_title,omitempty"`
|
|
Year int `json:"year,omitempty"`
|
|
ProviderHint string `json:"provider_hint,omitempty"`
|
|
Files []PlanFile `json:"files"`
|
|
Confidence float64 `json:"confidence"`
|
|
Notes string `json:"notes,omitempty"`
|
|
}
|
|
|
|
// PreParse — черновой разбор имени релиза (go-ptn).
|
|
type PreParse struct {
|
|
Title string
|
|
Year int
|
|
Season int
|
|
Episode int
|
|
Quality string
|
|
}
|
|
|
|
// Decision — решение модели уверенности.
|
|
type Decision struct {
|
|
Auto bool // авто-раскладка без review (в Ф2 всегда false)
|
|
Reasons []string // причины ухода в review / предупреждения валидации
|
|
}
|
|
|
|
// Match — подтверждение распознавания базой метаданных.
|
|
type Match struct {
|
|
Provider string // "tmdb" | "tvdb"
|
|
ProviderID string // официальный id
|
|
Title string // каноническое название
|
|
Year int // каноничный год
|
|
SeasonEpisodeCounts map[int]int // число серий по сезонам (для сериала)
|
|
}
|
|
|
|
// Result — итог распознавания.
|
|
type Result struct {
|
|
Plan Plan
|
|
PreParse PreParse
|
|
Decision Decision
|
|
Match *Match // подтверждённый единичный матч (nil — нет)
|
|
Candidates []metadata.Candidate // кандидаты базы для ручного выбора в review
|
|
Attempts int // сколько вызовов LLM понадобилось (вкл. ретраи)
|
|
Raw string // сырой ответ LLM последней попытки
|
|
}
|
|
|
|
// LLM — нужная recognize часть провайдера.
|
|
type LLM interface {
|
|
Complete(ctx context.Context, req llm.Request) (llm.Response, error)
|
|
}
|
|
|
|
// Config — параметры распознавания.
|
|
type Config struct {
|
|
MaxRetries int // переразбор ответа со схемой-в-промпте ([llm].max_retries)
|
|
MaxTokens int // лимит ответа модели (0 — дефолт)
|
|
MaxFiles int // усечение списка файлов в промпте (0 — дефолт)
|
|
AutoThreshold float64 // порог уверенности для авто (0 — дефолт 0.85)
|
|
}
|
|
|
|
const (
|
|
defaultMaxTokens = 4000
|
|
defaultMaxFiles = 100
|
|
defaultAutoThreshold = 0.85
|
|
)
|
|
|
|
// Recognizer — реализация распознавания.
|
|
type Recognizer struct {
|
|
llm LLM
|
|
providers []metadata.Provider
|
|
maxRetry int
|
|
maxTokens int
|
|
maxFiles int
|
|
threshold float64
|
|
log *slog.Logger
|
|
}
|
|
|
|
// New собирает распознаватель. providers — включённые базы метаданных
|
|
// (пусто → сверки нет, авто-раскладка не делается).
|
|
func New(provider LLM, providers []metadata.Provider, cfg Config, log *slog.Logger) *Recognizer {
|
|
maxTokens := cfg.MaxTokens
|
|
if maxTokens <= 0 {
|
|
maxTokens = defaultMaxTokens
|
|
}
|
|
maxFiles := cfg.MaxFiles
|
|
if maxFiles <= 0 {
|
|
maxFiles = defaultMaxFiles
|
|
}
|
|
retries := cfg.MaxRetries
|
|
if retries < 0 {
|
|
retries = 0
|
|
}
|
|
threshold := cfg.AutoThreshold
|
|
if threshold <= 0 {
|
|
threshold = defaultAutoThreshold
|
|
}
|
|
return &Recognizer{
|
|
llm: provider,
|
|
providers: providers,
|
|
maxRetry: retries,
|
|
maxTokens: maxTokens,
|
|
maxFiles: maxFiles,
|
|
threshold: threshold,
|
|
log: log,
|
|
}
|
|
}
|
|
|
|
// Recognize прогоняет конвейер. Транспортная ошибка LLM возвращается как
|
|
// error (наверху решат retry/failed). Неразобранный после ретраев ответ —
|
|
// не ошибка, а Result с решением review (см. recognition.md).
|
|
func (r *Recognizer) Recognize(ctx context.Context, in Input) (Result, error) {
|
|
pre := preParse(in.Name)
|
|
msgs := buildMessages(in, pre, r.maxFiles)
|
|
|
|
temp := 0.0
|
|
var raw string
|
|
var plan Plan
|
|
var parseErr error
|
|
attempts := 0
|
|
|
|
for attempt := 0; attempt <= r.maxRetry; attempt++ {
|
|
attempts++
|
|
resp, err := r.llm.Complete(ctx, llm.Request{
|
|
Messages: msgs,
|
|
JSONMode: true,
|
|
Temperature: &temp,
|
|
MaxTokens: r.maxTokens,
|
|
})
|
|
if err != nil {
|
|
return Result{}, fmt.Errorf("recognize: llm complete: %w", err)
|
|
}
|
|
raw = resp.Content
|
|
|
|
plan, parseErr = parsePlan(raw, in)
|
|
if parseErr == nil {
|
|
break
|
|
}
|
|
r.log.Warn("recognize: unparsed llm response",
|
|
"attempt", attempts, "err", parseErr)
|
|
// Просим модель исправиться, повторяя схему и ошибку.
|
|
msgs = append(msgs,
|
|
llm.Message{Role: llm.RoleAssistant, Content: raw},
|
|
llm.Message{Role: llm.RoleUser, Content: correctionMessage(parseErr, in, r.maxFiles)})
|
|
}
|
|
|
|
if parseErr != nil {
|
|
return Result{
|
|
PreParse: pre,
|
|
Attempts: attempts,
|
|
Raw: raw,
|
|
Decision: Decision{
|
|
Auto: false,
|
|
Reasons: []string{"ответ LLM не разобран после " + itoa(attempts) + " попыток: " + parseErr.Error()},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// Сверка с базой: подтверждаем id + каноническое имя; при матче имя/год
|
|
// в плане заменяем на каноничные. Кандидаты копим для ручного выбора в
|
|
// review, когда единичного сильного матча нет.
|
|
match, candidates := r.matchMetadata(ctx, plan)
|
|
if match != nil {
|
|
plan.Title = match.Title
|
|
if match.Year != 0 {
|
|
plan.Year = match.Year
|
|
}
|
|
}
|
|
|
|
dec := decide(plan, pre, match, len(r.providers) > 0, r.threshold)
|
|
r.log.Info("recognize: done",
|
|
"type", plan.Type, "title", plan.Title, "year", plan.Year,
|
|
"files", len(plan.Files), "attempts", attempts,
|
|
"matched", match != nil, "candidates", len(candidates),
|
|
"auto", dec.Auto, "reasons", len(dec.Reasons))
|
|
return Result{
|
|
Plan: plan,
|
|
PreParse: pre,
|
|
Decision: dec,
|
|
Match: match,
|
|
Candidates: candidates,
|
|
Attempts: attempts,
|
|
Raw: raw,
|
|
}, nil
|
|
}
|