File size: 4,997 Bytes
915668d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import asyncio
from typing import List, Dict, Any
from .mcp_client import MCPClient
class ToolRegistry:
"""Manages connections to all required MCP servers and their tools."""
def __init__(self, server_urls: List[str]):
self.servers: Dict[str, MCPClient] = {url: MCPClient(url) for url in server_urls}
self.tools: Dict[str, Dict[str, Any]] = {}
async def discover_tools(self):
"""Discovers all available tools from all connected MCP servers."""
discovery_tasks = [client.list_tools() for client in self.servers.values()]
results = await asyncio.gather(*discovery_tasks, return_exceptions=True)
for i, client in enumerate(self.servers.values()):
server_tools = results[i]
if isinstance(server_tools, list):
for tool in server_tools:
self.tools[tool["name"]] = {"client": client, "description": tool["description"]}
async def execute(self, tool_name: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""Finds the correct MCP server and executes the tool."""
if tool_name not in self.tools:
raise ValueError(f"Tool '{tool_name}' not found in registry.")
tool_info = self.tools[tool_name]
return await tool_info["client"].execute_tool(tool_name, params)
async def close_all(self):
"""Closes all client connections."""
await asyncio.gather(*(client.close() for client in self.servers.values()))
class AIAgent:
"""A mock AI agent that generates a plan based on a goal."""
def __init__(self, model_name: str = "mock-planner"):
self.model = model_name
async def generate_plan(self, goal: str, available_tools: List[str]) -> List[Dict[str, Any]]:
"""
Generates a step-by-step plan.
In a real application, this would involve a call to a powerful LLM.
Here, we use a hardcoded plan for demonstration.
"""
await asyncio.sleep(1) # Simulate LLM thinking time
# This is a mock plan. A real LLM would generate this dynamically.
plan = [
{"step": 1, "thought": "I need a place to store the code. I'll use the `create_repo` tool.", "tool": "create_repo", "params": {"name": "my-awesome-blog", "private": False}},
{"step": 2, "thought": "Now I need to scaffold a new Next.js application inside a secure sandbox.", "tool": "execute_shell", "params": {"command": "npx create-next-app@latest my-awesome-blog --yes"}},
{"step": 3, "thought": "The project is created. I should commit the initial code.", "tool": "execute_shell", "params": {"command": "cd my-awesome-blog && git init && git add . && git commit -m 'Initial commit'"}},
{"step": 4, "thought": "The plan is complete. I will report success.", "tool": "report_success", "params": {"message": "Blog project scaffolded successfully."}}
]
return plan
class ForgeApp:
"""The main orchestrator for the Forge application."""
def __init__(self, goal: str, mcp_server_urls: List[str]):
self.goal = goal
self.planner = AIAgent()
self.tool_registry = ToolRegistry(server_urls=mcp_server_urls)
async def run(self):
"""
Runs the agent and yields status updates as a generator.
"""
yield "🚀 **Starting Forge... Initializing systems.**"
await self.tool_registry.discover_tools()
yield f"✅ **Tool Discovery Complete.** Found {len(self.tool_registry.tools)} tools."
available_tool_names = list(self.tool_registry.tools.keys())
yield f"🧠 **Generating a plan for your goal:** '{self.goal}'"
plan = await self.planner.generate_plan(self.goal, available_tool_names)
yield "📝 **Plan Generated!** Starting execution..."
for task in plan:
yield f"\n**[Step {task['step']}]** 🤔 **Thought:** {task['thought']}"
if task["tool"] == "report_success":
yield f"🎉 **Final Result:** {task['params']['message']}"
break
try:
yield f"🛠️ **Action:** Executing tool `{task['tool']}` with params: `{task['params']}`"
result = await self.tool_registry.execute(task["tool"], task["params"])
if result.get("status") == "error":
yield f"❌ **Error:** {result.get('result', 'Unknown error')}"
yield "🛑 **Execution Halted due to error.**"
break
else:
yield f"✅ **Observation:** {result.get('result', 'Tool executed successfully.')}"
except Exception as e:
yield f"❌ **Critical Error executing step {task['step']}:** {e}"
yield "🛑 **Execution Halted due to critical error.**"
break
await self.tool_registry.close_all()
yield "\n🏁 **Forge execution finished.**" |