universal-mcp 0.1.24rc11__py3-none-any.whl → 0.1.24rc13__py3-none-any.whl

universal_mcp/agentr/registry.py

@@ -31,6 +31,8 @@ class AgentrRegistry(ToolRegistry):
             raise ValueError("Client is not initialized")
         try:
             all_apps = self.client.list_all_apps()
+            all_apps = [{"id": app["id"], "name": app["name"], "description": app["description"]} for app in all_apps]
+            # logger.debug(f"All apps: {all_apps}")
             return all_apps
         except Exception as e:
             logger.error(f"Error fetching apps from AgentR: {e}")
@@ -180,3 +182,7 @@ class AgentrRegistry(ToolRegistry):
     async def call_tool(self, tool_name: str, tool_args: dict[str, Any]) -> dict[str, Any]:
         """Call a tool with the given name and arguments."""
         return await self.tool_manager.call_tool(tool_name, tool_args)
+
+    async def list_connected_apps(self) -> list[str]:
+        """List all apps that the user has connected."""
+        return self.client.list_my_connections()

universal_mcp/agents/__init__.py

@@ -1,6 +1,9 @@
 from universal_mcp.agents.autoagent import AutoAgent
 from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.bigtool import BigToolAgent
+from universal_mcp.agents.builder import BuilderAgent
+from universal_mcp.agents.planner import PlannerAgent
 from universal_mcp.agents.react import ReactAgent
 from universal_mcp.agents.simple import SimpleAgent
 
-__all__ = ["BaseAgent", "ReactAgent", "SimpleAgent", "AutoAgent"]
+__all__ = ["BaseAgent", "ReactAgent", "SimpleAgent", "AutoAgent", "BigToolAgent", "PlannerAgent", "BuilderAgent"]

universal_mcp/agents/autoagent/__init__.py

@@ -1,6 +1,5 @@
 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
@@ -17,7 +16,7 @@ class AutoAgent(BaseAgent):
         **kwargs,
     ):
         super().__init__(name, instructions, model, memory, **kwargs)
-        self.tool_registry = registry or AgentrRegistry()
+        self.tool_registry = registry
 
     async def _build_graph(self):
         builder = await build_graph(self.tool_registry, self.instructions)

universal_mcp/agents/autoagent/__main__.py

@@ -1,5 +1,7 @@
 import asyncio
 
+from loguru import logger
+
 from universal_mcp.agentr.registry import AgentrRegistry
 from universal_mcp.agents.autoagent import AutoAgent
 
@@ -9,12 +11,13 @@ async def main():
         name="autoagent",
         instructions="You are a helpful assistant that can use tools to help the user.",
         model="azure/gpt-4.1",
-        tool_registry=AgentrRegistry(),
-    )
-    await agent.invoke(
-        user_input="Please send the email from google-mail to manoj@agentr.dev, with subject hello and body hello from auto",
-        thread_id="12345",
+        registry=AgentrRegistry(),
     )
+    async for event in agent.stream(
+        user_input="Send an email to manoj@agentr.dev",
+        thread_id="test123",
+    ):
+        logger.info(event.content)
     # from loguru import logger; logger.debug(result)
 
 

universal_mcp/agents/autoagent/graph.py

@@ -45,7 +45,7 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
     ):
         system_prompt = SYSTEM_PROMPT
         app_ids = await tool_registry.list_all_apps()
-        connections = tool_registry.client.list_my_connections()
+        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]
@@ -61,7 +61,6 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
 
         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)
@@ -80,10 +79,10 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
 
     def tool_router(state: State):
         last_message = state["messages"][-1]
-        if isinstance(last_message, ToolMessage):
-            return "agent"
-        else:
+        if isinstance(last_message, ToolMessage) and last_message.name == ask_user.name:
             return END
+        else:
+            return "agent"
 
     async def tool_node(state: State):
         outputs = []
