115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
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)
|