Harshil Patel commited on
Commit
6643e6e
·
1 Parent(s): b3c76c7

Add fire agent tool

Browse files
Files changed (1) hide show
  1. tools/fire_agent.py +63 -0
tools/fire_agent.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import importlib
2
+
3
+ __all__ = ['FireAgent']
4
+
5
+ class FireAgent():
6
+ dependencies = ["ollama==0.4.7",
7
+ "pydantic==2.11.1",
8
+ "pydantic_core==2.33.0"]
9
+
10
+ inputSchema = {
11
+ "name": "FireAgent",
12
+ "description": "Fires an AI agent for you.",
13
+ "parameters": {
14
+ "type": "object",
15
+ "properties":{
16
+ "agent_name": {
17
+ "type": "string",
18
+ "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.",
19
+ },
20
+ },
21
+ "required": ["agent_name"],
22
+ }
23
+ }
24
+
25
+ def does_agent_exist(self, agent_name):
26
+ ollama = importlib.import_module("ollama")
27
+ all_agents = [a.model for a in ollama.list().models]
28
+ if agent_name in all_agents or f'{agent_name}:latest' in all_agents:
29
+ return True
30
+
31
+ return False
32
+
33
+ def run(self, **kwargs):
34
+ print("Running Agent Creator")
35
+ agent_name= kwargs.get("agent_name")
36
+ ollama = importlib.import_module("ollama")
37
+ json = importlib.import_module("json")
38
+
39
+ if not self.does_agent_exist(agent_name):
40
+ return {
41
+ "status": "error",
42
+ "message": "Agent does not exists",
43
+ "output": None
44
+ }
45
+ ollama_response = ollama.delete(agent_name)
46
+
47
+ with open("./models/models.json", "r", encoding="utf8") as f:
48
+ models = f.read()
49
+ models = json.loads(models)
50
+ del models[agent_name]
51
+ with open("./models/models.json", "w", encoding="utf8") as f:
52
+ f.write(json.dumps(models, indent=4))
53
+
54
+ if "success" in ollama_response["status"]:
55
+ return {
56
+ "status": "success",
57
+ "message": "Agent successfully fired",
58
+ }
59
+ else:
60
+ return {
61
+ "status": "error",
62
+ "message": "Agent firing failed",
63
+ }