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

universal_mcp/agents/bigtool/prompts.py

@@ -1,15 +1,15 @@
 """Default prompts used by the agent."""
 
-SYSTEM_PROMPT = """You are a helpful AI assistant.
+SYSTEM_PROMPT = """You are {name}, 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.
+
+{instructions}
 """
 
 SELECT_TOOL_PROMPT = """You are an AI assistant that helps the user perform tasks using various apps (each app has multiple tools).

universal_mcp/agents/bigtool2/__init__.py

@@ -1,9 +1,10 @@
 from langgraph.checkpoint.base import BaseCheckpointSaver
+from universal_mcp.logger import logger
+from universal_mcp.tools.registry import ToolRegistry
 
 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.agents.utils import initialize_ui_tools
 
 from .graph import build_graph
 from .prompts import SYSTEM_PROMPT
@@ -19,15 +20,22 @@ class BigToolAgent2(BaseAgent):
         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)
+        super().__init__(name, instructions, model, memory, **kwargs)
 
         self.registry = registry
         self.llm = load_chat_model(self.model)
         self.recursion_limit = kwargs.get("recursion_limit", 10)
+        self.ui_tools = initialize_ui_tools()
+
+        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(
+            name=self.name,
+            instructions=f"**User Instructions:**\n{self.instructions}",
+        )
 
     async def _build_graph(self):
         """Build the bigtool agent graph using the existing create_agent function."""
@@ -36,6 +44,8 @@ class BigToolAgent2(BaseAgent):
             graph_builder = build_graph(
                 tool_registry=self.registry,
                 llm=self.llm,
+                system_prompt=self._build_system_message(),
+                ui_tools=self.ui_tools,
             )
 
             compiled_graph = graph_builder.compile(checkpointer=self.memory)

universal_mcp/agents/bigtool2/__main__.py

@@ -1,9 +1,10 @@
 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
 
 
 async def main():
@@ -13,11 +14,11 @@ async def main():
         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)
+    await agent.ainit()
+    output = await agent.invoke(
+        user_input="Send an email to manoj@agentr.dev"
+    )
+    logger.info(messages_to_list(output["messages"]))
 
 
 if __name__ == "__main__":

universal_mcp/agents/bigtool2/agent.py

@@ -1,6 +1,8 @@
-from universal_mcp.agents.bigtool2 import BigToolAgent2
 from universal_mcp.agentr.registry import AgentrRegistry
 
+from universal_mcp.agents.bigtool2 import BigToolAgent2
+
+
 async def agent():
     agent_object = await BigToolAgent2(
         name="BigTool Agent 2",
@@ -8,4 +10,4 @@ async def agent():
         model="anthropic/claude-4-sonnet-20250514",
         registry=AgentrRegistry(),
     )._build_graph()
-    return agent_object
+    return agent_object

universal_mcp/agents/bigtool2/graph.py

@@ -1,26 +1,21 @@
 import json
 from datetime import UTC, datetime
-from typing import Literal, TypedDict, cast
+from typing import Literal, 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
 
+from universal_mcp.agents.bigtool2.state import State
 
 
 def build_graph(
-    tool_registry: ToolRegistry,
-    llm: BaseChatModel
+    tool_registry: ToolRegistry, llm: BaseChatModel, system_prompt: str, ui_tools: list
 ):
     @tool
     async def search_tools(queries: list[str]) -> str:
@@ -33,12 +28,18 @@ def build_graph(
             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]
+            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:
                 tools_list = await tool_registry.search_tools(task_query, limit=40)
-                tool_candidates = [f"{tool['id']}: {tool['description']}" for tool in tools_list]
+                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:
@@ -49,65 +50,95 @@ def build_graph(
                         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"
+                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"]]:
+    async def call_model(
+        state: State,
+    ) -> 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 = 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)
+                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")
+            model_with_tools = model.bind_tools(
+                [search_tools, load_tools, *selected_tools, *ui_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.")
+                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.")
                     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"])
+                    tool_msg = ToolMessage(
+                        "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})
+                    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)
+                        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]})
+                        return Command(
+                            goto="call_tools", update={"messages": [response]}
+                        )
                     except Exception as e:
-                        logger.error(f"Unexpected tool call: {tool_call['name']}. Error: {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
@@ -120,12 +151,14 @@ 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) -> 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}")
@@ -136,10 +169,16 @@ 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']}")
+            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"])
+                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(
@@ -158,9 +197,12 @@ def build_graph(
                         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)
+    builder = StateGraph(State)
 
     builder.add_node(call_model)
     builder.add_node(select_tools)

universal_mcp/agents/bigtool2/prompts.py

@@ -8,5 +8,5 @@ SYSTEM_PROMPT = """You are a helpful AI assistant.
 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}
+{instructions}
 """

universal_mcp/agents/bigtoolcache/__init__.py

@@ -1,9 +1,9 @@
 from langgraph.checkpoint.base import BaseCheckpointSaver
+from universal_mcp.logger import logger
+from universal_mcp.tools.registry import ToolRegistry
 
 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
@@ -27,7 +27,9 @@ class BigToolAgentCache(BaseAgent):
         self.llm = load_chat_model(self.model)
         self.recursion_limit = kwargs.get("recursion_limit", 10)
 
-        logger.info(f"BigToolAgent '{self.name}' initialized with model '{self.model}'.")
+        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."""
@@ -42,7 +44,9 @@ class BigToolAgentCache(BaseAgent):
             logger.info("Graph built and compiled successfully.")
             return compiled_graph
         except Exception as e:
-            logger.error(f"Error building graph for BigToolAgentCache '{self.name}': {e}")
+            logger.error(
+                f"Error building graph for BigToolAgentCache '{self.name}': {e}"
+            )
             raise
 
     @property

universal_mcp/agents/bigtoolcache/__main__.py

@@ -1,8 +1,8 @@
 import asyncio
 
 from loguru import logger
-
 from universal_mcp.agentr.registry import AgentrRegistry
+
 from universal_mcp.agents.bigtoolcache import BigToolAgentCache
 
 

universal_mcp/agents/bigtoolcache/agent.py

@@ -1,11 +1,13 @@
-from universal_mcp.agents.bigtoolcache import BigToolAgentCache
 from universal_mcp.agentr.registry import AgentrRegistry
 
+from universal_mcp.agents.bigtoolcache import BigToolAgentCache
+
+
 async def agent():
     agent_object = await BigToolAgentCache(
-        name="BigTool Agent 2",
+        name="BigTool Agent Cache version",
         instructions="You are a helpful assistant that can use various tools to complete tasks.",
         model="anthropic/claude-4-sonnet-20250514",
         registry=AgentrRegistry(),
     )._build_graph()
-    return agent_object
+    return agent_object

universal_mcp/agents/bigtoolcache/context.py

@@ -30,4 +30,3 @@ class Context:
             "This is to prevent infinite recursion."
         },
     )
-