Update agents.py

agents.py

@@ -1,4 +1,4 @@
-# agents.py – Complete with full history and forced done
+# agents.py – FIXED: forced done, screenshot on error, truncated context
 
 import json
 import logging
@@ -15,7 +15,7 @@ class BaseAgent:
     def __init__(self, name: str, role: str, system_prompt: str, tools: ToolRegistry, state: Any):
         self.name = name
         self.role = role
-        self.system_prompt = system_prompt[:800]  # Keep essential
+        self.system_prompt = system_prompt[:400]  # Keep short
         self.tools = tools
         self.state = state
         self.history = []
@@ -25,20 +25,19 @@ class BaseAgent:
         self.plan_index = 0
 
     async def think(self, context: str) -> Dict:
-        # Build context with full history (last 3 steps, summarized)
+        # Build context with last 2 actions only, heavily truncated
         history_text = ""
         if self.history:
-            last_entries = self.history[-3:]
+            last_entries = self.history[-2:]
             history_text = "\n".join([
-                f"Step {i+1}: Action {entry['action'].get('tool')} -> Result: {self._summarize_result(entry['result'])}"
-                for i, entry in enumerate(last_entries)
+                f"Action: {entry['action'].get('tool')} -> Result: {self._summarize_result(entry['result'])}"
+                for entry in last_entries
             ])
         
-        full_context = f"Task: {self.state.task}\nHistory:\n{history_text}\nCurrent instruction: {context}"
-        if len(full_context) > 1200:
-            full_context = full_context[:1200] + "\n... (truncated)"
+        full_context = f"Task: {self.state.task}\nHistory:\n{history_text}\nInstruction: {context}"
+        if len(full_context) > 800:
+            full_context = full_context[:800] + "\n... (truncated)"
 
-        # Short system prompt
         short_prompt = self.system_prompt + """
 Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}}}
 """
@@ -49,6 +48,14 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
         raw, meta = await LLMClient.chat(messages, temperature=0.1, response_format="json")
         logger.info(f"Raw LLM response: {raw[:150]}...")
 
+        # If we already have a plan and all steps executed, force done
+        if self.plan and self.plan_index >= len(self.plan):
+            logger.info("Plan completed, forcing done")
+            return {
+                "thought": "Plan completed",
+                "action": {"tool": "done", "params": {"result": "Task completed by plan"}}
+            }
+
         json_data = self._extract_json(raw)
         if json_data:
             try:
@@ -64,15 +71,13 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
             except Exception as e:
                 logger.error(f"JSON parse error: {e}")
 
-        # If we have completed all steps in the plan, force done
-        if self.plan and self.plan_index >= len(self.plan):
-            logger.info("Plan completed, forcing done")
+        # Fallback: if we have a plan, execute next step; else navigate
+        if self.plan and self.plan_index < len(self.plan):
+            # We don't have a specific action for the next step, so fallback to navigate
             return {
-                "thought": "Plan completed",
-                "action": {"tool": "done", "params": {"result": "Task completed by plan"}}
+                "thought": "Fallback to navigate",
+                "action": {"tool": "navigate", "params": {"url": "https://example.com"}}
             }
