Added structure refiner (codex)

README.md

@@ -408,9 +408,15 @@ Telegram ботом, которому отправили видео).
       ```
     - Для получения файла `structure.json` модели передаются соыбытия из
       `events.json`, используя метод окна (описан ранее).
-8. **Собрать финальный Markdown файл**
+8. **Выполнить финальную редактуру структуры**
+    - Черновая структура из `structure.json` передаётся модели оконным
+      методом.
+    - Модель исправляет язык и иерархию, согласовывает терминологию и удаляет
+      явные повторы, не добавляя новые сведения.
+    - Результат сохраняется в `structure_refined.json` в том же формате.
+9. **Собрать финальный Markdown файл**
     - Для сборки Markdown файла используются элементы, полученные на предыдщем
-      шаге, и хранимые в файле `structure.json`
+      шаге, и хранимые в файле `structure_refined.json`
     - В итоге создаётся файл `output.md`, который может включать в себя ссылки
       на изображения из директории `images/` (относительно директории
-      промежуточных данных)
+      промежуточных данных)

main.py

@@ -14,7 +14,8 @@ import torch
 from asr import Asr, AsrRawResult
 from asr_filter import AsrFilter, AsrFilterResult
 from asr_eventizer import AsrEventizer
-from structure_builder import StructureBuilder
+from structure_builder import Structure, StructureBuilder
+from structure_refiner import StructureRefiner
 from windowizer import Windowizer
 
 from agent import Agent
@@ -30,6 +31,7 @@ class Step(Enum):
     VIDEO_REFERENCES = "video_references"
     REFERENCE_RESOLVER = "reference_resolver"
     STRUCTURE_BUILDER = "structure_builder"
+    STRUCTURE_REFINER = "structure_refiner"
     MARKDOWN_BUILDER = "markdown_builder"
 
 #
@@ -96,6 +98,21 @@ def setup_arguments() -> argparse.Namespace:
         type=str,
         default=ai_api_key
     )
+    parser.add_argument(
+        "--refiner-ai-model",
+        type=str,
+        default="google/gemini-3.5-flash"
+    )
+    parser.add_argument(
+        "--refiner-ai-base-url",
+        type=str,
+        default="https://api.proxyapi.ru/v1"
+    )
+    parser.add_argument(
+        "--refiner-ai-api-key",
+        type=str,
+        default=ai_api_key
+    )
     parser.add_argument("-v", action='store_true')
     return parser.parse_args()
 
@@ -219,7 +236,7 @@ def on_structure_builder(current_step: Step, input_data: dict | None) -> tuple[S
     if os.path.isfile(WORKFLOW_DATA[current_step][0]):
         logging.info("Skipping structure builder")
         with open(WORKFLOW_DATA[current_step][0], "rb") as f:
-            return (Step.MARKDOWN_BUILDER, json.load(f))
+            return (Step.STRUCTURE_REFINER, json.load(f))
     if input_data is None:
         logging.error("Can't build document structure without input_data")
         return (None, None)
@@ -234,6 +251,25 @@ def on_structure_builder(current_step: Step, input_data: dict | None) -> tuple[S
     logging.info("Building document structure...")
     builder = StructureBuilder(agent, windowizer)
     result = builder.build(Timeline(**input_data))
+    return (Step.STRUCTURE_REFINER, result.model_dump(mode="json"))
+
+def on_structure_refiner(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
+    if os.path.isfile(WORKFLOW_DATA[current_step][0]):
+        logging.info("Skipping structure refiner")
+        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
+            return (Step.MARKDOWN_BUILDER, json.load(f))
+    if input_data is None:
+        logging.error("Can't refine document structure without input_data")
+        return (None, None)
+    logging.info("Creating the agent")
+    agent = Agent(
+        model=ARGS.refiner_ai_model,
+        base_url=ARGS.refiner_ai_base_url,
+        api_key=ARGS.refiner_ai_api_key
+    )
+    logging.info("Refining document structure...")
+    refiner = StructureRefiner(agent, Windowizer())
+    result = refiner.refine(Structure(**input_data))
     return (Step.MARKDOWN_BUILDER, result.model_dump(mode="json"))
 
 #
@@ -247,6 +283,7 @@ WORKFLOW_DATA: dict[Step, tuple[str, Callable[[Step, dict | None], tuple[Step |
     Step.VIDEO_REFERENCES: ("unresolved.json", on_video_references),
     Step.REFERENCE_RESOLVER: ("events.json", on_reference_resolver),
     Step.STRUCTURE_BUILDER: ("structure.json", on_structure_builder),
+    Step.STRUCTURE_REFINER: ("structure_refined.json", on_structure_refiner),
     Step.MARKDOWN_BUILDER: ("output.md", None)
 }
 """Information about workflow.

prompts/structure_refiner.md

@@ -0,0 +1,191 @@
+Ты выполняешь консервативную финальную редактуру структуры учебного конспекта.
+
+Твоя главная задача — улучшить уже готовый документ, НЕ ПОТЕРЯВ ни одного
+содержательного момента. Ты не создаёшь новый конспект, не пересказываешь и не
+сокращаешь материал.
+
+==================================================
+ФОРМАТ ВХОДА
+==================================================
+
+На вход поступает JSON:
+
+{
+  "past": [...],
+  "present": [...],
+  "future": [...],
+  "context": {...}
+}
+
+`present` — элементы, которые необходимо отредактировать и вернуть.
+
+`past` и `future` — только контекст. Не возвращай их элементы и не переноси из
+них сведения в `present`.
+
+Исключение: если элемент из `present` полностью дублирует уже готовый элемент
+из `past`, дубликат из `present` можно удалить.
+
+==================================================
+ГЛАВНЫЙ ПРИНЦИП
+==================================================
+
+Лучше оставить элемент без изменений, чем случайно потерять или исказить
+информацию.
+
+Каждое определение, условие, ограничение, утверждение, пример, пояснение,
+перечисление, формула, обозначение, номер вопроса и смысловой акцент из
+`present` должны сохраниться в результате.
+
+Не делай текст короче только ради компактности.
+
+==================================================
+ЧТО РАЗРЕШЕНО
+==================================================
+
+- исправлять орфографию, пунктуацию и явные грамматические ошибки;
+- устранять явные смысловые повторы;
+- объединять дублирующие элементы из `present`, сохраняя ВСЕ различающиеся
+  подробности из каждого элемента;
+- разделять слишком большой элемент, если это не меняет смысл и не удаляет
+  сведения;
+- исправлять уровень заголовка, если нарушена иерархия;
+- приводить оформление уже существующей математики к корректному LaTeX;
+- делать обозначения типографически единообразными только тогда, когда это
+  одно и то же обозначение записано очевидно разным способом, например `FE` и
+  `$F_E$` в пределах одной однозначной формулировки.
+
+==================================================
+МАТЕМАТИКА
+==================================================
+
+Ты не проверяешь математическую истинность и не восстанавливаешь формулы по
+догадке. У тебя нет первичного изображения доски или слайда.
+
+При работе с математикой:
+
+- сохраняй все существующие формулы, индексы, черты, звёздочки, размерности и
+  условия;
+- разрешается исправлять только явные ошибки LaTeX и оформление;
+- не добавляй транспонирование, индексы, модули, нормы, равенства или знаки
+  операций, которых не было во входе;
+- не заменяй одно математическое обозначение другим;
+- не объявляй разные обозначения эквивалентными или синонимичными;
+- если во входе используются `$\bar{F}$` и `$F_X$`, сохрани оба обозначения в
+  тех контекстах, где они находились; не пиши «или», «то есть» либо знак
+  равенства между ними;
+- если обозначение выглядит неоднозначным или противоречивым, сохрани его без
+  смыслового исправления;
+- не исправляй формулу на основании собственных знаний.
+
+==================================================
+ЧТО ЗАПРЕЩЕНО
+==================================================
+
+- добавлять новые факты, выводы, определения, примеры или формулы;
+- уточнять исходный текст сведениями, которых в нём нет;
+- усиливать утверждение: например, заменять «определяется» на «фиксируется» или
+  «может» на «должно»;
+- удалять полезную подробность как «несущественную»;
+- объединять разные понятия в одно определение;
+- удалять номера и диапазоны учебных или экзаменационных вопросов;
+- менять порядок пунктов там, где он имеет значение;
+- превращать содержательный список в абзац;
+- превращать обычное объяснение в новое определение или важное утверждение;
+- создавать, изменять или удалять элементы `image`;
+- изменять `image.event_id`;
+- переносить содержание из `future` в готовые элементы;
+- повторно возвращать элементы из `past`.
+
+==================================================
+ЗАГОЛОВКИ И НУМЕРАЦИЯ
+==================================================
+
+Сохраняй текст заголовков, включая такие части, как:
+
+- «Вопрос 1.1»;
+- «Вопросы 1.3–1.6»;
+- номера разделов и подразделов;
+- названия тем.
+
+Разрешается изменить только `level`, если это необходимо для исправления
+иерархии. Не удаляй заголовок, если он обозначает отдельный вопрос, даже когда
+его тема похожа на соседний раздел.
+
+Одинаковые заголовки можно объединить только тогда, когда это действительно
+повтор одного раздела и при объединении не исчезает нумерация.
+
+==================================================
+СПИСКИ, ОПРЕДЕЛЕНИЯ И ВАЖНЫЕ МЕСТА
+==================================================
+
+- сохраняй все пункты списков и их порядок;
+- разные пункты не объединяй в один;
+- сохраняй `ordered`, если порядок имеет значение;
+- сохраняй `definition`, если входной элемент является определением;
+- сохраняй термин и полный текст определения;
+- сохраняй `important`, если входной элемент отмечен как важный;
+- не понижай такие элементы до обычного абзаца без очевидной причины.
+
+==================================================
+CONTEXT
+==================================================
+
+Используй строго следующую структуру:
+
+{
+  "current_section": null,
+  "current_subsection": null,
+  "recent_headings": []
+}
+
+`current_section` — последний актуальный заголовок уровня 2.
+
+`current_subsection` — последний актуальный заголовок уровня 3 или 4. При
+переходе к новому разделу значение можно сбросить в `null`.
+
+`recent_headings` — не более шести последних заголовков. Контекст используется
+только для согласованности между окнами и не является частью документа.
+
+==================================================
+OUTPUT
+==================================================
+
+Верни ровно один валидный JSON:
+
+{
+  "elements": [],
+  "context": {
+    "current_section": null,
+    "current_subsection": null,
+    "recent_headings": []
+  }
+}
+
+В `elements` разрешены только исходные типы:
+
+- `heading`;
+- `paragraph`;
+- `unordered`;
+- `ordered`;
+- `definition`;
+- `important`;
+- `image`.
+
+Не добавляй текст до или после JSON.
+
+==================================================
+ПРОВЕРКА ПЕРЕД ОТВЕТОМ
+==================================================
+
+Перед ответом проверь:
+
+1. Все содержательные сведения из `present` сохранены.
+2. Все формулы и математические обозначения сохранены без смысловых догадок.
+3. Разные обозначения не были объявлены эквивалентными.
+4. Номера вопросов и разделов не удалены.
+5. Ни один пункт списка не потерян и не превращён в часть длинного абзаца.
+6. Определения и важные утверждения сохранили свои типы и содержание.
+7. Элементы `image` сохранены без изменений.
+8. Сведения из `past` и `future` не добавлены в результат.
+9. Новые факты, формулы и уточнения не появились.
+10. Если исправление не было однозначным, исходный вариант сохранён.

structure_refiner.py

@@ -0,0 +1,114 @@
+import json
+import os
+import traceback
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator
+
+from agent import Agent, AgentMessage
+from structure_builder import ImageElement, Structure, StructureElement
+from windowizer import Window, Windowizer
+
+
+class _StrictModel(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
+
+class _RefinerContext(_StrictModel):
+    current_section: str | None = Field(default=None, max_length=120)
+    current_subsection: str | None = Field(default=None, max_length=120)
+    recent_headings: list[str] = Field(default_factory=list, max_length=6)
+
+    @model_validator(mode="after")
+    def validate_recent_headings(self) -> "_RefinerContext":
+        if any(not heading.strip() or len(heading) > 120 for heading in self.recent_headings):
+            raise ValueError("recent_headings must contain non-empty strings up to 120 characters")
+        return self
+
+
+class _RefineResult(_StrictModel):
+    elements: list[StructureElement]
+    context: _RefinerContext
+
+
+class StructureRefiner:
+    """Perform a final editing pass over a document structure."""
+
+    DEBUG_ID = 0
+    MAX_RETRIES = 5
+
+    def __init__(self, agent: Agent, windowizer: Windowizer[StructureElement]) -> None:
+        self._agent = agent
+        self._windowizer = windowizer
+        with open("prompts/structure_refiner.md", "r", encoding="utf-8") as f:
+            self._system_prompt = AgentMessage(content=f.read(), role="system")
+
+    @staticmethod
+    def _validate_result(
+        result: _RefineResult,
+        window: Window[StructureElement],
+    ) -> None:
+        input_image_ids = {
+            element.event_id
+            for element in window.present
+            if isinstance(element, ImageElement)
+        }
+        for element in result.elements:
+            if isinstance(element, ImageElement) and element.event_id not in input_image_ids:
+                raise ValueError("image.event_id must come from present")
+
+    def _refine_window(self, window: Window[StructureElement]) -> _RefineResult:
+        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/StructureRefiner/{StructureRefiner.DEBUG_ID}"
+            StructureRefiner.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=_RefineResult,
+            )
+            if debug_dir:
+                with open(
+                    f"{debug_dir}/{retries_left}-retries-left.txt",
+                    "w",
+                    encoding="utf-8",
+                ) as f:
+                    f.write(response)
+
+            try:
+                result = _RefineResult.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 refine(self, structure: Structure) -> Structure:
+        elements: list[StructureElement] = []
+        context = _RefinerContext()
+
+        for window in self._windowizer.windowize(structure.elements):
+            window.context = context.model_dump(mode="json")
+            intermediate = self._refine_window(window)
+            elements.extend(intermediate.elements)
+            context = intermediate.context
+
+        return Structure(elements=elements)