Implemented mode selection

Co-authored-by: Kunal Pai <kunpai@users.noreply.github.com>

main.py

@@ -1,7 +1,9 @@
+from typing import List
 import gradio as gr
 
 import base64
-from src.manager.manager import GeminiManager
+from src.manager.manager import GeminiManager, Mode
+from enum import Enum
 
 _logo_bytes = open("HASHIRU_LOGO.png", "rb").read()
 _logo_b64 = base64.b64encode(_logo_bytes).decode()
@@ -25,8 +27,6 @@ css = """
 
 
 def run_model(message, history):
-    print(f"User: {message}")
-    print(f"History: {history}")
     if 'text' in message:
         history.append({
             "role": "user",
@@ -38,33 +38,30 @@ def run_model(message, history):
                 "role": "user",
                 "content": (file,)
             })
-    print(f"History: {history}")
     yield "", history
     for messages in model_manager.run(history):
         yield "", messages
 
 
-def update_model(model_name):
-    print(f"Model changed to: {model_name}")
-    pass
+with gr.Blocks(css=css, fill_width=True, fill_height=True) as demo:
+    model_manager = GeminiManager(
+        gemini_model="gemini-2.0-flash", modes=[mode for mode in Mode])
 
+    def update_model(modeIndexes: List[int]):
+        modes = [Mode(i+1) for i in modeIndexes]
+        print(f"Selected modes: {modes}")
+        model_manager.set_modes(modes)
 
-with gr.Blocks(css=css, fill_width=True, fill_height=True) as demo:
-    model_manager = GeminiManager(gemini_model="gemini-2.0-flash")
-    
     with gr.Column(scale=1):
         with gr.Row(scale=0):
             gr.Markdown(_header_html)
             model_dropdown = gr.Dropdown(
-                choices=[
-                    "HASHIRU",
-                    "Static-HASHIRU",
-                    "Cloud-Only HASHIRU",
-                    "Local-Only HASHIRU",
-                    "No-Economy HASHIRU",
-                ],
-                value="HASHIRU",
+                choices=[mode.name for mode in Mode],
+                value=model_manager.get_current_modes,
                 interactive=True,
+                type="index",
+                multiselect=True,
+                label="Select Modes",
             )
 
             model_dropdown.change(

src/manager/agent_manager.py

@@ -16,12 +16,13 @@ from src.manager.budget_manager import BudgetManager
 MODEL_PATH = "./src/models/"
 MODEL_FILE_PATH = "./src/models/models.json"
 
+
 class Agent(ABC):
-    
-    def __init__(self, agent_name: str, 
-                 base_model: str, 
-                 system_prompt: str, 
-                 create_resource_cost: int, 
+
+    def __init__(self, agent_name: str,
+                 base_model: str,
+                 system_prompt: str,
+                 create_resource_cost: int,
                  invoke_resource_cost: int,
                  create_expense_cost: int = 0,
                  invoke_expense_cost: int = 0,):
@@ -33,22 +34,27 @@ class Agent(ABC):
         self.create_expense_cost = create_expense_cost
         self.invoke_expense_cost = invoke_expense_cost
         self.create_model()
-        
+
     @abstractmethod
     def create_model(self) -> None:
         """Create and Initialize agent"""
         pass
-    
+
     @abstractmethod
     def ask_agent(self, prompt: str) -> str:
         """ask agent a question"""
         pass
-    
+
     @abstractmethod
-    def delete_agent(self) ->None:
+    def delete_agent(self) -> None:
         """delete agent"""
         pass
-    
+
+    @abstractmethod
+    def get_type(self) -> None:
+        """get agent type"""
+        pass
+
     def get_costs(self):
         return {
             "create_resource_cost": self.create_resource_cost,
@@ -56,58 +62,68 @@ class Agent(ABC):
             "create_expense_cost": self.create_expense_cost,
             "invoke_expense_cost": self.invoke_expense_cost
         }
-    
+
+
 class OllamaAgent(Agent):
-    
+    type = "local"
+
     def create_model(self):
         ollama_response = ollama.create(
-            model = self.agent_name,
-            from_ = self.base_model,
-            system = self.system_prompt,
-            stream = False
+            model=self.agent_name,
+            from_=self.base_model,
+            system=self.system_prompt,
+            stream=False
         )
-    
+
     def ask_agent(self, prompt):
         output_assistant_response(f"Asked Agent {self.agent_name} a question")
         agent_response = ollama.chat(
             model=self.agent_name,
             messages=[{"role": "user", "content": prompt}],
         )
-        output_assistant_response(f"Agent {self.agent_name} answered with {agent_response.message.content}")
+        output_assistant_response(
+            f"Agent {self.agent_name} answered with {agent_response.message.content}")
         return agent_response.message.content
-    
+
     def delete_agent(self):
         ollama.delete(self.agent_name)
-    
+
+    def get_type(self):
+        return self.type
+
+
 class GeminiAgent(Agent):
-    def __init__(self, 
-                 agent_name: str, 
-                 base_model: str, 
-                 system_prompt: str, 
-                 create_resource_cost: int, 
+    type = "cloud"
+
+    def __init__(self,
+                 agent_name: str,
+                 base_model: str,
+                 system_prompt: str,
+                 create_resource_cost: int,
                  invoke_resource_cost: int,
                  create_expense_cost: int = 0,
                  invoke_expense_cost: int = 0,):
         load_dotenv()
         self.api_key = os.getenv("GEMINI_KEY")
         if not self.api_key:
-            raise ValueError("Google API key is required for Gemini models. Set GOOGLE_API_KEY environment variable or pass api_key parameter.")
-        
+            raise ValueError(
+                "Google API key is required for Gemini models. Set GOOGLE_API_KEY environment variable or pass api_key parameter.")
+
         # Initialize the Gemini API
         self.client = genai.Client(api_key=self.api_key)
-        
+
         # Call parent constructor after API setup
-        super().__init__(agent_name, 
-                         base_model, 
-                         system_prompt, 
-                         create_resource_cost, 
+        super().__init__(agent_name,
+                         base_model,
+                         system_prompt,
+                         create_resource_cost,
                          invoke_resource_cost,
                          create_expense_cost,
                          invoke_expense_cost)
 
     def create_model(self):
         self.messages = []
-    
+
     def ask_agent(self, prompt):
         response = self.client.models.generate_content(
             model=self.base_model,
@@ -117,11 +133,17 @@ class GeminiAgent(Agent):
             )
         )
         return response.text
-    
+
     def delete_agent(self):
         self.messages = []
 
+    def get_type(self):
+        return self.type
+
+
 class GroqAgent(Agent):
+    type = "cloud"
+
     def __init__(
         self,
         agent_name: str,
@@ -158,34 +180,65 @@ class GroqAgent(Agent):
         # No delete support for Groq
         pass
 
+    def get_type(self):
+        return self.type
+
+
 @singleton
 class AgentManager():
     budget_manager: BudgetManager = BudgetManager()
+    is_creation_enabled: bool = True
+    is_cloud_invocation_enabled: bool = True
+    is_local_invocation_enabled: bool = True
+
     def __init__(self):
         self._agents: Dict[str, Agent] = {}
-        self._agent_types ={
+        self._agent_types = {
             "ollama": OllamaAgent,
             "gemini": GeminiAgent,
             "groq": GroqAgent,
         }
-        
+
         self._load_agents()
-    
-    def create_agent(self, agent_name: str, 
-                     base_model: str, system_prompt: str, 
-                     description: str = "", create_resource_cost: float = 0, 
-                     invoke_resource_cost: float = 0, 
+
+    def set_creation_mode(self, status: bool):
+        self.is_creation_enabled = status
+        if status:
+            output_assistant_response("Agent creation mode is enabled.")
+        else:
+            output_assistant_response("Agent creation mode is disabled.")
+
+    def set_cloud_invocation_mode(self, status: bool):
+        self.is_cloud_invocation_enabled = status
+        if status:
+            output_assistant_response("Cloud invocation mode is enabled.")
+        else:
+            output_assistant_response("Cloud invocation mode is disabled.")
+
+    def set_local_invocation_mode(self, status: bool):
+        self.is_local_invocation_enabled = status
+        if status:
+            output_assistant_response("Local invocation mode is enabled.")
+        else:
+            output_assistant_response("Local invocation mode is disabled.")
+
+    def create_agent(self, agent_name: str,
+                     base_model: str, system_prompt: str,
+                     description: str = "", create_resource_cost: float = 0,
+                     invoke_resource_cost: float = 0,
                      create_expense_cost: float = 0,
                      invoke_expense_cost: float = 0,
-                 **additional_params) -> Tuple[Agent, int]:
-        
+                     **additional_params) -> Tuple[Agent, int]:
+        if not self.is_creation_enabled:
+            raise ValueError("Agent creation mode is disabled.")
+
         if agent_name in self._agents:
             raise ValueError(f"Agent {agent_name} already exists")
-        
+
         self._agents[agent_name] = self.create_agent_class(
-            agent_name, 
-            base_model, 
-            system_prompt, 
+            agent_name,
+            base_model,
+            system_prompt,
             description=description,
             create_resource_cost=create_resource_cost,
             invoke_resource_cost=invoke_resource_cost,
@@ -193,12 +246,12 @@ class AgentManager():
             invoke_expense_cost=invoke_expense_cost,
             **additional_params  # For any future parameters we might want to add
         )
-        
-        #save agent to file
+
+        # save agent to file
         self._save_agent(
-            agent_name, 
-            base_model, 
-            system_prompt, 
+            agent_name,
+            base_model,
+            system_prompt,
             description=description,
             create_resource_cost=create_resource_cost,
             invoke_resource_cost=invoke_resource_cost,
@@ -206,47 +259,47 @@ class AgentManager():
             invoke_expense_cost=invoke_expense_cost,
             **additional_params  # For any future parameters we might want to add
         )
-        return (self._agents[agent_name], 
-                self.budget_manager.get_current_remaining_resource_budget(), 
+        return (self._agents[agent_name],
+                self.budget_manager.get_current_remaining_resource_budget(),
                 self.budget_manager.get_current_remaining_expense_budget())
-    
-    def validate_budget(self, 
-                        resource_cost: float=0,
-                        expense_cost: float=0) -> None:
+
+    def validate_budget(self,
+                        resource_cost: float = 0,
+                        expense_cost: float = 0) -> None:
         if not self.budget_manager.can_spend_resource(resource_cost):
             raise ValueError(f"Do not have enough resource budget to create/use the agent. "
-                        +f"Creating/Using the agent costs {resource_cost} but only {self.budget_manager.get_current_remaining_resource_budget()} is remaining")
+                             + f"Creating/Using the agent costs {resource_cost} but only {self.budget_manager.get_current_remaining_resource_budget()} is remaining")
         if not self.budget_manager.can_spend_expense(expense_cost):
             raise ValueError(f"Do not have enough expense budget to create/use the agent. "
-                        +f"Creating/Using the agent costs {expense_cost} but only {self.budget_manager.get_current_remaining_expense_budget()} is remaining")
-        
-    def create_agent_class(self, 
+                             + f"Creating/Using the agent costs {expense_cost} but only {self.budget_manager.get_current_remaining_expense_budget()} is remaining")
+
+    def create_agent_class(self,
                            agent_name: str,
-                           base_model: str, 
-                           system_prompt: str, 
-                           description: str = "", 
-                           create_resource_cost: float = 0, 
+                           base_model: str,
+                           system_prompt: str,
+                           description: str = "",
+                           create_resource_cost: float = 0,
                            invoke_resource_cost: float = 0,
                            create_expense_cost: float = 0,
                            invoke_expense_cost: float = 0,
-                    **additional_params) -> Agent:
+                           **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}")
-        
-        created_agent = agent_class(agent_name, 
-                                    base_model, 
-                                    system_prompt, 
+
+        created_agent = agent_class(agent_name,
+                                    base_model,
+                                    system_prompt,
                                     create_resource_cost,
                                     invoke_resource_cost,
                                     create_expense_cost,
                                     invoke_expense_cost,)
-        
-        self.validate_budget(create_resource_cost, 
+
+        self.validate_budget(create_resource_cost,
                              create_expense_cost)
-        
+
         self.budget_manager.add_to_resource_budget(create_resource_cost)
         self.budget_manager.add_to_expense_budget(create_expense_cost)
         # create agent
@@ -257,14 +310,14 @@ class AgentManager():
         if agent_name not in self._agents:
             raise ValueError(f"Agent {agent_name} does not exists")
         return self._agents[agent_name]
-        
+
     def list_agents(self) -> dict:
         """Return agent information (name, description, costs)"""
         try:
             if os.path.exists(MODEL_FILE_PATH):
                 with open(MODEL_FILE_PATH, "r", encoding="utf8") as f:
                     full_models = json.loads(f.read())
-                    
+
                 # Create a simplified version with only the description and costs
                 simplified_agents = {}
                 for name, data in full_models.items():
@@ -282,47 +335,56 @@ class AgentManager():
         except Exception as e:
             output_assistant_response(f"Error listing agents: {e}")
             return {}
-    
+
     def delete_agent(self, agent_name: str) -> int:
-        agent = self.get_agent(agent_name)
-        
-        self.budget_manager.remove_from_resource_expense(agent.create_resource_cost)
+        agent: Agent = self.get_agent(agent_name)
+
+        self.budget_manager.remove_from_resource_expense(
+            agent.create_resource_cost)
         agent.delete_agent()
-        
+
         del self._agents[agent_name]
         try:
             if os.path.exists(MODEL_FILE_PATH):
                 with open(MODEL_FILE_PATH, "r", encoding="utf8") as f:
                     models = json.loads(f.read())
-                    
+
                 del models[agent_name]
                 with open(MODEL_FILE_PATH, "w", encoding="utf8") as f:
                     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_resource_budget(), 
+        return (self.budget_manager.get_current_remaining_resource_budget(),
                 self.budget_manager.get_current_remaining_expense_budget())
-    
-    def ask_agent(self, agent_name: str, prompt: str) -> Tuple[str,int]:
-        agent = self.get_agent(agent_name)
-        
-        self.validate_budget(agent.invoke_resource_cost, 
+
+    def ask_agent(self, agent_name: str, prompt: str) -> Tuple[str, int]:
+        agent: Agent = self.get_agent(agent_name)
+        print(agent.get_type())
+        print(agent_name)
+        print(self.is_local_invocation_enabled, self.is_cloud_invocation_enabled)
+        if not self.is_local_invocation_enabled and agent.get_type() == "local":
+            raise ValueError("Local invocation mode is disabled.")
+
+        if not self.is_cloud_invocation_enabled and agent.get_type() == "cloud":
+            raise ValueError("Cloud invocation mode is disabled.")
+
+        self.validate_budget(agent.invoke_resource_cost,
                              agent.invoke_expense_cost)
-        
+
         self.budget_manager.add_to_expense_budget(agent.invoke_expense_cost)
 
-        response = agent.ask_agent(prompt)        
-        return (response, 
-                self.budget_manager.get_current_remaining_resource_budget(), 
+        response = agent.ask_agent(prompt)
+        return (response,
+                self.budget_manager.get_current_remaining_resource_budget(),
                 self.budget_manager.get_current_remaining_expense_budget())
-    
-    def _save_agent(self, 
-                    agent_name: str, 
-                    base_model: str, 
-                    system_prompt: str, 
-                    description: str = "", 
-                    create_resource_cost: float = 0, 
-                    invoke_resource_cost: float = 0, 
+
+    def _save_agent(self,
+                    agent_name: str,
+                    base_model: str,
+                    system_prompt: str,
+                    description: str = "",
+                    create_resource_cost: float = 0,
+                    invoke_resource_cost: float = 0,
                     create_expense_cost: float = 0,
                     invoke_expense_cost: float = 0,
                     **additional_params) -> None:
@@ -330,14 +392,14 @@ class AgentManager():
         try:
             # Ensure the directory exists
             os.makedirs(MODEL_PATH, exist_ok=True)
-            
+
             # Read existing models file or create empty dict if it doesn't exist
             try:
                 with open(MODEL_FILE_PATH, "r", encoding="utf8") as f:
                     models = json.loads(f.read())
             except (FileNotFoundError, json.JSONDecodeError):
                 models = {}
-            
+
             # Update the models dict with the new agent
             models[agent_name] = {
                 "base_model": base_model,
@@ -348,19 +410,19 @@ class AgentManager():
                 "create_expense_cost": create_expense_cost,
                 "invoke_expense_cost": invoke_expense_cost,
             }
-            
+
             # Add any additional parameters that were passed
             for key, value in additional_params.items():
                 models[agent_name][key] = value
-            
+
             # Write the updated models back to the file
             with open(MODEL_FILE_PATH, "w", encoding="utf8") as f:
                 f.write(json.dumps(models, indent=4))
-                
+
         except Exception as e:
             output_assistant_response(f"Error saving agent {agent_name}: {e}")
 
-    def _get_agent_type(self, base_model)->str:
+    def _get_agent_type(self, base_model) -> str:
 
         if base_model == "llama3.2":
             return "ollama"
@@ -374,16 +436,16 @@ class AgentManager():
             return "groq"
         else:
             return "unknown"
-    
+
     def _load_agents(self) -> None:
         """Load agent configurations from disk"""
         try:
             if not os.path.exists(MODEL_FILE_PATH):
                 return
-                
+
             with open(MODEL_FILE_PATH, "r", encoding="utf8") as f:
                 models = json.loads(f.read())
-            
+
             for name, data in models.items():
                 if name in self._agents:
                     continue
@@ -395,13 +457,13 @@ class AgentManager():
                 invoke_expense_cost = data.get("invoke_expense_cost", 0)
                 model_type = self._get_agent_type(base_model)
                 manager_class = self._agent_types.get(model_type)
-                
+
                 if manager_class:
                     # Create the agent with the appropriate manager class
                     self._agents[name] = self.create_agent_class(
-                        name, 
-                        base_model, 
-                        system_prompt, 
+                        name,
+                        base_model,
+                        system_prompt,
                         description=data.get("description", ""),
                         create_resource_cost=create_resource_cost,
                         invoke_resource_cost=invoke_resource_cost,
@@ -410,7 +472,7 @@ class AgentManager():
                         **data.get("additional_params", {})
                     )
                     self._agents[name] = manager_class(
-                        name, 
+                        name,
                         base_model,
                         system_prompt,
                         create_resource_cost,
@@ -419,4 +481,4 @@ class AgentManager():
                         invoke_expense_cost,
                     )
         except Exception as e:
-            output_assistant_response(f"Error loading agents: {e}")
+            output_assistant_response(f"Error loading agents: {e}")

src/manager/budget_manager.py

@@ -9,11 +9,27 @@ class BudgetManager():
     total_expense_budget = 10
     current_expense = 0
     is_budget_initialized = False
+    is_resource_budget_enabled = True
+    is_expense_budget_enabled = True
     
     def __init__(self):
         if not self.is_budget_initialized:
             self.total_resource_budget = self.calculate_total_budget()
             self.is_budget_initialized = True
+    
+    def set_resource_budget_status(self, status: bool):
+        self.is_enabled = status
+        if status:
+            print("Budget manager is enabled.")
+        else:
+            print("Budget manager is disabled.")
+    
+    def set_expense_budget_status(self, status: bool):
+        self.is_expense_budget_enabled = status
+        if status:
+            print("Expense budget manager is enabled.")
+        else:
+            print("Expense budget manager is disabled.")
         
     def calculate_total_budget(self)-> int:
         total_mem = 0
@@ -43,14 +59,20 @@ class BudgetManager():
         return self.total_resource_budget - self.current_resource_usage
     
     def can_spend_resource(self, cost):
+        if not self.is_resource_budget_enabled:
+            return True
         return True if self.current_resource_usage + cost <= self.total_resource_budget else False
         
     def add_to_resource_budget(self, cost):
+        if not self.is_resource_budget_enabled:
+            return
         if not self.can_spend_resource(cost):
             raise Exception("No resource budget remaining")
         self.current_resource_usage += cost
     
     def remove_from_resource_expense(self, cost):
+        if not self.is_resource_budget_enabled:
+            return
         if self.current_resource_usage - cost < 0:
             raise Exception("Not enough resource budget to remove")
         self.current_resource_usage -= cost
@@ -65,9 +87,13 @@ class BudgetManager():
         return self.total_expense_budget - self.current_expense
     
     def can_spend_expense(self, cost):
+        if not self.is_expense_budget_enabled:
+            return True
         return True if self.current_expense + cost <= self.total_expense_budget else False
     
     def add_to_expense_budget(self, cost):
+        if not self.is_expense_budget_enabled:
+            return
         if not self.can_spend_expense(cost):
             raise Exception("No expense budget remaining")
         self.current_expense += cost

src/manager/manager.py

@@ -1,9 +1,13 @@
+from enum import Enum, auto
+from typing import List
 from google import genai
 from google.genai import types
 from google.genai.types import *
 import os
 from dotenv import load_dotenv
 import sys
+from src.manager.agent_manager import AgentManager
+from src.manager.budget_manager import BudgetManager
 from src.manager.tool_manager import ToolManager
 from src.manager.utils.suppress_outputs import suppress_output
 import logging
@@ -19,32 +23,59 @@ handler = logging.StreamHandler(sys.stdout)
 logger.addHandler(handler)
 
 
+class Mode(Enum):
+    ENABLE_AGENT_CREATION = auto()
+    ENABLE_LOCAL_AGENTS = auto()
+    ENABLE_CLOUD_AGENTS = auto()
+    ENABLE_TOOL_CREATION = auto()
+    ENABLE_TOOL_INVOCATION = auto()
+    ENABLE_RESOURCE_BUDGET = auto()
+    ENABLE_ECONOMY_BUDGET = auto()
+    ENABLE_MEMORY = auto()
+
+
 class GeminiManager:
-    def __init__(self, toolsLoader: ToolManager = None,
-                 system_prompt_file="./src/models/system4.prompt",
+    def __init__(self, system_prompt_file="./src/models/system4.prompt",
                  gemini_model="gemini-2.5-pro-exp-03-25",
-                 local_only=False, allow_tool_creation=True,
-                 cloud_only=False, use_economy=True,
-                 use_memory=True):
+                 modes: List[Mode] = []):
         load_dotenv()
-        self.toolsLoader: ToolManager = toolsLoader
-        if not toolsLoader:
-            self.toolsLoader: ToolManager = ToolManager()
+        self.budget_manager = BudgetManager()
+
+        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.use_memory = use_memory
+        self.agentManager: AgentManager = AgentManager()
 
         self.API_KEY = os.getenv("GEMINI_KEY")
         self.client = genai.Client(api_key=self.API_KEY)
-        self.toolsLoader.load_tools()
         self.model_name = gemini_model
-        self.memory_manager = MemoryManager() if use_memory else None
+        self.memory_manager = MemoryManager()
         with open(system_prompt_file, 'r', encoding="utf8") as f:
             self.system_prompt = f.read()
         self.messages = []
+        self.set_modes(modes)
+
+    def get_current_modes(self):
+        return [mode.name for mode in self.modes]
+
+    def set_modes(self, modes: List[Mode]):
+        self.modes = modes
+        self.budget_manager.set_resource_budget_status(
+            self.check_mode(Mode.ENABLE_RESOURCE_BUDGET))
+        self.budget_manager.set_expense_budget_status(
+            self.check_mode(Mode.ENABLE_ECONOMY_BUDGET))
+        self.toolsLoader.set_creation_mode(
+            self.check_mode(Mode.ENABLE_TOOL_CREATION))
+        self.toolsLoader.set_invocation_mode(
+            self.check_mode(Mode.ENABLE_TOOL_INVOCATION))
+        self.agentManager.set_creation_mode(
+            self.check_mode(Mode.ENABLE_AGENT_CREATION))
+        self.agentManager.set_local_invocation_mode(
+            self.check_mode(Mode.ENABLE_LOCAL_AGENTS))
+        self.agentManager.set_cloud_invocation_mode(
+            self.check_mode(Mode.ENABLE_CLOUD_AGENTS))
+
+    def check_mode(self, mode: Mode):
+        return mode in self.modes
 
     def generate_response(self, messages):
         tools = self.toolsLoader.getTools()
@@ -149,10 +180,12 @@ class GeminiManager:
                                     types.Part.from_text(text="Error uploading file: "+str(e)))
                             continue
                         else:
-                            parts = [types.Part.from_text(text=message.get("content", ""))]
+                            parts = [types.Part.from_text(
+                                text=message.get("content", ""))]
                     case "memories":
                         role = "user"
-                        parts = [types.Part.from_text(text="Relevant memories: "+message.get("content", ""))]
+                        parts = [types.Part.from_text(
+                            text="Relevant memories: "+message.get("content", ""))]
                     case "tool":
                         role = "tool"
                         formatted_history.append(
@@ -160,12 +193,14 @@ class GeminiManager:
                         continue
                     case "function_call":
                         role = "model"
+                        print(message)
                         formatted_history.append(
                             eval(message.get("content", "")))
                         continue
                     case _:
                         role = "model"
-                        parts = [types.Part.from_text(text=message.get("content", ""))]
+                        parts = [types.Part.from_text(
+                            text=message.get("content", ""))]
                 formatted_history.append(types.Content(
                     role=role,
                     parts=parts
@@ -173,8 +208,6 @@ class GeminiManager:
         return formatted_history
 
     def get_k_memories(self, query, k=5, threshold=0.0):
-        if not self.use_memory:
-            return []
         memories = MemoryManager().get_memories()
         for i in range(len(memories)):
             memories[i] = memories[i]['memory']
@@ -189,20 +222,24 @@ class GeminiManager:
         else:
             device = 'cpu'
         model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
-        doc_embeddings = model.encode(memories, convert_to_tensor=True, device=device)
-        query_embedding = model.encode(query, convert_to_tensor=True, device=device)
-        similarity_scores = model.similarity(query_embedding, doc_embeddings)[0]
+        doc_embeddings = model.encode(
+            memories, convert_to_tensor=True, device=device)
+        query_embedding = model.encode(
+            query, convert_to_tensor=True, device=device)
+        similarity_scores = model.similarity(
+            query_embedding, doc_embeddings)[0]
         scores, indices = torch.topk(similarity_scores, k=top_k)
         results = []
         for score, idx in zip(scores, indices):
             if score >= threshold:
                 results.append(memories[idx])
         return results
-    
+
     def run(self, messages):
         try:
-            if self.use_memory:
-                memories = self.get_k_memories(messages[-1]['content'], k=5, threshold=0.1)
+            if self.check_mode(Mode.ENABLE_MEMORY) and len(messages) > 0:
+                memories = self.get_k_memories(
+                    messages[-1]['content'], k=5, threshold=0.1)
                 if len(memories) > 0:
                     messages.append({
                         "role": "memories",
@@ -217,7 +254,7 @@ class GeminiManager:
         except Exception as e:
             pass
         yield from self.invoke_manager(messages)
-    
+
     def invoke_manager(self, messages):
         chat_history = self.format_chat_history(messages)
         logger.debug(f"Chat history: {chat_history}")
@@ -261,8 +298,8 @@ class GeminiManager:
         if response.function_calls:
             for call in self.handle_tool_calls(response):
                 yield messages + [call]
-                if (call.get("role") == "tool" 
-                    or (call.get("role") == "assistant" and call.get("metadata", {}).get("status") == "done")):
+                if (call.get("role") == "tool"
+                        or (call.get("role") == "assistant" and call.get("metadata", {}).get("status") == "done")):
                     messages.append(call)
             yield from self.invoke_manager(messages)
         yield messages

src/manager/tool_manager.py

@@ -55,14 +55,30 @@ class Tool:
 class ToolManager:
     toolsImported: List[Tool] = []
     budget_manager: BudgetManager = BudgetManager()
+    is_creation_enabled: bool = True
+    is_invocation_enabled: bool = True
 
     def __init__(self):
         self.load_tools()
         self._output_budgets()
     
+    def set_creation_mode(self, status: bool):
+        self.is_creation_enabled = status
+        if status:
+            output_assistant_response("Tool creation mode is enabled.")
+        else:
+            output_assistant_response("Tool creation mode is disabled.")
+    
+    def set_invocation_mode(self, status: bool):
+        self.is_invocation_enabled = status
+        if status:
+            output_assistant_response("Tool invocation mode is enabled.")
+        else:
+            output_assistant_response("Tool invocation mode is disabled.")
+    
     def _output_budgets(self):
         output_assistant_response(f"Resource budget Remaining: {self.budget_manager.get_current_remaining_resource_budget()}")
-        output_assistant_response(f"Expense budget Remaining: {self.budget_manager.get_current_remaining_resource_budget()}")
+        output_assistant_response(f"Expense budget Remaining: {self.budget_manager.get_current_remaining_expense_budget()}")
 
     def load_tools(self):
         newToolsImported = []
@@ -84,6 +100,11 @@ class ToolManager:
         self.toolsImported = newToolsImported
 
     def runTool(self, toolName, query):
+        if not self.is_invocation_enabled:
+            raise Exception("Tool invocation mode is disabled")
+        if toolName == "ToolCreator":
+            if not self.is_creation_enabled:
+                raise Exception("Tool creation mode is disabled")
         self._output_budgets()
         for tool in self.toolsImported:
             if tool.name == toolName: