universal-mcp-agents 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

universal_mcp/agents/builder.py

@@ -1,4 +1,5 @@
 import asyncio
+from collections import defaultdict
 from collections.abc import Sequence
 from typing import Annotated, TypedDict
 
@@ -15,7 +16,7 @@ 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.agents.utils import messages_to_list
-from collections import defaultdict
+
 
 class Agent(BaseModel):
     """Agent that can be created by the builder."""
@@ -24,9 +25,7 @@ class Agent(BaseModel):
     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
-    )
+    schedule: str | None = Field(description="The cron expression for the agent to run on.", default=None)
 
 
 class BuilderState(TypedDict):
@@ -79,9 +78,7 @@ Create an agent that feels thoughtfully designed, intelligent, and professionall
 """
 
 
-async def generate_agent(
-    llm: BaseChatModel, task: str, old_agent: Agent | None = None
-) -> Agent:
+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:
@@ -111,7 +108,7 @@ class BuilderAgent(BaseAgent):
     ):
         super().__init__(name, instructions, model, memory, **kwargs)
         self.registry = registry
-        self.llm: BaseChatModel = load_chat_model(model)
+        self.llm: BaseChatModel = load_chat_model(model, thinking=False)
 
     async def _create_agent(self, state: BuilderState):
         last_message = state["messages"][-1]
@@ -129,11 +126,7 @@ class BuilderAgent(BaseAgent):
         yield {
             "user_task": task,
             "generated_agent": generated_agent,
-            "messages": [
-                AIMessage(
-                    content=("I've designed an agent to help you with your task.")
-                )
-            ],
+            "messages": [AIMessage(content=("I've designed an agent to help you with your task."))],
         }
 
     async def _create_tool_config(self, state: BuilderState):
@@ -146,7 +139,7 @@ class BuilderAgent(BaseAgent):
             ]
         }
         tool_finder_graph = build_tool_node_graph(self.llm, self.registry)
-        
+
         initial_state = {
             "original_task": task,
             "messages": [HumanMessage(content=task)],
@@ -165,9 +158,7 @@ class BuilderAgent(BaseAgent):
                     apps_with_tools[app_id].extend(tool_ids)
 
             # Convert to a regular dict and remove any duplicate tool_ids for the same app
-            tool_config = {
-                app_id: list(set(tools)) for app_id, tools in apps_with_tools.items()
-            }
+            tool_config = {app_id: list(set(tools)) for app_id, tools in apps_with_tools.items()}
             final_message = "I have selected the necessary tools for the agent. The agent is ready!"
         else:
             # Handle the case where the graph failed to create a plan
@@ -175,9 +166,7 @@ class BuilderAgent(BaseAgent):
 
         yield {
             "tool_config": tool_config,
-            "messages": [
-                AIMessage(content=final_message)
-            ],
+            "messages": [AIMessage(content=final_message)],
         }
 
     async def _build_graph(self):
@@ -205,6 +194,7 @@ async def main():
         "Send a daily email to manoj@agentr.dev with daily agenda of the day",
     )
     from rich import print
+
     print(messages_to_list(result["messages"]))
     print(result["generated_agent"])
     print(result["tool_config"])

universal_mcp/agents/cli.py

@@ -20,7 +20,6 @@ def run(name: str = "react"):
     """Run the agent CLI"""
     import asyncio
 
-
     setup_logger(log_file=None, level="ERROR")
     client = AgentrClient()
     params = {
@@ -30,7 +29,7 @@ def run(name: str = "react"):
         "memory": MemorySaver(),
         "tools": {
             "google_mail": ["send_email"],
-        }
+        },
     }
     agent_cls = get_agent(name)
     agent = agent_cls(name, **params)

universal_mcp/agents/codeact/__init__.py

@@ -1,255 +1,3 @@
-import inspect
-import re
-from collections.abc import Awaitable, Callable, Sequence
-from typing import Any, TypeVar
+from .agent import CodeActAgent
 
-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
+__all__ = ["CodeActAgent"]

universal_mcp/agents/codeact/__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.codeact.agent import CodeActAgent
+from universal_mcp.agents.utils import messages_to_list
+
+
+async def main():
+    memory = MemorySaver()
+    agent = CodeActAgent(
+        "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/codeact/agent.py

@@ -0,0 +1,160 @@
+from collections.abc import Callable
+
+from langchain_core.messages import AIMessageChunk
+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 END, StateGraph
+from universal_mcp.logger import logger
+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.codeact.prompts import (
+    REFLECTION_PROMPT,
+    RETRY_PROMPT,
+    create_default_prompt,
+    make_safe_function_name,
+)
+from universal_mcp.agents.codeact.sandbox import eval_unsafe
+from universal_mcp.agents.codeact.state import CodeActState
+from universal_mcp.agents.codeact.utils import extract_and_combine_codeblocks
+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, instructions, model, memory, **kwargs)
+        self.model_instance = load_chat_model(model, thinking=False)
+        self.tools_config = tools or {}
+        self.registry = registry
+        self.eval_fn = eval_unsafe
+        self.reflection_prompt = REFLECTION_PROMPT
+        self.reflection_model = self.model_instance
+        self.max_reflections = 3
+        self.tools_context = {}
+        self.context = {}
+        self.sandbox_timeout = sandbox_timeout
+        self.processed_tools: list[StructuredTool | Callable] = []
+
+    async def _build_graph(self):
+        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)
+            self.processed_tools = [t if isinstance(t, StructuredTool) else create_tool(t) for t in exported_tools]
+
+        self.instructions = create_default_prompt(self.processed_tools, self.instructions)
+
+        for tool in self.processed_tools:
+            safe_name = make_safe_function_name(tool.name)
+            tool_callable = tool.coroutine if hasattr(tool, "coroutine") and tool.coroutine is not None else tool.func
+            self.tools_context[safe_name] = tool_callable
+        
+        self.context = {**self.context, **self.tools_context}
+
+        agent = StateGraph(CodeActState)
+        agent.add_node("call_model", self.call_model)
+        agent.add_node("sandbox", self.sandbox)
+
+        agent.set_entry_point("call_model")
+        agent.add_conditional_edges(
+            "call_model",
+            self.should_run_sandbox,
+            {
+                "sandbox": "sandbox",
+                END: END,
+            },
+        )
+        agent.add_edge("sandbox", "call_model")
+        return agent.compile(checkpointer=self.memory)
+
+    def should_run_sandbox(self, state: CodeActState) -> str:
+        last_message = state["messages"][-1]
+        if isinstance(last_message.content, str) and "TASK_COMPLETE" in last_message.content:
+            return END
+
+        if state.get("script"):
+            return "sandbox"
+        return END
+
+    def _extract_content(self, response: AIMessageChunk) -> str:
+        if isinstance(response.content, list):
+            content = " ".join([c.get("text", "") for c in response.content])
+        else:
+            content = response.content
+        return content
+
+    async def call_model(self, state: CodeActState) -> dict:
+        model = self.model_instance
+        reflection_model = self.reflection_model
+
+        messages = [{"role": "system", "content": self.instructions}] + state["messages"]
+
+        response = await model.ainvoke(messages)
+
+        text_content = self._extract_content(response)
+        if not isinstance(text_content, str):
+            raise ValueError(f"Content is not a string: {text_content}")
+        code = extract_and_combine_codeblocks(text_content)
+        logger.debug(f"Code: {code}")
+
+        if self.max_reflections > 0 and code:
+            reflection_count = 0
+            while reflection_count < self.max_reflections:
+                conversation_history = "\n".join(
+                    [
+                        f'<message role="{("user" if m.type == "human" else "assistant")}">\n{m.content}\n</message>'
+                        for m in state["messages"]
+                    ]
+                )
+                conversation_history += f'\n<message role="assistant">\n{response.content}\n</message>'
+
+                formatted_prompt = REFLECTION_PROMPT.format(conversation_history=conversation_history)
+
+                reflection_messages = [
+                    {"role": "system", "content": self.reflection_prompt},
+                    {"role": "user", "content": formatted_prompt},
+                ]
+                reflection_result = await reflection_model.ainvoke(reflection_messages)
+
+                if "NONE" in reflection_result.content:
+                    break
+
+                retry_prompt = RETRY_PROMPT.format(reflection_result=reflection_result.content)
+
+                regeneration_messages = [
+                    {"role": "system", "content": self.instructions},
+                    *state["messages"],
+                    {"role": "assistant", "content": response.content},
+                    {"role": "user", "content": retry_prompt},
+                ]
+                response = await model.ainvoke(regeneration_messages)
+
+                code = extract_and_combine_codeblocks(response.content)
+
+                if not code:
+                    break
+
+                reflection_count += 1
+
+        return {"messages": [response], "script": code}
+
+    async def sandbox(self, state: CodeActState) -> dict:
+        output, new_vars = await self.eval_fn(state["script"], self.context, timeout=self.sandbox_timeout)
+        self.context = {**self.context, **new_vars}
+        return {
+            "messages": [AIMessageChunk(content=output.strip())],
+            "script": None,
+        }

universal_mcp/agents/codeact/prompts.py

@@ -0,0 +1,91 @@
+import inspect
+import re
+from collections.abc import Sequence
+
+from langchain_core.tools import StructuredTool
+
+
+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: Sequence[StructuredTool],
+    base_prompt: str | None = None,
+):
+    """Create default prompt for the CodeAct agent."""
+    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'''\n{"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.
+
+Always use print() statements to explore data structures and function outputs. Simply returning values will not display them back to you for inspection. For example, use print(result) instead of just 'result'.
+
+As you don't know the output schema of the additional Python functions you have access to, start from exploring their contents before building a final solution.
+
+IMPORTANT CODING STRATEGY:
+1. Only write code up to the point where you make an API call/tool usage with an output
+2. Print the type/shape and a sample entry of this output, and using that knowledge proceed to write the further code
+
+This means:
+- Write code that makes the API call or tool usage
+- Print the result with type information: print(f"Type: {type(result)}")
+- Print the shape/structure: print(f"Shape/Keys: {result.keys() if isinstance(result, dict) else len(result) if isinstance(result, (list, tuple)) else 'N/A'}")
+- Print a sample entry: print(f"Sample: {result[0] if isinstance(result, (list, tuple)) and len(result) > 0 else result}")
+- Then, based on this knowledge, write the code to process/use this data
+
+Reminder: use Python code snippets to call tools
+
+When you have completely finished the task and provided the final answer, you MUST end your response with the exact phrase "TASK_COMPLETE".
+"""
+    return prompt
+
+
+REFLECTION_PROMPT = """
+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"
+"""
+
+RETRY_PROMPT = """
+I need you to completely regenerate your previous response based on this feedback:
+
+'''
+{reflection_result}
+'''
+
+DO NOT reference the feedback directly. Instead, provide a completely new response that addresses the issues.
+"""