main
py 33 lines 871 Bytes
Raw
1 from typing import Any
2
3 from helpers.extension import Extension, extensible
4 from helpers.mcp_handler import MCPConfig
5 from agent import Agent, LoopData
6
7
8 class MCPToolsPrompt(Extension):
9
10 async def execute(
11 self,
12 system_prompt: list[str] = [],
13 loop_data: LoopData = LoopData(),
14 **kwargs: Any,
15 ):
16 if not self.agent:
17 return
18 prompt = await build_prompt(self.agent)
19 if prompt:
20 system_prompt.append(prompt)
21
22
23 @extensible
24 async def build_prompt(agent: Agent) -> str:
25 mcp_config = MCPConfig.get_for_agent(agent)
26 if not mcp_config.servers:
27 return ""
28
29 pre_progress = agent.context.log.progress
30 agent.context.log.set_progress("Collecting MCP tools")
31 tools = mcp_config.get_tools_prompt(agent=agent)
32 agent.context.log.set_progress(pre_progress)
33 return tools