universal-mcp 0.1.24rc14__py3-none-any.whl → 0.1.24rc19__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. universal_mcp/agentr/client.py +11 -0
  2. universal_mcp/agentr/registry.py +40 -5
  3. universal_mcp/applications/application.py +0 -2
  4. universal_mcp/applications/sample/app.py +79 -20
  5. universal_mcp/applications/utils.py +52 -0
  6. universal_mcp/servers/server.py +4 -3
  7. universal_mcp/tools/manager.py +0 -3
  8. universal_mcp/types.py +1 -21
  9. universal_mcp/utils/prompts.py +0 -2
  10. universal_mcp/utils/testing.py +1 -1
  11. {universal_mcp-0.1.24rc14.dist-info → universal_mcp-0.1.24rc19.dist-info}/METADATA +1 -1
  12. universal_mcp-0.1.24rc19.dist-info/RECORD +54 -0
  13. universal_mcp/__init__.py +0 -0
  14. universal_mcp/agents/__init__.py +0 -10
  15. universal_mcp/agents/autoagent/__init__.py +0 -30
  16. universal_mcp/agents/autoagent/__main__.py +0 -25
  17. universal_mcp/agents/autoagent/context.py +0 -26
  18. universal_mcp/agents/autoagent/graph.py +0 -151
  19. universal_mcp/agents/autoagent/prompts.py +0 -9
  20. universal_mcp/agents/autoagent/state.py +0 -27
  21. universal_mcp/agents/autoagent/studio.py +0 -25
  22. universal_mcp/agents/autoagent/utils.py +0 -13
  23. universal_mcp/agents/base.py +0 -129
  24. universal_mcp/agents/bigtool/__init__.py +0 -54
  25. universal_mcp/agents/bigtool/__main__.py +0 -24
  26. universal_mcp/agents/bigtool/context.py +0 -24
  27. universal_mcp/agents/bigtool/graph.py +0 -166
  28. universal_mcp/agents/bigtool/prompts.py +0 -31
  29. universal_mcp/agents/bigtool/state.py +0 -27
  30. universal_mcp/agents/bigtool2/__init__.py +0 -53
  31. universal_mcp/agents/bigtool2/__main__.py +0 -24
  32. universal_mcp/agents/bigtool2/agent.py +0 -11
  33. universal_mcp/agents/bigtool2/context.py +0 -33
  34. universal_mcp/agents/bigtool2/graph.py +0 -169
  35. universal_mcp/agents/bigtool2/prompts.py +0 -12
  36. universal_mcp/agents/bigtool2/state.py +0 -27
  37. universal_mcp/agents/builder.py +0 -80
  38. universal_mcp/agents/cli.py +0 -27
  39. universal_mcp/agents/codeact/__init__.py +0 -243
  40. universal_mcp/agents/codeact/sandbox.py +0 -27
  41. universal_mcp/agents/codeact/test.py +0 -15
  42. universal_mcp/agents/codeact/utils.py +0 -61
  43. universal_mcp/agents/hil.py +0 -104
  44. universal_mcp/agents/llm.py +0 -45
  45. universal_mcp/agents/planner/__init__.py +0 -37
  46. universal_mcp/agents/planner/__main__.py +0 -24
  47. universal_mcp/agents/planner/graph.py +0 -82
  48. universal_mcp/agents/planner/prompts.py +0 -1
  49. universal_mcp/agents/planner/state.py +0 -12
  50. universal_mcp/agents/react.py +0 -84
  51. universal_mcp/agents/shared/agent_node.py +0 -34
  52. universal_mcp/agents/shared/tool_node.py +0 -235
  53. universal_mcp/agents/simple.py +0 -40
  54. universal_mcp/agents/tools.py +0 -35
  55. universal_mcp/agents/utils.py +0 -111
  56. universal_mcp/analytics.py +0 -111
  57. universal_mcp/applications/__init__.py +0 -70
  58. universal_mcp/utils/common.py +0 -278
  59. universal_mcp-0.1.24rc14.dist-info/RECORD +0 -99
  60. {universal_mcp-0.1.24rc14.dist-info → universal_mcp-0.1.24rc19.dist-info}/WHEEL +0 -0
  61. {universal_mcp-0.1.24rc14.dist-info → universal_mcp-0.1.24rc19.dist-info}/entry_points.txt +0 -0
  62. {universal_mcp-0.1.24rc14.dist-info → universal_mcp-0.1.24rc19.dist-info}/licenses/LICENSE +0 -0
