Update src/agents.py

src/agents.py

@@ -1,22 +1,22 @@
 import time
 from dataclasses import dataclass
-from typing import Dict, Any, Optional, Callable
+from typing import Dict, Any, Optional
 
 from .backends import LLMBackend, make_backend
 from .config import ModelSpec
-from .tasks import TaskContext
 
 
 @dataclass
 class AgentResult:
     role: str
+    provider: str
     model_id: str
-    backend: str
     prompt: str
     system: Optional[str]
     output: str
     elapsed_s: float
     params: Dict[str, Any]
+    used_image: bool
 
 
 class BaseAgent:
@@ -24,27 +24,32 @@ class BaseAgent:
 
     def __init__(self, spec: ModelSpec):
         self.spec = spec
-        self.backend: LLMBackend = make_backend(spec.backend, spec.model_id)
+        self.backend: LLMBackend = make_backend(spec.provider, spec.model_id)
 
-    def run(self, prompt: str) -> AgentResult:
+    def run(self, prompt: str, image_path: Optional[str] = None) -> AgentResult:
         t0 = time.time()
         params = {
             "temperature": self.spec.temperature,
             "max_new_tokens": self.spec.max_new_tokens,
             "top_p": self.spec.top_p,
-            "repetition_penalty": self.spec.repetition_penalty,
             **(self.spec.extra or {}),
         }
-        out = self.backend.generate(prompt, system=self.spec.system_prompt, params=params)
+        out = self.backend.generate(
+            prompt,
+            system=self.spec.system_prompt,
+            params=params,
+            image_path=image_path,
+        )
         return AgentResult(
             role=self.role,
+            provider=self.spec.provider,
             model_id=self.spec.model_id,
-            backend=self.spec.backend,
             prompt=prompt,
             system=self.spec.system_prompt,
             output=out,
             elapsed_s=time.time() - t0,
             params=params,
+            used_image=bool(image_path),
         )