helloparthshah commited on
Commit
167a572
·
1 Parent(s): 50fb0a0

Updated memory related tasks

Browse files
src/manager/manager.py CHANGED
@@ -158,6 +158,8 @@ class GeminiManager:
158
 
159
  def get_k_memories(self, query, k=5, threshold=0.0):
160
  memories = MemoryManager().get_memories()
 
 
161
  if len(memories) == 0:
162
  return []
163
  top_k = min(k, len(memories))
 
158
 
159
  def get_k_memories(self, query, k=5, threshold=0.0):
160
  memories = MemoryManager().get_memories()
161
+ for i in range(len(memories)):
162
+ memories[i] = memories[i]['memory']
163
  if len(memories) == 0:
164
  return []
165
  top_k = min(k, len(memories))
src/tools/default_tools/memory_manager.py CHANGED
@@ -9,22 +9,22 @@ class MemoryManager():
9
 
10
  inputSchema = {
11
  "name": "MemoryManager",
12
- "description": "Updates, retrieves, or deletes the memory of an AI agent. The memory is stored in a JSON file.",
13
  "parameters": {
14
  "type": "object",
15
  "properties":{
16
  "action": {
17
  "type": "string",
18
- "enum": ["add_memory", "get_memory", "delete_memory"],
19
- "description": "The action to perform: add_memory, get_memory, or delete_memory.",
20
  },
21
  "memory": {
22
  "type": "string",
23
  "description": "The memory to add. Required for 'add_memory' action.",
24
  },
25
- "index": {
26
- "type": "integer",
27
- "description": "The index of the memory to delete. Required for 'delete_memory' action.",
28
  },
29
  },
30
  "required": ["action"],
@@ -45,29 +45,40 @@ class MemoryManager():
45
  def update_memories(self, memories):
46
  # save the memory to src/data/memory.json
47
  with open("src/data/memory.json", "w") as f:
48
- json.dump(memories, f)
49
 
50
 
51
  def run(self, **kwargs):
52
  # save it to src/data/memory.json
53
  action = kwargs.get("action")
54
  memory = kwargs.get("memory")
55
- index = kwargs.get("index")
56
  memories = self.get_memories()
57
- if action == "get_memory":
58
  return {
59
  "status": "success",
60
  "message": "Memory retrieved successfully",
61
  "output": memories
62
  }
63
  if action == "add_memory":
64
- if memory is None:
65
  return {
66
  "status": "error",
67
- "message": "Memory is required for add_memory action",
68
  "output": None
69
  }
70
- memories.append(memory)
 
 
 
 
 
 
 
 
 
 
 
71
  self.update_memories(memories)
72
  return {
73
  "status": "success",
@@ -75,22 +86,24 @@ class MemoryManager():
75
  "output": None
76
  }
77
  if action == "delete_memory":
78
- if index is None:
79
  return {
80
  "status": "error",
81
- "message": "Index is required for delete_memory action",
82
  "output": None
83
  }
84
- if index < 0 or index >= len(memories):
85
- return {
86
- "status": "error",
87
- "message": "Index out of range",
88
- "output": None
89
- }
90
- memories.pop(index)
91
- self.update_memories(memories)
 
 
92
  return {
93
- "status": "success",
94
- "message": "Memory deleted successfully",
95
  "output": None
96
  }
 
9
 
10
  inputSchema = {
11
  "name": "MemoryManager",
12
+ "description": "Updates, retrieves, or deletes the memory of an AI agent.",
13
  "parameters": {
14
  "type": "object",
15
  "properties":{
16
  "action": {
17
  "type": "string",
18
+ "enum": ["add_memory", "get_all_memories", "delete_memory"],
19
+ "description": "The action to perform: add_memory, get_all_memories, or delete_memory.",
20
  },
21
  "memory": {
22
  "type": "string",
23
  "description": "The memory to add. Required for 'add_memory' action.",
24
  },
25
+ "key": {
26
+ "type": "string",
27
+ "description": "The key to delete or add memory."
28
  },
29
  },
30
  "required": ["action"],
 
45
  def update_memories(self, memories):
46
  # save the memory to src/data/memory.json
47
  with open("src/data/memory.json", "w") as f:
48
+ json.dump(memories, f, indent=4)
49
 
50
 
51
  def run(self, **kwargs):
52
  # save it to src/data/memory.json
53
  action = kwargs.get("action")
54
  memory = kwargs.get("memory")
55
+ key = kwargs.get("key")
56
  memories = self.get_memories()
57
+ if action == "get_all_memories":
58
  return {
59
  "status": "success",
60
  "message": "Memory retrieved successfully",
61
  "output": memories
62
  }
63
  if action == "add_memory":
64
+ if memory is None or key is None:
65
  return {
66
  "status": "error",
67
+ "message": "Memory and key are required for add_memory action",
68
  "output": None
69
  }
70
+ # check if the key already exists
71
+ for mem in memories:
72
+ if mem["key"] == key:
73
+ return {
74
+ "status": "error",
75
+ "message": f"Memory with key {key} already exists",
76
+ "output": None
77
+ }
78
+ memories.append({
79
+ "key": key,
80
+ "memory": memory
81
+ })
82
  self.update_memories(memories)
83
  return {
84
  "status": "success",
 
86
  "output": None
87
  }
88
  if action == "delete_memory":
89
+ if key is None:
90
  return {
91
  "status": "error",
92
+ "message": "Key is required for delete_memory action",
93
  "output": None
94
  }
95
+ # check if the key exists
96
+ for mem in memories:
97
+ if mem["key"] == key:
98
+ memories.remove(mem)
99
+ self.update_memories(memories)
100
+ return {
101
+ "status": "success",
102
+ "message": "Memory deleted successfully",
103
+ "output": None
104
+ }
105
  return {
106
+ "status": "error",
107
+ "message": f"Memory with key {key} not found",
108
  "output": None
109
  }