File size: 2,302 Bytes
6643e6e 8cf77a3 6643e6e 8cf77a3 6643e6e 8cf77a3 6643e6e 8cf77a3 6643e6e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import importlib
from CEO.budget_manager import BudgetManager
__all__ = ['FireAgent']
class FireAgent():
dependencies = ["ollama==0.4.7",
"pydantic==2.11.1",
"pydantic_core==2.33.0"]
inputSchema = {
"name": "FireAgent",
"description": "Fires an AI agent for you.",
"parameters": {
"type": "object",
"properties":{
"agent_name": {
"type": "string",
"description": "Name of the AI agent that is to be fired. This name cannot have spaces or special characters. It should be a single word.",
},
},
"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")
agent_name= kwargs.get("agent_name")
ollama = importlib.import_module("ollama")
json = importlib.import_module("json")
if not self.does_agent_exist(agent_name):
return {
"status": "error",
"message": "Agent does not exists",
"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))
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",
} |