Update agent/agent_core.py

agent/agent_core.py

@@ -1,11 +1,13 @@
 """
 LifeAdmin AI - Core Agent Logic
-Fully compatible with your UI and HuggingFace
-Includes:
-- process_files_to_rag()
-- manual_tool_call()
-- execute()
-- plan(), reflect(), memory, RAG, MCP
+Final stable version (HF / Gradio-compatible).
+Provides:
+ - plan()
+ - execute_task()
+ - reflect()
+ - execute() -> (final_answer, thoughts)
+ - process_files_to_rag()
+ - manual_tool_call()
 """
 
 import asyncio
@@ -14,6 +16,7 @@ import time
 from typing import List, Dict, Any, Optional
 from dataclasses import dataclass, asdict
 from enum import Enum
+from pathlib import Path
 
 from agent.mcp_client import MCPClient
 from agent.rag_engine import RAGEngine
@@ -21,10 +24,9 @@ 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"
@@ -35,7 +37,7 @@ class TaskStatus(Enum):
 @dataclass
 class AgentThought:
     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
@@ -58,178 +60,125 @@ class AgentTask:
     error: Optional[str] = None
 
 
-# ============================================================
-# MAIN AGENT
-# ============================================================
-
+# -------------------------
+# LifeAdminAgent
+# -------------------------
 class LifeAdminAgent:
-
     def __init__(self):
         self.mcp_client = MCPClient()
         self.rag_engine = RAGEngine()
         self.memory = MemoryStore()
         self.thoughts: List[AgentThought] = []
 
-    # ----------------------------------------------------
-    # RESET THOUGHTS
-    # ----------------------------------------------------
+        # ensure data directories exist
+        Path("data/uploads").mkdir(parents=True, exist_ok=True)
+        Path("data/outputs").mkdir(parents=True, exist_ok=True)
+
     def reset(self):
+        """Reset thoughts / context for a new request"""
         self.thoughts = []
 
-    # ----------------------------------------------------
-    # PROCESS FILES → TO RAG
-    # ----------------------------------------------------
-    async def process_files_to_rag(self, files: List[Dict]):
-        """
-        Expected format: [{ "path": "...", "name": "..." }]
-        Your UI calls this after uploads.
-        """
-
-        for f in files:
-            try:
-                await self.rag_engine.add_document(
-                    file_path=f["path"],
-                    metadata={"filename": f["name"]}
-                )
-
-                self.thoughts.append(AgentThought(
-                    step=len(self.thoughts) + 1,
-                    type="planning",
-                    content=f"Added to RAG: {f['name']}"
-                ))
-
-            except Exception as e:
-                self.thoughts.append(AgentThought(
-                    step=len(self.thoughts) + 1,
-                    type="planning",
-                    content=f"Failed indexing: {f['name']}, error={str(e)}"
-                ))
-
-        return True
-
-    # ----------------------------------------------------
-    # MANUAL TOOL CALL (Used in Manual Dashboard)
-    # ----------------------------------------------------
-    async def manual_tool_call(self, tool: str, args: Dict[str, Any]):
-        """
-        Your UI calls this for:
-        - OCR
-        - PDF extract
-        - email draft
-        - calendar event
-        - file tools etc.
-        """
-
+    # ---------------------
+    # 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="tool_call",
-            content=f"Manual call: {tool}",
-            tool_name=tool,
-            tool_args=args
-        ))
-
-        try:
-            result = await self.mcp_client.call_tool(tool, args)
-
-            self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
-                type="tool_call",
-                content="Manual tool execution succeeded",
-                tool_name=tool,
-                tool_result=result
-            ))
-
-            return result
-
-        except Exception as e:
-            self.thoughts.append(AgentThought(
-                step=len(self.thoughts) + 1,
-                type="tool_call",
-                content=f"Manual tool failed: {str(e)}",
-                tool_name=tool
-            ))
-            return {"error": str(e)}
-
-    # ----------------------------------------------------
-    # PLAN TASKS
-    # ----------------------------------------------------
-    async def plan(self, user_request: str, files: List[str] = None):
-        self.thoughts.append(AgentThought(
-            step=len(self.thoughts)+1,
             type="planning",
-            content=f"Analyzing user request: {user_request}"
+            content=f"Analyzing request: {user_request}"
         ))
 
