| from dataclasses import dataclass, asdict |
| from typing import Dict, Any, Optional |
|
|
| from .config import RunConfig, ModelSpec |
| from .agents import AnalyzerAgent, RefactorAgent, CriticAgent, AgentResult, BaseAgent |
| from .tasks import TaskContext, build_analyzer_prompt, build_refactor_prompt, build_critic_prompt |
| from .backends import HFInferenceAPIBackend |
|
|
| try: |
| from PIL import Image |
| except Exception: |
| Image = None |
|
|
|
|
| @dataclass |
| class RunTrace: |
| task: str |
| input_requirements: str |
| diagram_context: str |
| analyzer: AgentResult |
| refactor: AgentResult |
| critic: AgentResult |
|
|
| def to_dict(self) -> Dict[str, Any]: |
| return { |
| "task": self.task, |
| "input_requirements": self.input_requirements, |
| "diagram_context": self.diagram_context, |
| "analyzer": asdict(self.analyzer), |
| "refactor": asdict(self.refactor), |
| "critic": asdict(self.critic), |
| } |
|
|
|
|
| class Orchestrator: |
| def __init__(self, cfg: RunConfig, ctx: Optional[TaskContext] = None): |
| self.cfg = cfg |
| self.ctx = ctx or TaskContext() |
|
|
| self.analyzer = AnalyzerAgent(cfg.analyzer) |
| self.refactor = RefactorAgent(cfg.refactor) |
| self.critic = CriticAgent(cfg.critic) |
|
|
| |
| |
| self.ocr_backend = HFInferenceAPIBackend("microsoft/trocr-base-printed") |
| self.caption_backend = HFInferenceAPIBackend("Salesforce/blip-image-captioning-large") |
|
|
| def _extract_diagram_context(self, image: Optional["Image.Image"]) -> str: |
| if image is None: |
| return "" |
|
|
| parts = [] |
| try: |
| ocr = self.ocr_backend.image_to_text(image) |
| if ocr and ocr.strip(): |
| parts.append("OCR (текст на изображении):\n" + ocr.strip()) |
| except Exception as e: |
| parts.append(f"OCR: ошибка ({type(e).__name__})") |
|
|
| try: |
| cap = self.caption_backend.image_to_text(image) |
| if cap and cap.strip(): |
| parts.append("Описание изображения:\n" + cap.strip()) |
| except Exception as e: |
| parts.append(f"Caption: ошибка ({type(e).__name__})") |
|
|
| return "\n\n".join(parts).strip() |
|
|
| def run(self, requirements_text: str, image: Optional["Image.Image"] = None) -> RunTrace: |
| diagram_context = self._extract_diagram_context(image) |
|
|
| p1 = build_analyzer_prompt(requirements_text, diagram_context, self.ctx) |
| r1 = self.analyzer.run(p1) |
|
|
| p2 = build_refactor_prompt(requirements_text, r1.output, diagram_context, self.ctx) |
| r2 = self.refactor.run(p2) |
|
|
| p3 = build_critic_prompt(requirements_text, r2.output, diagram_context, self.ctx) |
| r3 = self.critic.run(p3) |
|
|
| return RunTrace( |
| task=self.ctx.name, |
| input_requirements=requirements_text, |
| diagram_context=diagram_context, |
| analyzer=r1, |
| refactor=r2, |
| critic=r3, |
| ) |
|
|
