universal-mcp-agents 0.1.9__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,3 +1,3 @@
 from .agent import CodeActAgent
 
-__all__ = ["CodeActAgent"]
+__all__ = ["CodeActAgent"]

universal_mcp/agents/codeact/__main__.py

@@ -1,23 +1,33 @@
 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="azure/gpt-4o",
-        tools={"google_mail": ["send_email"]},
+        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(
-        "Send an email to manoj@agentr.dev from my Gmail account with a subject 'testing codeact agent' and body 'This is a test of the codeact agent.'"
+        user_input="Fetch unsubscribe links from my Gmail inbox for promo emails I have received in the last 7 days"
     )
-    from rich import print
-
     print(messages_to_list(result["messages"]))
 
 

universal_mcp/agents/codeact/agent.py

@@ -1,25 +1,25 @@
-import inspect
-from typing import Callable, Union
+from collections.abc import Callable
 
-from langchain_core.language_models import BaseChatModel
-from langchain_core.tools import StructuredTool, tool as create_tool
+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 loguru import logger
+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.llm import load_chat_model
 from universal_mcp.agents.codeact.prompts import (
-    create_default_prompt,
-    make_safe_function_name,
     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):
@@ -31,51 +31,41 @@ class CodeActAgent(BaseAgent):
         memory: BaseCheckpointSaver | None = None,
         tools: ToolConfig | None = None,
         registry: ToolRegistry | None = None,
-        *,
-        reflection_prompt: str = None,
-        reflection_model: BaseChatModel = None,
-        max_reflections: int = 3,
+        sandbox_timeout: int = 20,
         **kwargs,
     ):
         super().__init__(name, instructions, model, memory, **kwargs)
-        self.model_instance = load_chat_model(model)
+        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 = reflection_model or self.model_instance
-        self.max_reflections = max_reflections if reflection_prompt else 0
+        self.reflection_prompt = REFLECTION_PROMPT
+        self.reflection_model = self.model_instance
+        self.max_reflections = 3
         self.tools_context = {}
-        self.processed_tools: list[Union[StructuredTool, Callable]] = []
+        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
-        )
+            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
-            )
+            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", lambda state, config: self.call_model(state, config))
+        agent.add_node("call_model", self.call_model)
         agent.add_node("sandbox", self.sandbox)
 
         agent.set_entry_point("call_model")
@@ -91,21 +81,34 @@ class CodeActAgent(BaseAgent):
         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 call_model(self, state: CodeActState, config: dict) -> dict:
-        context = config.get("context", {})
-        instructions = context.get("system_prompt", self.instructions)
+    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": instructions}] + state["messages"]
+        messages = [{"role": "system", "content": self.instructions}] + state["messages"]
 
-        response = model.invoke(messages)
+        response = await model.ainvoke(messages)
 
-        code = extract_and_combine_codeblocks(response.content)
+        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
@@ -116,34 +119,28 @@ class CodeActAgent(BaseAgent):
                         for m in state["messages"]
                     ]
                 )
-                conversation_history += (
-                    f'\n<message role="assistant">\n{response.content}\n</message>'
-                )
+                conversation_history += f'\n<message role="assistant">\n{response.content}\n</message>'
 
-                formatted_prompt = REFLECTION_PROMPT.format(
-                    conversation_history=conversation_history
-                )
+                formatted_prompt = REFLECTION_PROMPT.format(conversation_history=conversation_history)
 
                 reflection_messages = [
                     {"role": "system", "content": self.reflection_prompt},
                     {"role": "user", "content": formatted_prompt},
                 ]
-                reflection_result = reflection_model.invoke(reflection_messages)
+                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
-                )
+                retry_prompt = RETRY_PROMPT.format(reflection_result=reflection_result.content)
 
                 regeneration_messages = [
-                    {"role": "system", "content": instructions},
+                    {"role": "system", "content": self.instructions},
                     *state["messages"],
                     {"role": "assistant", "content": response.content},
                     {"role": "user", "content": retry_prompt},
                 ]
-                response = model.invoke(regeneration_messages)
+                response = await model.ainvoke(regeneration_messages)
 
                 code = extract_and_combine_codeblocks(response.content)
 
@@ -152,20 +149,12 @@ class CodeActAgent(BaseAgent):
 
                 reflection_count += 1
 
-        if code:
-            return {"messages": [response], "script": code}
-        else:
-            return {"messages": [response], "script": None}
+        return {"messages": [response], "script": code}
 
     async def sandbox(self, state: CodeActState) -> dict:
-        existing_context = state.get("context", {})
-        context = {**existing_context, **self.tools_context}
-        if inspect.iscoroutinefunction(self.eval_fn):
-            output, new_vars = await self.eval_fn(state["script"], context)
-        else:
-            output, new_vars = self.eval_fn(state["script"], context)
-        new_context = {**existing_context, **new_vars}
+        output, new_vars = await self.eval_fn(state["script"], self.context, timeout=self.sandbox_timeout)
+        self.context = {**self.context, **new_vars}
         return {
-            "messages": [{"role": "user", "content": output}],
-            "context": new_context,
-        }
+            "messages": [AIMessageChunk(content=output.strip())],
+            "script": None,
+        }

