agent/agent_core.py

"""
LifeAdmin AI - Core Agent Logic
Final stable version (No async-generators – fully HF compatible)
"""

import asyncio
import json
import time
from typing import List, Dict, Any, Optional
from dataclasses import dataclass, asdict
from enum import Enum

from agent.mcp_client import MCPClient
from agent.rag_engine import RAGEngine
from agent.memory import MemoryStore
from utils.llm_utils import get_llm_response


# =============================
# DATA MODELS
# =============================

class TaskStatus(Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class AgentThought:
    step: int
    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()


@dataclass
class AgentTask:
    id: str
    description: str
    tool: str
    args: Dict[str, Any]
    status: TaskStatus = TaskStatus.PENDING
    result: Optional[Any] = None
    error: Optional[str] = None


# =============================
# MAIN AGENT CLASS
# =============================

class LifeAdminAgent:

    def __init__(self):
        self.mcp_client = MCPClient()
        self.rag_engine = RAGEngine()
        self.memory = MemoryStore()
        self.thoughts: List[AgentThought] = []

    # -----------------------
    # UTIL
    # -----------------------

    def reset(self):
        self.thoughts = []

    # -----------------------
    # PLANNING
    # -----------------------

    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: {user_request}"
        ))

        tools = await self.mcp_client.list_tools()
        tool_desc = "\n".join([f"- {t['name']}: {t['description']}" for t in tools])

        rag_docs = []
        if user_request.strip():
            rag_docs = await self.rag_engine.search(user_request, k=3)

        rag_context = "\n".join([d["text"][:200] for d in rag_docs]) if rag_docs else "None"

        memory_context = self.memory.get_relevant_memories(user_request)

        prompt = f"""
You are a task planner.

REQUEST:
{user_request}

FILES: {files or []}

TOOLS:
{tool_desc}

RAG CONTEXT:
{rag_context}

MEMORY:
{memory_context}

Return ONLY JSON list:
[
  {{
    "id": "task1",
    "description": "Extract text",
    "tool": "ocr_extract_text",
    "args": {{"file_path": "x.pdf"}}
  }}
]
"""

        response = await get_llm_response(prompt, temperature=0.2)
        text = response.strip()

        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,
                type="planning",
                content="Planning failed (invalid JSON)"
            ))
            return []

        self.thoughts.append(AgentThought(
            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.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)
            task.result = result
            task.status = TaskStatus.COMPLETED

            self.thoughts.append(AgentThought(
                step=len(self.thoughts)+1,
                type="tool_call",
                content=f"✓ Completed",
                tool_name=task.tool,
                tool_result=result
            ))

        except Exception as e:
            task.status = TaskStatus.FAILED
            task.error = str(e)

            self.thoughts.append(AgentThought(
                step=len(self.thoughts)+1,
                type="tool_call",
                content=f"✗ Failed: {e}",
                tool_name=task.tool
            ))

        return task

    # -----------------------
    # REFLECTION
    # -----------------------

    async def reflect(self, tasks: List[AgentTask], original: str) -> str:

        self.thoughts.append(AgentThought(
            step=len(self.thoughts)+1,
            type="reflection",
            content="Summarizing results..."
        ))

        results = []
        for t in tasks:
            if t.status == TaskStatus.COMPLETED:
                results.append(f"✓ {t.description}: {str(t.result)[:200]}")
            else:
                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.
"""

        answer = await get_llm_response(prompt, temperature=0.4)

        self.thoughts.append(AgentThought(
            step=len(self.thoughts)+1,
            type="answer",
            content=answer
        ))

        self.memory.add_memory(
            content=f"Request: {original}\nAnswer: {answer}",
            metadata={"timestamp": time.time()}
        )

        return answer

    # -----------------------
    # MAIN EXECUTION (FULLY FIXED)
    # -----------------------

    async def execute(self, user_request: str, files: List[str] = None):
        """
        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:
            # return is allowed now
            return "Could not generate plan. Try rephrasing.", self.thoughts

        executed = []
        for t in tasks:
            executed.append(await self.execute_task(t))

        final_answer = await self.reflect(executed, user_request)
        return final_answer, self.thoughts

    # -----------------------
    # EXPORT THOUGHTS
    # -----------------------

    def get_thought_trace(self):
        return [asdict(t) for t in self.thoughts]