Move cost management from tool loader to agent manager

default_tools/agent_creater_tool.py

@@ -65,7 +65,7 @@ class AgentCreator():
 
         agent_manager = AgentManager()
         try:
-            agent_manager.create_agent(
+            _, remaining_budget = agent_manager.create_agent(
                 agent_name=agent_name,
                 base_model=base_model,
                 system_prompt=system_prompt,
@@ -83,5 +83,5 @@ class AgentCreator():
         return {
             "status": "success",
             "message": "Agent successfully created",
-            "cost": create_cost,
+            "remaining_budget": remaining_budget,
         }

default_tools/ask_agent_tool.py

@@ -34,10 +34,9 @@ class AskAgent():
         agent_name = kwargs.get("agent_name")
         prompt = kwargs.get("prompt")
         agent_manger = AgentManager()
-        budget_manager = BudgetManager()
 
         try:
-            agent = agent_manger.get_agent(agent_name=agent_name)
+            agent_response, remaining_budget = agent_manger.ask_agent(agent_name=agent_name, prompt=prompt)
         except ValueError as e:
             return {
                 "status": "error",
@@ -45,21 +44,10 @@ class AskAgent():
                 "output": None
             }
 
-        agent_costs = agent.get_costs()
-        agent_question_cost = agent_costs["invoke_cost"]
-        print("Agent question cost", agent_question_cost)
-        
-        if not budget_manager.can_spend(agent_question_cost):
-            return {
-                "status": "error",
-                "message": f"Do not have enough budget to ask the agent a question. Asking the agent costs {agent_question_cost} but only {budget_manager.get_current_remaining_budget()} is remaining",
-                "output": None
-            }
-
-        agent_response = agent.ask_agent(prompt=prompt)
         print("Agent response", agent_response)
         return {
             "status": "success",
             "message": "Agent has replied to the given prompt",
             "output": agent_response,
+            "remaining_budget": remaining_budget
         }

default_tools/fire_agent.py

@@ -28,26 +28,18 @@ class FireAgent():
         agent_name= kwargs.get("agent_name")
 
         agent_manager = AgentManager()
-        
-        agent = agent_manager.get_agent(agent_name=agent_name)
-        
-        agent_costs = agent.get_costs()
-        
+                
         try:
-            agent_manager.delete_agent(agent_name=agent_name)
+            remaining_budget = agent_manager.delete_agent(agent_name=agent_name)
         except ValueError as e:
             return {
                 "status": "error",
                 "message": f"Error occurred: {str(e)}",
                 "output": None
             }
-            
-        budget_manager = BudgetManager()
-        
-        budget_manager.add_to_expense(-1* int(agent_costs["create_cost"]))
 
         return {
             "status": "success",
             "message": "Agent successfully fired.",
-            "current_expense": budget_manager.get_current_expense()
+            "remaining_budget": remaining_budget
         }

src/agent_manager.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Type, Any, Optional
+from typing import Dict, Type, Any, Optional, Tuple
 import os
 import json
 import ollama
@@ -10,6 +10,8 @@ from google.genai import types
 from google.genai.types import *
 import os
 from dotenv import load_dotenv
+from src.budget_manager import BudgetManager
+
 class Agent(ABC):
     
     def __init__(self, agent_name: str, base_model: str, system_prompt: str, creation_cost: str, invoke_cost: str):
@@ -93,7 +95,7 @@ class GeminiAgent(Agent):
         self.messages = []
 @singleton
 class AgentManager():
-    
+    budget_manager = BudgetManager()
     def __init__(self):
         self._agents = {}
         self._agent_types ={
@@ -104,7 +106,7 @@ class AgentManager():
         self._load_agents()
     
     def create_agent(self, agent_name: str, base_model: str, system_prompt: str, description: str = "", create_cost: float = 0, invoke_cost: float = 0, 
-                 **additional_params) -> Agent:
+                 **additional_params) -> Tuple[Agent, int]:
         
         if agent_name in self._agents:
             raise ValueError(f"Agent {agent_name} already exists")
@@ -115,6 +117,11 @@ class AgentManager():
         if not agent_class:
             raise ValueError(f"Unsupported base model {base_model}")
         
+        if not self.budget_manager.can_spend(create_cost):
+            raise ValueError(f"Do not have enough budget to create the tool. "
+                        +f"Creating the tool costs {create_cost} but only {self.budget_manager.get_current_remaining_budget()} is remaining")
+        
+        self.budget_manager.add_to_expense(create_cost)
         # create agent
         self._agents[agent_name] = agent_class(agent_name, base_model, system_prompt, create_cost,invoke_cost )
         
@@ -129,7 +136,7 @@ class AgentManager():
             **additional_params  # For any future parameters we might want to add
         )
         
-        return self._agents[agent_name]
+        return (self._agents[agent_name], self.budget_manager.get_current_remaining_budget())
     
     def get_agent(self, agent_name: str) -> Agent:
         """Get existing agent by name"""
@@ -159,11 +166,14 @@ class AgentManager():
             output_assistant_response(f"Error listing agents: {e}")
             return {}
     
-    def delete_agent(self, agent_name: str) -> None:
+    def delete_agent(self, agent_name: str) -> int:
         if agent_name not in self._agents:
             raise ValueError(f"Agent {agent_name} does not exist")
+        
+        self.budget_manager.add_to_expense(-1 * self._agents[agent_name].creation_cost)
         self._agents[agent_name].delete_agent()
         
+        
         del self._agents[agent_name]
         try:
             if os.path.exists("./models/models.json"):
@@ -175,7 +185,17 @@ class AgentManager():
                     f.write(json.dumps(models, indent=4))
         except Exception as e:
             output_assistant_response(f"Error deleting agent: {e}")
+        return self.budget_manager.get_current_remaining_budget()
     
+    def ask_agent(self, agent_name: str, prompt: str) -> Tuple[str,int]:
+        agent = self._agents[agent_name]
+        
+        if not self.budget_manager.can_spend(agent.invoke_cost):
+            raise ValueError(f"Do not have enough budget to ask Agent {agent_name} a question. "
+                        +f"Asking Agent {agent_name} costs {agent.invoke_cost} but only {self.budget_manager.get_current_remaining_budget()} is remaining")
+        
+        response = agent.ask_agent(prompt)        
+        return (response, self.budget_manager.get_current_remaining_budget())
     
     def _save_agent(self, agent_name: str, base_model: str, system_prompt: str, 
                     description: str = "", create_cost: float = 0, invoke_cost: float = 0, 

src/tool_loader.py

@@ -20,7 +20,7 @@ TOOLS_DIRECTORIES = [os.path.abspath("./default_tools"), os.path.abspath("./tool
 class Tool:
     def __init__(self, toolClass):
         suppress_output(self.load_tool)(toolClass)
-
+        
     def load_tool(self, toolClass):
         self.tool = toolClass()
         self.inputSchema = self.tool.inputSchema
@@ -81,7 +81,6 @@ class ToolLoader:
         output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
         for tool in self.toolsImported:
             if tool.name == toolName:
-                self.update_budget(query, tool.inputSchema)
                 return tool.run(query)
         output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
         return {
@@ -90,28 +89,6 @@ class ToolLoader:
             "output": None
         }
     
-    def update_budget(self, query, inputSchema):
-        if "creates" in inputSchema:
-            selector = inputSchema["creates"]["selector"]
-            if selector in query:
-                create_cost = inputSchema["creates"]["types"][query[selector]]["create_cost"]
-                if not self.budget_manager.can_spend(create_cost):
-                    return {
-                        "status": "error",
-                        "message": f"Do not have enough budget to create the tool. "
-                        +f"Creating the tool costs {create_cost} but only {self.budget_manager.get_current_remaining_budget()} is remaining",
-                        "output": None
-                    }
-                self.budget_manager.add_to_expense(create_cost)
-        if "invoke_cost" in inputSchema:
-            invoke_cost = inputSchema["invoke_cost"]
-            if not self.budget_manager.can_spend(invoke_cost):
-                return {
-                    "status": "error",
-                    "message": f"Do not have enough budget to invoke the tool. "
-                    +f"Invoking the tool costs {invoke_cost} but only {self.budget_manager.get_current_remaining_budget()} is remaining",
-                    "output": None
-                }
 
     def getTools(self):
         toolsList = []