universal-mcp-agents 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

universal_mcp/agents/__init__.py

@@ -1,18 +1,15 @@
-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.bigtool2 import BigToolAgent2
+from universal_mcp.agents.bigtoolcache import BigToolAgentCache
 from universal_mcp.agents.builder import BuilderAgent
+from universal_mcp.agents.codeact import CodeActAgent
 from universal_mcp.agents.planner import PlannerAgent
 from universal_mcp.agents.react import ReactAgent
 from universal_mcp.agents.simple import SimpleAgent
-from universal_mcp.agents.codeact import CodeActAgent
 
 
 def get_agent(agent_name: str):
-    if agent_name == "auto":
-        return AutoAgent
-    elif agent_name == "react":
+    if agent_name == "react":
         return ReactAgent
     elif agent_name == "simple":
         return SimpleAgent
@@ -20,14 +17,17 @@ def get_agent(agent_name: str):
         return BuilderAgent
     elif agent_name == "planner":
         return PlannerAgent
-    elif agent_name == "bigtool":
-        return BigToolAgent
     elif agent_name == "bigtool2":
         return BigToolAgent2
+    elif agent_name == "bigtoolcache":
+        return BigToolAgentCache
     elif agent_name == "codeact":
         return CodeActAgent
     else:
-        raise ValueError(f"Unknown agent: {agent_name}. Possible values: auto, react, simple, builder, planner, bigtool, bigtool2, codeact")
+        raise ValueError(
+            f"Unknown agent: {agent_name}. Possible values:  react, simple, builder, planner, bigtoolcache, bigtool2, codeact"
+        )
+
 
 __all__ = [
     "BaseAgent",

universal_mcp/agents/base.py

@@ -36,7 +36,7 @@ class BaseAgent:
     async def _build_graph(self) -> StateGraph:
         raise NotImplementedError("Subclasses must implement this method")
 
-    async def stream(self, thread_id: str, user_input: str, metadata: dict = None):
+    async def stream(self, user_input: str, thread_id: str = str(uuid4()), metadata: dict = None):
         await self.ainit()
         aggregate = None
 
@@ -53,7 +53,7 @@ class BaseAgent:
             "metadata": run_metadata,
         }
 
