agents/code_agent.py

import sys
import io
import traceback
from agents.base_agent import BaseAgent, AgentResult

_BLOCKED_PATTERNS = [
    "import os", "import sys", "import subprocess", "__import__",
    "open(", "exec(", "eval(", "shutil", "socket", "requests",
]

_SAFE_BUILTINS = {
    "__builtins__": {
        "print": print, "range": range, "len": len, "list": list,
        "dict": dict, "str": str, "int": int, "float": float,
        "bool": bool, "tuple": tuple, "set": set, "sum": sum,
        "max": max, "min": min, "sorted": sorted, "enumerate": enumerate,
        "zip": zip, "map": map, "filter": filter, "abs": abs,
        "round": round, "isinstance": isinstance, "type": type,
    }
}


class CodeAgent(BaseAgent):
    def __init__(self):
        super().__init__(
            name="CodeAgent",
            role="expert software engineer who writes clean, functional Python code",
        )

    def _sandbox_exec(self, code: str) -> tuple:
        for pattern in _BLOCKED_PATTERNS:
            if pattern in code:
                return f"Blocked: '{pattern}' is not permitted in sandboxed execution.", False

        old_stdout = sys.stdout
        sys.stdout = io.StringIO()
        try:
            exec(code, dict(_SAFE_BUILTINS))
            output = sys.stdout.getvalue()
            return output or "Executed successfully (no printed output).", True
        except Exception:
            return traceback.format_exc(), False
        finally:
            sys.stdout = old_stdout

    def act(self, task: str, **kwargs) -> AgentResult:
        thinking = self.think(task)

        system = (
            f"You are {self.name}. Write Python code to complete the task.\n"
            "Return ONLY the raw Python code — no markdown fences, no explanations.\n"
            "The code should print its results using print()."
        )
        user = f"Task: {task}\n\nPython code:"
        raw_code = self._call_llm(system, user)

        # Strip markdown fences if present
        code = raw_code.strip()
        if code.startswith("```"):
            lines = code.split("\n")
            code = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])

        exec_output, success = self._sandbox_exec(code)

        output = f"**Generated Code:**\n```python\n{code}\n```\n\n**Output:**\n```\n{exec_output}\n```"

        return AgentResult(
            agent_name=self.name,
            task=task,
            thinking=thinking,
            output=output,
            tool_calls=[{"tool": "sandbox_exec", "success": success}],
            success=success,
            quality_score=0.9 if success else 0.4,
        )