Update agent/agent_core.py

agent/agent_core.py

@@ -1,6 +1,6 @@
 """
 LifeAdmin AI - Core Agent Logic
-Autonomous planning, tool orchestration, and execution
+Final stable version (No async-generators – fully HF compatible)
 """
 
 import asyncio
@@ -16,9 +16,9 @@ from agent.memory import MemoryStore
 from utils.llm_utils import get_llm_response
 
 
-# -------------------------
-# ENUMS & MODELS
-# -------------------------
+# =============================
+# DATA MODELS
+# =============================
 
 class TaskStatus(Enum):
     PENDING = "pending"
@@ -29,9 +29,8 @@ class TaskStatus(Enum):
 
 @dataclass
 class AgentThought:
-    """Represents a thought/step in agent reasoning"""
     step: int
-    type: str            # planning | tool_call | reflection | answer
+    type: str              # 'planning', 'tool_call', 'reflection', 'answer'
     content: str
     tool_name: Optional[str] = None
     tool_args: Optional[Dict] = None
@@ -45,7 +44,6 @@ class AgentThought:
 
 @dataclass
 class AgentTask:
-    """Represents an atomic MCP operation"""
     id: str
     description: str
     tool: str
@@ -55,9 +53,9 @@ class AgentTask:
     error: Optional[str] = None
 
 
-# -------------------------
+# =============================
 # MAIN AGENT CLASS
-# -------------------------
+# =============================
 
 class LifeAdminAgent:
 
@@ -66,58 +64,47 @@ class LifeAdminAgent:
         self.rag_engine = RAGEngine()
         self.memory = MemoryStore()
         self.thoughts: List[AgentThought] = []
-        self.current_context = {}
 
-    # -------------------------------------------
-    # RESET
-    # -------------------------------------------
+    # -----------------------
+    # UTIL
+    # -----------------------
 
-    def reset_context(self):
+    def reset(self):
         self.thoughts = []
-        self.current_context = {}
 
-    # -------------------------------------------
-    # PLANNING PHASE
-    # -------------------------------------------
+    # -----------------------
+    # PLANNING
+    # -----------------------
 
-    async def plan(self, user_request: str, available_files: List[str] = None) -> List[AgentTask]:
+    async def plan(self, user_request: str, files: List[str] = None) -> List[AgentTask]:
 
         self.thoughts.append(AgentThought(
             step=len(self.thoughts) + 1,
             type="planning",
-            content=f"Analyzing request: {user_request}"
+            content=f"Analyzing: {user_request}"
         ))
 
-        # List tools available through MCP
         tools = await self.mcp_client.list_tools()
-        tool_descriptions = "\n".join([
-            f"- {tool['name']}: {tool.get('description', '')}" for tool in tools
-        ])
+        tool_desc = "\n".join([f"- {t['name']}: {t['description']}" for t in tools])
 
-        # RAG context
-        relevant_docs = []
+        rag_docs = []
         if user_request.strip():
-            relevant_docs = await self.rag_engine.search(user_request, k=3)
+            rag_docs = await self.rag_engine.search(user_request, k=3)
 
-        rag_context = "\n".join(
-            [doc["text"][:200] for doc in relevant_docs]
-        ) if relevant_docs else "No relevant documents"
+        rag_context = "\n".join([d["text"][:200] for d in rag_docs]) if rag_docs else "None"
 
-        # Memory
         memory_context = self.memory.get_relevant_memories(user_request)
 
-        # Build plan prompt
-        planning_prompt = f"""
-You are an autonomous assistant. Create a JSON task plan.
+        prompt = f"""
+You are a task planner.
 
-USER REQUEST:
+REQUEST:
 {user_request}
 
-AVAILABLE FILES:
-{', '.join(available_files) if available_files else 'None'}
+FILES: {files or []}
 
-AVAILABLE TOOLS:
-{tool_descriptions}
+TOOLS:
+{tool_desc}
 
 RAG CONTEXT:
 {rag_context}
@@ -125,66 +112,54 @@ RAG CONTEXT:
 MEMORY:
 {memory_context}
 
