Rename planning_agent.py to tools_agent.py

planning_agent.py → tools_agent.py

@@ -1,41 +1,32 @@
 from typing import Dict, List, Any
-from dataclasses import dataclass
-from datetime import datetime
+from tool_registry import Tool
+import openai
 import os
 import json
-import requests
 
 
-@dataclass
-class Interaction:
-    """Record of a single interaction with the agent"""
-    timestamp: datetime
-    query: str
-    plan: Dict[str, Any]
-
 class Agent:
-    def __init__(self, model: str = "Qwen/Qwen2.5-Coder-32B-Instruct"):
-        """Initialize Agent with empty interaction history."""
-        self.interactions: List[Interaction] = []  # Working memory
-        self.model = model
-
-    def _query_llm(self, messages):
-        headers = {
-            "Content-Type": "application/json"
-        }
-        data = {
-            "model": self.model,
-            "messages": messages,
-            "max_tokens": 150
-        }
-        response = requests.post("https://api-inference.huggingface.co/v1/chat/completions", headers=headers, data=json.dumps(data))
-        print("Original response ", response.json())
-        print("\nOriginal response type", type(json.loads(response.choices[0].message.content)))
-        # final_response = response.json()['choices'][0]['message']['content'].strip()
-        # print("LLM Response ", final_response)
+    def __init__(self, client):
+        """Initialize Agent with empty tool registry."""
+        self.client = client
+        self.tools: Dict[str, Tool] = {}
+    
+    def add_tool(self, tool: Tool) -> None:
+        """Register a new tool with the agent."""
+        self.tools[tool.name] = tool
     
-        return response.choices[0].message.content
+    def get_available_tools(self) -> List[str]:
+        """Get list of available tool descriptions."""
+        return [f"{tool.name}: {tool.description}" for tool in self.tools.values()]
     
