Update agent/agent_core.py

agent/agent_core.py

@@ -1,13 +1,9 @@
 """
 LifeAdmin AI - Core Agent Logic
-Final stable version (HF / Gradio-compatible).
-Provides:
- - plan()
- - execute_task()
- - reflect()
- - execute() -> (final_answer, thoughts)
- - process_files_to_rag()
- - manual_tool_call()
+Final stable version (No async-generators – fully HF compatible)
+Includes async helpers used by the UI:
+- process_files_to_rag(files)
+- manual_tool_call(tool_name, args)
 """
 
 import asyncio
@@ -24,9 +20,10 @@ from agent.memory import MemoryStore
 from utils.llm_utils import get_llm_response
 
 
-# -------------------------
-# Data models
-# -------------------------
+# =============================
+# DATA MODELS
+# =============================
+
 class TaskStatus(Enum):
     PENDING = "pending"
     IN_PROGRESS = "in_progress"
@@ -60,124 +57,111 @@ class AgentTask:
     error: Optional[str] = None
 
 
-# -------------------------
-# LifeAdminAgent
-# -------------------------
+# =============================
+# MAIN AGENT CLASS
+# =============================
+
 class LifeAdminAgent:
+
     def __init__(self):
         self.mcp_client = MCPClient()
         self.rag_engine = RAGEngine()
         self.memory = MemoryStore()
         self.thoughts: List[AgentThought] = []
 
-        # ensure data directories exist
-        Path("data/uploads").mkdir(parents=True, exist_ok=True)
-        Path("data/outputs").mkdir(parents=True, exist_ok=True)
+    # -----------------------
+    # UTIL
+    # -----------------------
 
     def reset(self):
-        """Reset thoughts / context for a new request"""
         self.thoughts = []
 
-    # ---------------------
-    # Planning
-    # ---------------------
+    # -----------------------
+    # PLANNING
+    # -----------------------
+
     async def plan(self, user_request: str, files: List[str] = None) -> List[AgentTask]:
-        """Create an execution plan (list of AgentTask) using LLM + RAG + memory"""
+
         self.thoughts.append(AgentThought(
             step=len(self.thoughts) + 1,
             type="planning",
-            content=f"Analyzing request: {user_request}"
+            content=f"Analyzing: {user_request}"
         ))
 
-        # list tools
-        try:
-            tools = await self.mcp_client.list_tools()
-        except Exception:
-            tools = []
-        tool_desc = "\n".join([f"- {t['name']}: {t.get('description','')}" for t in tools]) if tools else "No tool metadata available"
+        tools = await self.mcp_client.list_tools()
+        tool_desc = "\n".join([f"- {t['name']}: {t['description']}" for t in tools])
 
-        # RAG search
         rag_docs = []
-        if user_request and user_request.strip():
-            try:
-                rag_docs = await self.rag_engine.search(user_request, k=3)
-            except Exception:
-                rag_docs = []
+        if user_request.strip():
+            rag_docs = await self.rag_engine.search(user_request, k=3)
 
-        rag_context = "\n".join([d.get("text","")[:250] for d in rag_docs]) if rag_docs else "No relevant docs"
+        rag_context = "\n".join([d.get("text", "")[:200] for d in rag_docs]) if rag_docs else "None"
 
-        memory_context = self.memory.get_relevant_memories(user_request) if self.memory else "No memory"
+        memory_context = self.memory.get_relevant_memories(user_request)
 
-        planning_prompt = f"""
-You are an autonomous life admin assistant. Produce a JSON array of tasks (no extra text).
-User request: {user_request}
-Available files: {files or []}
-Available tools:
+        prompt = f"""
+You are a task planner.
+REQUEST:
+{user_request}
+FILES: {files or []}
+TOOLS:
 {tool_desc}
-RAG context:
+RAG CONTEXT:
 {rag_context}
-Memory:
+MEMORY:
 {memory_context}
