Update src/forge_agent.py

src/forge_agent.py

@@ -1,5 +1,7 @@
 import asyncio
+import json
 from typing import List, Dict, Any
+from huggingface_hub import AsyncInferenceClient
 from .mcp_client import MCPClient
 
 class ToolRegistry:
@@ -30,33 +32,67 @@ class ToolRegistry:
         """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"):
+class HuggingFaceAgent:
+    """An AI agent that uses a Hugging Face model to generate plans."""
+    def __init__(self, hf_token: str, model_name: str = "mistralai/Mixtral-8x7B-Instruct-v0.1"):
         self.model = model_name
+        self.client = AsyncInferenceClient(model=model_name, token=hf_token)
 
-    async def generate_plan(self, goal: str, available_tools: List[str]) -> List[Dict[str, Any]]:
+    def _construct_prompt(self, goal: str, available_tools: List[Dict[str, Any]], previous_steps: List = None, error: str = None) -> str:
+        """Constructs the detailed prompt for the LLM."""
+        tools_json_string = json.dumps(available_tools, indent=2)
+        
+        prompt = f"""You are Forge, an autonomous AI agent. Your task is to create a step-by-step plan to achieve a goal.
+You must respond with a valid JSON array of objects, where each object represents a step in the plan.
+Each step must have 'step', 'thought', 'tool', and 'params' keys.
+The final step must always use the 'report_success' tool.
+
+Available Tools:
+{tools_json_string}
+
+Goal: "{goal}"
+"""
+        if previous_steps:
+            prompt += f"\nYou have already completed these steps:\n{json.dumps(previous_steps, indent=2)}\n"
+        if error:
+            prompt += f"\nAn error occurred during the last step: {error}\nAnalyze the error and create a new, corrected plan to achieve the original goal. Start the new plan from the current state."
+
+        prompt += "\nGenerate the JSON plan now:"
+        return prompt
+
+    async def _invoke_llm(self, prompt: str) -> List[Dict[str, Any]]:
+        """Invokes the LLM and parses the JSON response."""
+        try:
+            response = await self.client.text_generation(prompt, max_new_tokens=1024)
+            # The response might contain the JSON within backticks or other text.
+            json_response_str = response.strip().split('```json')[-1].split('```')[0].strip()
+            plan = json.loads(json_response_str)
+            if isinstance(plan, list):
+                return plan
+            else:
+                raise ValueError("LLM did not return a JSON list.")
+        except (json.JSONDecodeError, ValueError, IndexError) as e:
+            print(f"Error parsing LLM response: {e}\nRaw response:\n{response}")
+            # Fallback or re-try logic could be added here
+            return [{"step": 1, "thought": "Failed to generate a plan due to a parsing error.", "tool": "report_failure", "params": {"message": f"LLM response parsing failed: {e}"}}]
+
+    async def generate_plan(self, goal: str, available_tools: List[Dict[str, Any]]) -> 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
+        prompt = self._construct_prompt(goal, available_tools)
+        return await self._invoke_llm(prompt)
+
+    async def regenerate_plan_on_error(self, goal: str, available_tools: List[Dict[str, Any]], completed_steps: List, error_message: str) -> List[Dict[str, Any]]:
+        """Generates a new plan after an error occurred."""
+        prompt = self._construct_prompt(goal, available_tools, previous_steps=completed_steps, error=error_message)
+        return await self._invoke_llm(prompt)
 
 class ForgeApp:
     """The main orchestrator for the Forge application."""
-    def __init__(self, goal: str, mcp_server_urls: List[str]):
+    def __init__(self, goal: str, mcp_server_urls: List[str], hf_token: str):
         self.goal = goal
-        self.planner = AIAgent()
+        self.planner = HuggingFaceAgent(hf_token=hf_token)
         self.tool_registry = ToolRegistry(server_urls=mcp_server_urls)
 
     async def run(self):
@@ -66,34 +102,46 @@ class ForgeApp:
         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())
+        
+        # Provide the LLM with full tool details, not just names
+        available_tools_details = [{"name": name, "description": data["description"]} for name, data in self.tool_registry.tools.items()]
+        
         yield f"🧠 **Generating a plan for your goal:** '{self.goal}'"
-        plan = await self.planner.generate_plan(self.goal, available_tool_names)
+        plan = await self.planner.generate_plan(self.goal, available_tools_details)
         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
+        completed_steps = []
+        while plan:
+            task = plan.pop(0)
+            yield f"\n**[Step {task.get('step', '?')}]** 🤔 **Thought:** {task.get('thought', 'N/A')}"
+
+            tool_name = task.get("tool")
+            if tool_name in ["report_success", "report_failure"]:
+                emoji = "🎉" if tool_name == "report_success" else "🛑"
+                yield f"{emoji} **Final Result:** {task.get('params', {}).get('message', 'N/A')}"
+                plan = [] # End execution
+                continue
 
             try:
-                yield f"🛠️ **Action:** Executing tool `{task['tool']}` with params: `{task['params']}`"
-                result = await self.tool_registry.execute(task["tool"], task["params"])
+                yield f"🛠️ **Action:** Executing tool `{tool_name}` with params: `{task.get('params', {})}`"
+                result = await self.tool_registry.execute(tool_name, task.get("params", {}))
                 
                 if result.get("status") == "error":
-                    yield f"❌ **Error:** {result.get('result', 'Unknown error')}"
-                    yield "🛑 **Execution Halted due to error.**"
-                    break
+                    error_message = result.get('result', 'Unknown error')
+                    yield f"❌ **Error:** {error_message}"
+                    yield "🧠 **Agent is re-evaluating the plan based on the error...**"
+                    completed_steps.append({"step": task, "outcome": "error", "details": error_message})
+                    plan = await self.planner.regenerate_plan_on_error(self.goal, available_tools_details, completed_steps, error_message)
+                    yield "📝 **New Plan Generated!** Resuming execution..."
                 else:
-                    yield f"✅ **Observation:** {result.get('result', 'Tool executed successfully.')}"
+                    observation = result.get('result', 'Tool executed successfully.')
+                    yield f"✅ **Observation:** {observation}"
+                    completed_steps.append({"step": task, "outcome": "success", "details": observation})
 
             except Exception as e:
-                yield f"❌ **Critical Error executing step {task['step']}:** {e}"
+                yield f"❌ **Critical Error executing step {task.get('step', '?')}:** {e}"
                 yield "🛑 **Execution Halted due to critical error.**"
-                break
+                plan = [] # End execution
         
         await self.tool_registry.close_all()
         yield "\n🏁 **Forge execution finished.**"