universal-mcp-agents 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

universal_mcp/agents/bigtoolcache/graph.py

@@ -1,6 +1,6 @@
 import json
 from datetime import UTC, datetime
-from typing import Any, Literal, TypedDict, cast
+from typing import Literal, TypedDict, cast
 
 from langchain_anthropic import ChatAnthropic
 from langchain_core.language_models import BaseChatModel
@@ -15,94 +15,107 @@ from universal_mcp.agents.bigtoolcache.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.bigtoolcache.prompts import TOOLS_LIST
 
 
 
-def build_graph(
-    tool_registry: ToolRegistry,
-    llm: BaseChatModel
-):
+class ToolSelectionOutput(TypedDict):
+    connected_tool_ids: list[str]
+    unconnected_tool_ids: list[str]
+
+
+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)
-                for tool in tools_list:
-                    app = tool["id"].split("__")[0]
-                    if app not in app_tools:
-                        if len(app_tools.keys()) >= 10:
-                            break
-                        app_tools[app] = {}
-                    if len(app_tools[app]) < 3:
-                        if tool["id"] not in app_tools[app]:
-                            app_tools[app][tool["id"]] = tool["description"]
-            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}: {app_tools[app][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)
-    
+        messages = [
+            {
+                "role": "system",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": TOOLS_LIST,
+                        "cache_control": {"type": "ephemeral", "ttl": "1h"},
+                    }
+                ],
+            },
+            {"role": "user", "content": str(queries)},
+        ]
+        response = llm.with_structured_output(ToolSelectionOutput).invoke(messages)
+        response_text = f"Connected tools: {response['connected_tool_ids']}\nUnconnected tools: {response['unconnected_tool_ids']}"
+        return response_text
+
     @tool
     async def load_tools(tool_ids: list[str]) -> list[dict[str, Any]]:
-        """Load the tools for the given tool ids. Returns the tool name, description, parameters schema, and output schema."""
+        """Load the tools for the given tool ids. Returns the tool name, description, parameters schema, and output schema. A tool id is made up using the app_id and the tool_name, attached by double underscore (__). e.g. google_mail__send_email"""
         temp_manager = tool_registry.tool_manager
         temp_manager.clear_tools()
         await tool_registry.export_tools(tool_ids, format=ToolFormat.NATIVE)
         tool_details = []
         for tool_id in tool_ids:
             tool = temp_manager.get_tool(tool_id)
-            tool_details.append({
-                "name": tool.name,
-                "description": tool.description,
-                "parameters_schema": tool.parameters,
-                "output_schema": tool.output_schema,
-            })
+            tool_details.append(
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters_schema": tool.parameters,
+                    "output_schema": tool.output_schema,
+                }
+            )
         return tool_details
-    
+
     @tool
     async def call_tool(tool_id: str, tool_args: dict[str, Any]) -> Any:
         """Call the tool with the given id and arguments."""
         return await tool_registry.call_tool(tool_id, tool_args)
 
-
-    async def call_model(state: State, runtime: Runtime[Context]) -> Command[Literal["select_tools", "call_tools"]]:
+    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"]]
+            system_message = runtime.context.system_prompt.format(
+                system_time=datetime.now(tz=UTC).isoformat()
+            )
+            messages = [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": system_message,
+                            "cache_control": {"type": "ephemeral", "ttl": "1h"},
+                        }
+                    ],
+                },
+                *state["messages"],
+            ]
 
             model = llm
 
             if isinstance(model, ChatAnthropic):
                 model_with_tools = model.bind_tools(
-                    [search_tools, load_tools, call_tool], tool_choice="auto", cache_control={"type": "ephemeral"}
+                    [search_tools, load_tools, call_tool],
+                    tool_choice="auto",
+                    cache_control={"type": "ephemeral", "ttl": "1h"},
                 )
             else:
