universal-mcp 0.1.24rc6__py3-none-any.whl → 0.1.24rc8__py3-none-any.whl

universal_mcp/agents/auto.py

@@ -12,13 +12,12 @@ from pydantic import BaseModel
 from typing_extensions import TypedDict
 
 from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.llm import load_chat_model
 from universal_mcp.tools import ToolManager
 from universal_mcp.tools.adapters import ToolFormat
 from universal_mcp.tools.registry import ToolRegistry
 
-from .base import BaseAgent
-from .llm import get_llm
-
 # Auto Agent
 # Working
 # 1. For every message, and given list of tools, figure out if external tools are needed
@@ -61,9 +60,9 @@ class AutoAgent(BaseAgent):
     def __init__(self, name: str, instructions: str, model: str, app_registry: ToolRegistry):
         super().__init__(name, instructions, model)
         self.app_registry = app_registry
-        self.llm_tools = get_llm(model, tags=["tools"])
-        self.llm_choice = get_llm(model, tags=["choice"])
-        self.llm_quiet = get_llm(model, tags=["quiet"])
+        self.llm_tools = load_chat_model(model, tags=["tools"])
+        self.llm_choice = load_chat_model(model, tags=["choice"])
+        self.llm_quiet = load_chat_model(model, tags=["quiet"])
         self.tool_manager = ToolManager()
 
         self.task_analysis_prompt = """You are a task analysis expert. Given a task description and available apps, determine:
@@ -522,7 +521,7 @@ Be friendly and concise, but list each set of apps clearly. Do not return any ot
             return result
 
         # Get all available apps from platform manager
-        available_apps = await self.app_registry.list_apps()
+        available_apps =  self.app_registry.list_apps()
 
         logger.info(f"Found {len(available_apps)} available apps")
 
@@ -563,10 +562,10 @@ if __name__ == "__main__":
 
     # Create platform manager
     app_registry = AgentrRegistry(api_key=agentr_api_key)
-    want_instructions = input("Do you want to add a system prompt/instructions? (Y/N)")
+    want_instructions = input("Do you want to add a system prompt/instructions? (Y/N): ")
     instructions = "" if want_instructions.upper() == "N" else input("Enter your instructions/system prompt: ")
 
-    agent = AutoAgent("Auto Agent", instructions, "gpt-4.1", app_registry=app_registry)
+    agent = AutoAgent("Auto Agent", instructions, "azure/gpt-4.1", app_registry=app_registry)
 
     print("AutoAgent created successfully!")
     print(f"Agent name: {agent.name}")

universal_mcp/agents/autoagent/__init__.py

@@ -0,0 +1,31 @@
+from langgraph.checkpoint.base import BaseCheckpointSaver
+
+from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.agents.autoagent.graph import build_graph
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.tools.registry import ToolRegistry
+
+
+class AutoAgent(BaseAgent):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        memory: BaseCheckpointSaver | None = None,
+        registry: ToolRegistry | None = None,
+        **kwargs,
+    ):
+        super().__init__(name, instructions, model, memory, **kwargs)
+        self.tool_registry = registry or AgentrRegistry()
+
+    async def _build_graph(self):
+        builder = await build_graph(self.tool_registry, self.instructions)
+        return builder.compile(checkpointer=self.memory)
+
+    @property
+    def graph(self):
+        return self._graph
+
+
+__all__ = ["AutoAgent"]

universal_mcp/agents/autoagent/__main__.py

@@ -0,0 +1,21 @@
+import asyncio
+
+from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.agents.autoagent import AutoAgent
+
+
+async def main():
+    agent = AutoAgent(
+        name="autoagent",
+        instructions="You are a helpful assistant that can use tools to help the user.",
+        model="anthropic/claude-4-sonnet-20250514",
+        tool_registry=AgentrRegistry(),
+    )
+    result = await agent.invoke(
+        user_input="Send an email to Manoj from my google mail account, manoj@agentr.dev, with the subject 'Hello from auto agent' and the body 'testing'"
+    )
+    print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

universal_mcp/agents/autoagent/context.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass, field
+from typing import Annotated
+
+from universal_mcp.agents.autoagent.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: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = 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."
+        },
+    )

universal_mcp/agents/autoagent/graph.py

@@ -0,0 +1,148 @@
+import json
+from datetime import UTC, datetime
+from typing import cast
+
+from langchain_core.messages import AIMessage, ToolMessage
+from langchain_core.tools import tool
+from langgraph.graph import END, START, StateGraph
+from langgraph.runtime import Runtime
+
+from universal_mcp.agents.autoagent.context import Context
+from universal_mcp.agents.autoagent.prompts import SYSTEM_PROMPT
+from universal_mcp.agents.autoagent.state import State
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolFormat
+
+
+async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
+    @tool()
+    async def search_tools(query: str, app_ids: list[str] | None = None) -> list[str]:
+        """Retrieve tools using a search query and app id. Use multiple times if you require tools for different tasks."""
+        tools_list = []
+        if app_ids is not None:
+            for app_id in app_ids:
+                tools_list.extend(await tool_registry.search_tools(query, limit=10, app_id=app_id))
+        else:
+            tools_list = await tool_registry.search_tools(query, limit=10)
+        tools_list = [f"{tool['id']}: {tool['description']}" for tool in tools_list]
+        return tools_list
+
+    @tool()
+    async def ask_user(question: str) -> str:
+        """Ask the user a question. Use this tool to ask the user for any missing information for performing a task, or when you have multiple apps to choose from for performing a task."""
+        full_question = question
+        return f"ASKING_USER: {full_question}"
+    
+    @tool()
+    async def load_tools(tools: list[str]) -> list[str]:
+        """Choose the tools you want to use by passing their tool ids.  Loads the tools for the chosen tools and returns the tool ids."""
+        return tools
+    
+
+    async def call_model(
+        state: State,
+        runtime: Runtime[Context],
+    ):
+        system_prompt = runtime.context.system_prompt if runtime.context.system_prompt else SYSTEM_PROMPT
+        app_ids = await tool_registry.list_all_apps()
+        app_id_descriptions = "\n".join([f"{app['id']}: {app['description']}" for app in app_ids])
+        print(app_id_descriptions)
+        system_prompt = system_prompt.format(system_time=datetime.now(tz=UTC).isoformat(), app_ids=app_id_descriptions)
+        
+        messages = [{"role": "system", "content": system_prompt + "\n" + instructions}, *state["messages"]]
+        model = load_chat_model(runtime.context.model)
+        # Load tools from tool registry
+        loaded_tools = await tool_registry.export_tools(tools=state["selected_tool_ids"], format=ToolFormat.LANGCHAIN)
+        model_with_tools = model.bind_tools([search_tools, ask_user, load_tools, *loaded_tools], tool_choice="auto")
+        response_raw = model_with_tools.invoke(messages)
+        token_usage = state.get("token_usage", {})
+        for key in ["input_tokens", "output_tokens", "total_tokens"]:
+            if key in token_usage:
+                token_usage[key] += response_raw.usage_metadata[key]
+            else:
+                token_usage[key] = response_raw.usage_metadata[key]
+        print(response_raw.usage_metadata)
+        print(token_usage)
+        response = cast(AIMessage, response_raw) 
+        return {"messages": [response], "token_usage": token_usage}
+
+    # Define the conditional edge that determines whether to continue or not
+    def should_continue(state: State):
+        messages = state["messages"]
+        last_message = messages[-1]
+        # If there is no function call, then we finish
+        if not last_message.tool_calls:
+            return END
+        else:
+            return "tools"
+
+    def tool_router(state: State):
+        last_message = state["messages"][-1]
+        if isinstance(last_message, ToolMessage):
+            return "agent"
+        else:
+            return END
+
+    async def tool_node(state: State):
+        outputs = []
+        tool_ids = state["selected_tool_ids"]
+        for tool_call in state["messages"][-1].tool_calls:
+            if tool_call["name"] == ask_user.name:
+                outputs.append(
+                    ToolMessage(
+                        content=json.dumps(
+                            "The user has been asked the question, and the run will wait for the user's response."
+                        ),
+                        name=tool_call["name"],
+                        tool_call_id=tool_call["id"],
+                    )
+                )
+                ai_message = AIMessage(content=tool_call["args"]["question"])
+                outputs.append(ai_message)
+            elif tool_call["name"] == search_tools.name:
+                tools = await search_tools.ainvoke(tool_call["args"])
+                outputs.append(
+                    ToolMessage(
+                        content=json.dumps(tools)+"\n\nUse the load_tools tool to load the tools you want to use.",
+                        name=tool_call["name"],
+                        tool_call_id=tool_call["id"],
+                    )
+                )
+                
+            elif tool_call["name"] == load_tools.name:
+                tool_ids = await load_tools.ainvoke(tool_call["args"])
+                print(tool_ids)
+                outputs.append(
+                    ToolMessage(
+                        content=json.dumps(tool_ids),
+                        name=tool_call["name"],
+                        tool_call_id=tool_call["id"],
+                    )
+                )
+            else:
+                await tool_registry.export_tools([tool_call["name"]], ToolFormat.LANGCHAIN)
+                tool_result = await tool_registry.call_tool(tool_call["name"], tool_call["args"])
+                outputs.append(
+                    ToolMessage(
+                        content=json.dumps(tool_result),
+                        name=tool_call["name"],
+                        tool_call_id=tool_call["id"],
+                    )
+                )
+        return {"messages": outputs, "selected_tool_ids": tool_ids}
+
+                
+
+
+
+    builder = StateGraph(State, context_schema=Context)
+
+    builder.add_node("agent", call_model)
+    builder.add_node("tools", tool_node)
+
+    builder.add_edge(START, "agent")
+    builder.add_conditional_edges("agent", should_continue)
+    builder.add_conditional_edges("tools", tool_router)
+
+    return builder

universal_mcp/agents/autoagent/prompts.py

@@ -0,0 +1,8 @@
+"""Default prompts used by the agent."""
+
+SYSTEM_PROMPT = """You are a helpful AI assistant. When you lack tools for any task you should use the `search_tools` function to unlock relevant tools. Whenever you need to ask the user for any information, or choose between multiple different applications, you can ask the user using the `ask_user` function.
+
+System time: {system_time}
+App IDs: {app_ids}
+Note that when multiple apps seem relevant for a task, you should ask the user to choose the app.
+"""

universal_mcp/agents/autoagent/state.py

@@ -0,0 +1,28 @@
+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]
+    token_usage: dict[str, int]

universal_mcp/agents/autoagent/studio.py

@@ -0,0 +1,25 @@
+import asyncio
+
+from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.agents.autoagent import build_graph
+from universal_mcp.tools import ToolManager
+
+tool_registry = AgentrRegistry()
+tool_manager = ToolManager()
+
+
+
+async def main():
+    instructions = """
+    You are a helpful assistant that can use tools to help the user. If a task requires multiple steps, you should perform separate different searches for different actions. Prefer completing one action before searching for another.
+    """
+    graph = await build_graph(tool_registry, instructions=instructions)
+    return graph
+
+graph = asyncio.run(main())
+
+
+
+
+
+

universal_mcp/agents/autoagent/utils.py

@@ -0,0 +1,13 @@
+from langchain_core.messages import BaseMessage
+
+
+def get_message_text(msg: BaseMessage) -> str:
+    """Get the text content of a message."""
+    content = msg.content
+    if isinstance(content, str):
+        return content
+    elif isinstance(content, dict):
+        return content.get("text", "")
+    else:
+        txts = [c if isinstance(c, str) else (c.get("text") or "") for c in content]
+        return "".join(txts).strip()

universal_mcp/agents/base.py

@@ -7,7 +7,6 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.types import Command
 
-from .llm import get_llm
 from .utils import RichCLI
 
 
@@ -18,15 +17,19 @@ class BaseAgent:
         self.model = model
         self.memory = memory or MemorySaver()
         self._graph = None
-        self.llm = get_llm(model)
+        self._initialized = False
         self.cli = RichCLI()
 
+    async def ainit(self):
+        if not self._initialized:
+            self._graph = await self._build_graph()
+            self._initialized = True
+
     async def _build_graph(self):
         raise NotImplementedError("Subclasses must implement this method")
 
     async def stream(self, thread_id: str, user_input: str):
-        if self._graph is None:
-            self._graph = await self._build_graph()
+        await self.ainit()
         async for event, _ in self._graph.astream(
             {"messages": [{"role": "user", "content": user_input}]},
             config={"configurable": {"thread_id": thread_id}},
@@ -36,24 +39,24 @@ class BaseAgent:
             yield event
 
     async def stream_interactive(self, thread_id: str, user_input: str):
+        await self.ainit()
         with self.cli.display_agent_response_streaming(self.name) as stream_updater:
-            async for event in self.astream(thread_id, user_input):
+            async for event in self.stream(thread_id, user_input):
                 stream_updater.update(event.content)
 
-    async def run(self, user_input: str, thread_id: str = str(uuid4())):
+    async def invoke(self, user_input: str, thread_id: str = str(uuid4())):
         """Run the agent"""
-        if not self._graph:
-            self._graph = await self._build_graph()
+        await self.ainit()
         return await self._graph.ainvoke(
             {"messages": [{"role": "user", "content": user_input}]},
             config={"configurable": {"thread_id": thread_id}},
+            context={"system_prompt": self.instructions, "model": self.model},
         )
 
     async def run_interactive(self, thread_id: str = str(uuid4())):
         """Main application loop"""
 
-        if not self._graph:
-            self._graph = await self._build_graph()
+        await self.ainit()
         # Display welcome
         self.cli.display_welcome(self.name)
 

universal_mcp/agents/codeact/test.py

@@ -1,10 +1,10 @@
 from universal_mcp.agentr import Agentr
 from universal_mcp.agents.codeact import create_codeact
 from universal_mcp.agents.codeact.sandbox import eval_unsafe
-from universal_mcp.agents.llm import get_llm
+from universal_mcp.agents.llm import load_chat_model
 from universal_mcp.tools.adapters import ToolFormat
 
-model = get_llm("gpt-4.1")
+model = load_chat_model("gpt-4.1")
 
 agentr = Agentr()
 agentr.load_tools(["google-mail_send_email"])

universal_mcp/agents/hil.py

@@ -7,7 +7,7 @@ from langgraph.graph.message import add_messages
 from langgraph.types import Interrupt, interrupt
 
 from .base import BaseAgent
-from .llm import get_llm
+from .llm import load_chat_model
 
 
 class State(TypedDict):
@@ -70,7 +70,7 @@ def handle_interrupt(interrupt: Interrupt) -> str | bool:
 class HilAgent(BaseAgent):
     def __init__(self, name: str, instructions: str, model: str):
         super().__init__(name, instructions, model)
-        self.llm = get_llm(model)
+        self.llm = load_chat_model(model)
         self._graph = self._build_graph()
 
     def chatbot(self, state: State):

universal_mcp/agents/llm.py

@@ -1,10 +1,28 @@
+from langchain_anthropic import ChatAnthropic
+from langchain_core.language_models import BaseChatModel
+from langchain_google_vertexai.model_garden import ChatAnthropicVertex
 from langchain_openai import AzureChatOpenAI
 
 
-def get_llm(model: str, tags: list[str] | None = None):
-    return AzureChatOpenAI(model=model, api_version="2024-12-01-preview", azure_deployment=model, tags=tags)
+def load_chat_model(fully_specified_name: str, tags: list[str] | None = None) -> BaseChatModel:
+    """Load a chat model from a fully specified name.
+
+    Args:
+        fully_specified_name (str): String in the format 'provider/model'.
+    """
+    provider, model = fully_specified_name.split("/", maxsplit=1)
+    if provider == "google_anthropic_vertex":
+        return ChatAnthropicVertex(model=model, temperature=0.2, location="asia-east1", tags=tags)
+    elif provider == "anthropic":
+        return ChatAnthropic(
+            model=model, temperature=1, thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096, tags=tags
+        )  # pyright: ignore[reportCallIssue]
+    elif provider == "azure":
+        return AzureChatOpenAI(model=model, api_version="2024-12-01-preview", azure_deployment=model, tags=tags)
+    else:
+        raise ValueError(f"Unsupported provider: {provider}")
 
 
 if __name__ == "__main__":
-    llm = get_llm("gpt-4.1")
+    llm = load_chat_model("azure/gpt-4.1")
     print(llm.invoke("Hello, world!"))

universal_mcp/agents/react.py

@@ -2,9 +2,12 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.prebuilt import create_react_agent
 from loguru import logger
 
+from universal_mcp.agentr.registry import AgentrRegistry
 from universal_mcp.agents.base import BaseAgent
-from universal_mcp.agents.tools import load_agentr_tools, load_mcp_tools
-from universal_mcp.types import ToolConfig
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.tools import load_mcp_tools
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolConfig, ToolFormat
 
 
 class ReactAgent(BaseAgent):
@@ -15,18 +18,29 @@ class ReactAgent(BaseAgent):
         model: str,
         memory: BaseCheckpointSaver | None = None,
         tools: ToolConfig | None = None,
+        registry: ToolRegistry | None = None,
         max_iterations: int = 10,
         **kwargs,
     ):
         super().__init__(name, instructions, model, memory, **kwargs)
+        self.llm = load_chat_model(model)
         self.tools = tools
         self.max_iterations = max_iterations
+        self.registry = registry
 
     async def _build_graph(self):
         if self.tools:
             config = self.tools.model_dump(exclude_none=True)
-            agentr_tools = await load_agentr_tools(config["agentrServers"]) if config.get("agentrServers") else []
+            if config.get("agentrServers") and not self.registry:
+                raise ValueError("Agentr servers are configured but no registry is provided")
+            agentr_tools = (
+                await self.registry.export_tools(self.tools, ToolFormat.LANGCHAIN)
+                if config.get("agentrServers")
+                else []
+            )
+            print(agentr_tools)
             mcp_tools = await load_mcp_tools(config["mcpServers"]) if config.get("mcpServers") else []
+            print(mcp_tools)
             tools = agentr_tools + mcp_tools
         else:
             tools = []
@@ -39,7 +53,7 @@ class ReactAgent(BaseAgent):
         )
 
     def _build_system_message(self) -> str:
-        system_message = f"""You are {self.name}. {self.instructions}
+        system_message = f"""You are {self.name}.
 
 You have access to various tools that can help you answer questions and complete tasks. When you need to use a tool:
 
@@ -47,7 +61,10 @@ You have access to various tools that can help you answer questions and complete
 2. Call the appropriate tool with the right parameters
 3. Use the tool results to provide a comprehensive answer
 
-Always explain your reasoning and be thorough in your responses. If you need to use multiple tools to answer a question completely, do so."""
+Always explain your reasoning and be thorough in your responses. If you need to use multiple tools to answer a question completely, do so.
+
+{self.instructions}
+"""
         return system_message
 
 
@@ -57,8 +74,11 @@ if __name__ == "__main__":
     agent = ReactAgent(
         "Universal React Agent",
         instructions="",
-        model="gpt-4o",
+        model="azure/gpt-4o",
         tools=ToolConfig(agentrServers={"google-mail": {"tools": ["send_email"]}}),
+        registry=AgentrRegistry(),
+    )
+    result = asyncio.run(
+        agent.invoke(user_input="Send an email with the subject 'testing react agent' to manoj@agentr.dev")
     )
-    result = asyncio.run(agent.run(user_input="Send an email with the subject 'Hello' to john.doe@example.com"))
     print(result["messages"][-1].content)

universal_mcp/agents/simple.py

@@ -1,12 +1,13 @@
 import asyncio
 from typing import Annotated
 
+from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph import END, START, StateGraph
 from langgraph.graph.message import add_messages
 from typing_extensions import TypedDict
 
-from .base import BaseAgent
-from .llm import get_llm
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.llm import load_chat_model
 
 
 class State(TypedDict):
@@ -14,27 +15,26 @@ class State(TypedDict):
 
 
 class SimpleAgent(BaseAgent):
-    def __init__(self, name: str, instructions: str, model: str):
-        super().__init__(name, instructions, model)
-        self.llm = get_llm(model)
-        self._graph = self._build_graph()
+    def __init__(self, name: str, instructions: str, model: str, memory: BaseCheckpointSaver = None, **kwargs):
+        super().__init__(name, instructions, model, memory, **kwargs)
+        self.llm = load_chat_model(model)
 
-    def _build_graph(self):
+    async def _build_graph(self):
         graph_builder = StateGraph(State)
 
-        def chatbot(state: State):
-            return {"messages": [self.llm.invoke(state["messages"])]}
+        async def chatbot(state: State):
+            messages = [
+                {"role": "system", "content": self.instructions},
+                *state["messages"],
+            ]
+            return {"messages": [await self.llm.ainvoke(messages)]}
 
         graph_builder.add_node("chatbot", chatbot)
         graph_builder.add_edge(START, "chatbot")
         graph_builder.add_edge("chatbot", END)
         return graph_builder.compile(checkpointer=self.memory)
 
-    @property
-    def graph(self):
-        return self._graph
-
 
 if __name__ == "__main__":
-    agent = SimpleAgent("Simple Agent", "You are a helpful assistant", "openrouter/auto")
+    agent = SimpleAgent("Simple Agent", "You are a helpful assistant", "azure/gpt-4o")
     asyncio.run(agent.run_interactive())