diff --git a/README.md b/README.md index fd7c1ca..77fa117 100644 --- a/README.md +++ b/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/` (относительно директории - промежуточных данных) \ No newline at end of file + промежуточных данных) diff --git a/main.py b/main.py index 79979e6..38ae47c 100644 --- a/main.py +++ b/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. diff --git a/prompts/structure_refiner.md b/prompts/structure_refiner.md new file mode 100644 index 0000000..1f115dc --- /dev/null +++ b/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. Если исправление не было однозначным, исходный вариант сохранён. diff --git a/structure_refiner.py b/structure_refiner.py new file mode 100644 index 0000000..77ee476 --- /dev/null +++ b/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)