Update agent/agent_core.py

agent/agent_core.py

@@ -16,6 +16,10 @@ from agent.memory import MemoryStore
 from utils.llm_utils import get_llm_response
 
 
+# -------------------------
+# ENUMS & MODELS
+# -------------------------
+
 class TaskStatus(Enum):
     PENDING = "pending"
     IN_PROGRESS = "in_progress"
@@ -27,13 +31,13 @@ class TaskStatus(Enum):
 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
     tool_result: Optional[Any] = None
     timestamp: float = None
-    
+
     def __post_init__(self):
         if self.timestamp is None:
             self.timestamp = time.time()
@@ -41,7 +45,7 @@ class AgentThought:
 
 @dataclass
 class AgentTask:
-    """Represents a task to be executed"""
+    """Represents an atomic MCP operation"""
     id: str
     description: str
     tool: str
@@ -51,328 +55,283 @@ class AgentTask:
     error: Optional[str] = None
 
 
+# -------------------------
+# MAIN AGENT CLASS
+# -------------------------
+
 class LifeAdminAgent:
-    """Main autonomous agent with planning, tool calling, and reflection"""
-    
+
     def __init__(self):
         self.mcp_client = MCPClient()
         self.rag_engine = RAGEngine()
         self.memory = MemoryStore()
         self.thoughts: List[AgentThought] = []
         self.current_context = {}
-        
+
+    # -------------------------------------------
+    # RESET
+    # -------------------------------------------
+
     def reset_context(self):
-        """Reset agent context for new task"""
         self.thoughts = []
         self.current_context = {}
-    
+
+    # -------------------------------------------
+    # PLANNING PHASE
+    # -------------------------------------------
+
     async def plan(self, user_request: str, available_files: List[str] = None) -> List[AgentTask]:
-        """
-        Create execution plan from user request
-        
-        Args:
-            user_request: Natural language request from user
-            available_files: List of uploaded files
-            
-        Returns:
-            List of tasks to execute
-        """
+
         self.thoughts.append(AgentThought(
             step=len(self.thoughts) + 1,
-            type='planning',
+            type="planning",
             content=f"Analyzing request: {user_request}"
         ))
-        
-        # Get available tools
+
+        # 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
+            f"- {tool['name']}: {tool.get('description', '')}" for tool in tools
         ])
-        
-        # Search RAG for relevant context
+
+        # RAG context
         relevant_docs = []
-        if user_request:
+        if user_request.strip():
             relevant_docs = await self.rag_engine.search(user_request, k=3)
-        
-        context = "\n".join([doc['text'][:200] for doc in relevant_docs]) if relevant_docs else "No previous documents"
-        
-        # Get memory
+
+        rag_context = "\n".join(
+            [doc["text"][:200] for doc in relevant_docs]
+        ) if relevant_docs else "No relevant documents"
+
+        # Memory
         memory_context = self.memory.get_relevant_memories(user_request)
-        
-        # Create planning prompt
-        planning_prompt = f"""You are an autonomous life admin agent. Create a step-by-step execution plan.
 
-USER REQUEST: {user_request}
+        # Build plan prompt
+        planning_prompt = f"""
+You are an autonomous assistant. Create a JSON task plan.
+
+USER REQUEST:
+{user_request}
 
-AVAILABLE FILES: {', '.join(available_files) if available_files else 'None'}
+AVAILABLE FILES:
+{', '.join(available_files) if available_files else 'None'}
 
 AVAILABLE TOOLS:
 {tool_descriptions}
 
-RELEVANT CONTEXT:
-{context}
+RAG CONTEXT:
+{rag_context}
 
 MEMORY:
 {memory_context}
 