-
-        logger.warning("Using fallback navigate")
         return {
             "thought": "Fallback",
             "action": {"tool": "navigate", "params": {"url": "https://example.com"}}
@@ -80,8 +85,8 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
 
     def _summarize_result(self, result):
         if isinstance(result, dict):
-            return result.get("result", "") or result.get("data", "") or str(result)[:100]
-        return str(result)[:100]
+            return result.get("result", "") or result.get("data", "") or str(result)[:80]
+        return str(result)[:80]
 
     def _extract_json(self, text: str) -> str:
         start = text.find('{')
@@ -103,8 +108,8 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
         if tool_name == "done":
             self.done = True
             self.result = params.get("result", "Task completed.")
+            # Always capture a final screenshot
             try:
-                await asyncio.sleep(0.2)
                 screenshot = await self.tools("screenshot")
                 return {"result": self.result, "screenshot": screenshot.get("result", "")}
             except:
@@ -112,16 +117,24 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
 
         try:
             result = await self.tools(tool_name, **params)
-            await asyncio.sleep(0.1)
-            screenshot = await self.tools("screenshot")
-            result["screenshot"] = screenshot.get("result", "")
+            # Capture screenshot after action
+            try:
+                screenshot = await self.tools("screenshot")
+                result["screenshot"] = screenshot.get("result", "")
+            except:
+                result["screenshot"] = ""
             # Truncate large data
-            if "data" in result and isinstance(result["data"], str) and len(result["data"]) > 300:
-                result["data"] = result["data"][:300] + "... (truncated)"
+            if "data" in result and isinstance(result["data"], str) and len(result["data"]) > 200:
+                result["data"] = result["data"][:200] + "... (truncated)"
             return result
         except Exception as e:
             logger.error(f"Action failed: {e}")
-            return {"error": str(e)}
+            # Even on error, try to capture screenshot
+            try:
+                screenshot = await self.tools("screenshot")
+                return {"error": str(e), "screenshot": screenshot.get("result", "")}
+            except:
+                return {"error": str(e)}
 
     async def run(self, context: str) -> Dict:
         self.history = []
@@ -130,10 +143,9 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
         self.plan = []
         self.plan_index = 0
         steps = 0
-        max_steps = 10
+        max_steps = 8  # Reduce to speed up
 
         while not self.done and steps < max_steps:
-            # Build context with current step info
             data = await self.think(context)
             thought = data.get("thought", "")
             action = data.get("action", {"tool": "done", "params": {"result": "No action"}})
@@ -149,9 +161,13 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
             if self.done:
                 break
 
-        # If loop ended without done, set a default result
-        if not self.done:
-            self.result = self.history[-1].get("result", {}).get("result", "Task incomplete") if self.history else "No action taken"
+        # If loop ended without done, mark done with last result
+        if not self.done and self.history:
+            self.result = self.history[-1].get("result", {}).get("result", "Task incomplete")
+        elif not self.done:
+            self.result = "No actions executed"
+        self.done = True  # Force completion to avoid infinite loop
+
         return {
             "agent": self.name,
             "result": self.result,
@@ -170,42 +186,33 @@ Output ONLY JSON. Format: {"thought":"...","action":{"tool":"...","params":{...}
         }
 
 
-# ---------- Specialist Agents with clear step-by-step instructions ----------
+# ---------- Specialist Agents ----------
 class ResearchAgent(BaseAgent):
     def __init__(self, tools: ToolRegistry, state: Any):
-        prompt = """
-You are a Research Agent. Task: Go to wikipedia and search for AI.
-Follow these steps exactly:
-1. Navigate to https://www.wikipedia.org
-2. Fill #searchInput with "AI"
-3. Click button[type='submit']
-4. Extract text from #firstHeading
-5. Done with result: "Article title is: " + title.
-Only output one action per JSON.
-"""
+        prompt = "You are a Research Agent. Task: Go to wikipedia and search for AI. Steps: navigate, fill search, click, extract title, done."
         super().__init__("Research Agent", "Gather", prompt, tools, state)
 
 class SalesAgent(BaseAgent):
     def __init__(self, tools: ToolRegistry, state: Any):
-        prompt = "You are a Sales Agent. Generate proposals and emails. Output JSON actions only."
+        prompt = "You are a Sales Agent. Generate proposals and emails."
         super().__init__("Sales Agent", "Sales", prompt, tools, state)
 
 class SupportAgent(BaseAgent):
     def __init__(self, tools: ToolRegistry, state: Any):
-        prompt = "You are a Support Agent. Search knowledge base and generate FAQs. Output JSON actions only."
+        prompt = "You are a Support Agent. Search knowledge base and generate FAQs."
         super().__init__("Support Agent", "Support", prompt, tools, state)
 
 class DataAgent(BaseAgent):
     def __init__(self, tools: ToolRegistry, state: Any):
-        prompt = "You are a Data Agent. Analyse data and create plots. Output JSON actions only."
+        prompt = "You are a Data Agent. Analyse data and create plots."
         super().__init__("Data Agent", "Data", prompt, tools, state)
 
 class CodeAgent(BaseAgent):
     def __init__(self, tools: ToolRegistry, state: Any):
-        prompt = "You are a Code Agent. Write and execute Python code. Output JSON actions only."
+        prompt = "You are a Code Agent. Write and execute Python code."
         super().__init__("Code Agent", "Code", prompt, tools, state)
 
 class ContentAgent(BaseAgent):
     def __init__(self, tools: ToolRegistry, state: Any):
-        prompt = "You are a Content Agent. Summarise and translate text. Output JSON actions only."
+        prompt = "You are a Content Agent. Summarise and translate text."
         super().__init__("Content Agent", "Content", prompt, tools, state)