-        async for event, metadata in self._graph.astream(
+        async for event, meta in self._graph.astream(
             {"messages": [{"role": "user", "content": user_input}]},
             config=run_config,
             context={"system_prompt": self.instructions, "model": self.model},
@@ -62,21 +62,23 @@ class BaseAgent:
         ):
             # Only forward assistant token chunks that are not tool-related.
             type_ = type(event)
-            tags = metadata.get("tags", []) if isinstance(metadata, dict) else []
+            tags = meta.get("tags", []) if isinstance(meta, dict) else []
             is_quiet = isinstance(tags, list) and ("quiet" in tags)
             if is_quiet:
                 continue
             # Handle different types of messages
-            if type_ in (AIMessage, AIMessageChunk):
+            if type_ == AIMessageChunk:
                 # Accumulate billing and aggregate message
                 aggregate = event if aggregate is None else aggregate + event
             # Ignore intermeddite finish messages
             if "finish_reason" in event.response_metadata:
                 # Got LLM finish reason ignore it
-                logger.error(f"Finish event: {event}, reason: {event.response_metadata['finish_reason']}, Metadata: {metadata}")
+                logger.debug(
+                    f"Finish event: {event}, reason: {event.response_metadata['finish_reason']}, Metadata: {meta}"
+                )
                 pass
             else:
-                logger.debug(f"Event: {event}, Metadata: {metadata}")
+                logger.debug(f"Event: {event}, Metadata: {meta}")
                 yield event
         # Send a final finished message
         # The last event would be finish
@@ -97,9 +99,7 @@ class BaseAgent:
                 else:
                     stream_updater.update(event.content, type_="text")
 
-    async def invoke(
-        self, user_input: str, thread_id: str = str(uuid4()), metadata: dict = None
-    ):
+    async def invoke(self, user_input: str, thread_id: str = str(uuid4()), metadata: dict = None):
         """Run the agent"""
         await self.ainit()
 
@@ -133,9 +133,7 @@ class BaseAgent:
         # Main loop
         while True:
             try:
-                state = self._graph.get_state(
-                    config={"configurable": {"thread_id": thread_id}}
-                )
+                state = self._graph.get_state(config={"configurable": {"thread_id": thread_id}})
                 if state.interrupts:
                     value = self.cli.handle_interrupt(state.interrupts[0])
                     self._graph.invoke(
@@ -152,9 +150,7 @@ class BaseAgent:
                 if user_input.startswith("/"):
                     command = user_input.lower().lstrip("/")
                     if command == "about":
-                        self.cli.display_info(
-                            f"Agent is {self.name}. {self.instructions}"
-                        )
+                        self.cli.display_info(f"Agent is {self.name}. {self.instructions}")
                         continue
                     elif command in {"exit", "quit", "q"}:
                         self.cli.display_info("Goodbye! 👋")
@@ -165,9 +161,7 @@ class BaseAgent:
                         thread_id = str(uuid4())
                         continue
                     elif command == "help":
-                        self.cli.display_info(
-                            "Available commands: /about, /exit, /quit, /q, /reset"
-                        )
+                        self.cli.display_info("Available commands: /about, /exit, /quit, /q, /reset")
                         continue
                     else:
                         self.cli.display_error(f"Unknown command: {command}")
@@ -181,6 +175,7 @@ class BaseAgent:
                 break
             except Exception as e:
                 import traceback
+
                 traceback.print_exc()
                 self.cli.display_error(f"An error occurred: {str(e)}")
                 break

universal_mcp/agents/bigtool2/__init__.py

@@ -1,12 +1,13 @@
 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 universal_mcp.types import ToolConfig, ToolFormat
 
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.llm import load_chat_model
+
 from .graph import build_graph
+from .meta_tools import create_meta_tools
 from .prompts import SYSTEM_PROMPT
 
 
@@ -31,9 +32,7 @@ class BigToolAgent2(BaseAgent):
             # 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_table"]
 
-        logger.info(
-            f"BigToolAgent '{self.name}' initialized with model '{self.model}'."
-        )
+        logger.info(f"BigToolAgent '{self.name}' initialized with model '{self.model}'.")
 
     def _build_system_message(self):
         return SYSTEM_PROMPT.format(
@@ -65,4 +64,4 @@ class BigToolAgent2(BaseAgent):
         return self._graph
 
 
-__all__ = ["BigToolAgent2"]
+__all__ = ["BigToolAgent2", "create_meta_tools"]

universal_mcp/agents/bigtool2/__main__.py

@@ -1,8 +1,8 @@
 import asyncio
 
 from loguru import logger
-
 from universal_mcp.agentr.registry import AgentrRegistry
+
 from universal_mcp.agents.bigtool2 import BigToolAgent2
 from universal_mcp.agents.utils import messages_to_list
 
@@ -15,9 +15,7 @@ async def main():
         registry=AgentrRegistry(),
     )
     await agent.ainit()
-    output = await agent.invoke(
-        user_input="Send an email to manoj@agentr.dev"
-    )
+    output = await agent.invoke(user_input="Send an email to manoj@agentr.dev")
     logger.info(messages_to_list(output["messages"]))
 
 

universal_mcp/agents/bigtool2/agent.py

@@ -1,4 +1,5 @@
 from universal_mcp.agentr.registry import AgentrRegistry
+
 from universal_mcp.agents.bigtool2 import BigToolAgent2
 
 

universal_mcp/agents/bigtool2/graph.py

@@ -1,134 +1,30 @@
 import json
 from datetime import UTC, datetime
 from typing import Literal, cast
-import asyncio
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, ToolMessage
-from langchain_core.tools import BaseTool, tool
+from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
 from langgraph.types import Command
-
-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
 
+from universal_mcp.agents.bigtool2.meta_tools import create_meta_tools
+from universal_mcp.agents.bigtool2.state import State
 
-def build_graph(
-    tool_registry: ToolRegistry, llm: BaseChatModel, system_prompt: str, default_tools: list[BaseTool]
-):
-    @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
-            ]
-            [
-                app["id"] for app in app_ids if app["id"] not in connection_ids
-            ]
-            app_tools = {}
-            for task_query in queries:
-                apps_list = await tool_registry.search_apps(task_query, limit=5)
-                tools_list = []
-                for app in apps_list:
-                    tools_list.extend(await tool_registry.search_tools(task_query, limit=5, app_id=app["id"]))
-                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:
-                        app_tools[app] = []
-                    if len(app_tools[app]) < 5:
-                        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 valid tool ids after loading.
-        Tool ids are of form 'appid__toolid'. Example: 'google_mail__send_email'
-        """
-        correct, incorrect = [], []
-        app_tool_list: dict[str, list[str]] = {}
-
-        # Group tool_ids by app for fewer registry calls
-        app_to_tools: dict[str, list[str]] = {}
-        for tool_id in tool_ids:
-            if "__" not in tool_id:
-                incorrect.append(tool_id)
-                continue
-            app, tool = tool_id.split("__", 1)
-            app_to_tools.setdefault(app, []).append((tool_id, tool))
-
-        # Fetch all apps concurrently
-        async def fetch_tools(app: str):
-            try:
-                tools_dict = await tool_registry.list_tools(app)
-                return app, {tool_unit["name"] for tool_unit in tools_dict}
-            except Exception as e:
-                return app, None
-
-        results = await asyncio.gather(*(fetch_tools(app) for app in app_to_tools))
-
-        # Build map of available tools per app
-        for app, tools in results:
-            if tools is not None:
-                app_tool_list[app] = tools
-
-        # Validate tool_ids
-        for app, tool_entries in app_to_tools.items():
-            available = app_tool_list.get(app)
-            if available is None:
-                incorrect.extend(tool_id for tool_id, _ in tool_entries)
-                continue
-            for tool_id, tool in tool_entries:
-                if tool in available:
-                    correct.append(tool_id)
-                else:
-                    incorrect.append(tool_id)
-
-        return correct
-
-    @tool
-    async def web_search(query: str) -> str:
-        """Search the web for the given query. Returns the search results. Do not use for app-specific searches (for example, reddit or linkedin searches should be done using the app's tools)"""
-        tool = await tool_registry.export_tools(
-            ["exa__search_with_filters"], ToolFormat.LANGCHAIN
-        )
-        response = await tool_registry.call_tool("exa__search_with_filters", {"query": query, "contents": {"summary": True}})
-        return response
 
+def build_graph(tool_registry: ToolRegistry, llm: BaseChatModel, system_prompt: str, default_tools: list[BaseTool]):
+    # Instantiate meta tools (search, load, web_search)
+    search_tools, load_tools, web_search = create_meta_tools(tool_registry)
 
     async def call_model(
         state: State,
     ) -> Command[Literal["select_tools", "call_tools"]]:
         logger.info("Calling model...")
         try:
-            system_message = system_prompt.format(
-                system_time=datetime.now(tz=UTC).isoformat()
-            )
+            system_message = system_prompt.format(system_time=datetime.now(tz=UTC).isoformat())
             messages = [
                 {"role": "system", "content": system_message},
                 *state["messages"],
@@ -161,68 +57,7 @@ def build_graph(
             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.")
-                    selected_tool_ids = await load_tools.ainvoke(tool_call["args"])
-                    tool_msg = ToolMessage(
-                        f"Loaded tools- {selected_tool_ids}", tool_call_id=tool_call["id"]
-                    )
-                    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"] == web_search.name:
-                    logger.info(f"Tool '{tool_call['name']}' is a web search tool. Proceeding to call.")
-                    web_search_result = await web_search.ainvoke(input=tool_call["args"])
-                    tool_msg = ToolMessage(
-                        f"Web search result: {web_search_result}", tool_call_id=tool_call["id"]
-                    )
-                    return Command(goto="call_model", update={"messages": [response, tool_msg]})
-
-                elif "ui_tools" in tool_call["name"]:
-                    logger.info(f"Tool '{tool_call['name']}' is a UI tool. Proceeding to call.")
-                    ui_tool_result = await ui_tools_dict[tool_call["name"]].ainvoke(input=tool_call["args"])
-                    tool_msg = ToolMessage(
-                        f"UI tool result: {ui_tool_result}", tool_call_id=tool_call["id"]
-                    )
-                    return Command(goto="call_model", update={"messages": [response, tool_msg]})
-
-
-                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']}")
+                logger.info(f"Model responded with {len(response.tool_calls)} tool calls.")
                 return Command(goto="call_tools", update={"messages": [response]})
             else:
                 logger.info("Model responded with a message, ending execution.")
@@ -237,7 +72,8 @@ def build_graph(
             tool_call = state["messages"][-1].tool_calls[0]
             searched_tools = await search_tools.ainvoke(input=tool_call["args"])
             tool_msg = ToolMessage(
-                f"Available tool_ids: {searched_tools}. Call load_tools to select the required tools only.", tool_call_id=tool_call["id"]
+                f"Available tool_ids: {searched_tools}. Call load_tools to select the required tools only.",
+                tool_call_id=tool_call["id"],
             )
             return Command(goto="call_model", update={"messages": [tool_msg]})
         except Exception as e:
@@ -249,17 +85,45 @@ def build_graph(
         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.")
+                # Handle special tools internally (no export needed)
+                if tool_call["name"] == search_tools.name:
+                    search_result = await search_tools.ainvoke(input=tool_call["args"])
+                    outputs.append(
+                        ToolMessage(
+                            content=search_result,
+                            name=tool_call["name"],
+                            tool_call_id=tool_call["id"],
+                        )
+                    )
+                    continue
+
+                if tool_call["name"] == load_tools.name:
+                    selected_tool_ids = await load_tools.ainvoke(tool_call["args"])
+                    outputs.append(
+                        ToolMessage(
+                            content=json.dumps(f"Loaded tools- {selected_tool_ids}"),
+                            name=tool_call["name"],
+                            tool_call_id=tool_call["id"],
+                        )
+                    )
+                    recent_tool_ids = selected_tool_ids
+                    continue
+
+                if tool_call["name"] == web_search.name:
+                    web_search_result = await web_search.ainvoke(input=tool_call["args"])
+                    outputs.append(
+                        ToolMessage(
+                            content=json.dumps(f"Web search result: {web_search_result}"),
+                            name=tool_call["name"],
+                            tool_call_id=tool_call["id"],
+                        )
+                    )
+                    continue
+
+                # For other tools: export and call via registry
+                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),

universal_mcp/agents/bigtool2/meta_tools.py

@@ -0,0 +1,120 @@
+import asyncio
+from typing import Any
+
+from langchain_core.tools import BaseTool, tool
+from universal_mcp.logger import logger
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolFormat
+
+
+def create_meta_tools(tool_registry: ToolRegistry) -> list[BaseTool]:
+    @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 = ""
+
+            async def fetch_app_and_connection_metadata():
+                return await asyncio.gather(
+                    tool_registry.list_all_apps(),
+                    tool_registry.list_connected_apps(),
+                )
+
+            app_ids, connections = await fetch_app_and_connection_metadata()
+            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]
+            app_tools: dict[str, list[str]] = {}
+
+            async def find_tools_for_app(task_query: str, app_id: str) -> list[dict[str, Any]]:
+                return await tool_registry.search_tools(task_query, limit=5, app_id=app_id)
+
+            async def find_tools_for_query(task_query: str) -> list[str]:
+                apps_list = await tool_registry.search_apps(task_query, limit=5)
+                per_app_tool_lists = await asyncio.gather(
+                    *(find_tools_for_app(task_query, app_entry["id"]) for app_entry in apps_list)
+                )
+                tools_flat = [tool for sublist in per_app_tool_lists for tool in sublist]
+                return [f"{tool['id']}: {tool['description']}" for tool in tools_flat]
+
+            # Run all queries concurrently
+            query_results = await asyncio.gather(*(find_tools_for_query(q) for q in queries))
+
+            # Aggregate per-app with cap of 5 per app across all queries
+            for tool_desc in [tool for result in query_results for tool in result]:
+                app = tool_desc.split("__")[0]
+                if app not in app_tools:
+                    app_tools[app] = []
+                if len(app_tools[app]) < 5 and tool_desc not in app_tools[app]:
+                    app_tools[app].append(tool_desc)
+            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 valid tool ids after loading.
+        Tool ids are of form 'appid__toolid'. Example: 'google_mail__send_email'
+        """
+        correct, incorrect = [], []
+        app_tool_list: dict[str, list[str]] = {}
+
+        # Group tool_ids by app for fewer registry calls
+        app_to_tools: dict[str, list[str]] = {}
+        for tool_id in tool_ids:
+            if "__" not in tool_id:
+                incorrect.append(tool_id)
+                continue
+            app, tool = tool_id.split("__", 1)
+            app_to_tools.setdefault(app, []).append((tool_id, tool))
+
+        # Fetch all apps concurrently
+        async def fetch_tools(app: str):
+            try:
+                tools_dict = await tool_registry.list_tools(app)
+                return app, {tool_unit["name"] for tool_unit in tools_dict}
+            except Exception:
+                return app, None
+
+        results = await asyncio.gather(*(fetch_tools(app) for app in app_to_tools))
+
+        # Build map of available tools per app
+        for app, tools in results:
+            if tools is not None:
+                app_tool_list[app] = tools
+
+        # Validate tool_ids
+        for app, tool_entries in app_to_tools.items():
+            available = app_tool_list.get(app)
+            if available is None:
+                incorrect.extend(tool_id for tool_id, _ in tool_entries)
+                continue
+            for tool_id, tool in tool_entries:
+                if tool in available:
+                    correct.append(tool_id)
+                else:
+                    incorrect.append(tool_id)
+
+        return correct
+
+    @tool
+    async def web_search(query: str) -> str:
+        """Search the web for the given query. Returns the search results. Do not use for app-specific searches (for example, reddit or linkedin searches should be done using the app's tools)"""
+        await tool_registry.export_tools(["exa__search_with_filters"], ToolFormat.LANGCHAIN)
+        response = await tool_registry.call_tool(
+            "exa__search_with_filters", {"query": query, "contents": {"summary": True}}
+        )
+        return response
+
+    return [search_tools, load_tools, web_search]