-Create a JSON plan with tasks. Each task should have:
-- id: unique identifier
-- description: what this task does
-- tool: which tool to use
-- args: arguments for the tool (as a dict)
-
-Return ONLY valid JSON array of tasks, no other text.
-
-Example format:
+Return ONLY valid JSON list of tasks like:
 [
   {{
-    "id": "task_1",
-    "description": "Extract text from document",
+    "id": "t1",
+    "description": "Extract text",
     "tool": "ocr_extract_text",
-    "args": {{"file_path": "document.pdf", "language": "en"}}
+    "args": {{"file_path": "invoice.pdf"}}
   }}
 ]
 """
-        
+
         self.thoughts.append(AgentThought(
             step=len(self.thoughts) + 1,
-            type='planning',
-            content="Creating execution plan with LLM..."
+            type="planning",
+            content="Generating plan with LLM..."
         ))
-        
+
         try:
-            plan_response = await get_llm_response(planning_prompt, temperature=0.3)
-            
-            # Extract JSON from response
-            plan_text = plan_response.strip()
-            if '```json' in plan_text:
-                plan_text = plan_text.split('```json')[1].split('```')[0].strip()
-            elif '```' in plan_text:
-                plan_text = plan_text.split('```')[1].split('```')[0].strip()
-            
-            tasks_data = json.loads(plan_text)
-            
+            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()
+
+            plan_json = json.loads(txt)
+
             tasks = [
-                AgentTask(**{**task, 'status': TaskStatus.PENDING})
-                for task in tasks_data
+                AgentTask(**task, status=TaskStatus.PENDING)
+                for task in plan_json
             ]
-            
+
             self.thoughts.append(AgentThought(
                 step=len(self.thoughts) + 1,
-                type='planning',
-                content=f"Created plan with {len(tasks)} tasks"
+                type="planning",
+                content=f"Plan created: {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)}"
+                type="planning",
+                content=f"Planning failed: {e}"
             ))
             return []
-    
+
+    # -------------------------------------------
+    # TOOL EXECUTION PHASE
+    # -------------------------------------------
+
     async def execute_task(self, task: AgentTask) -> AgentTask:
-        """Execute a single task using MCP tools"""
-        
+
         self.thoughts.append(AgentThought(
             step=len(self.thoughts) + 1,
-            type='tool_call',
-            content=f"Executing: {task.description}",
+            type="tool_call",
+            content=f"Executing task: {task.description}",
             tool_name=task.tool,
             tool_args=task.args
         ))
-        
+
         task.status = TaskStatus.IN_PROGRESS
-        
+
         try:
-            # Call MCP tool
             result = await self.mcp_client.call_tool(task.tool, task.args)
-            
-            task.result = result
+
             task.status = TaskStatus.COMPLETED
-            
+            task.result = result
+
             self.thoughts.append(AgentThought(
                 step=len(self.thoughts) + 1,
-                type='tool_call',
+                type="tool_call",
                 content=f"✓ Completed: {task.description}",
                 tool_name=task.tool,
                 tool_result=result
             ))
-            
-            return task
-            
+
         except Exception as e:
-            task.error = str(e)
             task.status = TaskStatus.FAILED
-            
+            task.error = str(e)
+
             self.thoughts.append(AgentThought(
                 step=len(self.thoughts) + 1,
-                type='tool_call',
-                content=f"✗ Failed: {task.description} - {str(e)}",
+                type="tool_call",
+                content=f"✗ Failed: {task.description} — {e}",
                 tool_name=task.tool
             ))
-            
-            return task
-    
-    async def reflect(self, tasks: List[AgentTask], original_request: str) -> str:
-        """
-        Reflect on execution results and create final answer
-        
-        Args:
-            tasks: Executed tasks
-            original_request: Original user request
-            
-        Returns:
-            Final answer string
-        """
+
+        return task
+
+    # -------------------------------------------
+    # REFLECTION PHASE
+    # -------------------------------------------
+
+    async def reflect(self, tasks: List[AgentTask], request: str) -> str:
+
         self.thoughts.append(AgentThought(
             step=len(self.thoughts) + 1,
-            type='reflection',
-            content="Analyzing results and creating response..."
+            type="reflection",
+            content="Analyzing final results..."
         ))
-        
-        # Compile results
-        results_summary = []
-        for task in tasks:
-            if task.status == TaskStatus.COMPLETED:
-                results_summary.append(f"✓ {task.description}: {str(task.result)[:200]}")
+
+        results_string = []
+        for t in tasks:
+            if t.status == TaskStatus.COMPLETED:
+                short = str(t.result)[:200]
+                results_string.append(f"✓ {t.description}: {short}")
             else:
-                results_summary.append(f"✗ {task.description}: {task.error}")
-        
-        reflection_prompt = f"""You are an autonomous life admin agent. Review the execution results and create a helpful response.
+                results_string.append(f"✗ {t.description}: {t.error}")
 
-ORIGINAL REQUEST: {original_request}
+        reflection_prompt = f"""
+Summarize the final results of the following tasks:
 
-EXECUTION RESULTS:
-{chr(10).join(results_summary)}
+REQUEST:
+{request}
 
-Provide a clear, helpful response to the user about what was accomplished. Be specific about:
-1. What tasks were completed successfully
-2. What outputs were created (files, calendar events, etc.)
-3. Any issues encountered
-4. Next steps if applicable
+RESULTS:
+{chr(10).join(results_string)}
 
-Keep response concise but informative.
+Give a clear, helpful answer:
+- What succeeded
+- What failed
+- What files/events/emails were produced
+- Next steps
 """
-        
+
         try:
-            final_answer = await get_llm_response(reflection_prompt, temperature=0.7)
-            
+            answer = await get_llm_response(reflection_prompt, temperature=0.5)
+
             self.thoughts.append(AgentThought(
                 step=len(self.thoughts) + 1,
-                type='answer',
-                content=final_answer
+                type="answer",
+                content=answer
             ))
