Added structure refiner (codex)

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.