+    def use_tool(self, tool_name: str, **kwargs: Any) -> str:
+        """Execute a specific tool with given arguments."""
+        if tool_name not in self.tools:
+            raise ValueError(f"Tool '{tool_name}' not found. Available tools: {list(self.tools.keys())}")
+        
+        tool = self.tools[tool_name]
+        return tool.func(**kwargs)
+
     def create_system_prompt(self) -> str:
         """Create the system prompt for the LLM with available tools."""
         tools_json = {
@@ -43,34 +34,26 @@ class Agent:
             "capabilities": [
                 "Using provided tools to help users when necessary",
                 "Responding directly without tools for questions that don't require tool usage",
-                "Planning efficient tool usage sequences",
-                "If asked by the user, reflecting on the plan and suggesting changes if needed"
+                "Planning efficient tool usage sequences"
             ],
             "instructions": [
                 "Use tools only when they are necessary for the task",
                 "If a query can be answered directly, respond with a simple message instead of using tools",
-                "When tools are needed, plan their usage efficiently to minimize tool calls",
-                "If asked by the user, reflect on the plan and suggest changes if needed"
+                "When tools are needed, plan their usage efficiently to minimize tool calls"
             ],
             "tools": [
                 {
-                    "name": "convert_currency",
-                    "description": "Converts currency using latest exchange rates.",
+                    "name": tool.name,
+                    "description": tool.description,
                     "parameters": {
-                        "amount": {
-                            "type": "float",
-                            "description": "Amount to convert"
-                        },
-                        "from_currency": {
-                            "type": "str",
-                            "description": "Source currency code (e.g., USD)"
-                        },
-                        "to_currency": {
-                            "type": "str",
-                            "description": "Target currency code (e.g., EUR)"
+                        name: {
+                            "type": info["type"],
+                            "description": info["description"]
                         }
+                        for name, info in tool.parameters.items()
                     }
                 }
+                for tool in self.tools.values()
             ],
             "response_format": {
                 "type": "json",
@@ -170,129 +153,67 @@ class Agent:
         
         return f"""You are an AI assistant that helps users by providing direct answers or using tools when necessary.
 Configuration, instructions, and available tools are provided in JSON format below:
+
 {json.dumps(tools_json, indent=2)}
+
 Always respond with a JSON object following the response_format schema above. 
 Remember to use tools only when they are actually needed for the task."""
 
     def plan(self, user_query: str) -> Dict:
-        """Use LLM to create a plan and store it in memory."""
+        """Use LLM to create a plan for tool usage."""
         messages = [
             {"role": "system", "content": self.create_system_prompt()},
             {"role": "user", "content": user_query}
         ]
-
-        response = self._query_llm(messages=messages)
-        
-        try:
-            plan = json.loads(response)
-            # Store the interaction immediately after planning
-            interaction = Interaction(
-                timestamp=datetime.now(),
-                query=user_query,
-                plan=plan
-            )
-            self.interactions.append(interaction)
-            return plan
-        except json.JSONDecodeError:
-            raise ValueError("Failed to parse LLM response as JSON")
-
-    def reflect_on_plan(self) -> Dict[str, Any]:
-        """Reflect on the most recent plan using interaction history."""
-        if not self.interactions:
-            return {"reflection": "No plan to reflect on", "requires_changes": False}
-        
-        latest_interaction = self.interactions[-1]
-        
-        reflection_prompt = {
-            "task": "reflection",
-            "context": {
-                "user_query": latest_interaction.query,
-                "generated_plan": latest_interaction.plan
-            },
-            "instructions": [
-                "Review the generated plan for potential improvements",
-                "Consider if the chosen tools are appropriate",
-                "Verify tool parameters are correct",
-                "Check if the plan is efficient",
-                "Determine if tools are actually needed"
-            ],
-            "response_format": {
-                "type": "json",
-                "schema": {
-                    "requires_changes": {
-                        "type": "boolean",
-                        "description": "whether the plan needs modifications"
-                    },
-                    "reflection": {
-                        "type": "string",
-                        "description": "explanation of what changes are needed or why no changes are needed"
-                    },
-                    "suggestions": {
-                        "type": "array",
-                        "items": {"type": "string"},
-                        "description": "specific suggestions for improvements",
-                        "optional": True
-                    }
-                }
-            }
-        }
-        
-        messages = [
-            {"role": "system", "content": self.create_system_prompt()},
-            {"role": "user", "content": json.dumps(reflection_prompt, indent=2)}
-        ]
         
-        response = self._query_llm(messages=messages)
+        response = self.client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=messages,
+            temperature=0
+        )
         
         try:
-            return json.loads(response)
+            return json.loads(response.choices[0].message.content)
         except json.JSONDecodeError:
-            return {"reflection": response.choices[0].message.content}
+            raise ValueError("Failed to parse LLM response as JSON")
 
     def execute(self, user_query: str) -> str:
-        """Execute the full pipeline: plan, reflect, and potentially replan."""
+        """Execute the full pipeline: plan and execute tools."""
         try:
-            # Create initial plan (this also stores it in memory)
-            initial_plan = self.plan(user_query)
-            
-            # Reflect on the plan using memory
-            reflection = self.reflect_on_plan()
+            plan = self.plan(user_query)
             
-            # Check if reflection suggests changes
-            if reflection.get("requires_changes", False):
-                # Generate new plan based on reflection
-                messages = [
-                    {"role": "system", "content": self.create_system_prompt()},
-                    {"role": "user", "content": user_query},
-                    {"role": "assistant", "content": json.dumps(initial_plan)},
-                    {"role": "user", "content": f"Please revise the plan based on this feedback: {json.dumps(reflection)}"}
-                ]
-                
-                response = self._query_llm(messages=messages)
-                
-                try:
-                    final_plan = json.loads(response)
-                except json.JSONDecodeError:
-                    final_plan = initial_plan  # Fallback to initial plan if parsing fails
-            else:
-                final_plan = initial_plan
+            if not plan.get("requires_tools", True):
+                return plan["direct_response"]
             
-            # Update the stored interaction with all information
-            self.interactions[-1].plan = {
-                "initial_plan": initial_plan,
-                "reflection": reflection,
-                "final_plan": final_plan
-            }
+            # Execute each tool in sequence
+            results = []
+            for tool_call in plan["tool_calls"]:
+                tool_name = tool_call["tool"]
+                tool_args = tool_call["args"]
+                result = self.use_tool(tool_name, **tool_args)
+                results.append(result)
             
-            # Return the appropriate response
-            if final_plan.get("requires_tools", True):
-                return f"""Initial Thought: {initial_plan['thought']}
-Initial Plan: {'. '.join(initial_plan['plan'])}
-Reflection: {reflection.get('reflection', 'No improvements suggested')}
-Final Plan: {'. '.join(final_plan['plan'])}"""
-            else:
-                return f"""Response: {final_plan['direct_response']}
-Reflection: {reflection.get('reflection', 'No improvements suggested')}"""
+            # Combine results
+            return f"""Thought: {plan['thought']}
+Plan: {'. '.join(plan['plan'])}
+Results: {'. '.join(results)}"""
             
         except Exception as e:
             return f"Error executing plan: {str(e)}"
+
+def main():
+    from tools import convert_currency
+    
+    agent = Agent()
+    agent.add_tool(convert_currency)
+    
+    query_list = ["I am traveling to Japan from Serbia, I have 1500 of local currency, how much of Japaese currency will I be able to get?",
+                  "How are you doing?"]
+    
+    for query in query_list:
+        print(f"\nQuery: {query}")
+        result = agent.execute(query)
+        print(result)
+
+if __name__ == "__main__":
+    main()