universal-mcp-agents 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

universal_mcp/agents/codeact/sandbox.py

@@ -1,39 +1,85 @@
 import asyncio
-import builtins
 import contextlib
 import io
+from collections.abc import Callable
 from typing import Any
 
 from loguru import logger
 
+from .models import SandboxOutput
 
-async def eval_unsafe(code: str, _locals: dict[str, Any], timeout: int = 10) -> tuple[str, dict[str, Any]]:
+# Define a whitelist of safe built-in functions
+SAFE_BUILTINS = {
+    "abs": abs,
+    "all": all,
+    "any": any,
+    "bool": bool,
+    "callable": callable,
+    "chr": chr,
+    "dict": dict,
+    "divmod": divmod,
+    "enumerate": enumerate,
+    "filter": filter,
+    "float": float,
+    "getattr": getattr,
+    "hasattr": hasattr,
+    "hash": hash,
+    "id": id,
+    "int": int,
+    "isinstance": isinstance,
+    "iter": iter,
+    "len": len,
+    "list": list,
+    "max": max,
+    "min": min,
+    "next": next,
+    "ord": ord,
+    "pow": pow,
+    "print": print,
+    "range": range,
+    "repr": repr,
+    "reversed": reversed,
+    "round": round,
+    "set": set,
+    "slice": slice,
+    "sorted": sorted,
+    "str": str,
+    "sum": sum,
+    "tuple": tuple,
+    "type": type,
+    "zip": zip,
+}
+
+
+async def eval_unsafe(
+    code: str, _locals: dict[str, Callable], timeout: int = 10
+) -> tuple[SandboxOutput, dict[str, Any]]:
     """Executes a string of Python code in a sandboxed environment."""
-    # Store original keys before execution
     original_keys = set(_locals.keys())
-    result = f"Executing code...\n{code}\n\nOutput:\n"
-    result += "=" * 50 + "\n"
+    execution_context = _locals.copy()
+    execution_context["__builtins__"] = __builtins__  # TODO: Use SAFE_BUILTINS instead of __builtins__
+
+    stdout_capture = io.StringIO()
+    output = SandboxOutput(stdout="")
+
     try:
         logger.debug(f"Executing code with timeout {timeout}")
-        with contextlib.redirect_stdout(io.StringIO()) as f:
-            # Execute the code in the provided locals context
-            # This should define an async function `main`
-            exec(code, builtins.__dict__, _locals)
-
-            if "main" in _locals and asyncio.iscoroutinefunction(_locals["main"]):
-                # Run the main async function
-                await asyncio.wait_for(_locals["main"](), timeout=timeout)
+        with contextlib.redirect_stdout(stdout_capture):
+            exec(code, execution_context)
+
+            if "main" in execution_context and asyncio.iscoroutinefunction(execution_context["main"]):
+                return_val = await asyncio.wait_for(execution_context["main"](), timeout=timeout)
+                output.return_value = return_val
             else:
-                result += "\nError: No `async def main()` function found in the script."
+                output.error = "No `async def main()` function found in the script."
+
+        output.stdout = stdout_capture.getvalue()
 
-        output = f.getvalue()
-        result += output
-        if not output:
-            result += "<code ran, no output printed to stdout>"
     except Exception as e:
-        result += f"Error during execution: {repr(e)}"
+        output.error = f"{type(e).__name__}: {e}"
+        output.stdout = stdout_capture.getvalue()
+
+    new_keys = set(execution_context.keys()) - original_keys - {"__builtins__"}
+    new_vars = {key: execution_context[key] for key in new_keys}
 
-    # Determine new variables created during execution
-    new_keys = set(_locals.keys()) - original_keys
-    new_vars = {key: _locals[key] for key in new_keys}
-    return result, new_vars
+    return output, new_vars

universal_mcp/agents/codeact/state.py

@@ -7,3 +7,5 @@ class CodeActState(MessagesState):
 
     script: str | None = Field(default=None, description="The Python code script to be executed.")
     sandbox_output: str | None = Field(default=None, description="The output of the Python code script execution.")
+    syntax_error: str | None = Field(default=None, description="The syntax error from the last script validation.")
+    task_complete: bool = Field(default=False, description="Whether the task is complete.")

universal_mcp/agents/codeact0/__init__.py

@@ -0,0 +1,3 @@
+from .agent import CodeActAgent
+
+__all__ = ["CodeActAgent"]

universal_mcp/agents/codeact0/__main__.py

