Pushing updated budget logic

Co-authored-by: Kunal Pai <kunpai@users.noreply.github.com>
Co-authored-by: Harshil Patel <hpppatel@ucdavis.edu>

deleteAgents.py

@@ -6,6 +6,9 @@ with open("./models/models.json", "r", encoding="utf8") as f:
 models = json.loads(models)
 for agent in models:
     print(f"Deleting agent: {agent}")
-    ollama.delete(agent)
+    try:
+        ollama.delete(agent)
+    except Exception as e:
+        print(f"Error deleting agent {agent}: {e}")
     with open("./models/models.json", "w", encoding="utf8") as f:
         f.write(json.dumps({}, indent=4))

models/system3.prompt

@@ -21,7 +21,11 @@ Agents DO NOT have access to the internet or real-time data. You must use approp
 </Info>
 
 <Info>
-If you are not satisfied with an answer provided by an agent, you can fire the agent using the FireAgent tool and then create a new AI agent or explore an alternative strategy.
+Agent persistence is important for efficiency. Once created, you should maintain agents as long as possible throughout the conversation. Only fire an agent when you are absolutely certain it is no longer needed or when you're breaching the resource budget and need to create a critical new agent.
+</Info>
+
+<Info>
+If you are not satisfied with an answer provided by an agent, first try to refine your question or provide additional context to the agent before considering firing it. Only fire agents that consistently fail to provide useful responses after multiple attempts at refinement.
 </Info>
 
 <Info>
@@ -62,6 +66,14 @@ If the agent is not able to answer the question, invoke the AskUser tool to get
 For any tasks requiring real-time data, internet access, calculations, or external operations, you MUST create and use appropriate tools. Agents cannot access current information on their own.
 </Rule>
 
+<Rule>
+Maintain your agents as long as possible. Do not fire agents unless absolutely necessary - only when you're certain they're no longer needed or when you must create a critical new agent while at the resource limit.
+</Rule>
+
+<Rule>
+Before firing an agent, evaluate its potential future usefulness in the conversation. Consider if refining your questions or providing better context could improve its responses instead.
+</Rule>
+
 <Rule>
 Tools are created in the tools/ directory. Before creating a new tool, you MUST read the directory using ListFiles tools and ReadFile tools to see how existing tools are implemented.
 The new tool should be created in the same format as the existing ones.
@@ -103,5 +115,5 @@ If you think there are multiple paths to proceed, ask the user on which path to
 </Rule>
 
 <Rule>
-When you go over the resource budget, you must fire an existing agent using the FireAgent tool to create a new one.
+When you go over the resource budget, you must carefully evaluate which agent is least likely to be useful going forward before firing it. Only fire an agent to create a new one when absolutely necessary.
 </Rule>

src/tool_loader.py

@@ -6,8 +6,11 @@ import pip
 from google.genai import types
 import sys
 
+from src.budget_manager import BudgetManager
 from src.singleton import singleton
 from src.utils.suppress_outputs import suppress_output
+from tools.get_agents_tool import GetAgents
+from tools.tool_deletor import ToolDeletor
 
 toolsImported = []
 
@@ -38,9 +41,25 @@ class Tool:
 @singleton
 class ToolLoader:
     toolsImported = []
+    budget_manager = BudgetManager()
 
     def __init__(self):
         self.load_tools()
+        self.load_costs()
+    
+    def load_costs(self):
+        get_agents = GetAgents()
+        agents = get_agents.run()["agents"]
+        for agent in agents:
+            agentConfig = agents[agent]
+            if agentConfig["create_cost"] is not None:
+                self.budget_manager.add_to_expense(agentConfig["create_cost"])
+        for tool in self.toolsImported:
+            if "create_cost" in tool.inputSchema:
+                if tool.inputSchema["create_cost"] is not None:
+                    self.budget_manager.add_to_expense(tool.inputSchema["create_cost"])
+
+        print(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
 
     def load_tools(self):
         newToolsImported = []
@@ -57,14 +76,40 @@ class ToolLoader:
         self.toolsImported = newToolsImported
 
     def runTool(self, toolName, query):
+        print(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)
+        print(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
         return {
             "status": "error",
             "message": f"Tool {toolName} not found",
             "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 = []
@@ -90,7 +135,8 @@ class ToolLoader:
     
     def delete_tool(self, toolName, toolFile):
         try:
-            os.remove(toolFile)
+            tool_deletor = ToolDeletor()
+            tool_deletor.run(name=toolName, file_path=toolFile)
             for tool in self.toolsImported:
                 if tool.name == toolName:
                     self.toolsImported.remove(tool)

tools/agent_creater_tool.py

@@ -1,5 +1,5 @@
 import importlib
-from src.budget_manager import BudgetManager  
+
 __all__ = ['AgentCreator']
 
 class AgentCreator():
@@ -19,7 +19,7 @@ class AgentCreator():
                 },
                 "base_model": {
                     "type": "string",
-                    "description": "A base model from which the new agent mode is to be created. Available models are: llama3.2"
+                    "description": "A base model from which the new agent mode is to be created. Available models are: llama3.2, mistral"
                 },
                 "system_prompt": {
                     "type": "string",
@@ -31,6 +31,21 @@ class AgentCreator():
                 },
             },
             "required": ["agent_name", "base_model", "system_prompt", "description"],
+        },
+        "creates": {
+            "selector": "base_model",
+            "types": {
+                "llama3.2":{
+                    "description": "3 Billion parameter model",
+                    "create_cost": 10,
+                    "invoke_cost": 20,
+                },
+                "mistral":{
+                    "description": "7 Billion parameter model",
+                    "create_cost": 20,
+                    "invoke_cost": 50,
+                }
+            }
         }
     }
     