-            
-            # Store in memory
+
+            # Write to memory
             self.memory.add_memory(
-                f"Request: {original_request}\nResult: {final_answer}",
-                metadata={'type': 'task_completion', 'timestamp': time.time()}
+                f"Request: {request}\nAnswer: {answer}",
+                metadata={"type": "task_completion", "timestamp": time.time()}
             )
-            
-            return final_answer
-            
+
+            return answer
+
         except Exception as e:
-            error_msg = f"Reflection failed: {str(e)}"
+            errmsg = f"Reflection failed: {e}"
+
             self.thoughts.append(AgentThought(
                 step=len(self.thoughts) + 1,
-                type='answer',
-                content=error_msg
+                type="answer",
+                content=errmsg
             ))
-            return error_msg
-    
-    async def execute(self, user_request: str, files: List[str] = None, stream_thoughts: bool = False):
+
+            return errmsg
+
+    # -------------------------------------------
+    # STREAMING EXECUTION LOOP (FIXED)
+    # -------------------------------------------
+
+    async def execute(self, request: str, files: List[str] = None, stream_thoughts=False):
         """
-        Main execution loop - plan, execute, reflect
-        
-        Args:
-            user_request: User's natural language request
-            files: Uploaded files to process
-            stream_thoughts: Whether to yield thoughts as they happen
-            
-        Yields:
-            Thoughts if stream_thoughts=True
-            
-        Returns:
-            Final answer and complete thought trace
+        If stream_thoughts=True → yields AgentThought objects
+        If stream_thoughts=False → returns (answer, thoughts)
         """
+
         self.reset_context()
-        
-        # Phase 1: Planning
-        if stream_thoughts:
-            yield self.thoughts[-1] if self.thoughts else None
-        
-        tasks = await self.plan(user_request, files)
-        
+
+        # --- PLANNING ---
+        tasks = await self.plan(request, files)
         if stream_thoughts:
-            for thought in self.thoughts[-2:]:  # Last 2 planning thoughts
-                yield thought
-        
+            for th in self.thoughts:
+                yield th
+
         if not tasks:
-            error_thought = AgentThought(
+            # DO NOT return a value — async generator cannot return a value
+            thought = AgentThought(
                 step=len(self.thoughts) + 1,
-                type='answer',
-                content="Could not create execution plan. Please rephrase your request."
+                type="answer",
+                content="Could not create plan. Try rephrasing."
             )
-            self.thoughts.append(error_thought)
-            return error_thought.content, self.thoughts
-        
-        # Phase 2: Execution
-        executed_tasks = []
-        for task in tasks:
-            executed_task = await self.execute_task(task)
-            executed_tasks.append(executed_task)
-            
+            self.thoughts.append(thought)
             if stream_thoughts:
-                yield self.thoughts[-1]  # Latest thought
-        
-        # Phase 3: Reflection
-        final_answer = await self.reflect(executed_tasks, user_request)
-        
+                yield thought
+            return
+
+        # --- EXECUTION ---
+        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]  # Final answer thought
-        
+            yield self.thoughts[-1]
+            return
+
+        # If NOT streaming: return normal output
         return final_answer, self.thoughts
-    
+
+    # -------------------------------------------
+    # UTILITY
+    # -------------------------------------------
+
     def get_thought_trace(self) -> List[Dict]:
-        """Get formatted thought trace for UI display"""
-        return [asdict(thought) for thought in self.thoughts]
-    
-    async def process_files_to_rag(self, files: List[Dict[str, str]]):
-        """Process uploaded files and add to RAG engine"""
-        for file_info in files:
-            try:
-                # Extract text based on file type
-                if file_info['path'].endswith('.pdf'):
-                    from utils.pdf_utils import extract_text_from_pdf
-                    text = extract_text_from_pdf(file_info['path'])
-                elif file_info['path'].endswith(('.png', '.jpg', '.jpeg')):
-                    # Use OCR tool
-                    result = await self.mcp_client.call_tool(
-                        'ocr_extract_text',
-                        {'file_path': file_info['path'], 'language': 'en'}
-                    )
-                    text = result.get('text', '')
-                else:
-                    with open(file_info['path'], 'r', encoding='utf-8') as f:
-                        text = f.read()
-                
-                # Add to RAG
-                await self.rag_engine.add_document(
-                    text=text,
-                    metadata={'filename': file_info['name'], 'path': file_info['path']}
-                )
-                
-            except Exception as e:
-                print(f"Error processing {file_info['name']}: {e}")
-    
-    async def manual_tool_call(self, tool_name: str, args: Dict[str, Any]) -> Any:
-        """Direct tool call for manual mode"""
-        try:
-            result = await self.mcp_client.call_tool(tool_name, args)
-            return {'success': True, 'result': result}
-        except Exception as e:
-            return {'success': False, 'error': str(e)}
+        return [asdict(t) for t in self.thoughts]