-                model_with_tools = model.bind_tools([search_tools, load_tools, call_tool], tool_choice="auto")
+                model_with_tools = model.bind_tools(
+                    [search_tools, load_tools, call_tool], 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.")
+                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')")
+                    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.")
@@ -110,10 +123,18 @@ def build_graph(
                 elif tool_call["name"] == load_tools.name:
                     logger.info("Model requested to load tools.")
                     tool_details = await load_tools.ainvoke(input=tool_call["args"])
-                    tool_msg = ToolMessage(f"Loaded tools. {tool_details}", tool_call_id=tool_call["id"])
+                    tool_msg = ToolMessage(
+                        f"Loaded tools. {tool_details}", 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})
+                    return Command(
+                        goto="call_model",
+                        update={
+                            "messages": [response, tool_msg],
+                            "selected_tool_ids": selected_tool_ids,
+                        },
+                    )
                 elif tool_call["name"] == call_tool.name:
                     logger.info("Model requested to call tool.")
                     return Command(goto="call_tools", update={"messages": [response]})
@@ -125,12 +146,16 @@ def build_graph(
             logger.error(f"Error in call_model: {e}")
             raise
 
-    async def select_tools(state: State, runtime: Runtime[Context]) -> Command[Literal["call_model"]]:
+    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"])
+            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}")
@@ -146,7 +171,9 @@ def build_graph(
         logger.info(f"Executing tool: {tool_id} with args: {tool_args}")
         try:
             await tool_registry.export_tools([tool_id], ToolFormat.LANGCHAIN)
-            tool_result = await call_tool.ainvoke(input={"tool_id": tool_id, "tool_args": tool_args})
+            tool_result = await call_tool.ainvoke(
+                input={"tool_id": tool_id, "tool_args": tool_args}
+            )
             logger.info(f"Tool '{tool_id}' executed successfully.")
             outputs.append(
                 ToolMessage(
@@ -154,18 +181,21 @@ def build_graph(
                     name=tool_id,
                     tool_call_id=tool_call["id"],
                 )
-            )
-            recent_tool_ids.append(tool_id)
-        except Exception as e:
-            logger.error(f"Error executing tool '{tool_id}': {e}")
-            outputs.append(
-                ToolMessage(
-                    content=json.dumps("Error: " + str(e)),
-                    name=tool_id,
-                    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})
+        return Command(
+            goto="call_model",
+            update={"messages": outputs, "selected_tool_ids": recent_tool_ids},
+        )
 
     builder = StateGraph(State, context_schema=Context)
 

universal_mcp/agents/bigtoolcache/prompts.py

@@ -1,5 +1,24 @@
 """Default prompts used by the agent."""
 
+from pathlib import Path
+
+
+def load_tools_from_file():
+    """Load tools from the generated text file."""
+    # Get the directory where this file is located
+    current_dir = Path(__file__).parent
+
+    tools_file = current_dir / "tools_important.txt"
+    if not tools_file.exists():
+        tools_file = current_dir / "tools_all.txt"
+
+    if tools_file.exists():
+        with open(tools_file, encoding="utf-8") as f:
+            return f.read()
+    else:
+        return "No tools file found. Please run tool_retrieve.py to generate the tools list."
+
+
 SYSTEM_PROMPT = """You are a helpful AI assistant.
 
 **Core Directives:**
@@ -11,3 +30,12 @@ SYSTEM_PROMPT = """You are a helpful AI assistant.
 
 System time: {system_time}
 """
+
+
+TOOLS_LIST = f""" This is the list of all the tools available to you:
+{load_tools_from_file()}
+
+You will be provided a list of queries (which may be similar or different from each other). Your job is to select the relavent tools for the user's request. 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 return both apps tools. If the query specifically asks you to use an app that is not connected, return the tools for that app, they can still be connected by the user.
+
+You have to return the tool_ids by constructing the tool_id from the app_id and the tool_name, attached by double underscore (__). e.g. google_mail__send_email
+"""