QOL updates and refactoring. Also fixed the tool/agent budgeting

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

bench/benchmarking_connections.py

@@ -72,7 +72,7 @@ def benchmark_nyt_connections(num_samples=20, categories=None):
         if categories and sample["category"] not in categories:
             continue
         print(f"Sample {i}: {sample['contest']}")
-        prompt = f"Given the following words, group them into 4 categories of 4 words each:\n{' '.join(sample['words'])}\n\n"
+        prompt = f"Given the following words, group them into 4 categories of 4 words each:\n{' '.join(sample['words'])}\n\n Once you've solved it, final output should be in the following format Group 1: word1, word2, word3, word4\nGroup 2: ..."
         start_time = time.time()
         response = client.predict(messages=[{"role": "user", "content": prompt}], api_name="/run")
         end_time = time.time()

default_tools/fire_agent.py

@@ -1,4 +1,3 @@
-from src.budget_manager import BudgetManager  
 from src.agent_manager import AgentManager
 
 __all__ = ['FireAgent']

default_tools/get_budget.py

@@ -0,0 +1,35 @@
+
+__all__ = ['GetBudget']
+
+from src.budget_manager import BudgetManager
+
+
+class GetBudget():
+    dependencies = []
+
+    inputSchema = {
+        "name": "GetBudget",
+        "description": "Retrieves the current budget available.",
+        "parameters": {
+            "type": "object",
+            "properties":{},
+            "required": [],
+        },
+    }
+
+
+    def run(self, **kwargs):
+        budget_manager = BudgetManager()
+        total_budget = budget_manager.get_total_budget()
+        current_expense = budget_manager.get_current_expense()
+        current_remaining_budget = budget_manager.get_current_remaining_budget()
+        return {
+            "status": "success",
+            "message": "Budget retrieved successfully",
+            "output": {
+                "total_budget": total_budget,
+                "current_expense": current_expense,
+                "current_remaining_budget": current_remaining_budget
+            }
+        }
+        

main.py

@@ -1,6 +1,6 @@
 from google.genai import types
 from src.manager import GeminiManager
-from src.tool_loader import ToolLoader
+from src.tool_manager import ToolManager
 import gradio as gr
 import time
 import base64
@@ -22,11 +22,7 @@ _header_html = f"""
 """
 
 if __name__ == "__main__":
-    # Define the tool metadata for orchestration.
-    # Load the tools using the ToolLoader class.
-    tool_loader = ToolLoader()
-
-    model_manager = GeminiManager(toolsLoader=tool_loader, gemini_model="gemini-2.0-flash")
+    model_manager = GeminiManager(gemini_model="gemini-2.0-flash")
     
     def user_message(msg: str, history: list) -> tuple[str, list]:
         """Adds user message to chat history"""
@@ -54,7 +50,7 @@ if __name__ == "__main__":
     css = """
     #title-row { background: #2c2c2c; border-radius: 8px; padding: 8px; }
     """
-    with gr.Blocks(css=css) as demo:
+    with gr.Blocks(css=css, fill_width=True, fill_height=True) as demo:
         with gr.Row():
             gr.HTML(_header_html)
             model_dropdown = gr.Dropdown(
@@ -67,7 +63,7 @@ if __name__ == "__main__":
                     ],
                     value="HASHIRU",
                     # label="HASHIRU",
-                    scale=4,
+                    scale=2,
                     interactive=True,
             )
 

src/agent_manager.py

@@ -95,9 +95,9 @@ class GeminiAgent(Agent):
         self.messages = []
 @singleton
 class AgentManager():
-    budget_manager = BudgetManager()
+    budget_manager: BudgetManager = BudgetManager()
     def __init__(self):