-
-Return ONLY valid JSON array of tasks. Each task must contain:
- - id (string)
- - description (string)
- - tool (one of the tool names)
- - args (a JSON object)
-
-Example:
+Return ONLY JSON list:
 [
   {{
-    "id": "task_1",
-    "description": "Extract text from invoice.pdf",
+    "id": "task1",
+    "description": "Extract text",
     "tool": "ocr_extract_text",
-    "args": {{"file_path": "data/uploads/invoice.pdf", "language": "en"}}
+    "args": {{"file_path": "x.pdf"}}
   }}
 ]
 """
-        self.thoughts.append(AgentThought(
-            step=len(self.thoughts) + 1,
-            type="planning",
-            content="Asking LLM to create a plan..."
-        ))
 
-        try:
-            plan_text = await get_llm_response(planning_prompt, temperature=0.2)
-            plan_text = plan_text.strip()
+        response = await get_llm_response(prompt, temperature=0.2)
+        text = response.strip()
 
-            # try to extract JSON if wrapped in code fences
-            if "```json" in plan_text:
-                plan_text = plan_text.split("```json", 1)[1].split("```", 1)[0].strip()
-            elif "```" in plan_text:
-                plan_text = plan_text.split("```", 1)[1].split("```", 1)[0].strip()
+        if "```json" in text:
+            text = text.split("```json")[1].split("```")[0].strip()
+        elif "```" in text:
+            text = text.split("```")[1].split("```")[0].strip()
 
-            tasks_data = json.loads(plan_text)
-            tasks = [AgentTask(**t) for t in tasks_data]
-            self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
-                type="planning",
-                content=f"Plan created with {len(tasks)} tasks."
-            ))
-            return tasks
-        except Exception as e:
+        try:
+            plan_data = json.loads(text)
+            tasks = [AgentTask(**t) for t in plan_data]
+        except Exception:
             self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
+                step=len(self.thoughts)+1,
                 type="planning",
-                content=f"Planning failed: {str(e)}"
+                content="Planning failed (invalid JSON)"
             ))
             return []
 
-    # ---------------------
-    # Execution of a single task
-    # ---------------------
-    async def execute_task(self, task: AgentTask) -> AgentTask:
         self.thoughts.append(AgentThought(
-            step=len(self.thoughts) + 1,
+            step=len(self.thoughts)+1,
+            type="planning",
+            content=f"Created {len(tasks)} tasks"
+        ))
+
+        return tasks
+
+    # -----------------------
+    # EXECUTION
+    # -----------------------
+
+    async def execute_task(self, task: AgentTask):
+
+        self.thoughts.append(AgentThought(
+            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
         ))
+
         task.status = TaskStatus.IN_PROGRESS
 
         try:
@@ -186,160 +170,153 @@ Example:
             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
             ))
-            return task
+
         except Exception as e:
             task.status = TaskStatus.FAILED
             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} - {str(e)}",
+                content=f"✗ Failed: {e}",
                 tool_name=task.tool
             ))
-            return task
 
-    # ---------------------
-    # Reflection / final answer
-    # ---------------------
-    async def reflect(self, tasks: List[AgentTask], original_request: str) -> str:
+        return task
+
+    # -----------------------
+    # REFLECTION
+    # -----------------------
+
+    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="Synthesizing results..."
+            content="Summarizing results..."
         ))
 
-        summary_lines = []
+        results = []
         for t in tasks:
             if t.status == TaskStatus.COMPLETED:
-                summary_lines.append(f"✓ {t.description}: {str(t.result)[:300]}")
+                results.append(f"✓ {t.description}: {str(t.result)[:200]}")
             else:
-                summary_lines.append(f"✗ {t.description}: {t.error}")
-
-        reflection_prompt = f"""
-You are the agent summarizing execution results.
-Original request: {original_request}
-Execution summary:
-{chr(10).join(summary_lines)}
-
-Write a clear, friendly reply telling the user what was done, outputs created, any errors, and next steps.
+                results.append(f"✗ {t.description}: {t.error}")
+
+        prompt = f"""
+Provide a helpful summary for the user.
+REQUEST:
+{original}
+RESULTS:
+{chr(10).join(results)}
+Write a clear, friendly answer.
 """
 
-        try:
-            answer = await get_llm_response(reflection_prompt, temperature=0.5)
-        except Exception as e:
-            answer = f"Reflection failed: {str(e)}"
+        answer = await get_llm_response(prompt, temperature=0.4)
 
         self.thoughts.append(AgentThought(
-            step=len(self.thoughts) + 1,
+            step=len(self.thoughts)+1,
             type="answer",
             content=answer
         ))
 
-        # store short memory
+        # store memory
         try:
             self.memory.add_memory(
-                content=f"Request: {original_request}\nResult: {answer}",
-                memory_type="task_completion",
+                content=f"Request: {original}\nAnswer: {answer}",
                 metadata={"timestamp": time.time()}
             )
         except Exception:
+            # don't break on memory errors
             pass
 
         return answer
 