-Return ONLY valid JSON list of tasks like:
+Return ONLY JSON list:
 [
   {{
-    "id": "t1",
+    "id": "task1",
     "description": "Extract text",
     "tool": "ocr_extract_text",
-    "args": {{"file_path": "invoice.pdf"}}
+    "args": {{"file_path": "x.pdf"}}
   }}
 ]
 """
 
-        self.thoughts.append(AgentThought(
-            step=len(self.thoughts) + 1,
-            type="planning",
-            content="Generating plan with LLM..."
-        ))
-
-        try:
-            raw = await get_llm_response(planning_prompt, temperature=0.2)
-            txt = raw.strip()
-
-            # Remove markdown wrappers
-            if "```json" in txt:
-                txt = txt.split("```json")[1].split("```")[0].strip()
-            elif "```" in txt:
-                txt = txt.split("```")[1].split("```")[0].strip()
+        response = await get_llm_response(prompt, temperature=0.2)
+        text = response.strip()
 
-            plan_json = json.loads(txt)
-
-            tasks = [
-                AgentTask(**task, status=TaskStatus.PENDING)
-                for task in plan_json
-            ]
+        if "```json" in text:
+            text = text.split("```json")[1].split("```")[0].strip()
+        elif "```" in text:
+            text = text.split("```")[1].split("```")[0].strip()
 
+        try:
+            plan_data = json.loads(text)
+            tasks = [AgentTask(**t) for t in plan_data]
+        except:
             self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
+                step=len(self.thoughts)+1,
                 type="planning",
-                content=f"Plan created: {len(tasks)} tasks"
+                content="Planning failed (invalid JSON)"
             ))
+            return []
 
-            return tasks
+        self.thoughts.append(AgentThought(
+            step=len(self.thoughts)+1,
+            type="planning",
+            content=f"Created {len(tasks)} tasks"
+        ))
 
-        except Exception as e:
-            self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
-                type="planning",
-                content=f"Planning failed: {e}"
-            ))
-            return []
+        return tasks
 
-    # -------------------------------------------
-    # TOOL EXECUTION PHASE
-    # -------------------------------------------
+    # -----------------------
+    # EXECUTION
+    # -----------------------
 
-    async def execute_task(self, task: AgentTask) -> AgentTask:
+    async def execute_task(self, task: AgentTask):
 
         self.thoughts.append(AgentThought(
-            step=len(self.thoughts) + 1,
+            step=len(self.thoughts)+1,
             type="tool_call",
-            content=f"Executing task: {task.description}",
+            content=f"Executing: {task.description}",
             tool_name=task.tool,
             tool_args=task.args
         ))
@@ -193,14 +168,13 @@ Return ONLY valid JSON list of tasks like:
 
         try:
             result = await self.mcp_client.call_tool(task.tool, task.args)
-
-            task.status = TaskStatus.COMPLETED
             task.result = result
+            task.status = TaskStatus.COMPLETED
 
             self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
+                step=len(self.thoughts)+1,
                 type="tool_call",
-                content=f"✓ Completed: {task.description}",
+                content=f"✓ Completed",
                 tool_name=task.tool,
                 tool_result=result
             ))
@@ -210,128 +184,87 @@ Return ONLY valid JSON list of tasks like:
             task.error = str(e)
 
             self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
+                step=len(self.thoughts)+1,
                 type="tool_call",
-                content=f"✗ Failed: {task.description} — {e}",
+                content=f"✗ Failed: {e}",
                 tool_name=task.tool
             ))
 
         return task
 
-    # -------------------------------------------
-    # REFLECTION PHASE
-    # -------------------------------------------
+    # -----------------------
+    # REFLECTION
+    # -----------------------
 
-    async def reflect(self, tasks: List[AgentTask], request: str) -> str:
+    async def reflect(self, tasks: List[AgentTask], original: str) -> str:
 
         self.thoughts.append(AgentThought(
-            step=len(self.thoughts) + 1,
+            step=len(self.thoughts)+1,
             type="reflection",
-            content="Analyzing final results..."
+            content="Summarizing results..."
         ))
 
-        results_string = []
+        results = []
         for t in tasks:
             if t.status == TaskStatus.COMPLETED:
-                short = str(t.result)[:200]
-                results_string.append(f"✓ {t.description}: {short}")
+                results.append(f"✓ {t.description}: {str(t.result)[:200]}")
             else:
-                results_string.append(f"✗ {t.description}: {t.error}")
+                results.append(f"✗ {t.description}: {t.error}")
 
-        reflection_prompt = f"""
-Summarize the final results of the following tasks:
+        prompt = f"""
+Provide a helpful summary for the user.
 
 REQUEST:
-{request}
+{original}
 
 RESULTS:
-{chr(10).join(results_string)}
+{chr(10).join(results)}
 
-Give a clear, helpful answer:
-- What succeeded
-- What failed
-- What files/events/emails were produced
-- Next steps
+Write a clear, friendly answer.
 """
 