-        self._agents = {}
+        self._agents: Dict[str, Agent] = {}
         self._agent_types ={
             "ollama": OllamaAgent,
             "gemini": GeminiAgent
@@ -111,19 +111,15 @@ class AgentManager():
         if agent_name in self._agents:
             raise ValueError(f"Agent {agent_name} already exists")
         
-        agent_type = self._get_agent_type(base_model)
-        agent_class = self._agent_types.get(agent_type)
-        
-        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 )
+        self._agents[agent_name] = self.create_agent_class(
+            agent_name, 
+            base_model, 
+            system_prompt, 
+            description=description,
+            create_cost=create_cost,
+            invoke_cost=invoke_cost,
+            **additional_params  # For any future parameters we might want to add
+        )
         
         #save agent to file
         self._save_agent(
@@ -135,9 +131,28 @@ class AgentManager():
             invoke_cost=invoke_cost,
             **additional_params  # For any future parameters we might want to add
         )
-        
         return (self._agents[agent_name], self.budget_manager.get_current_remaining_budget())
     
+    def validate_budget(self, amount: float) -> None:
+        if not self.budget_manager.can_spend(amount):
+            raise ValueError(f"Do not have enough budget to create the tool. "
+                        +f"Creating the tool costs {amount} but only {self.budget_manager.get_current_remaining_budget()} is remaining")
+        
+    def create_agent_class(self, agent_name: str, base_model: str, system_prompt: str, description: str = "", create_cost: float = 0, invoke_cost: float = 0,
+                    **additional_params) -> Agent:
+        agent_type = self._get_agent_type(base_model)
+        agent_class = self._agent_types.get(agent_type)
+        
+        if not agent_class:
+            raise ValueError(f"Unsupported base model {base_model}")
+        
+        self.validate_budget(create_cost)
+        
+        self.budget_manager.add_to_expense(create_cost)
+        # create agent
+        return agent_class(agent_name, base_model, system_prompt, create_cost,invoke_cost)
+        
+
     def get_agent(self, agent_name: str) -> Agent:
         """Get existing agent by name"""
         if agent_name not in self._agents:
@@ -167,12 +182,10 @@ class AgentManager():
             return {}
     
     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()
+        agent = self.get_agent(agent_name)
         
+        self.budget_manager.remove_from_expense(agent.creation_cost)
+        agent.delete_agent()
         
         del self._agents[agent_name]
         try:
@@ -188,12 +201,10 @@ class AgentManager():
         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")
+        agent = self.get_agent(agent_name)
         
+        self.validate_budget(agent.invoke_cost)
+
         response = agent.ask_agent(prompt)        
         return (response, self.budget_manager.get_current_remaining_budget())
     
@@ -264,6 +275,15 @@ class AgentManager():
                 
                 if manager_class:
                     # Create the agent with the appropriate manager class