-        tools = await self.mcp_client.list_tools()
-        tool_desc = "\n".join([f"{t['name']}: {t['description']}" for t in tools])
-
-        rag_docs = await self.rag_engine.search(user_request, k=3)
-        rag_context = "\n".join([d["text"][:300] for d in rag_docs]) or "None"
+        # 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"
 
-        mem_context = self.memory.get_relevant_memories(user_request)
+        # 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 = []
 
-        prompt = f"""
-You are a task planner. USER REQUEST:
-{user_request}
+        rag_context = "\n".join([d.get("text","")[:250] for d in rag_docs]) if rag_docs else "No relevant docs"
 
-FILES: {files or []}
+        memory_context = self.memory.get_relevant_memories(user_request) if self.memory else "No memory"
 
-TOOLS AVAILABLE:
+        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:
 {tool_desc}
-
-RAG CONTEXT:
+RAG context:
 {rag_context}
+Memory:
+{memory_context}
 
-MEMORY:
-{mem_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)
 
-Return JSON ONLY:
+Example:
 [
   {{
-    "id": "task1",
-    "description": "Extract something",
+    "id": "task_1",
+    "description": "Extract text from invoice.pdf",
     "tool": "ocr_extract_text",
-    "args": {{"file_path": "x.pdf"}}
+    "args": {{"file_path": "data/uploads/invoice.pdf", "language": "en"}}
   }}
 ]
 """
-
-        resp = await get_llm_response(prompt, temperature=0.2)
-        txt = resp.strip()
-
-        if "```json" in txt:
-            txt = txt.split("```json")[1].split("```")[0].strip()
+        self.thoughts.append(AgentThought(
+            step=len(self.thoughts) + 1,
+            type="planning",
+            content="Asking LLM to create a plan..."
+        ))
 
         try:
-            parsed = json.loads(txt)
-            tasks = [AgentTask(**t) for t in parsed]
+            plan_text = await get_llm_response(planning_prompt, temperature=0.2)
+            plan_text = plan_text.strip()
 
-        except Exception:
+            # 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()
+
+            tasks_data = json.loads(plan_text)
+            tasks = [AgentTask(**t) for t in tasks_data]
             self.thoughts.append(AgentThought(
-                step=len(self.thoughts)+1,
+                step=len(self.thoughts) + 1,
                 type="planning",
-                content="Planning failed – invalid JSON returned by LLM."
+                content=f"Plan created with {len(tasks)} tasks."
+            ))
+            return tasks
+        except Exception as e:
+            self.thoughts.append(AgentThought(
+                step=len(self.thoughts) + 1,
+                type="planning",
+                content=f"Planning failed: {str(e)}"
             ))
             return []
 
+    # ---------------------
+    # Execution of a single task
+    # ---------------------
+    async def execute_task(self, task: AgentTask) -> AgentTask:
         self.thoughts.append(AgentThought(
-            step=len(self.thoughts)+1,
-            type="planning",
-            content=f"Created {len(tasks)} tasks"
-        ))
-        return tasks
-
-    # ----------------------------------------------------
-    # EXECUTE A SINGLE TASK
-    # ----------------------------------------------------
-    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.description}",
+            content=f"Executing task: {task.description}",
             tool_name=task.tool,
             tool_args=task.args
         ))
+        task.status = TaskStatus.IN_PROGRESS
 
         try:
             result = await self.mcp_client.call_tool(task.tool, task.args)
@@ -237,86 +186,160 @@ Return JSON ONLY:
             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: {task.description}",
                 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: {task.description} - {str(e)}",
                 tool_name=task.tool
             ))
+            return task
 