universal_mcp/agents/codeact/prompts.py

@@ -1,8 +1,8 @@
 import inspect
 import re
-from typing import Optional, Sequence
+from collections.abc import Sequence
 
-from langchain_core.tools import StructuredTool, tool as create_tool
+from langchain_core.tools import StructuredTool
 
 
 def make_safe_function_name(name: str) -> str:
@@ -20,7 +20,7 @@ def make_safe_function_name(name: str) -> str:
 
 def create_default_prompt(
     tools: Sequence[StructuredTool],
-    base_prompt: Optional[str] = None,
+    base_prompt: str | None = None,
 ):
     """Create default prompt for the CodeAct agent."""
     prompt = f"{base_prompt}\n\n" if base_prompt else ""
@@ -32,11 +32,7 @@ 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
-        )
+        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
@@ -44,7 +40,7 @@ In addition to the Python Standard Library, you can use the following functions:
         # Add appropriate function definition
         prompt += f'''\n{"async " if is_async else ""}def {safe_name}{str(inspect.signature(tool_callable))}:
     """{tool.description}"""
-    ... 
+    ...
 '''
 
     prompt += """
@@ -66,7 +62,10 @@ This means:
 - 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"""
+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
 
 

universal_mcp/agents/codeact/sandbox.py

@@ -5,7 +5,7 @@ import io
 from typing import Any
 
 
-async def eval_unsafe(code: str, _locals: dict[str, Any]) -> tuple[str, dict[str, Any]]:
+async def eval_unsafe(code: str, _locals: dict[str, Any], timeout: int = 10) -> tuple[str, dict[str, Any]]:
     """
     Execute code in a non-blocking way and return the output and changed variables.
     """
