Add agent.py — core agent with router, planner, executor, memory, safety

multeclaw/agent.py

@@ -0,0 +1,505 @@
+"""
+Multeclaw Agent Core — reasoning loops, tool execution, planning, memory, repair, and safety.
+This is the brain of the system.
+"""
+
+import json
+import math
+import time
+import traceback
+import subprocess
+import tempfile
+import os
+from typing import Generator, Optional, Any
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+
+from multeclaw.config import AgentConfig, BUILT_IN_TOOLS, SYSTEM_PROMPTS
+from multeclaw.llm_client import MultiModelClient, LLMResponse
+
+
+# ─── Types ─────────────────────────────────────────────────────────────────────
+class TaskStatus(str, Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+@dataclass
+class Step:
+    id: int
+    description: str
+    status: TaskStatus = TaskStatus.PENDING
+    result: str = ""
+    error: str = ""
+
+
+@dataclass
+class AgentMemory:
+    """Short-term conversation + long-term preferences."""
+    conversation: list[dict] = field(default_factory=list)
+    task_plans: list[list[Step]] = field(default_factory=list)
+    tool_results: list[dict] = field(default_factory=list)
+    user_preferences: dict = field(default_factory=dict)
+    session_notes: list[str] = field(default_factory=list)
+
+    def add_message(self, role: str, content: str):
+        self.conversation.append({
+            "role": role,
+            "content": content,
+            "timestamp": datetime.now().isoformat(),
+        })
+
+    def get_messages(self, limit: int = 50) -> list[dict]:
+        """Return recent messages in OpenAI format."""
+        recent = self.conversation[-limit:]
+        return [{"role": m["role"], "content": m["content"]} for m in recent]
+
+    def add_tool_result(self, tool_name: str, args: dict, result: str, success: bool):
+        self.tool_results.append({
+            "tool": tool_name,
+            "args": args,
+            "result": result[:2000],
+            "success": success,
+            "timestamp": datetime.now().isoformat(),
+        })
+
+    def clear(self):
+        self.conversation.clear()
+        self.task_plans.clear()
+        self.tool_results.clear()
+        self.session_notes.clear()
+
+
+# ─── Tool Executor ─────────────────────────────────────────────────────────────
+class ToolExecutor:
+    """Sandboxed execution of built-in tools with safety checks."""
+
+    SAFE_MATH_NAMES = {
+        "abs": abs, "round": round, "min": min, "max": max, "sum": sum,
+        "pow": pow, "int": int, "float": float, "len": len,
+        "sqrt": math.sqrt, "log": math.log, "log2": math.log2, "log10": math.log10,
+        "sin": math.sin, "cos": math.cos, "tan": math.tan,
+        "pi": math.pi, "e": math.e, "inf": math.inf,
+        "ceil": math.ceil, "floor": math.floor, "factorial": math.factorial,
+    }
+
+    def execute(self, tool_name: str, arguments: dict) -> dict:
+        """Execute a tool and return {result, success, error}."""
+        try:
+            if tool_name == "calculator":
+                return self._calculator(arguments["expression"])
+            elif tool_name == "code_executor":
+                return self._code_executor(arguments["code"])
+            elif tool_name == "file_reader":
+                return self._file_reader(arguments["path"])
+            elif tool_name == "file_writer":
+                return self._file_writer(arguments["path"], arguments["content"])
+            elif tool_name == "web_search":
+                return {"result": "🔍 Web search is not available in this environment. Reason about the question using your knowledge.", "success": True}
+            else:
+                return {"result": "", "success": False, "error": f"Unknown tool: {tool_name}"}
+        except Exception as e:
+            return {"result": "", "success": False, "error": f"{type(e).__name__}: {str(e)}"}
+
+    def _calculator(self, expression: str) -> dict:
+        """Safe math evaluation."""
+        blocked = ["import", "exec", "eval", "open", "__", "os.", "sys.", "subprocess"]
+        expr_lower = expression.lower()
+        for b in blocked:
+            if b in expr_lower:
+                return {"result": "", "success": False, "error": f"Blocked unsafe expression containing '{b}'"}
+        try:
+            result = eval(expression, {"__builtins__": {}}, self.SAFE_MATH_NAMES)
+            return {"result": str(result), "success": True}
+        except Exception as e:
+            return {"result": "", "success": False, "error": f"Math error: {str(e)}"}
+
+    def _code_executor(self, code: str) -> dict:
+        """Execute Python code in a subprocess with timeout."""
+        blocked = ["os.system", "subprocess.call", "subprocess.run", "shutil.rmtree",
+                    "os.remove", "os.unlink", "os.rmdir"]
+        for b in blocked:
+            if b in code:
+                return {"result": "", "success": False, "error": f"Blocked unsafe code containing '{b}'"}
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write(code)
+            f.flush()
+            try:
+                result = subprocess.run(
+                    ["python", f.name],
+                    capture_output=True, text=True, timeout=30,
+                    env={**os.environ, "PYTHONDONTWRITEBYTECODE": "1"},
+                )
+                output = result.stdout
+                if result.stderr:
+                    output += f"\n[stderr]: {result.stderr}"
+                return {"result": output[:5000], "success": result.returncode == 0,
+                        "error": result.stderr[:1000] if result.returncode != 0 else ""}
+            except subprocess.TimeoutExpired:
+                return {"result": "", "success": False, "error": "Code execution timed out (30s limit)"}
+            finally:
+                os.unlink(f.name)
+
+    def _file_reader(self, path: str) -> dict:
+        """Read a file safely."""
+        path = os.path.abspath(path)
+        if not os.path.exists(path):
+            return {"result": "", "success": False, "error": f"File not found: {path}"}
+        if os.path.getsize(path) > 1_000_000:
+            return {"result": "", "success": False, "error": "File too large (>1MB)"}
+        try:
+            with open(path, "r", encoding="utf-8", errors="replace") as f:
+                content = f.read()
+            return {"result": content, "success": True}
+        except Exception as e:
+            return {"result": "", "success": False, "error": str(e)}
+
+    def _file_writer(self, path: str, content: str) -> dict:
+        """Write a file safely."""
+        path = os.path.abspath(path)
+        try:
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            with open(path, "w", encoding="utf-8") as f:
+                f.write(content)
+            return {"result": f"Written {len(content)} bytes to {path}", "success": True}
+        except Exception as e:
+            return {"result": "", "success": False, "error": str(e)}
+
+
+# ─── Safety Layer ──────────────────────────────────────────────────────────────
+class SafetyLayer:
+    """Content filtering and request safety checks."""
+
+    BLOCKED_PATTERNS = [
+        "make a bomb", "how to hack", "generate malware", "create a virus",
+        "steal credentials", "phishing email", "exploit vulnerability",
+    ]
+
+    @classmethod
+    def check_input(cls, text: str) -> tuple[bool, str]:
+        """Returns (is_safe, reason)."""
+        text_lower = text.lower()
+        for pattern in cls.BLOCKED_PATTERNS:
+            if pattern in text_lower:
+                return False, f"Request blocked: contains disallowed content pattern."
+        return True, ""
+
+    @classmethod
+    def check_output(cls, text: str) -> tuple[bool, str]:
+        """Basic output safety check."""
+        return True, ""
+
+
+# ─── Router ────────────────────────────────────────────────────────────────────
+class TaskType(str, Enum):
+    DIRECT = "direct"
+    TOOL_ASSISTED = "tool"
+    MULTI_STEP = "multi_step"
+    CODE = "code"
+    ANALYSIS = "analysis"
+
+
+class TaskRouter:
+    """Classifies incoming requests to determine the execution path."""
+
+    CODE_SIGNALS = ["write code", "implement", "function", "class", "script",
+                    "python", "javascript", "debug", "fix this code", "refactor"]
+    TOOL_SIGNALS = ["calculate", "compute", "search", "look up", "find", "read file",
+                    "write file", "execute", "run"]
+    MULTI_STEP_SIGNALS = ["step by step", "plan", "analyze and", "research",
+                          "compare", "build", "create a complete", "design"]
+    ANALYSIS_SIGNALS = ["analyze", "evaluate", "assess", "review", "audit",
+                        "summarize this data", "what are the trends"]
+
+    @classmethod
+    def classify(cls, message: str) -> TaskType:
+        msg_lower = message.lower()
+        scores = {
+            TaskType.CODE: sum(1 for s in cls.CODE_SIGNALS if s in msg_lower),
+            TaskType.TOOL_ASSISTED: sum(1 for s in cls.TOOL_SIGNALS if s in msg_lower),
+            TaskType.MULTI_STEP: sum(1 for s in cls.MULTI_STEP_SIGNALS if s in msg_lower),
+            TaskType.ANALYSIS: sum(1 for s in cls.ANALYSIS_SIGNALS if s in msg_lower),
+        }
+        best = max(scores, key=scores.get)
+        if scores[best] > 0:
+            return best
+        return TaskType.DIRECT
+
+
+# ─── Planner ───────────────────────────────────────────────────────────────────
+class Planner:
+    """Decomposes complex tasks into executable steps using the LLM."""
+
+    PLANNING_PROMPT = """You are a task planner. Break down the user's request into 2-6 concrete, actionable steps.
+Return ONLY a JSON array of step descriptions. Example:
+["Research the topic", "Outline the structure", "Write the first draft", "Review and refine"]
+
+User request: {request}
+
+Steps (JSON array only):"""
+
+    @staticmethod
+    def create_plan_from_llm(client: MultiModelClient, model_name: str, request: str) -> list[Step]:
+        """Use LLM to generate a plan."""
+        try:
+            resp = client.complete(
+                model_name=model_name,
+                messages=[{"role": "user", "content": Planner.PLANNING_PROMPT.format(request=request)}],
+                temperature=0.3,
+                max_tokens=500,
+            )
+            if resp.error:
+                return [Step(id=1, description=request)]
+            text = resp.content.strip()
+            if "```" in text:
+                text = text.split("```")[1]
+                if text.startswith("json"):
+                    text = text[4:]
+                text = text.strip()
+            steps_list = json.loads(text)
+            return [Step(id=i + 1, description=s) for i, s in enumerate(steps_list)]
+        except Exception:
+            return [Step(id=1, description=request)]
+
+    @staticmethod
+    def create_simple_plan(request: str) -> list[Step]:
+        return [Step(id=1, description=request)]
+
+
+# ─── Agent Core ────────────────────────────────────────────────────────────────
+class MulteclawAgent:
+    """
+    The main agent orchestrator. Runs the reasoning loop:
+    Router → Planner → Executor → Verifier → Reporter
+    with repair on failure.
+    """
+
+    def __init__(self):
+        self.client = MultiModelClient()
+        self.memory = AgentMemory()
+        self.tools = ToolExecutor()
+        self.config = AgentConfig()
+        self.logs: list[str] = []
+
+    def _log(self, msg: str):
+        entry = f"[{datetime.now().strftime('%H:%M:%S')}] {msg}"
+        self.logs.append(entry)
+
+    def get_logs(self, n: int = 50) -> str:
+        return "\n".join(self.logs[-n:])
+
+    # ─── Main Entry: Streaming Chat ────────────────────────────────────────
+    def chat_stream(
+        self,
+        message: str,
+        history: list[dict],
+        model_name: str,
+        system_prompt: str = "",
+        temperature: float = 0.7,
+        max_tokens: int = 4096,
+        enable_tools: bool = True,
+        enable_planning: bool = True,
+    ) -> Generator[str, None, None]:
+        """
+        Main chat entry point. Yields partial response text for streaming.
+        Handles the full agent loop: safety → routing → planning → execution → verification.
+        """
+        # 1. Safety check
+        is_safe, reason = SafetyLayer.check_input(message)
+        if not is_safe:
+            yield f"⚠️ {reason}"
+            return
+
+        # 2. Update memory
+        self.memory.add_message("user", message)
+        self._log(f"User: {message[:80]}...")
+
+        # 3. Route the task
+        task_type = TaskRouter.classify(message)
+        self._log(f"Router → {task_type.value}")
+
+        # 4. Build context
+        all_messages = self._build_context(history, message)
+
+        # 5. Execute based on task type
+        if task_type == TaskType.MULTI_STEP and enable_planning:
+            yield from self._execute_multi_step(model_name, all_messages, system_prompt, temperature, max_tokens, message)
+        elif task_type == TaskType.TOOL_ASSISTED and enable_tools:
+            yield from self._execute_with_tools(model_name, all_messages, system_prompt, temperature, max_tokens)
+        else:
+            yield from self._execute_direct(model_name, all_messages, system_prompt, temperature, max_tokens)
+
+    # ─── Direct Execution ──────────────────────────────────────────────────
+    def _execute_direct(self, model_name, messages, system_prompt, temperature, max_tokens) -> Generator[str, None, None]:
+        """Simple streaming completion."""
+        self._log("Executing: direct stream")
+        full_response = ""
+        for chunk in self.client.stream(model_name, messages, system_prompt, temperature, max_tokens):
+            full_response += chunk
+            yield chunk
+        self.memory.add_message("assistant", full_response)
+        self._log(f"Response: {len(full_response)} chars")
+
+    # ─── Tool-Assisted Execution ───────────────────────────────────────────
+    def _execute_with_tools(self, model_name, messages, system_prompt, temperature, max_tokens) -> Generator[str, None, None]:
+        """
+        Agentic tool loop:
+        1. Ask LLM what tool to call (if any)
+        2. Execute the tool
+        3. Feed result back to LLM
+        4. Repeat until LLM gives a final answer
+        """
+        self._log("Executing: tool-assisted loop")
+
+        tool_prompt = system_prompt + "\n\nYou have access to these tools:\n"
+        for name, tool in BUILT_IN_TOOLS.items():
+            tool_prompt += f"- **{name}**: {tool['description']}\n"
+        tool_prompt += (
+            "\nTo use a tool, respond with a JSON block:\n"
+            '```json\n{"tool": "tool_name", "arguments": {"arg1": "value1"}}\n```\n'
+            "After receiving the result, provide your final answer.\n"
+            "If no tool is needed, just answer directly."
+        )
+
+        current_messages = list(messages)
+        attempts = 0
+
+        while attempts < self.config.max_tool_retries:
+            attempts += 1
+            self._log(f"Tool loop iteration {attempts}")
+
+            resp = self.client.complete(model_name, current_messages, tool_prompt, temperature, max_tokens)
+
+            if resp.error:
+                yield f"❌ Error: {resp.error}"
+                return
+
+            content = resp.content
+            tool_call = self._extract_tool_call(content)
+
+            if tool_call:
+                tool_name = tool_call["tool"]
+                tool_args = tool_call["arguments"]
+                self._log(f"Tool call: {tool_name}({json.dumps(tool_args)[:60]})")
+
+                yield f"🔧 Using **{tool_name}**"
+                if tool_name == "calculator":
+                    yield f" → `{tool_args.get('expression', '')}`\n"
+                elif tool_name == "code_executor":
+                    yield f"\n```python\n{tool_args.get('code', '')[:200]}\n```\n"
+                else:
+                    yield f" → `{json.dumps(tool_args)[:100]}`\n"
+
+                result = self.tools.execute(tool_name, tool_args)
+                self.memory.add_tool_result(tool_name, tool_args, result.get("result", ""), result.get("success", False))
+
+                if result["success"]:
+                    yield f"✅ Result: `{result['result'][:200]}`\n\n"
+                    self._log(f"Tool success: {result['result'][:60]}")
+                else:
+                    yield f"⚠️ Error: {result.get('error', 'Unknown error')}\n\n"
+                    self._log(f"Tool error: {result.get('error', '')[:60]}")
+
+                current_messages.append({"role": "assistant", "content": content})
+                current_messages.append({"role": "user", "content": f"Tool result for {tool_name}: {result['result'][:2000]}"})
+            else:
+                yield content
+                self.memory.add_message("assistant", content)
+                return
+
+        yield "\n\n⚠️ Reached maximum tool iterations. Here's what I have so far."
+
+    # ─── Multi-Step Execution ──────────────────────────────────────────────
+    def _execute_multi_step(self, model_name, messages, system_prompt, temperature, max_tokens, original_request) -> Generator[str, None, None]:
+        """
+        Planning loop:
+        1. Generate a plan
+        2. Execute each step
+        3. Combine results
+        """
+        self._log("Executing: multi-step planning")
+
+        yield "📋 **Planning...**\n\n"
+        plan = Planner.create_plan_from_llm(self.client, model_name, original_request)
+        self.memory.task_plans.append(plan)
+
+        for step in plan:
+            yield f"  {step.id}. {step.description}\n"
+        yield "\n---\n\n"
+
+        step_results = []
+        for step in plan:
+            step.status = TaskStatus.RUNNING
+            yield f"**Step {step.id}: {step.description}**\n"
+            self._log(f"Step {step.id}: {step.description}")
+
+            step_context = f"""You are executing step {step.id} of a plan.
+Original request: {original_request}
+Current step: {step.description}
+Previous step results: {json.dumps(step_results[-3:]) if step_results else 'None yet'}
+
+Execute this step thoroughly. Be specific and detailed."""
+
+            step_messages = messages + [{"role": "user", "content": step_context}]
+
+            step_response = ""
+            for chunk in self.client.stream(model_name, step_messages, system_prompt, temperature, max_tokens):
+                step_response += chunk
+                yield chunk
+
+            step.status = TaskStatus.COMPLETED
+            step.result = step_response
+            step_results.append({"step": step.id, "description": step.description, "result": step_response[:500]})
+            yield "\n\n"
+
+        self.memory.add_message("assistant", f"[Multi-step task completed with {len(plan)} steps]")
+        self._log(f"Plan completed: {len(plan)} steps")
+
+    # ─── Helpers ────────────────���──────────────────────────────────────────
+    def _build_context(self, history: list[dict], current_message: str) -> list[dict]:
+        """Build the message context from history + current message."""
+        messages = []
+        for msg in history:
+            messages.append({"role": msg["role"], "content": msg["content"]})
+        messages.append({"role": "user", "content": current_message})
+        return messages
+
+    @staticmethod
+    def _extract_tool_call(text: str) -> Optional[dict]:
+        """Extract a JSON tool call from LLM response text."""
+        if "```json" in text:
+            try:
+                json_str = text.split("```json")[1].split("```")[0].strip()
+                obj = json.loads(json_str)
+                if "tool" in obj and "arguments" in obj:
+                    return obj
+            except (IndexError, json.JSONDecodeError):
+                pass
+
+        if '"tool"' in text and '"arguments"' in text:
+            start = text.find("{")
+            if start >= 0:
+                depth = 0
+                for i in range(start, len(text)):
+                    if text[i] == "{":
+                        depth += 1
+                    elif text[i] == "}":
+                        depth -= 1
+                        if depth == 0:
+                            try:
+                                obj = json.loads(text[start:i + 1])
+                                if "tool" in obj and "arguments" in obj:
+                                    return obj
+                            except json.JSONDecodeError:
+                                pass
+                            break
+        return None
+
+    def clear_memory(self):
+        self.memory.clear()
+        self.logs.clear()
+        self._log("Memory and logs cleared")