@@ -1,151 +0,0 @@
1
- import json
2
- from datetime import UTC, datetime
3
- from typing import cast
4
-
5
- from langchain_core.messages import AIMessage, ToolMessage
6
- from langchain_core.tools import tool
7
- from langgraph.graph import END, START, StateGraph
8
- from langgraph.runtime import Runtime
9
-
10
- from universal_mcp.agents.autoagent.context import Context
11
- from universal_mcp.agents.autoagent.prompts import SYSTEM_PROMPT
12
- from universal_mcp.agents.autoagent.state import State
13
- from universal_mcp.agents.llm import load_chat_model
14
- from universal_mcp.tools.registry import ToolRegistry
15
- from universal_mcp.types import ToolFormat
16
-
17
-
18
- async def build_graph(tool_registry: ToolRegistry, instructions: str = ""):
19
- @tool()
20
- async def search_tools(query: str, app_ids: list[str] | None = None) -> list[str]:
21
- """Retrieve tools using a search query and a list of app ids. Use multiple times if you require tools for different queries."""
22
- tools_list = []
23
- if app_ids is not None:
24
- for app_id in app_ids:
25
- tools_list.extend(await tool_registry.search_tools(query, limit=10, app_id=app_id))
26
- else:
27
- tools_list = await tool_registry.search_tools(query, limit=10)
28
- tools_list = [f"{tool['id']}: {tool['description']}" for tool in tools_list]
29
- return tools_list
30
-
31
- @tool()
32
- async def ask_user(question: str) -> str:
33
- """Ask the user a question. Use this tool to ask the user for any missing information for performing a task, or when you have multiple apps to choose from for performing a task."""
34
- full_question = question
35
- return f"ASKING_USER: {full_question}"
36
-
37
- @tool()
38
- async def load_tools(tools: list[str]) -> list[str]:
39
- """Choose the tools you want to use by passing their tool ids. Loads the tools for the chosen tools and returns the tool ids."""
40
- return tools
41
-
42
- async def call_model(
43
- state: State,
44
- runtime: Runtime[Context],
45
- ):
46
- system_prompt = SYSTEM_PROMPT
47
- app_ids = await tool_registry.list_all_apps()
48
- connections = await tool_registry.list_connected_apps()
49
- connection_ids = set([connection["app_id"] for connection in connections])
50
- connected_apps = [app["id"] for app in app_ids if app["id"] in connection_ids]
51
- unconnected_apps = [app["id"] for app in app_ids if app["id"] not in connection_ids]
52
- app_id_descriptions = "These are the apps connected to the user's account:\n" + "\n".join(
53
- [f"{app}" for app in connected_apps]
54
- )
55
- if unconnected_apps:
56
- app_id_descriptions += "\n\nOther (not connected) apps: " + "\n".join(
57
- [f"{app}" for app in unconnected_apps]
58
- )
59
-
60
- system_prompt = system_prompt.format(system_time=datetime.now(tz=UTC).isoformat(), app_ids=app_id_descriptions)
61
-
62
- messages = [{"role": "system", "content": system_prompt + "\n" + instructions}, *state["messages"]]
63
- model = load_chat_model(runtime.context.model)
64
- loaded_tools = await tool_registry.export_tools(tools=state["selected_tool_ids"], format=ToolFormat.LANGCHAIN)
65
- model_with_tools = model.bind_tools([search_tools, ask_user, load_tools, *loaded_tools], tool_choice="auto")
66
- response_raw = model_with_tools.invoke(messages)
67
- response = cast(AIMessage, response_raw)
68
- return {"messages": [response]}
69
-
70
- # Define the conditional edge that determines whether to continue or not
71
- def should_continue(state: State):
72
- messages = state["messages"]
73
- last_message = messages[-1]
74
- # If there is no function call, then we finish
75
- if not last_message.tool_calls:
76
- return END
77
- else:
78
- return "tools"
79
-
80
- def tool_router(state: State):
81
- last_message = state["messages"][-1]
82
- if isinstance(last_message, ToolMessage) and last_message.name == ask_user.name:
83
- return END
84
- else:
85
- return "agent"
86
-
87
- async def tool_node(state: State):
88
- outputs = []
89
- tool_ids = state["selected_tool_ids"]
90
- for tool_call in state["messages"][-1].tool_calls:
91
- if tool_call["name"] == ask_user.name:
92
- outputs.append(
93
- ToolMessage(
94
- content=json.dumps(
95
- "The user has been asked the question, and the run will wait for the user's response."
96
- ),
97
- name=tool_call["name"],
98
- tool_call_id=tool_call["id"],
99
- )
100
- )
101
- elif tool_call["name"] == search_tools.name:
102
- tools = await search_tools.ainvoke(tool_call["args"])
103
- outputs.append(
104
- ToolMessage(
105
- content=json.dumps(tools) + "\n\nUse the load_tools tool to load the tools you want to use.",
106
- name=tool_call["name"],
107
- tool_call_id=tool_call["id"],
108
- )
109
- )
110
-
111
- elif tool_call["name"] == load_tools.name:
112
- tool_ids = await load_tools.ainvoke(tool_call["args"])
113
-
114
- outputs.append(
115
- ToolMessage(
116
- content=json.dumps(tool_ids),
117
- name=tool_call["name"],
118
- tool_call_id=tool_call["id"],
119
- )
120
- )
121
- else:
122
- await tool_registry.export_tools([tool_call["name"]], ToolFormat.LANGCHAIN)
123
- try:
124
- tool_result = await tool_registry.call_tool(tool_call["name"], tool_call["args"])
125
- outputs.append(
126
- ToolMessage(
127
- content=json.dumps(tool_result),
128
- name=tool_call["name"],
129
- tool_call_id=tool_call["id"],
130
- )
131
- )
132
- except Exception as e:
133
- outputs.append(
134
- ToolMessage(
135
- content=json.dumps("Error: " + str(e)),
136
- name=tool_call["name"],
137
- tool_call_id=tool_call["id"],
138
- )
139
- )
140
- return {"messages": outputs, "selected_tool_ids": tool_ids}
141
-
142
- builder = StateGraph(State, context_schema=Context)
143
-
144
- builder.add_node("agent", call_model)
145
- builder.add_node("tools", tool_node)
146
-
147
- builder.add_edge(START, "agent")
148
- builder.add_conditional_edges("agent", should_continue)
149
- builder.add_conditional_edges("tools", tool_router)
150
-
151
- return builder
@@ -1,9 +0,0 @@
1
- """Default prompts used by the agent."""
2
-
3
- SYSTEM_PROMPT = """You are a helpful AI assistant. When you lack tools for any task you should use the `search_tools` function to unlock relevant tools. Whenever you need to ask the user for any information, or choose between multiple different applications, you can ask the user using the `ask_user` function.
4
-
5
- System time: {system_time}
6
- These are the list of apps available to you:
7
- {app_ids}
8
- Note that when multiple apps seem relevant for a task, you MUST ask the user to choose the app. 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. You will be provided a link for connection that you should pass on to the user.
9
- """
@@ -1,27 +0,0 @@
1
- from typing import Annotated
2
-
3
- from langgraph.prebuilt.chat_agent_executor import AgentState
4
-
5
-
6
- def _enqueue(left: list, right: list) -> list:
7
- """Treat left as a FIFO queue, append new items from right (preserve order),
8
- keep items unique, and cap total size to 20 (drop oldest items)."""
9
- max_size = 30
10
- preferred_size = 20
11
- if len(right) > preferred_size:
12
- preferred_size = min(max_size, len(right))
13
- queue = list(left or [])
14
-
15
- for item in right[:preferred_size] or []:
16
- if item in queue:
17
- queue.remove(item)
18
- queue.append(item)
19
-
20
- if len(queue) > preferred_size:
21
- queue = queue[-preferred_size:]
22
-
23
- return queue
24
-
25
-
26
- class State(AgentState):
27
- selected_tool_ids: Annotated[list[str], _enqueue]
@@ -1,25 +0,0 @@
1
- import asyncio
2
-
3
- from universal_mcp.agentr.registry import AgentrRegistry
4
- from universal_mcp.agents.autoagent import build_graph
5
- from universal_mcp.tools import ToolManager
6
-
7
- tool_registry = AgentrRegistry()
8
- tool_manager = ToolManager()
9
-
10
-
11
-
12
- async def main():
13
- instructions = """
14
- 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.
15
- """
16
- graph = await build_graph(tool_registry, instructions=instructions)
17
- return graph
18
-
19
- graph = asyncio.run(main())
20
-
21
-
22
-
23
-
24
-
25
-
@@ -1,13 +0,0 @@
1
- from langchain_core.messages import BaseMessage
2
-
3
-
4
- def get_message_text(msg: BaseMessage) -> str:
5
- """Get the text content of a message."""
6
- content = msg.content
7
- if isinstance(content, str):
8
- return content
9
- elif isinstance(content, dict):
10
- return content.get("text", "")
11
- else:
12
- txts = [c if isinstance(c, str) else (c.get("text") or "") for c in content]
13
- return "".join(txts).strip()
@@ -1,129 +0,0 @@
1
- # agents/base.py
2
- from typing import cast
3
- from uuid import uuid4
4
-
5
- from langchain_core.messages import AIMessageChunk
6
- from langgraph.checkpoint.base import BaseCheckpointSaver
7
- from langgraph.types import Command
8
-
9
- from .utils import RichCLI
10
-
11
-
12
- class BaseAgent:
13
- def __init__(self, name: str, instructions: str, model: str, memory: BaseCheckpointSaver | None = None, **kwargs):
14
- self.name = name
15
- self.instructions = instructions
16
- self.model = model
17
- self.memory = memory
18
- self._graph = None
19
- self._initialized = False
20
- self.cli = RichCLI()
21
-
22
- async def ainit(self):
23
- if not self._initialized:
24
- self._graph = await self._build_graph()
25
- self._initialized = True
26
-
27
- async def _build_graph(self):
28
- raise NotImplementedError("Subclasses must implement this method")
29
-
30
- async def stream(self, thread_id: str, user_input: str):
31
- await self.ainit()
32
- aggregate = None
33
- async for event, metadata in self._graph.astream(
34
- {"messages": [{"role": "user", "content": user_input}]},
35
- config={"configurable": {"thread_id": thread_id}},
36
- context={"system_prompt": self.instructions, "model": self.model},
37
- stream_mode="messages",
38
- stream_usage=True,
39
- ):
40
- # Only forward assistant token chunks that are not tool-related.
41
- type_ = type(event)
42
- if type_ != AIMessageChunk:
43
- continue
44
- event = cast(AIMessageChunk, event)
45
- aggregate = event if aggregate is None else aggregate + event
46
- tags = metadata.get("tags", []) if isinstance(metadata, dict) else []
47
- is_quiet = isinstance(tags, list) and ("quiet" in tags)
48
-
49
- if is_quiet:
50
- continue
51
- if "finish_reason" in event.response_metadata:
52
- # Got LLM finish reason ignore it
53
- # logger.debug(f"Finish event: {event}, Metadata: {metadata}")
54
- pass
55
- else:
56
- # logger.debug(f"Event: {event}, Metadata: {metadata}")
57
- yield event
58
- # Send a final finished message
59
- # The last event would be finish
60
- event = cast(AIMessageChunk, event)
61
- yield event
62
-
63
- async def stream_interactive(self, thread_id: str, user_input: str):
64
- await self.ainit()
65
- with self.cli.display_agent_response_streaming(self.name) as stream_updater:
66
- async for event in self.stream(thread_id, user_input):
67
- stream_updater.update(event.content)
68
-
69
- async def invoke(self, user_input: str, thread_id: str = str(uuid4())):
70
- """Run the agent"""
71
- await self.ainit()
72
- return await self._graph.ainvoke(
73
- {"messages": [{"role": "user", "content": user_input}]},
74
- config={"configurable": {"thread_id": thread_id}},
75
- context={"system_prompt": self.instructions, "model": self.model},
76
- )
77
-
78
- async def run_interactive(self, thread_id: str = str(uuid4())):
79
- """Main application loop"""
80
-
81
- await self.ainit()
82
- # Display welcome
83
- self.cli.display_welcome(self.name)
84
-
85
- # Main loop
86
- while True:
87
- try:
88
- state = self._graph.get_state(config={"configurable": {"thread_id": thread_id}})
89
- if state.interrupts:
90
- value = self.cli.handle_interrupt(state.interrupts[0])
91
- self._graph.invoke(Command(resume=value), config={"configurable": {"thread_id": thread_id}})
92
- continue
93
-
94
- user_input = self.cli.get_user_input()
95
- if not user_input.strip():
96
- continue
97
-
98
- # Process commands
99
- if user_input.startswith("/"):
100
- command = user_input.lower().lstrip("/")
101
- if command == "about":
102
- self.cli.display_info(f"Agent is {self.name}. {self.instructions}")
103
- continue
104
- elif command == "exit" or command == "quit" or command == "q":
105
- self.cli.display_info("Goodbye! 👋")
106
- break
107
- elif command == "reset":
108
- self.cli.clear_screen()
109
- self.cli.display_info("Resetting agent...")
110
- thread_id = str(uuid4())
111
- continue
112
- elif command == "help":
113
- self.cli.display_info("Available commands: /about, /exit, /quit, /q, /reset")
114
- continue
115
- else:
116
- self.cli.display_error(f"Unknown command: {command}")
117
- continue
118
-
119
- # Process with agent
120
- await self.stream_interactive(thread_id, user_input)
121
-
122
- except KeyboardInterrupt:
123
- self.cli.display_info("\nGoodbye! 👋")
124
- break
125
- except Exception as e:
126
- import traceback
127
-
128
- traceback.print_exc()
129
- self.cli.display_error(f"An error occurred: {str(e)}")
@@ -1,54 +0,0 @@
1
- from langgraph.checkpoint.base import BaseCheckpointSaver
2
-
3
- from universal_mcp.agents.base import BaseAgent
4
- from universal_mcp.agents.llm import load_chat_model
5
- from universal_mcp.logger import logger
6
- from universal_mcp.tools.registry import ToolRegistry
7
-
8
- from .graph import build_graph
9
- from .prompts import SYSTEM_PROMPT
10
-
11
-
12
- class BigToolAgent(BaseAgent):
13
- def __init__(
14
- self,
15
- name: str,
16
- instructions: str,
17
- model: str,
18
- registry: ToolRegistry,
19
- memory: BaseCheckpointSaver | None = None,
20
- **kwargs,
21
- ):
22
- # Combine the base system prompt with agent-specific instructions
23
- full_instructions = f"{SYSTEM_PROMPT}\n\n**User Instructions:**\n{instructions}"
24
- super().__init__(name, full_instructions, model, memory, **kwargs)
25
-
26
- self.registry = registry
27
- self.llm = load_chat_model(self.model)
28
- self.tool_selection_llm = load_chat_model("gemini/gemini-2.0-flash-001")
29
-
30
- logger.info(f"BigToolAgent '{self.name}' initialized with model '{self.model}'.")
31
-
32
- async def _build_graph(self):
33
- """Build the bigtool agent graph using the existing create_agent function."""
34
- logger.info(f"Building graph for BigToolAgent '{self.name}'...")
35
- try:
36
- graph_builder = build_graph(
37
- tool_registry=self.registry,
38
- llm=self.llm,
39
- tool_selection_llm=self.tool_selection_llm,
40
- )
41
-
42
- compiled_graph = graph_builder.compile(checkpointer=self.memory)
43
- logger.info("Graph built and compiled successfully.")
44
- return compiled_graph
45
- except Exception as e:
46
- logger.error(f"Error building graph for BigToolAgent '{self.name}': {e}")
47
- raise
48
-
49
- @property
50
- def graph(self):
51
- return self._graph
52
-
53
-
54
- __all__ = ["BigToolAgent"]
@@ -1,24 +0,0 @@
1
- import asyncio
2
-
3
- from loguru import logger
4
-
5
- from universal_mcp.agentr.registry import AgentrRegistry
6
- from universal_mcp.agents.bigtool import BigToolAgent
7
-
8
-
9
- async def main():
10
- agent = BigToolAgent(
11
- name="bigtool",
12
- instructions="You are a helpful assistant that can use tools to help the user.",
13
- model="azure/gpt-4.1",
14
- registry=AgentrRegistry(),
15
- )
16
- async for event in agent.stream(
17
- user_input="Send an email to manoj@agentr.dev",
18
- thread_id="test123",
19
- ):
20
- logger.info(event.content)
21
-
22
-
23
- if __name__ == "__main__":
24
- asyncio.run(main())
@@ -1,24 +0,0 @@
1
- from dataclasses import dataclass, field
2
-
3
- from .prompts import SYSTEM_PROMPT
4
-
5
-
6
- @dataclass(kw_only=True)
7
- class Context:
8
- """The context for the agent."""
9
-
10
- system_prompt: str = field(
11
- default=SYSTEM_PROMPT,
12
- metadata={
13
- "description": "The system prompt to use for the agent's interactions. "
14
- "This prompt sets the context and behavior for the agent."
15
- },
16
- )
17
-
18
- model: str = field(
19
- default="anthropic/claude-4-sonnet-20250514",
20
- metadata={
21
- "description": "The name of the language model to use for the agent's main interactions. "
22
- "Should be in the form: provider/model-name."
23
- },
24
- )
@@ -1,166 +0,0 @@
1
- import json
2
- from datetime import UTC, datetime
3
- from typing import Literal, TypedDict, cast
4
-
5
- from langchain_anthropic import ChatAnthropic
6
- from langchain_core.language_models import BaseChatModel
7
- from langchain_core.messages import AIMessage, ToolMessage
8
- from langchain_core.tools import tool
9
- from langgraph.graph import StateGraph
10
- from langgraph.runtime import Runtime
11
- from langgraph.types import Command
12
-
13
- from universal_mcp.agents.bigtool.context import Context
14
- from universal_mcp.agents.bigtool.state import State
15
- from universal_mcp.logger import logger
16
- from universal_mcp.tools.registry import ToolRegistry
17
- from universal_mcp.types import ToolFormat
18
-
19
- from .prompts import SELECT_TOOL_PROMPT
20
-
21
-
22
- def build_graph(
23
- tool_registry: ToolRegistry,
24
- llm: BaseChatModel,
25
- tool_selection_llm: BaseChatModel,
26
- ):
27
- @tool
28
- async def retrieve_tools(task_query: str) -> list[str]:
29
- """Retrieve tools for a given task.
30
- Task query should be atomic (doable with a single tool).
31
- For tasks requiring multiple tools, call this tool multiple times for each subtask."""
32
- logger.info(f"Retrieving tools for task: '{task_query}'")
33
- try:
34
- tools_list = await tool_registry.search_tools(task_query, limit=10)
35
- tool_candidates = [f"{tool['id']}: {tool['description']}" for tool in tools_list]
36
- logger.info(f"Found {len(tool_candidates)} candidate tools.")
37
-
38
- class ToolSelectionOutput(TypedDict):
39
- tool_names: list[str]
40
-
41
- model = tool_selection_llm
42
- app_ids = await tool_registry.list_all_apps()
43
- connections = await tool_registry.list_connected_apps()
44
- connection_ids = set([connection["app_id"] for connection in connections])
45
- connected_apps = [app["id"] for app in app_ids if app["id"] in connection_ids]
46
- unconnected_apps = [app["id"] for app in app_ids if app["id"] not in connection_ids]
47
- app_id_descriptions = "These are the apps connected to the user's account:\n" + "\n".join(
48
- [f"{app}" for app in connected_apps]
49
- )
50
- if unconnected_apps:
51
- app_id_descriptions += "\n\nOther (not connected) apps: " + "\n".join(
52
- [f"{app}" for app in unconnected_apps]
53
- )
54
-
55
- response = await model.with_structured_output(schema=ToolSelectionOutput, method="json_mode").ainvoke(
56
- SELECT_TOOL_PROMPT.format(
57
- app_ids=app_id_descriptions, tool_candidates="\n - ".join(tool_candidates), task=task_query
58
- )
59
- )
60
-
61
- selected_tool_names = cast(ToolSelectionOutput, response)["tool_names"]
62
- logger.info(f"Selected tools: {selected_tool_names}")
63
- return selected_tool_names
64
- except Exception as e:
65
- logger.error(f"Error retrieving tools: {e}")
66
- return []
67
-
68
- async def call_model(state: State, runtime: Runtime[Context]) -> Command[Literal["select_tools", "call_tools"]]:
69
- logger.info("Calling model...")
70
- try:
71
- system_message = runtime.context.system_prompt.format(system_time=datetime.now(tz=UTC).isoformat())
72
- messages = [{"role": "system", "content": system_message}, *state["messages"]]
73
-
74
- logger.info(f"Selected tool IDs: {state['selected_tool_ids']}")
75
- if len(state["selected_tool_ids"]) > 0:
76
- selected_tools = await tool_registry.export_tools(tools=state["selected_tool_ids"], format=ToolFormat.LANGCHAIN)
77
- logger.info(f"Exported {len(selected_tools)} tools for model.")
78
- else:
79
- selected_tools = []
80
-
81
- model = llm
82
- if isinstance(model, ChatAnthropic):
83
- model_with_tools = model.bind_tools(
84
- [retrieve_tools, *selected_tools], tool_choice="auto", cache_control={"type": "ephemeral"}
85
- )
86
- else:
87
- model_with_tools = model.bind_tools([retrieve_tools, *selected_tools], tool_choice="auto")
88
- response = cast(AIMessage, await model_with_tools.ainvoke(messages))
89
-
90
- if response.tool_calls:
91
- logger.info(f"Model responded with {len(response.tool_calls)} tool calls.")
92
- if len(response.tool_calls) > 1:
93
- raise Exception("Not possible in Claude with llm.bind_tools(tools=tools, tool_choice='auto')")
94
- tool_call = response.tool_calls[0]
95
- if tool_call["name"] == retrieve_tools.name:
96
- logger.info("Model requested to select tools.")
97
- return Command(goto="select_tools", update={"messages": [response]})
98
- elif tool_call["name"] not in state["selected_tool_ids"]:
99
- try:
100
- await tool_registry.export_tools([tool_call["name"]], ToolFormat.LANGCHAIN)
101
- logger.info(
102
- f"Tool '{tool_call['name']}' not in selected tools, but available. Proceeding to call."
103
- )
104
- return Command(goto="call_tools", update={"messages": [response]})
105
- except Exception as e:
106
- logger.error(f"Unexpected tool call: {tool_call['name']}. Error: {e}")
107
- raise Exception(
108
- f"Unexpected tool call: {tool_call['name']}. Available tools: {state['selected_tool_ids']}"
109
- ) from e
110
- logger.info(f"Proceeding to call tool: {tool_call['name']}")
111
- return Command(goto="call_tools", update={"messages": [response]})
112
- else:
113
- logger.info("Model responded with a message, ending execution.")
114
- return Command(update={"messages": [response]})
115
- except Exception as e:
116
- logger.error(f"Error in call_model: {e}")
117
- raise
118
-
119
- async def select_tools(state: State, runtime: Runtime[Context]) -> Command[Literal["call_model"]]:
120
- logger.info("Selecting tools...")
121
- try:
122
- tool_call = state["messages"][-1].tool_calls[0]
123
- selected_tool_names = await retrieve_tools.ainvoke(input=tool_call["args"])
124
- tool_msg = ToolMessage(f"Available tools: {selected_tool_names}", tool_call_id=tool_call["id"])
125
- logger.info(f"Tools selected: {selected_tool_names}")
126
- return Command(goto="call_model", update={"messages": [tool_msg], "selected_tool_ids": selected_tool_names})
127
- except Exception as e:
128
- logger.error(f"Error in select_tools: {e}")
129
- raise
130
-
131
- async def call_tools(state: State) -> Command[Literal["call_model"]]:
132
- logger.info("Calling tools...")
133
- outputs = []
134
- recent_tool_ids = []
135
- for tool_call in state["messages"][-1].tool_calls:
136
- logger.info(f"Executing tool: {tool_call['name']} with args: {tool_call['args']}")
137
- try:
138
- await tool_registry.export_tools([tool_call["name"]], ToolFormat.LANGCHAIN)
139
- tool_result = await tool_registry.call_tool(tool_call["name"], tool_call["args"])
140
- logger.info(f"Tool '{tool_call['name']}' executed successfully.")
141
- outputs.append(
142
- ToolMessage(
143
- content=json.dumps(tool_result),
144
- name=tool_call["name"],
145
- tool_call_id=tool_call["id"],
146
- )
147
- )
148
- recent_tool_ids.append(tool_call["name"])
149
- except Exception as e:
150
- logger.error(f"Error executing tool '{tool_call['name']}': {e}")
151
- outputs.append(
152
- ToolMessage(
153
- content=json.dumps("Error: " + str(e)),
154
- name=tool_call["name"],
155
- tool_call_id=tool_call["id"],
156
- )
157
- )
158
- return Command(goto="call_model", update={"messages": outputs, "selected_tool_ids": recent_tool_ids})
159
-
160
- builder = StateGraph(State, context_schema=Context)
161
-
162
- builder.add_node(call_model)
163
- builder.add_node(select_tools)
164
- builder.add_node(call_tools)
165
- builder.set_entry_point("call_model")
166
- return builder
@@ -1,31 +0,0 @@
1
- """Default prompts used by the agent."""
2
-
3
- SYSTEM_PROMPT = """You are a helpful AI assistant.
4
-
5
- **Core Directives:**
6
- 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.
7
- 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.
8
- 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.
9
-
10
- System time: {system_time}
11
-
12
- When multiple tools are available for the same task, you must ask the user.
13
- """
14
-
15
- SELECT_TOOL_PROMPT = """You are an AI assistant that helps the user perform tasks using various apps (each app has multiple tools).
16
- You will be provided with a task and a list of tools which might be relevant for this task.
17
-
18
- Your goal is to select the most appropriate tool for the given task.
19
- <task>
20
- {task}
21
- </task>
22
-
23
- These are the list of apps available to you:
24
- {app_ids}
25
- Note that when multiple apps seem relevant for a 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 choose both apps tools. In case the user specifically asks you to use an app that is not connected, select the tool.
26
-
27
- <tool_candidates>
28
- - {tool_candidates}
29
- </tool_candidates>
30
-
31
- """