-        return task
-
-    # ----------------------------------------------------
-    # REFLECT / SUMMARIZE
-    # ----------------------------------------------------
-    async def reflect(self, tasks: List[AgentTask], original: str):
+    # ---------------------
+    # Reflection / final answer
+    # ---------------------
+    async def reflect(self, tasks: List[AgentTask], original_request: str) -> str:
         self.thoughts.append(AgentThought(
-            step=len(self.thoughts)+1,
+            step=len(self.thoughts) + 1,
             type="reflection",
-            content="Summarizing results…"
+            content="Synthesizing results..."
         ))
 
-        results_txt = "\n".join([
-            f"✓ {t.description}: {str(t.result)[:200]}"
-            if t.status == TaskStatus.COMPLETED
-            else f"✗ {t.description}: {t.error}"
-            for t in tasks
-        ])
-
-        prompt = f"""
-Summarize results clearly for a user.
-
-REQUEST:
-{original}
-
-RESULTS:
-{results_txt}
+        summary_lines = []
+        for t in tasks:
+            if t.status == TaskStatus.COMPLETED:
+                summary_lines.append(f"✓ {t.description}: {str(t.result)[:300]}")
+            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.
 """
 
-        answer = await get_llm_response(prompt, temperature=0.4)
+        try:
+            answer = await get_llm_response(reflection_prompt, temperature=0.5)
+        except Exception as e:
+            answer = f"Reflection failed: {str(e)}"
 
         self.thoughts.append(AgentThought(
-            step=len(self.thoughts)+1,
+            step=len(self.thoughts) + 1,
             type="answer",
             content=answer
         ))
 
-        self.memory.add_memory(
-            content=f"Request: {original}\nAnswer: {answer}",
-            metadata={"timestamp": time.time()}
-        )
+        # store short memory
+        try:
+            self.memory.add_memory(
+                content=f"Request: {original_request}\nResult: {answer}",
+                memory_type="task_completion",
+                metadata={"timestamp": time.time()}
+            )
+        except Exception:
+            pass
+
         return answer
 
-    # ----------------------------------------------------
-    # MAIN EXECUTION (Used by Voice Mode)
-    # ----------------------------------------------------
-    async def execute(self, user_request: str, files: List[str] = None):
+    # ---------------------
+    # Main execute (no streaming)
+    # ---------------------
+    async def execute(self, user_request: str, files: List[str] = None) -> (str, List[AgentThought]):
+        """
+        Run plan -> execute each task -> reflect
+        Returns: (final_answer, list_of_thoughts)
+        """
         self.reset()
 
         tasks = await self.plan(user_request, files)
         if not tasks:
-            return "I could not create a plan. Try rephrasing.", self.thoughts
+            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
 
         executed = []
         for t in tasks:
-            executed.append(await self.execute_task(t))
+            executed_task = await self.execute_task(t)
+            executed.append(executed_task)
 
-        final = await self.reflect(executed, user_request)
-        return final, self.thoughts
+        final_answer = await self.reflect(executed, user_request)
+        return final_answer, self.thoughts
 
-    # ----------------------------------------------------
-    # EXPORT THOUGHT TRACE
-    # ----------------------------------------------------
-    def get_thought_trace(self):
+    # ---------------------
+    # Utility: provide thought trace for UI
+    # ---------------------
+    def get_thought_trace(self) -> List[Dict[str, Any]]:
         return [asdict(t) for t in self.thoughts]
+
+    # ---------------------
+    # Add uploaded files into RAG index (helper used 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.
+        """
+        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
+                    try:
+                        result = await self.mcp_client.call_tool("ocr_extract_text", {"file_path": path, "language": "en"})
+                        text = result.get("text", "")
+                    except Exception:
+                        text = ""
+                else:
+                    # read plain text files
+                    try:
+                        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
+
+    # ---------------------
+    # 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}
+        """
+        try:
+            result = await self.mcp_client.call_tool(tool_name, args)
+            return {"success": True, "result": result, "error": None}
+        except Exception as e:
+            return {"success": False, "result": None, "error": str(e)}