2026-linux-sumka/video_references.py

import json
import os
import traceback
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator

from agent import Agent, AgentMessage
from utils import Timeline
from windowizer import Window, Windowizer


class _StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)


class VideoReferenceEvent(_StrictModel):
    id: int
    timestamp: float
    duration: float
    text: str


class UnresolvedReference(_StrictModel):
    ids: list[int] = Field(min_length=1)
    type: Literal["vis", "ocr"]
    text: str = Field(min_length=1)

    @model_validator(mode="after")
    def validate_ids(self) -> "UnresolvedReference":
        if len(self.ids) != len(set(self.ids)):
            raise ValueError("ids must not contain duplicates")
        return self


class UnresolvedReferences(_StrictModel):
    unresolved: list[UnresolvedReference]


class _ReferenceContext(_StrictModel):
    topic: str | None = Field(default=None, max_length=120)
    recent_references: list[str] = Field(default_factory=list, max_length=6)

    @model_validator(mode="after")
    def validate_recent_references(self) -> "_ReferenceContext":
        if any(
            not reference.strip() or len(reference) > 160
            for reference in self.recent_references
        ):
            raise ValueError(
                "recent_references must contain non-empty strings up to 160 characters"
            )
        return self


class _WindowResult(_StrictModel):
    unresolved: list[UnresolvedReference]
    context: _ReferenceContext


class VideoReferenceBuilder:
    """Find timeline regions that require visual or OCR processing."""

    DEBUG_ID = 0
    MAX_RETRIES = 5

    def __init__(
        self,
        agent: Agent,
        windowizer: Windowizer[VideoReferenceEvent],
    ) -> None:
        self._agent = agent
        self._windowizer = windowizer
        with open("prompts/video_references.md", "r", encoding="utf-8") as f:
            self._system_prompt = AgentMessage(content=f.read(), role="system")

    @staticmethod
    def _validate_result(
        result: _WindowResult,
        window: Window[VideoReferenceEvent],
    ) -> None:
        present_ids = [event.id for event in window.present]
        present_positions = {
            event_id: position for position, event_id in enumerate(present_ids)
        }
        seen: set[tuple[tuple[int, ...], str]] = set()

        for reference in result.unresolved:
            if any(event_id not in present_positions for event_id in reference.ids):
                raise ValueError("Reference ids must come from present")
            positions = [present_positions[event_id] for event_id in reference.ids]
            if positions != sorted(positions):
                raise ValueError("Reference ids must be in chronological order")

            key = (tuple(reference.ids), reference.type)
            if key in seen:
                raise ValueError("Duplicate reference in a single window")
            seen.add(key)

    def _build_window(
        self,
        window: Window[VideoReferenceEvent],
    ) -> _WindowResult:
        messages = [
            self._system_prompt,
            AgentMessage(
                content=json.dumps(
                    window.model_dump(mode="json"),
                    indent=2,
                    ensure_ascii=False,
                ),
                role="user",
            ),
        ]

        debug_dir = None
        if os.path.isdir("debug"):
            debug_dir = f"debug/VideoReferenceBuilder/{self.DEBUG_ID}"
            VideoReferenceBuilder.DEBUG_ID += 1
            os.makedirs(debug_dir, exist_ok=True)
            with open(f"{debug_dir}/request.txt", "w", encoding="utf-8") as f:
                f.write(messages[1].content)

        retries_left = self.MAX_RETRIES
        while retries_left > 0:
            retries_left -= 1
            response = self._agent.completion(
                messages=messages,
                response_format=_WindowResult,
            )
            if debug_dir:
                with open(
                    f"{debug_dir}/{retries_left}-retries-left.txt",
                    "w",
                    encoding="utf-8",
                ) as f:
                    f.write(response)

            try:
                result = _WindowResult.model_validate_json(response)
                self._validate_result(result, window)
                return result
            except (ValidationError, ValueError, TypeError):
                traceback.print_exc()

        raise RuntimeError("Agent has failed to provide valid schema too many times")

    def build(self, timeline: Timeline) -> UnresolvedReferences:
        events: list[VideoReferenceEvent] = []
        seen_ids: set[int] = set()
        for event in timeline.events:
            if event.type != "voice":
                raise ValueError("Video references must be built from voice events")
            if event.id in seen_ids:
                raise ValueError(f"Duplicate event ID: {event.id}")
            seen_ids.add(event.id)
            events.append(
                VideoReferenceEvent(
                    id=event.id,
                    timestamp=event.timestamp,
                    duration=event.duration,
                    text=event.text,
                )
            )

        unresolved: list[UnresolvedReference] = []
        context = _ReferenceContext()
        for window in self._windowizer.windowize(events):
            window.context = context.model_dump(mode="json")
            result = self._build_window(window)
            unresolved.extend(result.unresolved)
            context = result.context

        return UnresolvedReferences(unresolved=unresolved)