Update src/orchestrator.py

src/orchestrator.py

@@ -1,16 +1,22 @@
 from dataclasses import dataclass, asdict
 from typing import Dict, Any, Optional
 
-from .config import RunConfig
-from .agents import AnalyzerAgent, RefactorAgent, CriticAgent, AgentResult
+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
-    image_attached: bool
+    diagram_context: str
     analyzer: AgentResult
     refactor: AgentResult
     critic: AgentResult
@@ -19,7 +25,7 @@ class RunTrace:
         return {
             "task": self.task,
             "input_requirements": self.input_requirements,
-            "image_attached": self.image_attached,
+            "diagram_context": self.diagram_context,
             "analyzer": asdict(self.analyzer),
             "refactor": asdict(self.refactor),
             "critic": asdict(self.critic),
@@ -35,22 +41,48 @@ class Orchestrator:
         self.refactor = RefactorAgent(cfg.refactor)
         self.critic = CriticAgent(cfg.critic)
 
-    def run(self, requirements_text: str, image_path: Optional[str] = None) -> RunTrace:
-        has_image = bool(image_path)
+        # Dedicated models for diagram extraction (free/open on HF Inference API)
+        # OCR: TrOCR, Caption: BLIP
+        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, has_image, self.ctx)
-        r1 = self.analyzer.run(p1, image_path=image_path)
+        p1 = build_analyzer_prompt(requirements_text, diagram_context, self.ctx)
+        r1 = self.analyzer.run(p1)
 
-        p2 = build_refactor_prompt(requirements_text, r1.output, has_image, self.ctx)
-        r2 = self.refactor.run(p2, image_path=image_path)
+        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, has_image, self.ctx)
-        r3 = self.critic.run(p3, image_path=image_path)
+        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,
-            image_attached=has_image,
+            diagram_context=diagram_context,
             analyzer=r1,
             refactor=r2,
             critic=r3,