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

universal_mcp/agents/autoagent/graph.py

@@ -22,7 +22,9 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
         tools_list = []
         if app_ids is not None:
             for app_id in app_ids:
-                tools_list.extend(await tool_registry.search_tools(query, limit=10, app_id=app_id))
+                tools_list.extend(
+                    await tool_registry.search_tools(query, limit=10, app_id=app_id)
+                )
         else:
             tools_list = await tool_registry.search_tools(query, limit=10)
         tools_list = [f"{tool['id']}: {tool['description']}" for tool in tools_list]
@@ -48,21 +50,33 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
         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]
+        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]
             )
 
-        system_prompt = system_prompt.format(system_time=datetime.now(tz=UTC).isoformat(), app_ids=app_id_descriptions)
+        system_prompt = system_prompt.format(
+            system_time=datetime.now(tz=UTC).isoformat(), app_ids=app_id_descriptions
+        )
 
-        messages = [{"role": "system", "content": system_prompt + "\n" + instructions}, *state["messages"]]
+        messages = [
+            {"role": "system", "content": system_prompt + "\n" + instructions},
+            *state["messages"],
+        ]
         model = load_chat_model(runtime.context.model)
-        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")
+        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)
         response = cast(AIMessage, response_raw)
         return {"messages": [response]}
@@ -102,7 +116,8 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
                 tools = await search_tools.ainvoke(tool_call["args"])
                 outputs.append(
                     ToolMessage(
-                        content=json.dumps(tools) + "\n\nUse the load_tools tool to load the tools you want to use.",
+                        content=json.dumps(tools)
+                        + "\n\nUse the load_tools tool to load the tools you want to use.",
                         name=tool_call["name"],
                         tool_call_id=tool_call["id"],
                     )
@@ -119,9 +134,13 @@ async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
                     )
                 )
             else:
-                await tool_registry.export_tools([tool_call["name"]], ToolFormat.LANGCHAIN)
+                await tool_registry.export_tools(
+                    [tool_call["name"]], ToolFormat.LANGCHAIN
+                )
                 try:
-                    tool_result = await tool_registry.call_tool(tool_call["name"], tool_call["args"])
+                    tool_result = await tool_registry.call_tool(
+                        tool_call["name"], tool_call["args"]
+                    )
                     outputs.append(
                         ToolMessage(
                             content=json.dumps(tool_result),

universal_mcp/agents/autoagent/studio.py

@@ -8,7 +8,6 @@ tool_registry = AgentrRegistry()
 tool_manager = ToolManager()
 
 
-
 async def main():
     instructions = """
     You are a helpful assistant that can use tools to help the user. If a task requires multiple steps, you should perform separate different searches for different actions. Prefer completing one action before searching for another.
@@ -16,10 +15,5 @@ async def main():
     graph = await build_graph(tool_registry, instructions=instructions)
     return graph
 
-graph = asyncio.run(main())
-
-
-
-
-
 
+graph = asyncio.run(main())

universal_mcp/agents/base.py

@@ -10,7 +10,14 @@ from .utils import RichCLI
 
 
 class BaseAgent:
-    def __init__(self, name: str, instructions: str, model: str, memory: BaseCheckpointSaver | None = None, **kwargs):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        memory: BaseCheckpointSaver | None = None,
+        **kwargs,
+    ):
         self.name = name
         self.instructions = instructions
         self.model = model
@@ -27,12 +34,26 @@ class BaseAgent:
     async def _build_graph(self):
         raise NotImplementedError("Subclasses must implement this method")
 
-    async def stream(self, thread_id: str, user_input: str):
+    async def stream(self, thread_id: str, user_input: str, metadata: dict = None):
         await self.ainit()
         aggregate = None
+
+        run_metadata = {
+            "agent_name": self.name,
+            "is_background_run": False,  # Default to False
+        }
+
+        if metadata:
+            run_metadata.update(metadata)
+
+        run_config = {
+            "configurable": {"thread_id": thread_id},
+            "metadata": run_metadata,
+        }
+
         async for event, metadata in self._graph.astream(
             {"messages": [{"role": "user", "content": user_input}]},
-            config={"configurable": {"thread_id": thread_id}},
+            config=run_config,
             context={"system_prompt": self.instructions, "model": self.model},
             stream_mode="messages",
             stream_usage=True,
@@ -66,12 +87,28 @@ class BaseAgent:
             async for event in self.stream(thread_id, user_input):
                 stream_updater.update(event.content)
 
-    async def invoke(self, user_input: str, thread_id: str = str(uuid4())):
+    async def invoke(
+        self, user_input: str, thread_id: str = str(uuid4()), metadata: dict = None
+    ):
         """Run the agent"""
         await self.ainit()
+
+        run_metadata = {
+            "agent_name": self.name,
+            "is_background_run": False,  # Default to False
+        }
+
+        if metadata:
+            run_metadata.update(metadata)
+
+        run_config = {
+            "configurable": {"thread_id": thread_id},
+            "metadata": run_metadata,
+        }
+
         return await self._graph.ainvoke(
             {"messages": [{"role": "user", "content": user_input}]},
-            config={"configurable": {"thread_id": thread_id}},
+            config=run_config,
             context={"system_prompt": self.instructions, "model": self.model},
         )
 
@@ -85,10 +122,15 @@ 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(Command(resume=value), config={"configurable": {"thread_id": thread_id}})
+                    self._graph.invoke(
+                        Command(resume=value),
+                        config={"configurable": {"thread_id": thread_id}},
+                    )
                     continue
 
                 user_input = self.cli.get_user_input()
@@ -99,7 +141,9 @@ 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 == "exit" or command == "quit" or command == "q":
                         self.cli.display_info("Goodbye! 👋")
@@ -110,7 +154,9 @@ 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}")

universal_mcp/agents/bigtool/__init__.py

@@ -27,7 +27,9 @@ class BigToolAgent(BaseAgent):
         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}'.")
+        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."""

universal_mcp/agents/bigtool/graph.py

@@ -32,7 +32,9 @@ def build_graph(
         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]
+            tool_candidates = [
+                f"{tool['id']}: {tool['description']}" for tool in tools_list
+            ]
             logger.info(f"Found {len(tool_candidates)} candidate tools.")
 
             class ToolSelectionOutput(TypedDict):
@@ -42,19 +44,28 @@ 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]
-            app_id_descriptions = "These are the apps connected to the user's account:\n" + "\n".join(
-                [f"{app}" for app in connected_apps]
+            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(
+            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
+                    app_ids=app_id_descriptions,
+                    tool_candidates="\n - ".join(tool_candidates),
+                    task=task_query,
                 )
             )
 
@@ -65,15 +76,24 @@ def build_graph(
             logger.error(f"Error retrieving tools: {e}")
             return []
 
-    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": 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 = []
@@ -81,29 +101,43 @@ def build_graph(
             model = llm
             if isinstance(model, ChatAnthropic):
                 model_with_tools = model.bind_tools(
-                    [retrieve_tools, *selected_tools], tool_choice="auto", cache_control={"type": "ephemeral"}
+                    [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")
+                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.")
+                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"] == 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)
+                        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
@@ -116,14 +150,24 @@ 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]
             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"])
+            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})
+            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
@@ -133,10 +177,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(
@@ -155,7 +205,10 @@ 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)
 

universal_mcp/agents/bigtool2/__init__.py

@@ -27,7 +27,9 @@ class BigToolAgent2(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."""

universal_mcp/agents/bigtool2/agent.py

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

universal_mcp/agents/bigtool2/context.py

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

universal_mcp/agents/bigtool2/graph.py

@@ -17,11 +17,7 @@ from universal_mcp.tools.registry import ToolRegistry
 from universal_mcp.types import ToolFormat
 
 
-
-def build_graph(
-    tool_registry: ToolRegistry,
-    llm: BaseChatModel
-):
+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
@@ -33,12 +29,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
+            ]
+            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]
+                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 +51,94 @@ 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, 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": 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], 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(
+                        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})
+                    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,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}")
@@ -136,10 +171,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,7 +199,10 @@ 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)