@@ -0,0 +1,35 @@
+import asyncio
+
+from langgraph.checkpoint.memory import MemorySaver
+from rich import print
+from universal_mcp.agentr.registry import AgentrRegistry
+
+from universal_mcp.agents.codeact0.agent import CodeActAgent
+from universal_mcp.agents.utils import messages_to_list
+
+
+async def main():
+    memory = MemorySaver()
+    agent = CodeActAgent(
+        name="CodeAct Agent",
+        instructions="Be very concise in your answers.",
+        model="anthropic:claude-4-sonnet-20250514",
+        tools={"google_mail": ["list_messages"]},
+        registry=AgentrRegistry(),
+        memory=memory,
+    )
+    print("Starting agent...")
+    # await agent.ainit()
+    # await agent.run_interactive()
+    # async for event in agent.stream(
+    #     user_input="Fetch unsubscribe links from my Gmail inbox for promo emails I have received in the last 7 days"
+    # ):
+    #     print(event.content, end="")
+    result = await agent.invoke(
+        user_input="Fetch unsubscribe links from my Gmail inbox for promo emails I have received in the last 7 days"
+    )
+    print(messages_to_list(result["messages"]))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

universal_mcp/agents/codeact0/agent.py

@@ -0,0 +1,136 @@
+import inspect
+from collections.abc import Callable
+from typing import Literal, cast
+
+from langchain_core.messages import AIMessage, ToolMessage
+from langchain_core.tools import StructuredTool
+from langchain_core.tools import tool as create_tool
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import START, StateGraph
+from langgraph.types import Command, RetryPolicy
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolConfig, ToolFormat
+
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.codeact0.llm_tool import ai_classify, call_llm, data_extractor, smart_print
+from universal_mcp.agents.codeact0.prompts import (
+    create_default_prompt,
+)
+from universal_mcp.agents.codeact0.sandbox import eval_unsafe, execute_ipython_cell
+from universal_mcp.agents.codeact0.state import CodeActState
+from universal_mcp.agents.codeact0.utils import filter_retry_on, inject_context
+from universal_mcp.agents.llm import load_chat_model
+
+
+class CodeActAgent(BaseAgent):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        memory: BaseCheckpointSaver | None = None,
+        tools: ToolConfig | None = None,
+        registry: ToolRegistry | None = None,
+        sandbox_timeout: int = 20,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            instructions=instructions,
+            model=model,
+            memory=memory,
+            **kwargs,
+        )
+        self.model_instance = load_chat_model(model, thinking=True)
+        self.tools_config = tools or {}
+        self.registry = registry
+        self.eval_fn = eval_unsafe
+        self.sandbox_timeout = sandbox_timeout
+        self.processed_tools: list[StructuredTool | Callable] = []
+
+        # TODO(manoj): Use toolformat native instead of langchain
+        # TODO(manoj, later): Add better sandboxing
+        # Old Nishant TODO s:
+        # - Make codeact faster by calling upto API call (this done but should be tested)
+        # - Add support for async eval_fn
+        # - Throw Error if code snippet is too long (> 1000 characters) and suggest to split it into smaller parts
+        # - Multiple models from config
+
+    async def _build_graph(self):
+        exported_tools = []
+        if self.tools_config:
+            if not self.registry:
+                raise ValueError("Tools are configured but no registry is provided")
+            # Langchain tools are fine
+            exported_tools = await self.registry.export_tools(self.tools_config, ToolFormat.LANGCHAIN)
+        exported_tools.extend([smart_print, data_extractor, ai_classify, call_llm])
+        self.processed_tools = [t if isinstance(t, StructuredTool) else create_tool(t) for t in exported_tools]
+        self.instructions, self.tools_context = create_default_prompt(self.processed_tools, self.instructions)
+
+        def call_model(state: CodeActState) -> Command[Literal["sandbox"]]:
+            messages = [{"role": "system", "content": self.instructions}] + state["messages"]
+
+            # Run the model and potentially loop for reflection
+            model_with_tools = self.model_instance.bind_tools(tools=[execute_ipython_cell], tool_choice="auto")
+            response = cast(AIMessage, model_with_tools.invoke(messages))
+
+            if response.tool_calls:
+                if len(response.tool_calls) > 1:
+                    raise Exception("Not possible in Claude with llm.bind_tools(tools=tools, tool_choice='auto')")
+                if response.tool_calls[0]["name"] != "execute_ipython_cell":
+                    raise Exception(
+                        f"Unexpected tool call: {response.tool_calls[0]['name']}. Expected 'execute_ipython_cell'."
+                    )
+                if (
+                    response.tool_calls[0]["args"].get("snippet") is None
+                    or not response.tool_calls[0]["args"]["snippet"].strip()
+                ):
+                    raise Exception("Tool call 'execute_ipython_cell' requires a non-empty 'snippet' argument.")
+                return Command(goto="sandbox", update={"messages": [response]})
+            else:
+                return Command(update={"messages": [response]})
+
+        # If eval_fn is a async, we define async node function.
+        if inspect.iscoroutinefunction(self.eval_fn):
+            raise ValueError("eval_fn must be a synchronous function, not a coroutine.")
+            # async def sandbox(state: StateSchema):
+            #     existing_context = state.get("context", {})
+            #     context = {**existing_context, **tools_context}
+            #     # Execute the script in the sandbox
+            #     output, new_vars = await eval_fn(state["script"], context)
+            #     new_context = {**existing_context, **new_vars}
+            #     return {
+            #         "messages": [{"role": "user", "content": output}],
+            #         "context": new_context,
+            #     }
+        else:
+
+            def sandbox(state: CodeActState) -> Command[Literal["call_model"]]:
+                tool_call = state["messages"][-1].tool_calls[0]  # type: ignore
+                code = tool_call["args"]["snippet"]
+                previous_add_context = state.get("add_context", {})
+                add_context = inject_context(previous_add_context, self.tools_context)
+                existing_context = state.get("context", {})
+                context = {**existing_context, **add_context}
+                # Execute the script in the sandbox
+                output, new_context, new_add_context = self.eval_fn(code, context, previous_add_context)
+                return Command(
+                    goto="call_model",
+                    update={
+                        "messages": [
+                            ToolMessage(
+                                content=output,
+                                name=tool_call["name"],
+                                tool_call_id=tool_call["id"],
+                            )
+                        ],
+                        "context": new_context,
+                        "add_context": new_add_context,
+                    },
+                )
+
+        agent = StateGraph(state_schema=CodeActState)
+        agent.add_node(call_model, retry_policy=RetryPolicy(max_attempts=3, retry_on=filter_retry_on))
+        agent.add_node(sandbox)
+        agent.add_edge(START, "call_model")
+        return agent.compile(checkpointer=self.memory)