+                    self._agents[name] = self.create_agent_class(
+                        name, 
+                        base_model, 
+                        system_prompt, 
+                        description=data.get("description", ""),
+                        create_cost=creation_cost,
+                        invoke_cost=invoke_cost,
+                        **data.get("additional_params", {})
+                    )
                     self._agents[name] = manager_class(
                         name, 
                         base_model,

src/budget_manager.py

@@ -20,4 +20,9 @@ class BudgetManager():
     def add_to_expense(self, cost):
         if not self.can_spend(cost):
             raise Exception("No budget remaining")
-        self.current_expense += cost
+        self.current_expense += cost
+    
+    def remove_from_expense(self, cost):
+        if self.current_expense - cost < 0:
+            raise Exception("Cannot remove more than current expense")
+        self.current_expense -= cost

src/manager.py

@@ -4,7 +4,7 @@ from google.genai.types import *
 import os
 from dotenv import load_dotenv
 import sys
-from src.tool_loader import ToolLoader
+from src.tool_manager import ToolManager
 from src.utils.suppress_outputs import suppress_output
 import logging
 import gradio as gr
@@ -16,17 +16,29 @@ logger.addHandler(handler)
 
 
 class GeminiManager:
-    def __init__(self, toolsLoader: ToolLoader, system_prompt_file="./models/system3.prompt", gemini_model="gemini-2.5-pro-exp-03-25"):
+    def __init__(self, toolsLoader: ToolManager = None,
+                 system_prompt_file="./models/system3.prompt",
+                 gemini_model="gemini-2.5-pro-exp-03-25",
+                 local_only=False, allow_tool_creation=True,
+                 cloud_only=False, use_economy=True):
         load_dotenv()
+        self.toolsLoader: ToolManager = toolsLoader
+        if not toolsLoader:
+            self.toolsLoader: ToolManager = ToolManager()
+
+        self.local_only = local_only
+        self.allow_tool_creation = allow_tool_creation
+        self.cloud_only = cloud_only
+        self.use_economy = use_economy
+
         self.API_KEY = os.getenv("GEMINI_KEY")
         self.client = genai.Client(api_key=self.API_KEY)
-        self.toolsLoader: ToolLoader = toolsLoader
         self.toolsLoader.load_tools()
         self.model_name = gemini_model
         with open(system_prompt_file, 'r', encoding="utf8") as f:
             self.system_prompt = f.read()
         self.messages = []
-
+        
     def generate_response(self, messages):
         return self.client.models.generate_content(
             model=self.model_name,

src/{tool_loader.py → tool_manager.py}

@@ -2,6 +2,7 @@ import importlib
 import importlib.util
 import os
 import types
+from typing import List
 import pip
 from google.genai import types
 import sys
@@ -27,7 +28,17 @@ class Tool:
         self.name = self.inputSchema["name"]
         self.description = self.inputSchema["description"]
         self.dependencies = self.tool.dependencies
-        for package in self.tool.dependencies:
+        self.create_cost = None
+        self.invoke_cost = None
+        if "create_cost" in self.tool.inputSchema:
+            self.create_cost = self.tool.inputSchema["create_cost"]
+        if "invoke_cost" in self.tool.inputSchema:
+            self.invoke_cost = self.tool.inputSchema["invoke_cost"]
+        if self.dependencies:
+            self.install_dependencies()
+    
+    def install_dependencies(self):
+        for package in self.dependencies:
             try:
                 __import__(package.split('==')[0])
             except ImportError:
@@ -40,25 +51,12 @@ class Tool:
         return self.tool.run(**query)
 
 @singleton
-class ToolLoader:
-    toolsImported = []
-    budget_manager = BudgetManager()
+class ToolManager:
+    toolsImported: List[Tool] = []
+    budget_manager: BudgetManager = 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"])
 
         output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
 
@@ -75,12 +73,16 @@ class ToolLoader:
                     toolClass = getattr(foo, class_name)
                     toolObj = Tool(toolClass)
                     newToolsImported.append(toolObj)
+                    if toolObj.create_cost is not None:
+                        self.budget_manager.add_to_expense(toolObj.create_cost)
         self.toolsImported = newToolsImported
 
     def runTool(self, toolName, query):
         output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
         for tool in self.toolsImported:
             if tool.name == toolName:
+                if tool.invoke_cost is not None:
+                    self.budget_manager.add_to_expense(tool.invoke_cost)
                 return tool.run(query)
         output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
         return {
@@ -118,6 +120,9 @@ class ToolLoader:
             tool_deletor.run(name=toolName, file_path=toolFile)
             for tool in self.toolsImported:
                 if tool.name == toolName:
+                    # remove budget for the tool
+                    if tool.create_cost is not None:
+                        self.budget_manager.remove_from_expense(tool.create_cost)
                     self.toolsImported.remove(tool)
                     return {
                         "status": "success",
@@ -129,10 +134,4 @@ class ToolLoader:
                 "status": "error",
                 "message": f"Tool {toolName} not found",
                 "output": None
-            }
-
-# toolLoader = ToolLoader()
-
-# Example usage
-# print(toolLoader.getTools())
-# print(toolLoader.runTool("AgentCreator", {"agent_name": "Kunla","base_model":"llama3.2","system_prompt": "You love making the indian dish called Kulcha. You declare that in every conversation you have in a witty way." }))
+            }