@@ -29,7 +29,10 @@ async def eval_unsafe(code: str, _locals: dict[str, Any]) -> tuple[str, dict[str
             return f"Error during execution: {repr(e)}"
 
     # Run the synchronous exec in a separate thread to avoid blocking the event loop.
-    output = await asyncio.to_thread(sync_eval_in_thread)
+    try:
+        output = await asyncio.wait_for(asyncio.to_thread(sync_eval_in_thread), timeout=timeout)
+    except asyncio.TimeoutError:
+        output = f"Error: Code execution timed out after {timeout} seconds."
     result += output
 
     # Identify all variables that are not part of the original builtins

universal_mcp/agents/codeact/state.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 from langgraph.graph import MessagesState
 
@@ -6,7 +6,5 @@ from langgraph.graph import MessagesState
 class CodeActState(MessagesState):
     """State for CodeAct agent."""
 
-    script: Optional[str]
+    script: str | None
     """The Python code script to be executed."""
-    context: dict[str, Any]
-    """Dictionary containing the execution context with available tools and variables."""

universal_mcp/agents/codeact/utils.py

@@ -1,5 +1,7 @@
 import re
 
+from universal_mcp.logger import logger
+
 BACKTICK_PATTERN = r"(?:^|\n)```(.*?)(?:```(?:\n|$))"
 
 
@@ -37,7 +39,12 @@ def extract_and_combine_codeblocks(text: str) -> str:
     """
     # Find all code blocks in the text using regex
     # Pattern matches anything between triple backticks, with or without a language identifier
-    code_blocks = re.findall(BACKTICK_PATTERN, text, re.DOTALL)
+    try:
+        code_blocks = re.findall(BACKTICK_PATTERN, text, re.DOTALL)
+    except Exception as e:
+        logger.error(f"Error extracting code blocks: {e}")
+        logger.error(f"Text: {text}")
+        return ""
 
     if not code_blocks:
         return ""
@@ -46,15 +53,15 @@ def extract_and_combine_codeblocks(text: str) -> str:
     processed_blocks = []
     for block in code_blocks:
         # Strip leading and trailing whitespace
-        block = block.strip()
+        cleaned_block = block.strip()
 
         # If the first line looks like a language identifier, remove it
-        lines = block.split("\n")
+        lines = cleaned_block.split("\n")
         if lines and (not lines[0].strip() or " " not in lines[0].strip()):
             # First line is empty or likely a language identifier (no spaces)
-            block = "\n".join(lines[1:])
+            cleaned_block = "\n".join(lines[1:])
 
-        processed_blocks.append(block)
+        processed_blocks.append(cleaned_block)
 
     # Combine all codeblocks with newlines between them
     combined_code = "\n\n".join(processed_blocks)

universal_mcp/agents/hil.py

@@ -63,12 +63,7 @@ def handle_interrupt(interrupt: Interrupt) -> str | bool:
         value = input("Do you accept this? (y/n): " + interrupt.value["question"])
         return value.lower() in ["y", "yes"]
     elif interrupt_type == "choice":
-        value = input(
-            "Enter your choice: "
-            + interrupt.value["question"]
-            + " "
-            + ", ".join(interrupt.value["choices"])
-        )
+        value = input("Enter your choice: " + interrupt.value["question"] + " " + ", ".join(interrupt.value["choices"]))
         if value in interrupt.value["choices"]:
             return value
         else:

universal_mcp/agents/planner/__init__.py

@@ -26,9 +26,7 @@ class PlannerAgent(BaseAgent):
         self.executor_agent_cls = executor_agent_cls
 
     def _build_system_message(self):
-        return DEVELOPER_PROMPT.format(
-            name=self.name, instructions=self.instructions
-        )
+        return DEVELOPER_PROMPT.format(name=self.name, instructions=self.instructions)
 
     async def _build_graph(self):
         return build_graph(

universal_mcp/agents/planner/graph.py

@@ -18,9 +18,7 @@ def build_graph(llm, registry, instructions, model, executor_agent_cls):
         task = state["messages"][-1].content
         logger.info(f"Running tool finder for task: {task}")
         tool_finder_graph = build_tool_node_graph(llm, registry)
-        tool_finder_state = await tool_finder_graph.ainvoke(
-            {"original_task": task, "messages": state["messages"]}
-        )
+        tool_finder_state = await tool_finder_graph.ainvoke({"original_task": task, "messages": state["messages"]})
 
         if not tool_finder_state.get("apps_required"):
             logger.info("Tool finder determined no apps are required.")

universal_mcp/agents/react.py

@@ -4,6 +4,7 @@ from loguru import logger
 from universal_mcp.agentr.registry import AgentrRegistry
 from universal_mcp.tools.registry import ToolRegistry
 from universal_mcp.types import ToolConfig, ToolFormat
+from rich import print
 
 from universal_mcp.agents.base import BaseAgent
 from universal_mcp.agents.llm import load_chat_model
@@ -40,7 +41,18 @@ class ReactAgent(BaseAgent):
         self.llm = load_chat_model(model)
         self.tools = tools or {}
         if "ui" not in self.tools:
-            self.tools["ui"] = ["create_bar_chart", "create_line_chart", "create_pie_chart", "create_table", "http_get", "http_post", "http_put", "http_delete", "http_patch", "read_file"]
+            self.tools["ui"] = [
+                "create_bar_chart",
+                "create_line_chart",
+                "create_pie_chart",
+                "create_table",
+                "http_get",
+                "http_post",
+                "http_put",
+                "http_delete",
+                "http_patch",
+                "read_file",
+            ]
         self.max_iterations = max_iterations
         self.registry = registry
 
@@ -54,7 +66,6 @@ class ReactAgent(BaseAgent):
         else:
             tools = []
 
-
         logger.debug(f"Initialized ReactAgent: name={self.name}, model={self.model}")
         return create_react_agent(
             self.llm,
@@ -75,10 +86,7 @@ async def main():
         tools={"google-mail": ["send_email"]},
         registry=AgentrRegistry(),
     )
-    result = await agent.invoke(
-        "Send an email with the subject 'testing react agent' to manoj@agentr.dev"
-    )
-    from rich import print
+    result = await agent.invoke("Send an email with the subject 'testing react agent' to manoj@agentr.dev")
 
     print(messages_to_list(result["messages"]))
 

universal_mcp/agents/shared/prompts.py

@@ -6,7 +6,7 @@ You are an expert planner. Your goal is to consolidate a complex user request in
 2.  **Focus on Data Handoffs:** A good decomposition often involves one sub-task to *retrieve* information and a subsequent sub-task to *use* that information.
 3.  **Assume Internal Capabilities:** Do NOT create sub-tasks for abstract cognitive work like 'summarize' or 'analyze'.
 4.  **Simplify Single Actions:** If the user's task is already a single, simple action, the output should be a single sub-task that concisely describes that action. Do not make it identical to the user's input.
-5.  **General purpose sub tasks:** You also need to realise that these subtasks are going to be used to search for tools and apps. And the names and description of these tools and apps are going to be general in nature so the sub tasks should not be too specific. The task which you will get may be specific in nature but the sub taks must be general. 
+5.  **General purpose sub tasks:** You also need to realise that these subtasks are going to be used to search for tools and apps. And the names and description of these tools and apps are going to be general in nature so the sub tasks should not be too specific. The task which you will get may be specific in nature but the sub taks must be general.
 **--- EXAMPLES ---**
 
 **EXAMPLE 1:**
@@ -116,7 +116,7 @@ You are an AI assistant that selects the most appropriate tool(s) from a list to
 **INSTRUCTIONS:**
 1.  Carefully review the sub-task to understand the required action.
 2.  Examine the list of available tools and their descriptions.
-3.  Select the best tool ID that matches the sub-task. You are encouraged to select multiple tools if there are multiple tools with similar capabilties 
+3.  Select the best tool ID that matches the sub-task. You are encouraged to select multiple tools if there are multiple tools with similar capabilties
 or names. It is always good to have more tools than having insufficent tools.
 4.  If no tool is a good fit, return an empty list.
 5.  Only return the tool IDs.
@@ -129,4 +129,4 @@ or names. It is always good to have more tools than having insufficent tools.
 {tool_candidates}
 
 **YOUR SELECTED TOOL ID(s):**
-"""
+"""