universal_mcp/agents/codeact0/config.py

@@ -0,0 +1,77 @@
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+# Literal type for all available usecase filenames
+UseCaseName = Literal[
+    " ",
+    "1-unsubscribe",
+    "2-reddit",
+    "2.1-reddit",
+    "3-earnings",
+    "4-maps",
+    "4.1-maps",
+    "5-gmailreply",
+    "6-contract",
+    "7-overnight",
+    "8-sheets_chart",
+    "9-learning",
+    "10-reddit2",
+    "11-github",
+]
+
+
+class ContextSchema(BaseModel):
+    """The configuration for the agent."""
+
+    base_prompt: str = Field(
+        default=" ",
+        description="The base prompt to use for the agent's interactions. Leave blank if using a JSON prompt from the dropdown.",
+    )
+    model_provider: Annotated[
+        Literal[
+            "openai",
+            "anthropic",
+            "azure_openai",
+            "azure_ai",
+            "google_vertexai",
+            "google_genai",
+            "bedrock",
+            "bedrock_converse",
+            "cohere",
+            "fireworks",
+            "together",
+            "mistralai",
+            "huggingface",
+            "groq",
+            "ollama",
+            "google_anthropic_vertex",
+            "deepseek",
+            "ibm",
+            "nvidia",
+            "xai",
+            "perplexity",
+        ],
+        {"__template_metadata__": {"kind": "provider"}},
+    ] = Field(
+        default="anthropic",
+        description="The name of the model provider to use for the agent's main interactions. ",
+    )
+    model: Annotated[
+        Literal[
+            "claude-4-sonnet-20250514",
+            "claude-sonnet-4@20250514",
+        ],
+        {"__template_metadata__": {"kind": "llm"}},
+    ] = Field(
+        default="claude-4-sonnet-20250514",
+        description="The name of the language model to use for the agent's main interactions. ",
+    )
+    tool_names: list[str] = Field(
+        default=[],
+        description="The names of the tools to use for the agent's main interactions. Leave blank if using a JSON prompt from the dropdown.",
+    )
+    json_prompt_name: UseCaseName = Field(
+        default=" ",
+        description="The name of the JSON prompt to use for the agent's main interactions, instead of providing a base prompt and tool names. ",
+    )