| 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 |