Integrate AgentManager with tools

src/agent_manager.py

@@ -7,10 +7,12 @@ from src.singleton import singleton
 
 class Agent(ABC):
     
-    def __init__(self, agent_name: str, base_model: str, system_prompt: str):
+    def __init__(self, agent_name: str, base_model: str, system_prompt: str, creation_cost: str, invoke_cost: str):
         self.agent_name = agent_name
         self.base_model = base_model
         self.system_prompt = system_prompt
+        self.creation_cost = creation_cost
+        self.invoke_cost = invoke_cost
         self.create_model()
         
     @abstractmethod
@@ -28,6 +30,12 @@ class Agent(ABC):
         """delete agent"""
         pass
     
+    def get_costs(self):
+        return {
+            "create_cost": self.creation_cost,
+            "invoke_cost": self.invoke_cost
+        }
+    
 class OllamaAgent(Agent):
     
     def create_model(self):
@@ -56,7 +64,7 @@ class AgentManager():
     def __init__(self):
         self._agents = {}
         self._agent_types ={
-            ollama: OllamaAgent
+            "ollama": OllamaAgent
         }
         
         self._load_agents()
@@ -74,7 +82,7 @@ class AgentManager():
             raise ValueError(f"Unsupported base model {base_model}")
         
         # create agent
-        self._agents[agent_name] = agent_class(agent_name, base_model, system_prompt)
+        self._agents[agent_name] = agent_class(agent_name, base_model, system_prompt, create_cost,invoke_cost )
         
         #save agent to file
         self._save_agent(
@@ -91,8 +99,8 @@ class AgentManager():
     
     def get_agent(self, agent_name: str) -> Agent:
         """Get existing agent by name"""
-        if agent_name in self._agents:
-            raise ValueError(f"Agent {agent_name} already exists")
+        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:
@@ -171,7 +179,10 @@ class AgentManager():
             print(f"Error saving agent {agent_name}: {e}")
 
     def _get_agent_type(self, base_model)->str:
-        if base_model.startswith("ollama"):
+
+        if base_model == "llama3.2":
+            return "ollama"
+        elif base_model == "mistral":
             return "ollama"
         else:
             return "unknown"
@@ -186,18 +197,23 @@ class AgentManager():
                 models = json.loads(f.read())
             
             for name, data in models.items():
+                if name in self._agents:
+                    continue
                 base_model = data["base_model"]
                 system_prompt = data["system_prompt"]
-                
-                model_type = self._get_model_type(base_model)
-                manager_class = self._model_types.get(model_type)
+                creation_cost = data["create_cost"]
+                invoke_cost = data["invoke_cost"]
+                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] = manager_class(
                         name, 
                         base_model,
-                        system_prompt
+                        system_prompt,
+                        creation_cost,
+                        invoke_cost
                     )
         except Exception as e:
             print(f"Error loading agents: {e}")

tools/agent_creater_tool.py

@@ -1,5 +1,4 @@
-import importlib
-
+from src.agent_manager import AgentManager
 __all__ = ['AgentCreator']
 
 class AgentCreator():
@@ -48,65 +47,36 @@ class AgentCreator():
             }
         }
     }
-    
-    def does_agent_exist(self, agent_name):
-        ollama = importlib.import_module("ollama")
-        all_agents = [a.model for a in ollama.list().models]
-        if agent_name in all_agents or f'{agent_name}:latest' in all_agents:
-            return True 
-        
-        return False
+
 
     def run(self, **kwargs):
         print("Running Agent Creator")
         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")
+        description = kwargs.get("description")
+        create_cost = self.inputSchema["creates"]["types"][base_model]["create_cost"]
+        invoke_cost = self.inputSchema["creates"]["types"][base_model]["invoke_cost"]
 
-        if self.does_agent_exist(agent_name):
+        agent_manager = AgentManager()
+        try:
+            agent_manager.create_agent(
+                agent_name=agent_name,
+                base_model=base_model,
+                system_prompt=system_prompt,
+                description=description,
+                create_cost=create_cost,
+                invoke_cost=invoke_cost
+            )
+        except ValueError as e:
             return {
                 "status": "error",
-                "message": "Agent already exists",
+                "message": f"Error occurred: {str(e)}",
                 "output": None
             }
         
