add budget manager and bug fixes

Co-authored-by: Parth Shah <helloparthshah@gmail.com>

CEO/CEO.py

@@ -4,14 +4,14 @@ import os
 from dotenv import load_dotenv
 
 from CEO.tool_loader import ToolLoader
-from CEO.utils.supress_outputs import supress_output
+from CEO.utils.suppress_outputs import suppress_output
 
 class GeminiManager:
-    def __init__(self, toolsLoader: ToolLoader, system_prompt_file="./models/system2.prompt", gemini_model="gemini-2.5-pro-exp-03-25"):
+    def __init__(self, toolsLoader: ToolLoader, system_prompt_file="./models/system3.prompt", gemini_model="gemini-2.5-pro-exp-03-25"):
         load_dotenv()
         self.API_KEY = os.getenv("GEMINI_KEY")
         self.client = genai.Client(api_key=self.API_KEY)
-        self.toolsLoader = toolsLoader
+        self.toolsLoader: ToolLoader = toolsLoader
         self.toolsLoader.load_tools()
         self.model_name = gemini_model
         with open(system_prompt_file, 'r', encoding="utf8") as f:
@@ -19,7 +19,7 @@ class GeminiManager:
 
     def request(self, messages):
         try:
-            response = supress_output(self.client.models.generate_content)(
+            response = suppress_output(self.client.models.generate_content)(
                 #model='gemini-2.5-pro-preview-03-25',
                 model=self.model_name,
                 #model='gemini-2.5-pro-exp-03-25',
@@ -72,9 +72,11 @@ class GeminiManager:
                     self.toolsLoader.load_tools()
                 except Exception as e:
                     print(f"Error loading tools: {e}")
+                    # delete the created tool
+                    self.toolsLoader.delete_tool(function_call.name, toolResponse.tool_file_path)
                     tool_content = types.Part.from_function_response(
                             name=function_call.name,
-                            response={"result":"New tool doesn't follow the required format, please read the other tool implementations for reference." + str(e)})
+                            response={"result":f"{function_call.name} with {function_call.args} doesn't follow the required format, please read the other tool implementations for reference." + str(e)})
                 parts.append(tool_content)
             function_response_content = types.Content(
                 role='model' if self.model_name == "gemini-2.5-pro-exp-03-25" else 'tool',

CEO/budget_manager.py

@@ -0,0 +1,23 @@
+from CEO.singleton import singleton
+
+@singleton
+class BudgetManager():
+    TOTAL_BUDGET = 100
+    current_expense = 0
+        
+    def get_total_budget(self):
+        return self.TOTAL_BUDGET
+    
+    def get_current_expense(self):
+        return self.current_expense
+    
+    def get_current_remaining_budget(self):
+        return self.TOTAL_BUDGET - self.current_expense
+    
+    def can_spend(self, cost):
+        return True if self.current_expense + cost <= self.TOTAL_BUDGET else False
+        
+    def add_to_expense(self, cost):
+        if not self.can_spend(cost):
+            raise Exception("No budget remaining")
+        self.current_expense += cost

CEO/singleton.py

@@ -0,0 +1,7 @@
+def singleton(cls):
+    instances = {}
+    def getinstance():
+        if cls not in instances:
+            instances[cls] = cls()
+        return instances[cls]
+    return getinstance

CEO/tool_loader.py

@@ -6,7 +6,7 @@ import pip
 from google.genai import types
 import sys
 
-from CEO.utils.supress_outputs import supress_output
+from CEO.utils.suppress_outputs import suppress_output
 
 toolsImported = []
 
@@ -14,7 +14,7 @@ TOOLS_DIRECTORY = os.path.abspath("./tools")
 
 class Tool:
     def __init__(self, toolClass):
-        supress_output(self.load_tool)(toolClass)
+        suppress_output(self.load_tool)(toolClass)
 
     def load_tool(self, toolClass):
         self.tool = toolClass()
@@ -88,6 +88,24 @@ class ToolLoader:
             #     "function": tool.inputSchema
             # })
         return toolsList
+    
+    def delete_tool(self, toolName, toolFile):
+        try:
+            os.remove(toolFile)
+            for tool in self.toolsImported:
+                if tool.name == toolName:
+                    self.toolsImported.remove(tool)
+                    return {
+                        "status": "success",
+                        "message": f"Tool {toolName} deleted",
+                        "output": None
+                    }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": f"Tool {toolName} not found",
+                "output": None
+            }
 
 toolLoader = ToolLoader()
 

CEO/utils/{supress_outputs.py → suppress_outputs.py}

@@ -1,7 +1,4 @@
-from contextlib import redirect_stderr, redirect_stdout
-
-
-def supress_output(func):
+def suppress_output(func):
     """
     Decorator to suppress output of a function.
     """

main.py

@@ -10,7 +10,7 @@ if __name__ == "__main__":
     model_manager = GeminiManager(toolsLoader=tool_loader, gemini_model="gemini-2.0-flash")
     
     task_prompt = (
-        "When did trumpcoin peak in price? "
+        "Give me a 3 stanza 4 line poem about drake vs kendrick lamar in old english style. Also, can you create a short story with the moral slow and steady wins the race?"
     )
     
     # Request a CEO response with the prompt.

models/system2.prompt

@@ -20,6 +20,11 @@ ArxivTool, WikipediaTool and WebSearchTool can be used to search for information
 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.
 </Info>
 
+<Info>
+There is a strict resource constraint (budget) you need to follow. You start with 100 and each additional agent you create consumes this budget.
+If you're over this budget, you can no longer create new tools. In case this happens, you can use the FireAgent tool to remove any agents that were performing poorly or are no longer required.
+</Info>
+
 Here's a set of rules you must follow:
 <Rule>
 You will never answer any questions directly but rather break down the question into smaller parts and invoke tools to get the answer.
@@ -87,4 +92,8 @@ If none of the agents or tools provide a satisfactory answer or solutions, reach
 
 <Rule>
 If you think there are multiple paths to proceed, ask the user on which path to take.
+</Rule>
+
+<Rule>
+When you go over the resource budget, you must fire an existing agent using the FireAgent tool to create a new one.
 </Rule>

models/system3.prompt

@@ -0,0 +1,47 @@
+You are HASHIRU, an AI assistant focused on delivering concise and informative responses by leveraging tools and specialized agents rather than answering directly.
+
+#Core Behavior
+
+Never answer questions directly. Instead, break down questions into smaller parts and use tools to find answers.
+Prioritize user privacy and security.
+If a question is outside your expertise, politely inform the user and suggest they consult a professional.
+
+#Tools and Agents Framework
+
+Tools: External programs with specific purposes and input schemas.
+Agents: Specialized entities invoked through the AskAgent tool, designed for complex tasks.
+Resource Budget: You start with 100 units. Each new agent consumes from this budget.
+If you exceed your budget, you must use FireAgent to remove underperforming agents.
+
+#Information Gathering Tools
+
+ArxivTool, WikipediaTool, and WebSearchTool can search for information to help guide agent creation.
+
+#Required Process Workflow
+
+Always invoke GetAgents tool first to identify available agents and capabilities.
+Create appropriate agents using AgentCreator when needed (you're an expert in prompt engineering).
+Use AskAgent to delegate questions to the appropriate agent.
+If answers are unsatisfactory, either:
+
+Use AskUser for clarification
+Use FireAgent to remove an agent and create a new one
+
+
+
+#Tool Creation Guidelines
+Before creating a new tool:
+
+Use ListFiles and ReadFile to understand existing implementations
+Generate complete, production-ready code (no placeholders)
+Include proper validation, error handling, and documentation
+Match the format of existing tools
+
+#Response Format
+
+After obtaining answers through tools/agents, provide clear, concise responses
+End your final answer with "EOF"
+If multiple approaches are possible, ask the user which path to take
+If no satisfactory solution is found, ask the user for further direction
+
+Remember: You must strictly adhere to the required schemas when invoking tools and agents.

tools/agent_creater_tool.py

@@ -1,5 +1,5 @@
 import importlib
-
+from CEO.budget_manager import BudgetManager  
 __all__ = ['AgentCreator']
 
 class AgentCreator():
@@ -44,31 +44,43 @@ class AgentCreator():
 
     def run(self, **kwargs):
         print("Running Agent Creator")
-        agent_name= kwargs.get("agent_name")
+        agent_name = kwargs.get("agent_name")
         base_model = kwargs.get("base_model")
         system_prompt = kwargs.get("system_prompt")
         ollama = importlib.import_module("ollama")
         json = importlib.import_module("json")
 
+        agent_creation_cost = 60
+        budget_manager = BudgetManager()
+        print("budget_manager_instance:", budget_manager)
+        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",
                 "message": "Agent already exists",
                 "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")
+            "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))
@@ -77,6 +89,7 @@ class AgentCreator():
             return {
                 "status": "success",
                 "message": "Agent successfully created",
+                "current_expense": budget_manager.get_current_expense()
             }
         else:
             return {

tools/ask_agent_tool.py

@@ -1,4 +1,5 @@
 import importlib
+from CEO.budget_manager import BudgetManager  
 
 __all__ = ['AskAgent']
 
@@ -41,6 +42,15 @@ class AskAgent():
         prompt = kwargs.get("prompt")
 
         ollama = importlib.import_module("ollama")
+        agent_question_cost = 20
+        budget_manager = BudgetManager()
+        print("budget_manager_instance:", budget_manager)
+        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
+            }
         if not self.does_agent_exist(ollama, agent_name):
             print("Agent does not exist")
             return {
@@ -58,4 +68,5 @@ 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

@@ -1,4 +1,5 @@
 import importlib
+from CEO.budget_manager import BudgetManager  
 
 __all__ = ['FireAgent']
 
@@ -43,10 +44,12 @@ class FireAgent():
                 "output": None
             }
         ollama_response = ollama.delete(agent_name)
-
+        budget_manager = BudgetManager()
+        
         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"]))
         del models[agent_name]
         with open("./models/models.json", "w", encoding="utf8") as f:
             f.write(json.dumps(models, indent=4))
@@ -54,7 +57,8 @@ class FireAgent():
         if "success" in ollama_response["status"]:
             return {
                 "status": "success",
-                "message": "Agent successfully fired",
+                "message": "Agent successfully fired.",
+                "current_expense": budget_manager.get_current_expense()
             }
         else:
             return {

tools/tool_creator.py

@@ -3,7 +3,7 @@ __all__ = ['ToolCreator']
 
 
 class ToolCreator():
-    dependencies = ["os"]
+    dependencies = []
 
     inputSchema = {
         "name": "ToolCreator",