-        try:
-            answer = await get_llm_response(reflection_prompt, temperature=0.5)
-
-            self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
-                type="answer",
-                content=answer
-            ))
-
-            # Write to memory
-            self.memory.add_memory(
-                f"Request: {request}\nAnswer: {answer}",
-                metadata={"type": "task_completion", "timestamp": time.time()}
-            )
+        answer = await get_llm_response(prompt, temperature=0.4)
 
-            return answer
+        self.thoughts.append(AgentThought(
+            step=len(self.thoughts)+1,
+            type="answer",
+            content=answer
+        ))
 
-        except Exception as e:
-            errmsg = f"Reflection failed: {e}"
+        self.memory.add_memory(
+            content=f"Request: {original}\nAnswer: {answer}",
+            metadata={"timestamp": time.time()}
+        )
 
-            self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
-                type="answer",
-                content=errmsg
-            ))
+        return answer
 
-            return errmsg
+    # -----------------------
+    # MAIN EXECUTION (FULLY FIXED)
+    # -----------------------
 
-    # -------------------------------------------
-    # STREAMING EXECUTION LOOP (FIXED)
-    # -------------------------------------------
-
-    async def execute(self, request: str, files: List[str] = None, stream_thoughts=False):
+    async def execute(self, user_request: str, files: List[str] = None):
         """
-        If stream_thoughts=True → yields AgentThought objects
-        If stream_thoughts=False → returns (answer, thoughts)
+        A simple coroutine returning final_answer, thoughts
+        No yields → No async generator → No syntax errors
         """
 
-        self.reset_context()
-
-        # --- PLANNING ---
-        tasks = await self.plan(request, files)
-        if stream_thoughts:
-            for th in self.thoughts:
-                yield th
+        self.reset()
 
+        tasks = await self.plan(user_request, files)
         if not tasks:
-            # DO NOT return a value — async generator cannot return a value
-            thought = AgentThought(
-                step=len(self.thoughts) + 1,
-                type="answer",
-                content="Could not create plan. Try rephrasing."
-            )
-            self.thoughts.append(thought)
-            if stream_thoughts:
-                yield thought
-            return
-
-        # --- EXECUTION ---
+            # return is allowed now
+            return "Could not generate plan. Try rephrasing.", self.thoughts
+
         executed = []
         for t in tasks:
-            done = await self.execute_task(t)
-            executed.append(done)
-            if stream_thoughts:
-                yield self.thoughts[-1]
-
-        # --- REFLECTION ---
-        final_answer = await self.reflect(executed, request)
-        if stream_thoughts:
-            yield self.thoughts[-1]
-            return
-
-        # If NOT streaming: return normal output
+            executed.append(await self.execute_task(t))
+
+        final_answer = await self.reflect(executed, user_request)
         return final_answer, self.thoughts
 
-    # -------------------------------------------
-    # UTILITY
-    # -------------------------------------------
+    # -----------------------
+    # EXPORT THOUGHTS
+    # -----------------------
 
-    def get_thought_trace(self) -> List[Dict]:
+    def get_thought_trace(self):
         return [asdict(t) for t in self.thoughts]