@@ -99,8 +98,6 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
                         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(

universal_mcp/agents/autoagent/prompts.py

@@ -5,5 +5,5 @@ SYSTEM_PROMPT = """You are a helpful AI assistant. When you lack tools for any t
 System time: {system_time}
 These are the list of apps available to you:
 {app_ids}
-Note that when multiple apps seem relevant for a task, you MUST ask the user to choose the app. 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.
+Note that when multiple apps seem relevant for a task, you MUST ask the user to choose the app. 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. You will be provided a link for connection that you should pass on to the user.
 """

universal_mcp/agents/base.py

@@ -4,7 +4,6 @@ from uuid import uuid4
 
 from langchain_core.messages import AIMessageChunk
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.checkpoint.memory import MemorySaver
 from langgraph.types import Command
 
 from .utils import RichCLI
@@ -15,7 +14,7 @@ class BaseAgent:
         self.name = name
         self.instructions = instructions
         self.model = model
-        self.memory = memory or MemorySaver()
+        self.memory = memory
         self._graph = None
         self._initialized = False
         self.cli = RichCLI()
@@ -30,6 +29,7 @@ class BaseAgent:
 
     async def stream(self, thread_id: str, user_input: str):
         await self.ainit()
+        aggregate = None
         async for event, metadata in self._graph.astream(
             {"messages": [{"role": "user", "content": user_input}]},
             config={"configurable": {"thread_id": thread_id}},
@@ -38,25 +38,23 @@ class BaseAgent:
             stream_usage=True,
         ):
             # Only forward assistant token chunks that are not tool-related.
-            event = cast(AIMessageChunk, event)
-            if "finish_reason" in event.response_metadata:
-                # Got LLM finish reason ignore it
-                # logger.debug(f"Finish event: {event}")
+            type_ = type(event)
+            if type_ != AIMessageChunk:
                 continue
-
-            # Skip chunks that correspond to tool calls or tool execution phases
-            # - tool_call_chunks present => model is emitting tool call(s)
-            # - metadata tags or node names may indicate tool/quiet phases
+            event = cast(AIMessageChunk, event)
+            aggregate = event if aggregate is None else aggregate + event
             tags = metadata.get("tags", []) if isinstance(metadata, dict) else []
-
             is_quiet = isinstance(tags, list) and ("quiet" in tags)
 
             if is_quiet:
                 continue
-
-            # Emit only the token chunks for the final assistant message.
-            # logger.debug(f"Event: {event}, Metadata: {metadata}")
-            yield event
+            if "finish_reason" in event.response_metadata:
+                # Got LLM finish reason ignore it
+                # logger.debug(f"Finish event: {event}, Metadata: {metadata}")
+                pass
+            else:
+                # logger.debug(f"Event: {event}, Metadata: {metadata}")
+                yield event
         # Send a final finished message
         # The last event would be finish
         event = cast(AIMessageChunk, event)

universal_mcp/agents/bigtool/__init__.py

@@ -0,0 +1,54 @@
+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.logger import logger
+from universal_mcp.tools.registry import ToolRegistry
+
+from .graph import build_graph
+from .prompts import SYSTEM_PROMPT
+
+
+class BigToolAgent(BaseAgent):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        registry: ToolRegistry,
+        memory: BaseCheckpointSaver | None = None,
+        **kwargs,
+    ):
+        # Combine the base system prompt with agent-specific instructions
+        full_instructions = f"{SYSTEM_PROMPT}\n\n**User Instructions:**\n{instructions}"
+        super().__init__(name, full_instructions, model, memory, **kwargs)
+
+        self.registry = registry
+        self.llm = load_chat_model(self.model)
+        self.tool_selection_llm = load_chat_model("gemini/gemini-2.0-flash-001")
+
+        logger.info(f"BigToolAgent '{self.name}' initialized with model '{self.model}'.")
+
+    async def _build_graph(self):
+        """Build the bigtool agent graph using the existing create_agent function."""
+        logger.info(f"Building graph for BigToolAgent '{self.name}'...")
+        try:
+            graph_builder = build_graph(
+                tool_registry=self.registry,
+                llm=self.llm,
+                tool_selection_llm=self.tool_selection_llm,
+            )
+
+            compiled_graph = graph_builder.compile(checkpointer=self.memory)
+            logger.info("Graph built and compiled successfully.")
+            return compiled_graph
+        except Exception as e:
+            logger.error(f"Error building graph for BigToolAgent '{self.name}': {e}")
+            raise
+
+    @property
+    def graph(self):
+        return self._graph
+
+
+__all__ = ["BigToolAgent"]

universal_mcp/agents/bigtool/__main__.py

@@ -0,0 +1,24 @@
+import asyncio
+
+from loguru import logger
+
+from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.agents.bigtool import BigToolAgent
+
+
+async def main():
+    agent = BigToolAgent(
+        name="bigtool",
+        instructions="You are a helpful assistant that can use tools to help the user.",
+        model="azure/gpt-4.1",
+        registry=AgentrRegistry(),
+    )
+    async for event in agent.stream(
+        user_input="Send an email to manoj@agentr.dev",
+        thread_id="test123",
+    ):
+        logger.info(event.content)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

universal_mcp/agents/bigtool/context.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass, field
+
+from .prompts import SYSTEM_PROMPT
+
+
+@dataclass(kw_only=True)
+class Context:
+    """The context for the agent."""
+
+    system_prompt: str = field(
+        default=SYSTEM_PROMPT,
+        metadata={
+            "description": "The system prompt to use for the agent's interactions. "
+            "This prompt sets the context and behavior for the agent."
+        },
+    )
+
+    model: str = field(
+        default="anthropic/claude-4-sonnet-20250514",
+        metadata={
+            "description": "The name of the language model to use for the agent's main interactions. "
+            "Should be in the form: provider/model-name."
+        },
+    )