-        ollama_response = ollama.create(
-            model = agent_name,
-            from_ = base_model,
-            system = system_prompt,
-            stream = False
-        )
-
-        if "success" in ollama_response["status"]:
-            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"),
-                    "system_prompt": system_prompt,
-                    "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",
-                "message": "Agent creation failed",
-            }
+        return {
+            "status": "success",
+            "message": "Agent successfully created",
+            "cost": create_cost,
+        }

tools/ask_agent_tool.py

@@ -1,6 +1,5 @@
-import importlib
 from src.budget_manager import BudgetManager
-from tools.get_agents_tool import GetAgents
+from src.agent_manager import AgentManager
 
 __all__ = ['AskAgent']
 
@@ -29,51 +28,38 @@ class AskAgent():
         }
     }
 
-    def does_agent_exist(self, ollama, agent_name):
-        all_agents = [a.model for a in ollama.list().models]
-        if agent_name in all_agents or f'{agent_name}:latest' in all_agents:
-            return True
-
-        return False
-
     def run(self, **kwargs):
         print("Asking agent a question")
 
         agent_name = kwargs.get("agent_name")
         prompt = kwargs.get("prompt")
-
-        ollama = importlib.import_module("ollama")
-        
+        agent_manger = AgentManager()
         budget_manager = BudgetManager()
-        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):
+
+        try:
+            agent = agent_manger.get_agent(agent_name=agent_name)
+        except ValueError as e:
             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",
+                "message": f"Error occurred: {str(e)}",
                 "output": None
             }
-        if not self.does_agent_exist(ollama, agent_name):
-            print("Agent does not exist")
+
+        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": "Agent does not exists",
+                "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 = ollama.chat(
-            model=agent_name,
-            messages=[{"role": "user", "content": prompt}],
-        )
-        print("Agent response", agent_response.message.content)
+        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.message.content,
+            "output": agent_response,
         }

tools/fire_agent.py

@@ -1,5 +1,5 @@
-import importlib
 from src.budget_manager import BudgetManager  
+from src.agent_manager import AgentManager
 
 __all__ = ['FireAgent']
 
@@ -22,46 +22,32 @@ class FireAgent():
             "required": ["agent_name"],
         }
     }
-    
-    def does_agent_exist(self, agent_name):
-        ollama = importlib.import_module("ollama")
-        all_agents = [a.model for a in ollama.list().models]
-        if agent_name in all_agents or f'{agent_name}:latest' in all_agents:
-            return True 
-        
-        return False
 
     def run(self, **kwargs):
-        print("Running Agent Creator")
+        print("Running Fire Agent")
         agent_name= kwargs.get("agent_name")
-        ollama = importlib.import_module("ollama")
-        json = importlib.import_module("json")
 
-        if not self.does_agent_exist(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)
+        except ValueError as e:
             return {
                 "status": "error",
-                "message": "Agent does not exists",
+                "message": f"Error occurred: {str(e)}",
                 "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]["create_cost"]))
-        del models[agent_name]
-        with open("./models/models.json", "w", encoding="utf8") as f:
-            f.write(json.dumps(models, indent=4))
+        budget_manager.add_to_expense(-1* int(agent_costs["create_cost"]))
 
-        if "success" in ollama_response["status"]:
-            return {
-                "status": "success",
-                "message": "Agent successfully fired.",
-                "current_expense": budget_manager.get_current_expense()
-            }
-        else:
-            return {
-                "status": "error",
-                "message": "Agent firing failed",
-            }
+        return {
+            "status": "success",
+            "message": "Agent successfully fired.",
+            "current_expense": budget_manager.get_current_expense()
+        }

tools/get_agents_tool.py

@@ -1,4 +1,4 @@
-import json
+from src.agent_manager import AgentManager
 
 __all__ = ['GetAgents']
 
@@ -16,11 +16,12 @@ class GetAgents():
     }
 
     def run(self, **kwargs):
-        with open("./models/models.json", "r", encoding="utf8") as f:
-            models = f.read()
-        models = json.loads(models)
+        
+        agent_manger = AgentManager()
+        agents = agent_manger.list_agents()
+
         return {
             "status": "success",
             "message": "Agents list retrieved successfully",
-            "agents": models,
+            "agents": agents,
         }