Added video references search (codex)

README.md

@@ -312,6 +312,9 @@ Telegram ботом, которому отправили видео).
           следующем вызове.
     - После выполнения этого шага модель должна сформировать файл
       `unresolved.json`
+    - Модель и адрес API можно задать через `--video-ref-ai-model`,
+      `--video-ref-ai-base-url` и `--video-ref-ai-api-key`. По умолчанию
+      используется `google/gemini-3.1-flash-lite`.
 6. **Разрешить ссылки на видео**
     - После этого шага должен быть получен файл `events.json`, используя который
       будет построена структура итогового документа.

main.py

@@ -14,6 +14,7 @@ import torch
 from asr import Asr, AsrRawResult
 from asr_filter import AsrFilter, AsrFilterResult
 from asr_eventizer import AsrEventizer
+from video_references import VideoReferenceBuilder
 from structure_builder import Structure, StructureBuilder
 from structure_refiner import StructureRefiner
 from markdown_builder import MarkdownBuilder
@@ -84,6 +85,21 @@ def setup_arguments() -> argparse.Namespace:
         type=str,
         default=ai_api_key
     )
+    parser.add_argument(
+        "--video-ref-ai-model",
+        type=str,
+        default="google/gemini-3.1-flash-lite"
+    )
+    parser.add_argument(
+        "--video-ref-ai-base-url",
+        type=str,
+        default="https://api.proxyapi.ru/v1"
+    )
+    parser.add_argument(
+        "--video-ref-ai-api-key",
+        type=str,
+        default=ai_api_key
+    )
     parser.add_argument(
         "--structure-ai-model",
         type=str,
@@ -215,9 +231,19 @@ def on_video_references(current_step: Step, input_data: dict | None) -> tuple[St
     if not os.path.isfile("video.mp4"):
         logging.info("No video, skipping")
         return (Step.REFERENCE_RESOLVER, None)
-    # NOT IMPLEMENTED
-    logging.warning("Video references are not implemented yet")
-    return (Step.REFERENCE_RESOLVER, None)
+    if input_data is None:
+        logging.error("Can't build video references without input_data")
+        return (None, None)
+    logging.info("Creating the agent")
+    agent = Agent(
+        model=ARGS.video_ref_ai_model,
+        base_url=ARGS.video_ref_ai_base_url,
+        api_key=ARGS.video_ref_ai_api_key
+    )
+    logging.info("Building video references...")
+    builder = VideoReferenceBuilder(agent, Windowizer())
+    result = builder.build(Timeline(**input_data))
+    return (Step.REFERENCE_RESOLVER, result.model_dump(mode="json"))
 
 def on_reference_resolver(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
     # do not filter if output file exists
@@ -230,7 +256,7 @@ def on_reference_resolver(current_step: Step, input_data: dict | None) -> tuple[
         return (Step.STRUCTURE_BUILDER, input_data)
     # NOT IMPLEMENTED
     logging.warning("Reference resolver is not implemented yet")
-    return (Step.STRUCTURE_BUILDER, input_data)
+    return (None, None)
 
 def on_structure_builder(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
     # don't if done

prompts/video_references.md

@@ -0,0 +1,85 @@
+Ты определяешь, какие фрагменты видеолекции нужно дополнительно обработать как
+изображение или при помощи OCR.
+
+Ты НЕ анализируешь видео и НЕ создаёшь конспект. Ты только формируешь точные
+запросы для следующего этапа программы.
+
+На вход поступает JSON с полями `past`, `present`, `future` и `context`.
+Каждое событие содержит `id`, `timestamp`, `duration` и `text`.
+
+Обрабатывай только события из `present`. `past` и `future` используй только для
+понимания контекста. Никогда не добавляй их ID в результат.
+
+Создавай ссылку типа `vis`, когда изображение существенно помогает понять
+материал, например преподаватель обсуждает:
+
+- схему, график, диаграмму, чертёж или рисунок;
+- таблицу, расположение элементов или пространственную связь;
+- показанный объект или визуальную демонстрацию;
+- содержимое экрана, которое трудно полноценно передать словами.
+
+Создавай ссылку типа `ocr`, когда для конспекта нужен точный текст с экрана:
+
+- формула или математическое обозначение;
+- определение, заголовок, список или таблица;
+- фрагмент кода;
+- термин, имя или подпись, произнесённые неоднозначно.
+
+Для одного диапазона разрешается создать одновременно `vis` и `ocr`, если
+нужны и изображение, и точный текст.
+
+Не создавай ссылку:
+
+- для обычного изображения преподавателя;
+- для декоративного слайда;
+- если речь уже полностью и однозначно передаёт материал;
+- только из-за слов «запишите», «смотрите» или «на экране», если визуальный
+  источник не добавляет полезной информации;
+- на основании догадки о том, чего может не быть в видео.
+
+`ids` должны содержать только ID событий из `present`, во время которых нужно
+искать содержимое. ID указывай в хронологическом порядке и без повторов.
+
+`text` — короткая конкретная инструкция следующему этапу:
+
+- для `vis` опиши, какое полезное изображение нужно найти;
+- для `ocr` назови, какой именно текст или формулу нужно распознать;
+- не утверждай, что содержимое точно присутствует в кадре;
+- не придумывай внешний вид или текст.
+
+Используй контекст строго в таком формате:
+
+{
+  "topic": null,
+  "recent_references": []
+}
+
+`topic` — текущая тема, строка до 120 символов или `null`.
+
+`recent_references` — не более шести коротких описаний последних запросов. Они
+нужны только для предотвращения бессмысленных повторов.
+
+Верни ровно один валидный JSON:
+
+{
+  "unresolved": [
+    {
+      "ids": [1, 2],
+      "type": "vis",
+      "text": "Найти схему, которую преподаватель использует для объяснения устройства системы."
+    },
+    {
+      "ids": [3],
+      "type": "ocr",
+      "text": "Распознать точную запись формулы нормы рассогласования."
+    }
+  ],
+  "context": {
+    "topic": null,
+    "recent_references": []
+  }
+}
+
+Если полезные ссылки не нужны, верни пустой список `unresolved`.
+
+Не добавляй Markdown, комментарии или текст до и после JSON.

video_references.py

@@ -0,0 +1,174 @@
+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)