universal_mcp/agents/bigtool/graph.py

@@ -0,0 +1,166 @@
+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.bigtool.context import Context
+from universal_mcp.agents.bigtool.state import State
+from universal_mcp.logger import logger
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolFormat
+
+from .prompts import SELECT_TOOL_PROMPT
+
+
+def build_graph(
+    tool_registry: ToolRegistry,
+    llm: BaseChatModel,
+    tool_selection_llm: BaseChatModel,
+):
+    @tool
+    async def retrieve_tools(task_query: str) -> list[str]:
+        """Retrieve tools for a given task.
+        Task query should be atomic (doable with a single tool).
+        For tasks requiring multiple tools, call this tool multiple times for each subtask."""
+        logger.info(f"Retrieving tools for task: '{task_query}'")
+        try:
+            tools_list = await tool_registry.search_tools(task_query, limit=10)
+            tool_candidates = [f"{tool['id']}: {tool['description']}" for tool in tools_list]
+            logger.info(f"Found {len(tool_candidates)} candidate tools.")
+
+            class ToolSelectionOutput(TypedDict):
+                tool_names: list[str]
+
+            model = tool_selection_llm
+            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_id_descriptions = "These are the apps connected to the user's account:\n" + "\n".join(
+                [f"{app}" for app in connected_apps]
+            )
+            if unconnected_apps:
+                app_id_descriptions += "\n\nOther (not connected) apps: " + "\n".join(
+                    [f"{app}" for app in unconnected_apps]
+                )
+
+            response = await model.with_structured_output(schema=ToolSelectionOutput, method="json_mode").ainvoke(
+                SELECT_TOOL_PROMPT.format(
+                    app_ids=app_id_descriptions, tool_candidates="\n - ".join(tool_candidates), task=task_query
+                )
+            )
+
+            selected_tool_names = cast(ToolSelectionOutput, response)["tool_names"]
+            logger.info(f"Selected tools: {selected_tool_names}")
+            return selected_tool_names
+        except Exception as e:
+            logger.error(f"Error retrieving tools: {e}")
+            return []
+
+    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']}")
+            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.")
+
+            model = llm
+            if isinstance(model, ChatAnthropic):
+                model_with_tools = model.bind_tools(
+                    [retrieve_tools, *selected_tools], tool_choice="auto", cache_control={"type": "ephemeral"}
+                )
+            else:
+                model_with_tools = model.bind_tools([retrieve_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"] == retrieve_tools.name:
+                    logger.info("Model requested to select tools.")
+                    return Command(goto="select_tools", update={"messages": [response]})
+                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]
+            selected_tool_names = await retrieve_tools.ainvoke(input=tool_call["args"])
+            tool_msg = ToolMessage(f"Available tools: {selected_tool_names}", tool_call_id=tool_call["id"])
+            logger.info(f"Tools selected: {selected_tool_names}")
+            return Command(goto="call_model", update={"messages": [tool_msg], "selected_tool_ids": selected_tool_names})
+        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/bigtool/prompts.py

@@ -0,0 +1,31 @@
+"""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 `retrieve_tools`:** Before you can use any other tool, you MUST first call the `retrieve_tools` function to find the right tool for the user's request. This is your mandatory first action.
+3.  **Strictly Follow the Process:** Your only job in your first turn is to analyze the user's request and call `retrieve_tools` with a concise query describing the core task. Do not engage in conversation.
+
+System time: {system_time}
+
+When multiple tools are available for the same task, you must ask the user.
+"""
+
+SELECT_TOOL_PROMPT = """You are an AI assistant that helps the user perform tasks using various apps (each app has multiple tools).
+You will be provided with a task and a list of tools which might be relevant for this task.
+
+Your goal is to select the most appropriate tool for the given task.
+<task>
+{task}
+</task>
+
+These are the list of apps available to you:
+{app_ids}
+Note that when multiple apps seem relevant for a 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 choose both apps tools. In case the user specifically asks you to use an app that is not connected, select the tool.
+
+<tool_candidates>
+ - {tool_candidates}
+</tool_candidates>
+
+"""

universal_mcp/agents/bigtool/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,26 +1,37 @@
 from langchain_anthropic import ChatAnthropic
 from langchain_core.language_models import BaseChatModel
-from langchain_google_vertexai import ChatVertexAI
-from langchain_google_vertexai.model_garden import ChatAnthropicVertex
+from langchain_google_genai import ChatGoogleGenerativeAI
 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)
-    elif provider == "anthropic":
+    if provider == "anthropic":
         return ChatAnthropic(
-            model=model, temperature=1, thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096, tags=tags
+            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)
-    elif provider == "vertex":
-        return ChatVertexAI(model=model, temperature=0.1, tags=tags)
+        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}")
 
@@ -28,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!"))