universal-mcp-agents 0.1.2__py3-none-any.whl

universal_mcp/agents/bigtoolcache/context.py

@@ -0,0 +1,33 @@
+from dataclasses import dataclass, field
+
+from .prompts import SYSTEM_PROMPT
+
+
+@dataclass(kw_only=True)
+class Context:
+    """The context for the agent."""
+
+    system_prompt: str = field(
+        default=SYSTEM_PROMPT,
+        metadata={
+            "description": "The system prompt to use for the agent's interactions. "
+            "This prompt sets the context and behavior for the agent."
+        },
+    )
+
+    model: str = field(
+        default="anthropic/claude-4-sonnet-20250514",
+        metadata={
+            "description": "The name of the language model to use for the agent's main interactions. "
+            "Should be in the form: provider/model-name."
+        },
+    )
+
+    recursion_limit: int = field(
+        default=10,
+        metadata={
+            "description": "The maximum number of times the agent can call itself recursively. "
+            "This is to prevent infinite recursion."
+        },
+    )
+

universal_mcp/agents/bigtoolcache/graph.py

@@ -0,0 +1,176 @@
+import json
+from datetime import UTC, datetime
+from typing import Any, Literal, TypedDict, cast
+
+from langchain_anthropic import ChatAnthropic
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, ToolMessage
+from langchain_core.tools import tool
+from langgraph.graph import StateGraph
+from langgraph.runtime import Runtime
+from langgraph.types import Command
+
+from universal_mcp.agents.bigtoolcache.context import Context
+from universal_mcp.agents.bigtoolcache.state import State
+from universal_mcp.logger import logger
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolFormat
+
+
+
+def build_graph(
+    tool_registry: ToolRegistry,
+    llm: BaseChatModel
+):
+    @tool
+    async def search_tools(queries: list[str]) -> str:
+        """Search tools for a given list of queries
+        Each single query should be atomic (doable with a single tool).
+        For tasks requiring multiple tools, add separate queries for each subtask"""
+        logger.info(f"Searching tools for queries: '{queries}'")
+        try:
+            all_tool_candidates = ""
+            app_ids = await tool_registry.list_all_apps()
+            connections = await tool_registry.list_connected_apps()
+            connection_ids = set([connection["app_id"] for connection in connections])
+            connected_apps = [app["id"] for app in app_ids if app["id"] in connection_ids]
+            unconnected_apps = [app["id"] for app in app_ids if app["id"] not in connection_ids]
+            app_tools = {}
+            for task_query in queries:
+                tools_list = await tool_registry.search_tools(task_query, limit=40)
+                for tool in tools_list:
+                    app = tool["id"].split("__")[0]
+                    if app not in app_tools:
+                        if len(app_tools.keys()) >= 10:
+                            break
+                        app_tools[app] = {}
+                    if len(app_tools[app]) < 3:
+                        if tool["id"] not in app_tools[app]:
+                            app_tools[app][tool["id"]] = tool["description"]
+            for app in app_tools:
+                app_status = "connected" if app in connected_apps else "NOT connected"
+                all_tool_candidates += f"Tools from {app} (status: {app_status} by user):\n"
+                for tool in app_tools[app]:
+                    all_tool_candidates += f" - {tool}: {app_tools[app][tool]}\n"
+                all_tool_candidates += "\n"
+                
+            
+            return all_tool_candidates
+        except Exception as e:
+            logger.error(f"Error retrieving tools: {e}")
+            return "Error: " + str(e)
+    
+    @tool
+    async def load_tools(tool_ids: list[str]) -> list[dict[str, Any]]:
+        """Load the tools for the given tool ids. Returns the tool name, description, parameters schema, and output schema."""
+        temp_manager = tool_registry.tool_manager
+        temp_manager.clear_tools()
+        await tool_registry.export_tools(tool_ids, format=ToolFormat.NATIVE)
+        tool_details = []
+        for tool_id in tool_ids:
+            tool = temp_manager.get_tool(tool_id)
+            tool_details.append({
+                "name": tool.name,
+                "description": tool.description,
+                "parameters_schema": tool.parameters,
+                "output_schema": tool.output_schema,
+            })
+        return tool_details
+    
+    @tool
+    async def call_tool(tool_id: str, tool_args: dict[str, Any]) -> Any:
+        """Call the tool with the given id and arguments."""
+        return await tool_registry.call_tool(tool_id, tool_args)
+
+
+    async def call_model(state: State, runtime: Runtime[Context]) -> Command[Literal["select_tools", "call_tools"]]:
+        logger.info("Calling model...")
+        try:
+            system_message = runtime.context.system_prompt.format(system_time=datetime.now(tz=UTC).isoformat())
+            messages = [{"role": "system", "content": system_message}, *state["messages"]]
+
+            model = llm
+
+            if isinstance(model, ChatAnthropic):
+                model_with_tools = model.bind_tools(
+                    [search_tools, load_tools, call_tool], tool_choice="auto", cache_control={"type": "ephemeral"}
+                )
+            else:
+                model_with_tools = model.bind_tools([search_tools, load_tools, call_tool], tool_choice="auto")
+            response = cast(AIMessage, await model_with_tools.ainvoke(messages))
+
+            if response.tool_calls:
+                logger.info(f"Model responded with {len(response.tool_calls)} tool calls.")
+                if len(response.tool_calls) > 1:
+                    raise Exception("Not possible in Claude with llm.bind_tools(tools=tools, tool_choice='auto')")
+                tool_call = response.tool_calls[0]
+                if tool_call["name"] == search_tools.name:
+                    logger.info("Model requested to select tools.")
+                    return Command(goto="select_tools", update={"messages": [response]})
+                elif tool_call["name"] == load_tools.name:
+                    logger.info("Model requested to load tools.")
+                    tool_details = await load_tools.ainvoke(input=tool_call["args"])
+                    tool_msg = ToolMessage(f"Loaded tools. {tool_details}", tool_call_id=tool_call["id"])
+                    selected_tool_ids = tool_call["args"]["tool_ids"]
+                    logger.info(f"Loaded tools: {selected_tool_ids}")
+                    return Command(goto="call_model", update={ "messages": [response, tool_msg], "selected_tool_ids": selected_tool_ids})
+                elif tool_call["name"] == call_tool.name:
+                    logger.info("Model requested to call tool.")
+                    return Command(goto="call_tools", update={"messages": [response]})
+                return Command(goto="call_tools", update={"messages": [response]})
+            else:
+                logger.info("Model responded with a message, ending execution.")
+                return Command(update={"messages": [response]})
+        except Exception as e:
+            logger.error(f"Error in call_model: {e}")
+            raise
+
+    async def select_tools(state: State, runtime: Runtime[Context]) -> Command[Literal["call_model"]]:
+        logger.info("Selecting tools...")
+        try:
+            tool_call = state["messages"][-1].tool_calls[0]
+            searched_tools= await search_tools.ainvoke(input=tool_call["args"])
+            tool_msg = ToolMessage(f"Available tools: {searched_tools}", tool_call_id=tool_call["id"])
+            return Command(goto="call_model", update={"messages": [tool_msg]})
+        except Exception as e:
+            logger.error(f"Error in select_tools: {e}")
+            raise
+
+    async def call_tools(state: State) -> Command[Literal["call_model"]]:
+        logger.info("Calling tools...")
+        outputs = []
+        recent_tool_ids = []
+        tool_call = state["messages"][-1].tool_calls[0]
+        tool_id = tool_call["args"]["tool_id"]
+        tool_args = tool_call["args"]["tool_args"]
+        logger.info(f"Executing tool: {tool_id} with args: {tool_args}")
+        try:
+            await tool_registry.export_tools([tool_id], ToolFormat.LANGCHAIN)
+            tool_result = await call_tool.ainvoke(input={"tool_id": tool_id, "tool_args": tool_args})
+            logger.info(f"Tool '{tool_id}' executed successfully.")
+            outputs.append(
+                ToolMessage(
+                    content=json.dumps(tool_result),
+                    name=tool_id,
+                    tool_call_id=tool_call["id"],
+                )
+            )
+            recent_tool_ids.append(tool_id)
+        except Exception as e:
+            logger.error(f"Error executing tool '{tool_id}': {e}")
+            outputs.append(
+                ToolMessage(
+                    content=json.dumps("Error: " + str(e)),
+                    name=tool_id,
+                    tool_call_id=tool_call["id"],
+                )
+            )
+        return Command(goto="call_model", update={"messages": outputs, "selected_tool_ids": recent_tool_ids})
+
+    builder = StateGraph(State, context_schema=Context)
+
+    builder.add_node(call_model)
+    builder.add_node(select_tools)
+    builder.add_node(call_tools)
+    builder.set_entry_point("call_model")
+    return builder

universal_mcp/agents/bigtoolcache/prompts.py

@@ -0,0 +1,13 @@
+"""Default prompts used by the agent."""
+
+SYSTEM_PROMPT = """You are a helpful AI assistant.
+
+**Core Directives:**
+1.  **Always Use Tools for Tasks:** For any user request that requires an action (e.g., sending an email, searching for information, creating an event), you MUST use a tool. Do not answer from your own knowledge or refuse a task if a tool might exist for it.
+2.  **First Step is ALWAYS `search_tools`:** Before you can use any other tool, you MUST first call the `search_tools` function to find the right tools for the user's request. This is your mandatory first action. You must not use the same/similar query multiple times in the list. The list should have multiple queries only if the task has clearly different sub-tasks.
+3.  **Load Tools:** After looking at the output of `search_tools`, you MUST call the `load_tools` function to load only the tools you want to use. Use your judgement to eliminate irrelevant apps that came up just because of semantic similarity. However, sometimes, multiple apps might be relevant for the same task. Prefer connected apps over unconnected apps while breaking a tie. If more than one relevant app (or none of the relevant apps) are connected, you must ask the user to choose the app. In case the user asks you to use an app that is not connected, call the apps tools normally. The tool will return a link for connecting that you should pass on to the user.
+4.  **Call Tools:** After loading the tools, you MUST call the `call_tool` function to call the tools you want to use. You must call the tool with the correct arguments. You can only call the tool once you have loaded it.
+5.  **Strictly Follow the Process:** Your only job in your first turn is to analyze the user's request and call `search_tools` with a concise query describing the core task. Do not engage in conversation.
+
+System time: {system_time}
+"""

universal_mcp/agents/bigtoolcache/state.py

@@ -0,0 +1,27 @@
+from typing import Annotated
+
+from langgraph.prebuilt.chat_agent_executor import AgentState
+
+
+def _enqueue(left: list, right: list) -> list:
+    """Treat left as a FIFO queue, append new items from right (preserve order),
+    keep items unique, and cap total size to 20 (drop oldest items)."""
+    max_size = 30
+    preferred_size = 20
+    if len(right) > preferred_size:
+        preferred_size = min(max_size, len(right))
+    queue = list(left or [])
+
+    for item in right[:preferred_size] or []:
+        if item in queue:
+            queue.remove(item)
+        queue.append(item)
+
+    if len(queue) > preferred_size:
+        queue = queue[-preferred_size:]
+
+    return queue
+
+
+class State(AgentState):
+    selected_tool_ids: Annotated[list[str], _enqueue]

universal_mcp/agents/builder.py

@@ -0,0 +1,146 @@
+import asyncio
+from collections.abc import Sequence
+from typing import Annotated, TypedDict
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.message import add_messages
+from pydantic import BaseModel, Field
+
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.shared.tool_node import build_tool_node_graph
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolConfig
+
+
+class Agent(BaseModel):
+    """Agent that can be created by the builder."""
+
+    name: str = Field(description="Name of the agent.")
+    description: str = Field(description="A small description of the agent.")
+    expertise: str = Field(description="The expertise of the agent.")
+    instructions: str = Field(description="The instructions for the agent to follow.")
+    schedule: str | None = Field(description="The cron expression for the agent to run on.", default=None)
+
+
+class BuilderState(TypedDict):
+    user_task: str
+    generated_agent: Agent | None
+    tool_config: ToolConfig | None
+    messages: Annotated[Sequence[BaseMessage], add_messages]
+
+
+AGENT_BUILDER_INSTRUCTIONS = """
+You are an agent builder. Your goal is to create an agent that can accomplish the user's task.
+Your will be given a task and you need to generate an agent that can accomplish the task.
+The agent should have a name, role, instructions, and a model.
+- The name should be a short and descriptive name for the agent.
+- The description should be a small description of the agent. For example, research a stock and write a buy sell analysis report.
+- The expertise should be the expertise of the agent. For example, GTM Expert, SEO Expert, etc.
+- The instructions should be a detailed description of what the agent should do. This should include the input, the output, and the tool usage. The agent will be provided a set of tools, you can use that to give a more accurate response.
+- The model should be the model to use for the agent.
+- The reasoning should be a detailed explanation of why you are creating this agent with these parameters.
+- If the user specifies a schedule, you should also provide a cron expression for the agent to run on. The schedule should be in a proper cron expression and nothing more.
+"""
+
+
+async def generate_agent(llm: BaseChatModel, task: str, old_agent: Agent | None = None) -> Agent:
+    """Generates an agent from a task, optionally modifying an existing one."""
+    prompt_parts = [AGENT_BUILDER_INSTRUCTIONS]
+    if old_agent:
+        prompt_parts.append(
+            "\nThe user wants to modify the following agent design. "
+            "Incorporate their feedback into a new design.\n\n"
+            f"{old_agent.model_dump_json(indent=2)}"
+        )
+    else:
+        prompt_parts.append(f"\n\n**Task:** {task}")
+
+    prompt = "\n".join(prompt_parts)
+    structured_llm = llm.with_structured_output(Agent)
+    agent = await structured_llm.ainvoke(prompt)
+    return agent
+
+
+class BuilderAgent(BaseAgent):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        registry: ToolRegistry,
+        memory: BaseCheckpointSaver | None = None,
+        **kwargs,
+    ):
+        super().__init__(name, instructions, model, memory, **kwargs)
+        self.registry = registry
+        self.llm: BaseChatModel = load_chat_model(model)
+
+    async def _create_agent(self, state: BuilderState):
+        last_message = state["messages"][-1]
+        task = last_message.content
+        agent = state.get("generated_agent")
+
+        yield {
+            "messages": [
+                AIMessage(
+                    content="Thinking... I will now design an agent to handle your request.",
+                )
+            ],
+        }
+        generated_agent = await generate_agent(self.llm, task, agent)
+        yield {
+            "user_task": task,
+            "generated_agent": generated_agent,
+            "messages": [AIMessage(content=("I've designed an agent to help you with your task."))],
+        }
+
+    async def _create_tool_config(self, state: BuilderState):
+        task = state["user_task"]
+        yield {
+            "messages": [
+                AIMessage(
+                    content="Great! Now, I will select the appropriate tools for this agent. This may take a moment.",
+                )
+            ]
+        }
+        tool_finder_graph = build_tool_node_graph(self.llm, self.registry)
+        tool_config = await tool_finder_graph.ainvoke({"task": task, "messages": [HumanMessage(content=task)]})
+        tool_config = tool_config.get("apps_with_tools", {})
+        yield {
+            "tool_config": tool_config,
+            "messages": [AIMessage(content="I have selected the necessary tools for the agent. The agent is ready!")],
+        }
+
+    async def _build_graph(self):
+        builder = StateGraph(BuilderState)
+        builder.add_node("create_agent", self._create_agent)
+        builder.add_node("create_tool_config", self._create_tool_config)
+
+        builder.add_edge(START, "create_agent")
+        builder.add_edge("create_agent", "create_tool_config")
+        builder.add_edge("create_tool_config", END)
+        return builder.compile(checkpointer=self.memory)
+
+
+async def main():
+    from universal_mcp.agentr.registry import AgentrRegistry
+
+    registry = AgentrRegistry()
+    agent = BuilderAgent(
+        name="Builder Agent",
+        instructions="You are a builder agent that creates other agents.",
+        model="gemini/gemini-1.5-pro",
+        registry=registry,
+    )
+    result = await agent.invoke(
+        "Send a daily email to manoj@agentr.dev with daily agenda of the day",
+    )
+    print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

universal_mcp/agents/cli.py

@@ -0,0 +1,27 @@
+from typer import Typer
+
+from universal_mcp.agents import ReactAgent
+from universal_mcp.logger import setup_logger
+
+app = Typer()
+
+
+@app.command(
+    help="Run the agent CLI",
+    epilog="""
+    Example:
+    mcp client run --config client_config.json
+    """,
+)
+def run():
+    """Run the agent CLI"""
+    import asyncio
+
+    setup_logger(log_file=None, level="WARNING")
+
+    agent = ReactAgent("React Agent", "You are a helpful assistant", "openrouter/auto")
+    asyncio.run(agent.run_interactive())
+
+
+if __name__ == "__main__":
+    app()

universal_mcp/agents/codeact/__init__.py

@@ -0,0 +1,243 @@
+import inspect
+import re
+from collections.abc import Awaitable, Callable, Sequence
+from typing import Any, TypeVar
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.tools import StructuredTool
+from langchain_core.tools import tool as create_tool
+from langgraph.graph import END, START, MessagesState, StateGraph
+from langgraph.types import Command
+
+from .utils import extract_and_combine_codeblocks
+
+EvalFunction = Callable[[str, dict[str, Any]], tuple[str, dict[str, Any]]]
+EvalCoroutine = Callable[[str, dict[str, Any]], Awaitable[tuple[str, dict[str, Any]]]]
+
+
+class CodeActState(MessagesState):
+    """State for CodeAct agent."""
+
+    script: str | None
+    """The Python code script to be executed."""
+    context: dict[str, Any]
+    """Dictionary containing the execution context with available tools and variables."""
+
+
+StateSchema = TypeVar("StateSchema", bound=CodeActState)
+StateSchemaType = type[StateSchema]
+
+
+def make_safe_function_name(name: str) -> str:
+    """Convert a tool name to a valid Python function name."""
+    # Replace non-alphanumeric characters with underscores
+    safe_name = re.sub(r"[^a-zA-Z0-9_]", "_", name)
+    # Ensure the name doesn't start with a digit
+    if safe_name and safe_name[0].isdigit():
+        safe_name = f"tool_{safe_name}"
+    # Handle empty name edge case
+    if not safe_name:
+        safe_name = "unnamed_tool"
+    return safe_name
+
+
+def create_default_prompt(tools: list[StructuredTool], base_prompt: str | None = None):
+    """Create default prompt for the CodeAct agent."""
+    tools = [t if isinstance(t, StructuredTool) else create_tool(t) for t in tools]
+    prompt = f"{base_prompt}\n\n" if base_prompt else ""
+    prompt += """You will be given a task to perform. You should output either
+- a Python code snippet that provides the solution to the task, or a step towards the solution. Any output you want to extract from the code should be printed to the console. Code should be output in a fenced code block.
+- text to be shown directly to the user, if you want to ask for more information or provide the final answer.
+
+In addition to the Python Standard Library, you can use the following functions:
+"""
+
+    for tool in tools:
+        # Use coroutine if it exists, otherwise use func
+        tool_callable = tool.coroutine if hasattr(tool, "coroutine") and tool.coroutine is not None else tool.func
+        # Create a safe function name
+        safe_name = make_safe_function_name(tool.name)
+        # Determine if it's an async function
+        is_async = inspect.iscoroutinefunction(tool_callable)
+        # Add appropriate function definition
+        prompt += f'''
+{"async " if is_async else ""}def {safe_name}{str(inspect.signature(tool_callable))}:
+    """{tool.description}"""
+    ...
+'''
+
+    prompt += """
+
+Variables defined at the top level of previous code snippets can be referenced in your code.
+
+Reminder: use Python code snippets to call tools"""
+    return prompt
+
+
+def create_codeact(
+    model: BaseChatModel,
+    tools: Sequence[StructuredTool | Callable],
+    eval_fn: EvalFunction | EvalCoroutine,
+    *,
+    prompt: str | None = None,
+    reflection_prompt: str | None = None,
+    reflection_model: BaseChatModel | None = None,
+    max_reflections: int = 3,
+    state_schema: StateSchemaType = CodeActState,
+) -> StateGraph:
+    """Create a CodeAct agent.
+
+    Args:
+        model: The language model to use for generating code
+        tools: List of tools available to the agent. Can be passed as python functions or StructuredTool instances.
+        eval_fn: Function or coroutine that executes code in a sandbox. Takes code string and locals dict,
+            returns a tuple of (stdout output, new variables dict)
+        prompt: Optional custom system prompt. If None, uses default prompt.
+            To customize default prompt you can use `create_default_prompt` helper:
+            `create_default_prompt(tools, "You are a helpful assistant.")`
+        reflection_prompt: Optional prompt for reflection. If provided, will be used to evaluate responses.
+            If the reflection output contains "NONE", the response is considered valid, otherwise the
+            reflection output is passed back to the model for regeneration.
+        reflection_model: Optional model to use for reflection. If None, uses the same model as for generation.
+        max_reflections: Maximum number of reflection iterations (default: 3).
+        state_schema: The state schema to use for the agent.
+
+    Returns:
+        A StateGraph implementing the CodeAct architecture
+    """
+    tools = [t if isinstance(t, StructuredTool) else create_tool(t) for t in tools]
+
+    if prompt is None:
+        prompt = create_default_prompt(tools)
+
+    # If no reflection model is provided, use the main model
+    if reflection_model is None:
+        reflection_model = model
+
+    # Make tools available to the code sandbox - use safe names for keys
+    tools_context = {}
+    for tool in tools:
+        safe_name = make_safe_function_name(tool.name)
+        # Use coroutine if it exists, otherwise use func (same as in create_default_prompt)
+        tool_callable = tool.coroutine if hasattr(tool, "coroutine") and tool.coroutine is not None else tool.func
+        # Only use the safe name for consistency with the prompt
+        tools_context[safe_name] = tool_callable
+
+    def call_model(state: StateSchema) -> Command:
+        messages = [{"role": "system", "content": prompt}] + state["messages"]
+
+        # Run the model and potentially loop for reflection
+        response = model.invoke(messages)
+
+        # Extract and combine all code blocks
+        code = extract_and_combine_codeblocks(response.content)
+
+        # Loop for reflection if needed and if code is present
+        if reflection_prompt and code:
+            reflection_count = 0
+            while reflection_count < max_reflections:
+                # Format conversation history with XML-style tags
+                conversation_history = "\n".join(
+                    [
+                        f'<message role="{("user" if m.type == "human" else "assistant")}">\n{m.content}\n</message>'
+                        for m in state["messages"]
+                    ]
+                )
+
+                # Add the current response
+                conversation_history += f'\n<message role="assistant">\n{response.content}\n</message>'
+
+                # Create the reflection prompt with the tagged conversation history
+                formatted_prompt = f"""
+Review the assistant's latest code for as per the quality rules:
+
+<conversation_history>
+{conversation_history}
+</conversation_history>
+
+If you find ANY of these issues, describe the problem briefly and clearly.
+If NO issues are found, respond with EXACTLY: "NONE"
+"""
+
+                # Create messages for reflection with correct ordering
+                reflection_messages = [
+                    {"role": "system", "content": reflection_prompt},
+                    # Include the formatted reflection prompt as the final user message
+                    {"role": "user", "content": formatted_prompt},
+                ]
+                reflection_result = reflection_model.invoke(reflection_messages)
+
+                # Check if reflection passed
+                if "NONE" in reflection_result.content:
+                    # Reflection passed, exit loop
+                    break
+
+                # Reflection didn't pass, regenerate response
+                reflection_messages = [
+                    {"role": "system", "content": prompt},
+                    *state["messages"],
+                    {"role": "assistant", "content": response.content},
+                    {
+                        "role": "user",
+                        "content": f"""
+I need you to completely regenerate your previous response based on this feedback:
+
+'''
+{reflection_result.content}
+'''
+
+DO NOT reference the feedback directly. Instead, provide a completely new response that addresses the issues.
+""",
+                    },
+                ]
+                response = model.invoke(reflection_messages)
+
+                # Extract code from the new response
+                code = extract_and_combine_codeblocks(response.content)
+
+                # If no code in the new response, exit the reflection loop
+                if not code:
+                    break
+
+                # Increment reflection count
+                reflection_count += 1
+
+        # Return appropriate command with only the latest response
+        if code:
+            return Command(goto="sandbox", update={"messages": [response], "script": code})
+        else:
+            # no code block, end the loop and respond to the user
+            return Command(update={"messages": [response], "script": None})
+
+    # If eval_fn is a async, we define async node function.
+    if inspect.iscoroutinefunction(eval_fn):
+
+        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: StateSchema):
+            existing_context = state.get("context", {})
+            context = {**existing_context, **tools_context}
+            # Execute the script in the sandbox
+            output, new_vars = eval_fn(state["script"], context)
+            new_context = {**existing_context, **new_vars}
+            return {
+                "messages": [{"role": "user", "content": output}],
+                "context": new_context,
+            }
+
+    agent = StateGraph(state_schema)
+    agent.add_node(call_model, destinations=(END, "sandbox"))
+    agent.add_node(sandbox)
+    agent.add_edge(START, "call_model")
+    agent.add_edge("sandbox", "call_model")
+    return agent