universal-mcp 0.1.24rc12__py3-none-any.whl → 0.1.24rc14__py3-none-any.whl

universal_mcp/agents/bigtool2/graph.py

@@ -0,0 +1,169 @@
+import json
+from datetime import UTC, datetime
+from typing import 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.bigtool2.context import Context
+from universal_mcp.agents.bigtool2.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)
+                tool_candidates = [f"{tool['id']}: {tool['description']}" for tool in tools_list]
+                for tool in tool_candidates:
+                    app = tool.split("__")[0]
+                    if app not in app_tools:
+                        if len(app_tools.keys()) >= 10:
+                            break
+                        app_tools[app] = []
+                    if len(app_tools[app]) < 3:
+                        app_tools[app].append(tool)
+            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}\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[str]:
+        """Load the tools for the given tool ids. Returns the tool ids."""
+        return tool_ids
+
+
+    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"]]
+
+            logger.info(f"Selected tool IDs: {state['selected_tool_ids']}")
+            if len(state["selected_tool_ids"]) > 0:
+                selected_tools = await tool_registry.export_tools(tools=state["selected_tool_ids"], format=ToolFormat.LANGCHAIN)
+                logger.info(f"Exported {len(selected_tools)} tools for model.")
+            else:
+                selected_tools = []
+
+            model = llm
+
+            model_with_tools = model.bind_tools([search_tools, load_tools, *selected_tools], 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_msg = ToolMessage(f"Loaded tools.", 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"] not in state["selected_tool_ids"]:
+                    try:
+                        await tool_registry.export_tools([tool_call["name"]], ToolFormat.LANGCHAIN)
+                        logger.info(
+                            f"Tool '{tool_call['name']}' not in selected tools, but available. Proceeding to call."
+                        )
+                        return Command(goto="call_tools", update={"messages": [response]})
+                    except Exception as e:
+                        logger.error(f"Unexpected tool call: {tool_call['name']}. Error: {e}")
+                        raise Exception(
+                            f"Unexpected tool call: {tool_call['name']}. Available tools: {state['selected_tool_ids']}"
+                        ) from e
+                logger.info(f"Proceeding to call tool: {tool_call['name']}")
+                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 = []
+        for tool_call in state["messages"][-1].tool_calls:
+            logger.info(f"Executing tool: {tool_call['name']} with args: {tool_call['args']}")
+            try:
+                await tool_registry.export_tools([tool_call["name"]], ToolFormat.LANGCHAIN)
+                tool_result = await tool_registry.call_tool(tool_call["name"], tool_call["args"])
+                logger.info(f"Tool '{tool_call['name']}' executed successfully.")
+                outputs.append(
+                    ToolMessage(
+                        content=json.dumps(tool_result),
+                        name=tool_call["name"],
+                        tool_call_id=tool_call["id"],
+                    )
+                )
+                recent_tool_ids.append(tool_call["name"])
+            except Exception as e:
+                logger.error(f"Error executing tool '{tool_call['name']}': {e}")
+                outputs.append(
+                    ToolMessage(
+                        content=json.dumps("Error: " + str(e)),
+                        name=tool_call["name"],
+                        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/bigtool2/prompts.py

@@ -0,0 +1,12 @@
+"""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.
+3.  **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/bigtool2/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,80 @@
+import asyncio
+from collections.abc import Sequence
+from typing import Annotated, TypedDict
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import BaseMessage
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.message import add_messages
+
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.shared.agent_node import Agent, generate_agent
+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 BuilderState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], add_messages]
+    generated_agent: Agent | None
+    tool_config: ToolConfig | None
+
+
+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]
+        generated_agent = await generate_agent(self.llm, last_message.content)
+        return {"generated_agent": generated_agent}
+
+    async def _create_tool_config(self, state: BuilderState):
+        last_message = state["messages"][-1]
+        tool_finder_graph = build_tool_node_graph(self.llm, self.registry)
+        tool_config = await tool_finder_graph.ainvoke({"task": last_message.content, "messages": [last_message]})
+        tool_config = tool_config.get("apps_with_tools", {})
+        return {"tool_config": tool_config}
+
+    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()
+
+
+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.model_dump_json(indent=2))
+    # print(f"Agent: {result['generated_agent'].model_dump_json(indent=2)}")
+    # print(f"Tool Config: {result['tool_config'].model_dump_json(indent=2)}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

universal_mcp/agents/llm.py

@@ -1,32 +1,37 @@
-import os
 from langchain_anthropic import ChatAnthropic
 from langchain_core.language_models import BaseChatModel
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_google_vertexai import ChatVertexAI
-from langchain_google_vertexai.model_garden import ChatAnthropicVertex
 from langchain_openai import AzureChatOpenAI
 
 
-def load_chat_model(fully_specified_name: str, tags: list[str] | None = None) -> BaseChatModel:
+def load_chat_model(
+    fully_specified_name: str, temperature: float = 1.0, 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, stream_usage=True,)
-    elif provider == "anthropic":
+    if provider == "anthropic":
         return ChatAnthropic(
-            model=model, temperature=1, thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096, tags=tags, stream_usage=True,
+            model=model,
+            temperature=temperature,
+            thinking={"type": "enabled", "budget_tokens": 2048},
+            max_tokens=4096,
+            tags=tags,
+            stream_usage=True,
         )  # pyright: ignore[reportCallIssue]
     elif provider == "azure":
-        return AzureChatOpenAI(model=model, api_version="2024-12-01-preview", azure_deployment=model, tags=tags, stream_usage=True,)
-    elif provider == "gemini":
-        return ChatGoogleGenerativeAI(
-            model= model,
-            temperature=1.0,
-            max_retries=2
+        return AzureChatOpenAI(
+            model=model,
+            api_version="2024-12-01-preview",
+            azure_deployment=model,
+            temperature=temperature,
+            tags=tags,
+            stream_usage=True,
         )
+    elif provider == "gemini":
+        return ChatGoogleGenerativeAI(model=model, temperature=temperature)
     else:
         raise ValueError(f"Unsupported provider: {provider}")
 
@@ -34,5 +39,7 @@ def load_chat_model(fully_specified_name: str, tags: list[str] | None = None) ->
 if __name__ == "__main__":
     from loguru import logger
 
-    llm = load_chat_model("azure/gpt-4.1")
-    logger.info(llm.invoke("Hello, world!"))
+    models_to_test = ["azure/gpt-5-chat", "anthropic/claude-4-sonnet-20250514", "gemini/gemini-2.5-pro"]
+    for model in models_to_test:
+        llm = load_chat_model(model)
+        logger.info(llm.invoke("Hi!"))

universal_mcp/agents/planner/__init__.py

@@ -0,0 +1,37 @@
+from langgraph.checkpoint.base import BaseCheckpointSaver
+
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.react import ReactAgent
+from universal_mcp.tools.registry import ToolRegistry
+
+from .graph import build_graph
+
+
+class PlannerAgent(BaseAgent):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        registry: ToolRegistry,
+        memory: BaseCheckpointSaver | None = None,
+        executor_agent_cls: type[BaseAgent] = ReactAgent,
+        **kwargs,
+    ):
+        super().__init__(name, instructions, model, memory, **kwargs)
+        self.app_registry = registry
+        self.llm = load_chat_model(model)
+        self.executor_agent_cls = executor_agent_cls
+
+    async def _build_graph(self):
+        return build_graph(self.llm, self.app_registry, self.instructions, self.model, self.executor_agent_cls).compile(
+            checkpointer=self.memory
+        )
+
+    @property
+    def graph(self):
+        return self._graph
+
+
+__all__ = ["PlannerAgent"]

universal_mcp/agents/planner/__main__.py

@@ -0,0 +1,24 @@
+import asyncio
+
+from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.agents.planner import PlannerAgent
+
+
+async def main():
+    registry = AgentrRegistry()
+    agent = PlannerAgent(
+        name="planner-agent",
+        instructions="You are a helpful assistant.",
+        model="gemini/gemini-2.5-flash",
+        registry=registry,
+    )
+    from rich.console import Console
+
+    console = Console()
+    console.print("Starting agent...", style="yellow")
+    async for event in agent.stream(user_input="Send an email to manoj@agentr.dev'", thread_id="xyz"):
+        console.print(event.content, style="red")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

universal_mcp/agents/planner/graph.py

@@ -0,0 +1,82 @@
+from typing import Any
+
+from langchain_core.messages import AIMessage
+from langgraph.graph import END, START, StateGraph
+from loguru import logger
+
+from universal_mcp.agents.shared.tool_node import build_tool_node_graph
+from universal_mcp.types import AgentrToolConfig, ToolConfig
+
+from .state import State
+
+
+def build_graph(llm, registry, instructions, model, executor_agent_cls):
+    """Build the graph for the planner agent."""
+    graph_builder = StateGraph(State)
+
+    async def _tool_finder_node(state: State) -> dict[str, Any]:
+        """Runs the tool finder subgraph to identify necessary tools."""
+        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({"task": task, "messages": state["messages"]})
+
+        if not tool_finder_state.get("apps_required"):
+            logger.info("Tool finder determined no apps are required.")
+            return {"apps_with_tools": AgentrToolConfig(agentrServers={})}
+
+        apps_with_tools = tool_finder_state.get("apps_with_tools", AgentrToolConfig(agentrServers={}))
+        logger.info(f"Tool finder identified apps and tools: {apps_with_tools}")
+        return {"apps_with_tools": apps_with_tools, "task": task}
+
+    def _should_continue(state: State) -> str:
+        """Determines whether to continue to the executor or end."""
+        if state.get("apps_with_tools") and state["apps_with_tools"].agentrServers:
+            return "continue"
+        return "end"
+
+    async def _executor_node(state: State) -> dict[str, Any]:
+        """Executes the task with the identified tools."""
+        tool_config = state["apps_with_tools"]
+
+        logger.info(f"Preparing executor with tools: {tool_config}")
+        agent = executor_agent_cls(
+            name="executor-agent",
+            instructions=instructions,
+            model=model,
+            registry=registry,
+            tools=ToolConfig(agentrServers=tool_config.agentrServers),
+        )
+
+        await agent.ainit()
+        react_graph = agent._graph
+        logger.info("Invoking executor agent with tools.")
+        # We invoke the agent to make it run the tool
+        response = await react_graph.ainvoke({"messages": state["messages"]})
+
+        final_message = AIMessage(content=response["messages"][-1].content)
+        return {"messages": [final_message]}
+
+    async def _no_tools_node(state: State) -> dict[str, Any]:
+        """Handles tasks that don't require tools by invoking the LLM directly."""
+        logger.info("No tools required. Invoking LLM directly.")
+        response = await llm.ainvoke(state["messages"])
+        return {"messages": [response]}
+
+    graph_builder.add_node("tool_finder", _tool_finder_node)
+    graph_builder.add_node("executor", _executor_node)
+    graph_builder.add_node("no_tools_executor", _no_tools_node)
+
+    graph_builder.add_edge(START, "tool_finder")
+    graph_builder.add_conditional_edges(
+        "tool_finder",
+        _should_continue,
+        {
+            "continue": "executor",
+            "end": "no_tools_executor",
+        },
+    )
+    graph_builder.add_edge("executor", END)
+    graph_builder.add_edge("no_tools_executor", END)
+
+    return graph_builder

universal_mcp/agents/planner/prompts.py

@@ -0,0 +1 @@
+# Prompts for the planner agent

universal_mcp/agents/planner/state.py

@@ -0,0 +1,12 @@
+from typing import Annotated
+
+from langgraph.graph.message import add_messages
+from typing_extensions import TypedDict
+
+from universal_mcp.types import AgentrToolConfig
+
+
+class State(TypedDict):
+    messages: Annotated[list, add_messages]
+    task: str
+    apps_with_tools: AgentrToolConfig

universal_mcp/agents/shared/agent_node.py

@@ -0,0 +1,34 @@
+from pydantic import BaseModel, Field
+
+
+class Agent(BaseModel):
+    name: str = Field(description="The name of the agent")
+    description: str = Field(description="A small paragraph description of the agent")
+    expertise: str = Field(description="Agents expertise. Growth expert, SEO expert, etc")
+    instructions: str = Field(description="The instructions for the agent")
+    schedule: str = Field(
+        description="The schedule for the agent in crontab syntax (e.g., '0 9 * * *' for daily at 9 AM)"
+    )
+
+
+AGENT_PROMPT = """You are an AI assistant that creates autonomous agents based on user requests.
+
+Your task is to analyze the user's request and create a structured agent definition that includes:
+- A clear name for the agent
+- A concise description of what the agent does
+- The agent's area of expertise
+- Detailed instructions for executing the task
+- A cron schedule for when the agent should run
+- A list of apps/services the agent will need to use
+
+Be specific and actionable in your agent definitions. Consider the user's intent and create agents that can effectively accomplish their goals.
+
+<query>
+{query}
+</query>
+"""
+
+
+async def generate_agent(llm, query):
+    response = await llm.with_structured_output(Agent).ainvoke(input=AGENT_PROMPT.format(query=query))
+    return response