backend/agent/react_agent.py

# backend/agent/react_agent.py
import time
import json
import re
from dataclasses import dataclass
from typing import Any, Callable, Optional, AsyncGenerator

try:
    from google import genai
    from config import GEMINI_API_KEY
except ImportError:
    pass  # We'll assume these exist based on the core/brain.py implementation

@dataclass
class Tool:
    name: str
    description: str
    parameters: dict  # JSON Schema
    handler: Callable  # async function that executes the tool

@dataclass
class ReActStep:
    step_type: str  # "think" | "act" | "observe" | "reflect" | "final_answer"
    content: str
    tool_name: Optional[str] = None
    tool_input: Optional[dict] = None
    tool_output: Optional[Any] = None
    duration_ms: int = 0
    confidence: float = 1.0

class ReActAgent:
    def __init__(self, personality: str, tools: list[Tool], memory_client, language: str = "en"):
        self.personality = personality  # "jarvis" or "friday"
        self.language = language
        self.tools: dict[str, Tool] = {t.name: t for t in tools}
        self.memory = memory_client
        self.trace: list[ReActStep] = []
        self._client = None
        try:
            self._client = genai.Client(api_key=GEMINI_API_KEY)
        except Exception as e:
            import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")

    def _build_system_prompt(self) -> str:
        tools_str = json.dumps([{
            "name": t.name, 
            "description": t.description, 
            "parameters": t.parameters
        } for t in self.tools.values()], indent=2)

        try:
            from modules.assistant_identity import set_mode, get_identity_persona_prompt
            set_mode(self.personality.lower())
            identity_prompt = get_identity_persona_prompt()
        except ImportError:
            identity_prompt = f"You are {self.personality.upper()}, an autonomous agent."

        # ── Micro research knowledge block ────────────────────────────────────
        # Injects top 3 HIGH/CRITICAL notes only — titles only, no summaries.
        # Hard cap: 400 chars total. Zero impact on normal conversations.
        research_block = ""
        try:
            from backend.omega.research_engine import get_research_notes
            notes = get_research_notes(limit=20)
            top = [
                n for n in notes
                if n.get("importance") in ("HIGH", "CRITICAL")
            ][:3]
            if top:
                lines = [f"• {n['title']}" for n in top]
                block = "RECENT KNOWLEDGE (top findings):\n" + "\n".join(lines)
                research_block = "\n\n" + block[:400]  # hard cap 400 chars
        except Exception:
            pass  # Never block chat if research DB is unavailable
        # ─────────────────────────────────────────────────────────────────────

        return f"""{identity_prompt}
You operate in a ReAct (Reasoning and Acting) loop.
Respond in {self.language} unless the user explicitly asks for another language.{research_block}

AVAILABLE TOOLS:
{tools_str}

OUTPUT FORMAT:
You must ALWAYS respond with a SINGLE valid JSON object matching this schema. Do not include markdown codeblocks, just the raw JSON:
{{
  "step_type": "think" | "act" | "final_answer",
  "content": "Your internal monologue, reasoning, or final response",
  "tool_name": "name of the tool if step_type is 'act'",
  "tool_input": {{"param": "value"}} // if step_type is 'act'
}}

RULES:
1. Always 'think' before you 'act'.
2. If you need information from the system or world, use a tool ('act').
3. When you have enough information, output 'final_answer' with the content being your response.
4. Keep your personality intact during 'final_answer'.
"""

    def _build_prompt(self, user_input: str, memories: list, context: dict) -> str:
        prompt = f"USER REQUEST: {user_input}\n\n"
        if memories:
            prompt += "RELEVANT MEMORIES:\n"
            for m in memories:
                prompt += f"- {m}\n"
            prompt += "\n"
        if context:
            prompt += f"CONTEXT: {json.dumps(context)}\n\n"
        prompt += "BEGIN REASONING LOOP:\n"
        return prompt

    def _update_prompt(self, prompt: str, step: ReActStep, obs: ReActStep) -> str:
        if step.step_type == "think":
            prompt += f"\nTHOUGHT: {step.content}"
        elif step.step_type == "act":
            prompt += f"\nACTION: {step.tool_name}({json.dumps(step.tool_input)})\n"
            prompt += f"OBSERVATION: {obs.content}\n"
        return prompt

    async def _llm_step_stream(self, prompt: str, context: dict) -> AsyncGenerator[ReActStep, None]:
        t0 = time.time()
        try:
            provider = context.get("llm_provider", "OPENAI").upper()
            api_key = context.get("llm_key", "")
            
            if not api_key and provider != "OLLAMA":
                yield ReActStep("final_answer", f"Error: API Key missing for provider {provider}", duration_ms=int((time.time() - t0) * 1000))
                return

            system_prompt = self._build_system_prompt()
            
            from backend.services.connectors import stream_with_fallback
            
            raw = ""
            last_content_len = 0
            
            async for token in stream_with_fallback(system_prompt, prompt):
                raw += token
                
                # Flawless streaming parse loop: extract "content" as it streams
                match = re.search(r'"content"\s*:\s*"((?:[^"\\]|\\.)*)', raw)
                if match:
                    current_content = match.group(1)
                    try:
                        parsed_content = json.loads(f'"{current_content}"')
                        if len(parsed_content) > last_content_len:
                            delta = parsed_content[last_content_len:]
                            yield ReActStep("token", delta)
                            last_content_len = len(parsed_content)
                    except Exception as e:
                        import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")

            # Parse JSON
            try:
                if raw.startswith("```"):
                    raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw, flags=re.MULTILINE)
                data = json.loads(raw)
                yield ReActStep(
                    step_type=data.get("step_type", "think"),
                    content=data.get("content", ""),
                    tool_name=data.get("tool_name"),
                    tool_input=data.get("tool_input"),
                    duration_ms=int((time.time() - t0) * 1000)
                )
            except json.JSONDecodeError as e:
                yield ReActStep(
                    step_type="think",
                    content=f"Failed to parse JSON output: {e}. Raw: {raw}",
                    duration_ms=int((time.time() - t0) * 1000)
                )

        except Exception as e:
            logging.error(f"ReAct LLM Error: {e}")
            yield ReActStep("final_answer", f"Error querying LLM: {e}", duration_ms=int((time.time() - t0) * 1000))

    async def run(self, user_input: str, context: dict) -> AsyncGenerator[ReActStep, None]:
        # 1. Retrieve relevant memories
        memories = await self.memory.query(user_input, top_k=5) if hasattr(self.memory, 'query') else []
        # 2. Build prompt with personality, memories, tool descriptions, user input
        prompt = self._build_prompt(user_input, memories, context)
        
        # --- JARVIS 10X Emergency Playback Fast Path ---
        from backend.voice.emergency_playback import matches_emergency_replay_intent, get_live_audio_buffer, extract_duration, play_audio_immediately
        if matches_emergency_replay_intent(user_input):
            import logging
            logging.info("*** EMERGENCY REPLAY TRIGGERED ***")
            seconds = extract_duration(user_input, default=30)
            audio = await get_live_audio_buffer(seconds_back=seconds)
            await play_audio_immediately(audio)
            
            from backend.ws.agent_ws import ws_manager
            await ws_manager.broadcast({"event": "emergency_replay", "payload": {"seconds": seconds}})
            
            yield ReActStep(
                step_type="final_answer",
                content=f"Playing back the last {seconds} seconds of emergency audio, sir.",
                duration_ms=0
            )
            return

        # Check Easter Eggs first — they short-circuit the normal pipeline
        from backend.easter_eggs.engine import check_easter_eggs
        egg_response = await check_easter_eggs(
            user_input=user_input,
            ui_trigger=context.get("ui_trigger"),
            context={**context, "db_path": getattr(self, "db_path", context.get("db_path"))}
        )
        if egg_response:
            # Yield as a special step type that the frontend renders with Iron Man flair
            yield ReActStep(
                step_type="easter_egg",
                content=egg_response.get("text", ""),
                tool_output=egg_response,
            )
            return

        # 3. ReAct loop: max 10 iterations to prevent infinite loops
        for iteration in range(10):
            final_step = None
            async for partial_step in self._llm_step_stream(prompt, context):
                if partial_step.step_type == "token":
                    # We might not want to yield raw JSON to UI, but maybe we do.
                    # Or we yield it so UI can render streaming text.
                    pass # Don't yield tokens to `agent.run()`, we'll let agent_service handle it or we can yield it.
                    yield partial_step
                else:
                    final_step = partial_step
                    
            step = final_step
            self.trace.append(step)
            yield step  # stream final parsed step to UI
            
            if step.step_type == "final_answer":
                break
            
            if step.step_type == "think":
                prompt = self._update_prompt(prompt, step, None)
                continue

            if step.step_type == "act":
                tool = self.tools.get(step.tool_name)
                if not tool:
                    obs = ReActStep("observe", f"Tool '{step.tool_name}' not found", duration_ms=0)
                else:
                    t0 = time.time()
                    try:
                        result = await tool.handler(**(step.tool_input or {}))
                        obs = ReActStep("observe", str(result)[:2000], tool_output=result,
                                        duration_ms=int((time.time()-t0)*1000))
                    except Exception as e:
                        obs = ReActStep("observe", f"Tool error: {e}",
                                        duration_ms=int((time.time()-t0)*1000))
                self.trace.append(obs)
                yield obs
                prompt = self._update_prompt(prompt, step, obs)