-    # ---------------------
-    # Main execute (no streaming)
-    # ---------------------
-    async def execute(self, user_request: str, files: List[str] = None) -> (str, List[AgentThought]):
+    # -----------------------
+    # MAIN EXECUTION (FULLY FIXED)
+    # -----------------------
+
+    async def execute(self, user_request: str, files: List[str] = None):
         """
-        Run plan -> execute each task -> reflect
-        Returns: (final_answer, list_of_thoughts)
+        A simple coroutine returning final_answer, thoughts
+        No yields → No async generator → No syntax errors
         """
+
         self.reset()
 
         tasks = await self.plan(user_request, files)
         if not tasks:
-            err_msg = "Could not create an execution plan. Try rephrasing your request."
-            self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
-                type="answer",
-                content=err_msg
-            ))
-            return err_msg, self.thoughts
+            # return is allowed now
+            return "Could not generate plan. Try rephrasing.", self.thoughts
 
         executed = []
         for t in tasks:
-            executed_task = await self.execute_task(t)
-            executed.append(executed_task)
+            executed.append(await self.execute_task(t))
 
         final_answer = await self.reflect(executed, user_request)
         return final_answer, self.thoughts
 
-    # ---------------------
-    # Utility: provide thought trace for UI
-    # ---------------------
-    def get_thought_trace(self) -> List[Dict[str, Any]]:
+    # -----------------------
+    # EXPORT THOUGHTS
+    # -----------------------
+
+    def get_thought_trace(self):
         return [asdict(t) for t in self.thoughts]
 
-    # ---------------------
-    # Add uploaded files into RAG index (helper used by UI)
-    # ---------------------
+    # -----------------------
+    # Additional helpers expected by UI
+    # -----------------------
+
     async def process_files_to_rag(self, files: List[Dict[str, str]]):
         """
-        files: list of dicts {'path': <path>, 'name': <filename>}
-        Extract text using available local tools (pdf/text/ocr) and add to RAG.
+        Process uploaded files and add text to RAG.
+        files: List[{'path': '/abs/path', 'name': 'file.pdf'}]
         """
         for file_info in files:
-            path = file_info.get("path")
-            name = file_info.get("name", Path(path).name if path else "")
             try:
-                text = ""
-                if path and path.lower().endswith(".pdf"):
-                    # try utils.pdf_utils
-                    try:
-                        from utils.pdf_utils import extract_text_from_pdf
-                        text = extract_text_from_pdf(path)
-                    except Exception:
-                        text = ""
-                elif path and path.lower().endswith((".png", ".jpg", ".jpeg", ".tiff")):
-                    # use MCP OCR tool (via client) or local easyocr
+                path = file_info.get('path')
+                if not path:
+                    continue
+
+                # small heuristic on extension
+                if path.lower().endswith('.pdf'):
+                    from utils.pdf_utils import extract_text_from_pdf
+                    text = extract_text_from_pdf(path)
+                elif path.lower().endswith(('.png', '.jpg', '.jpeg')):
+                    # Use OCR tool via MCPClient (local fallback)
                     try:
-                        result = await self.mcp_client.call_tool("ocr_extract_text", {"file_path": path, "language": "en"})
-                        text = result.get("text", "")
+                        res = await self.mcp_client.call_tool('ocr_extract_text', {'file_path': path, 'language': 'en'})
+                        text = res.get('text', '')
                     except Exception:
+                        # Last-resort: empty text
                         text = ""
                 else:
-                    # read plain text files
+                    # treat as text file
                     try:
-                        with open(path, "r", encoding="utf-8") as f:
+                        with open(path, 'r', encoding='utf-8') as f:
                             text = f.read()
                     except Exception:
                         text = ""
 
-                if text and len(text.strip()) > 20:
-                    try:
-                        await self.rag_engine.add_document(text=text, metadata={"filename": name, "path": path})
-                    except Exception:
-                        pass
-            except Exception:
-                continue
+                # add to RAG
+                if text:
+                    await self.rag_engine.add_document(text=text, metadata={'filename': file_info.get('name'), 'path': path})
+            except Exception as e:
+                # log (print) but don't raise
+                print(f"Error processing {file_info.get('name')}: {e}")
 
-    # ---------------------
-    # Manual tool call wrapper for UI (guarantees consistent return shape)
-    # ---------------------
     async def manual_tool_call(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Calls an MCP tool (via MCPClient). Returns dict:
-        {'success': bool, 'result': <tool_result> or None, 'error': <err_msg> or None}
+        Direct tool call helper used by UI buttons. Returns dict with 'success' key.
         """
         try:
             result = await self.mcp_client.call_tool(tool_name, args)
-            return {"success": True, "result": result, "error": None}
+            return {'success': True, 'result': result}
         except Exception as e:
-            return {"success": False, "result": None, "error": str(e)}
+            return {'success': False, 'error': str(e)}