Update src/agents.py

src/agents.py

@@ -1,22 +1,22 @@
 import time
 from dataclasses import dataclass
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Callable
 
 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,32 +24,27 @@ class BaseAgent:
 
     def __init__(self, spec: ModelSpec):
         self.spec = spec
-        self.backend: LLMBackend = make_backend(spec.provider, spec.model_id)
+        self.backend: LLMBackend = make_backend(spec.backend, spec.model_id)
 
-    def run(self, prompt: str, image_path: Optional[str] = None) -> AgentResult:
+    def run(self, prompt: str) -> 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,
-            image_path=image_path,
-        )
+        out = self.backend.generate(prompt, system=self.spec.system_prompt, params=params)
         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),
         )
 
 
@@ -62,4 +57,4 @@ class RefactorAgent(BaseAgent):
 
 
 class CriticAgent(BaseAgent):
-    role = "critic"
+    role = "critic"