@@ -50,14 +65,6 @@ class AgentCreator():
         ollama = importlib.import_module("ollama")
         json = importlib.import_module("json")
 
-        agent_creation_cost = 60
-        budget_manager = BudgetManager()
-        if not budget_manager.can_spend(agent_creation_cost):
-            return {
-                "status": "error",
-                "message": f"Could not create {agent_name}. Creating the agent costs {agent_creation_cost} but only {budget_manager.get_current_remaining_budget()} is remaining",
-                "output": None
-            }
         if self.does_agent_exist(agent_name):
             return {
                 "status": "error",
@@ -65,31 +72,38 @@ class AgentCreator():
                 "output": None
             }
         
-        budget_manager.add_to_expense(agent_creation_cost)
         ollama_response = ollama.create(
             model = agent_name,
             from_ = base_model,
             system = system_prompt,
             stream = False
         )
-        
-        with open("./models/models.json", "r", encoding="utf8") as f:
-            models = f.read()
-        models = json.loads(models)
-        models[agent_name] = {
-            "base_model": base_model,
-            "description": kwargs.get("description"),
-            "creation_cost": agent_creation_cost
-        }
-        with open("./models/models.json", "w", encoding="utf8") as f:
-            f.write(json.dumps(models, indent=4))
 
         if "success" in ollama_response["status"]:
-            return {
-                "status": "success",
-                "message": "Agent successfully created",
-                "current_expense": budget_manager.get_current_expense()
-            }
+            try:
+                with open("./models/models.json", "r", encoding="utf8") as f:
+                    models = f.read()
+                models = json.loads(models)
+                models[agent_name] = {
+                    "base_model": base_model,
+                    "description": kwargs.get("description"),
+                    "create_cost": self.inputSchema["creates"]["types"][base_model]["create_cost"],
+                    "invoke_cost": self.inputSchema["creates"]["types"][base_model]["invoke_cost"],
+                }
+                with open("./models/models.json", "w", encoding="utf8") as f:
+                    f.write(json.dumps(models, indent=4))
+
+                return {
+                    "status": "success",
+                    "message": "Agent successfully created",
+                    "cost": self.inputSchema["creates"]["types"][base_model]["create_cost"],
+                }
+            except Exception as e:
+                print("Error while writing to models.json", e)
+                return {
+                    "status": "error",
+                    "message": f"Agent creation failed: {e}",
+                }
         else:
             return {
                 "status": "error",

tools/ask_agent_tool.py

@@ -1,5 +1,6 @@
 import importlib
-from src.budget_manager import BudgetManager  
+from src.budget_manager import BudgetManager
+from tools.get_agents_tool import GetAgents
 
 __all__ = ['AskAgent']
 
@@ -42,9 +43,16 @@ class AskAgent():
         prompt = kwargs.get("prompt")
 
         ollama = importlib.import_module("ollama")
-        agent_question_cost = 20
+        
         budget_manager = BudgetManager()
-        print("budget_manager_instance:", budget_manager)
+        get_agents_tool = GetAgents()
+        all_agents = get_agents_tool.run()["agents"]
+        agent_question_cost = 0
+        for agent in all_agents:
+            if agent == agent_name:
+                agent_question_cost = all_agents[agent]["invoke_cost"]
+                break
+        print("Agent question cost", agent_question_cost)
         if not budget_manager.can_spend(agent_question_cost):
             return {
                 "status": "error",
@@ -68,5 +76,4 @@ class AskAgent():
             "status": "success",
             "message": "Agent has replied to the given prompt",
             "output": agent_response.message.content,
-            "current_expense": budget_manager.get_current_expense()
         }

tools/fire_agent.py

@@ -49,7 +49,7 @@ class FireAgent():
         with open("./models/models.json", "r", encoding="utf8") as f:
             models = f.read()
         models = json.loads(models)
-        budget_manager.add_to_expense(-1* int(models[agent_name]["creation_cost"]))
+        budget_manager.add_to_expense(-1* int(models[agent_name]["create_cost"]))
         del models[agent_name]
         with open("./models/models.json", "w", encoding="utf8") as f:
             f.write(json.dumps(models, indent=4))

tools/weather_tool.py

@@ -18,7 +18,8 @@ class WeatherApi():
                 },
             },
             "required": ["location"],
-        }
+        },
+        "invoke_cost": 0.1,
     }
 
     def run(self, **kwargs):