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

universal_mcp/agents/__init__.py

@@ -7,6 +7,25 @@ from universal_mcp.agents.planner import PlannerAgent
 from universal_mcp.agents.react import ReactAgent
 from universal_mcp.agents.simple import SimpleAgent
 
+
+def get_agent(agent_name: str):
+    if agent_name == "auto":
+        return AutoAgent
+    elif agent_name == "react":
+        return ReactAgent
+    elif agent_name == "simple":
+        return SimpleAgent
+    elif agent_name == "builder":
+        return BuilderAgent
+    elif agent_name == "planner":
+        return PlannerAgent
+    elif agent_name == "bigtool":
+        return BigToolAgent
+    elif agent_name == "bigtool2":
+        return BigToolAgent2
+    else:
+        raise ValueError(f"Unknown agent: {agent_name}. Possible values: auto, react, simple, builder, planner, bigtool, bigtool2")
+
 __all__ = [
     "BaseAgent",
     "ReactAgent",

universal_mcp/agents/autoagent/__init__.py

@@ -1,8 +1,8 @@
 from langgraph.checkpoint.base import BaseCheckpointSaver
+from universal_mcp.tools.registry import ToolRegistry
 
 from universal_mcp.agents.autoagent.graph import build_graph
 from universal_mcp.agents.base import BaseAgent
-from universal_mcp.tools.registry import ToolRegistry
 
 
 class AutoAgent(BaseAgent):

universal_mcp/agents/autoagent/__main__.py

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

universal_mcp/agents/autoagent/graph.py

@@ -6,13 +6,13 @@ from langchain_core.messages import AIMessage, ToolMessage
 from langchain_core.tools import tool
 from langgraph.graph import END, START, StateGraph
 from langgraph.runtime import Runtime
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolFormat
 
 from universal_mcp.agents.autoagent.context import Context
 from universal_mcp.agents.autoagent.prompts import SYSTEM_PROMPT
 from universal_mcp.agents.autoagent.state import State
 from universal_mcp.agents.llm import load_chat_model
-from universal_mcp.tools.registry import ToolRegistry
-from universal_mcp.types import ToolFormat
 
 
 async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
@@ -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

@@ -1,14 +1,14 @@
 import asyncio
 
 from universal_mcp.agentr.registry import AgentrRegistry
-from universal_mcp.agents.autoagent import build_graph
 from universal_mcp.tools import ToolManager
 
+from universal_mcp.agents.autoagent import build_graph
+
 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 +16,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

@@ -2,15 +2,24 @@
 from typing import cast
 from uuid import uuid4
 
-from langchain_core.messages import AIMessageChunk
+from langchain_core.messages import AIMessage, AIMessageChunk
 from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import StateGraph
 from langgraph.types import Command
+from universal_mcp.logger import logger
 
 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
@@ -24,56 +33,96 @@ class BaseAgent:
             self._graph = await self._build_graph()
             self._initialized = True
 
-    async def _build_graph(self):
+    async def _build_graph(self) -> StateGraph:
         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,
         ):
             # Only forward assistant token chunks that are not tool-related.
             type_ = type(event)
-            if type_ != AIMessageChunk:
-                continue
-            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
+            # Handle different types of messages
+            if type_ in (AIMessage, 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.debug(f"Finish event: {event}, Metadata: {metadata}")
+                logger.debug(f"Finish event: {event}, Metadata: {metadata}")
                 pass
             else:
-                # logger.debug(f"Event: {event}, Metadata: {metadata}")
+                logger.debug(f"Event: {event}, Metadata: {metadata}")
                 yield event
         # Send a final finished message
         # The last event would be finish
         event = cast(AIMessageChunk, event)
+        event.usage_metadata = aggregate.usage_metadata
+        logger.debug(f"Usage metadata: {event.usage_metadata}")
         yield event
 
     async def stream_interactive(self, thread_id: str, user_input: str):
         await self.ainit()
         with self.cli.display_agent_response_streaming(self.name) as stream_updater:
             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())):
+                if isinstance(event.content, list):
+                    thinking_content = "".join([c.get("thinking", "") for c in event.content])
+                    stream_updater.update(thinking_content, type_="thinking")
+                    content = "".join([c.get("text", "") for c in event.content])
+                    stream_updater.update(content, type_="text")
+                else:
+                    stream_updater.update(event.content, type_="text")
+
+    async def invoke(
+        self, user_input: str, thread_id: str = str(uuid4()), metadata: dict = None
+    ):
         """Run the agent"""
         await self.ainit()
-        return await self._graph.ainvoke(
+
+        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,
+        }
+
+        result = 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},
         )
+        return result
 
     async def run_interactive(self, thread_id: str = str(uuid4())):
         """Main application loop"""
@@ -85,10 +134,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,9 +153,11 @@ 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":
+                    elif command in {"exit", "quit", "q"}:
                         self.cli.display_info("Goodbye! 👋")
                         break
                     elif command == "reset":
@@ -110,7 +166,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}")
@@ -124,6 +182,6 @@ 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/bigtool/__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
@@ -19,15 +19,19 @@ class BigToolAgent(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.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}'."
+        )
+
+    def _build_system_message(self):
+        return SYSTEM_PROMPT.format(
+            name=self.name,
+            instructions=self.instructions,
+        )
 
     async def _build_graph(self):
         """Build the bigtool agent graph using the existing create_agent function."""
@@ -36,7 +40,7 @@ class BigToolAgent(BaseAgent):
             graph_builder = build_graph(
                 tool_registry=self.registry,
                 llm=self.llm,
-                tool_selection_llm=self.tool_selection_llm,
+                system_prompt=self._build_system_message(),
             )
 
             compiled_graph = graph_builder.compile(checkpointer=self.memory)

universal_mcp/agents/bigtool/__main__.py

@@ -1,9 +1,10 @@
 import asyncio
 
 from loguru import logger
-
 from universal_mcp.agentr.registry import AgentrRegistry
+
 from universal_mcp.agents.bigtool import BigToolAgent
+from universal_mcp.agents.utils import messages_to_list
 
 
 async def main():
@@ -13,12 +14,10 @@ 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__":
     asyncio.run(main())

universal_mcp/agents/bigtool/graph.py

@@ -1,28 +1,24 @@
 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 universal_mcp.agents.bigtool.state import State
+
 from .prompts import SELECT_TOOL_PROMPT
 
 
 def build_graph(
     tool_registry: ToolRegistry,
     llm: BaseChatModel,
-    tool_selection_llm: BaseChatModel,
+    system_prompt: str,
 ):
     @tool
     async def retrieve_tools(task_query: str) -> list[str]:
@@ -32,29 +28,40 @@ 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):
                 tool_names: list[str]
 
-            model = tool_selection_llm
+            model = 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]
+            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,45 +72,63 @@ 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
+    ) -> 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"]]
+            messages = [
+                {"role": "system", "content": system_prompt},
+                *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
-            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))
+            model_with_tools = llm.bind_tools(
+                [retrieve_tools, *selected_tools], tool_choice="auto"
+            )
+
+
+            response = await model_with_tools.ainvoke(messages)
+            cast(AIMessage, response)
+            logger.debug(f"Response: {response}")
+
 
             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 +141,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
+    ) -> 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 +168,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,9 +196,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)