main
py 280 lines 11.1 KB
Raw
1 import numpy as np
2
3 from helpers.api import ApiHandler, Request, Response
4 from helpers import files
5 from helpers.localization import Localization
6 from models import ModelConfig, ModelType
7 from langchain_core.documents import Document
8 from agent import AgentContext
9
10 from plugins._memory.helpers.memory import Memory, get_existing_memory_subdirs, get_context_memory_subdir
11
12
13 class MemoryDashboard(ApiHandler):
14
15 async def process(self, input: dict, request: Request) -> dict | Response:
16 try:
17 action = input.get("action", "search")
18 if action == "get_memory_subdirs":
19 return await self._get_memory_subdirs()
20 elif action == "get_current_memory_subdir":
21 return await self._get_current_memory_subdir(input)
22 elif action == "search":
23 return await self._search_memories(input)
24 elif action == "delete":
25 return await self._delete_memory(input)
26 elif action == "bulk_delete":
27 return await self._bulk_delete_memories(input)
28 elif action == "update":
29 return await self._update_memory(input)
30 else:
31 return {
32 "success": False,
33 "error": f"Unknown action: {action}",
34 "memories": [],
35 "total_count": 0,
36 }
37
38 except Exception as e:
39 return {"success": False, "error": str(e), "memories": [], "total_count": 0}
40
41 async def _delete_memory(self, input: dict) -> dict:
42 """Delete a memory by ID from the specified subdirectory."""
43 try:
44 memory_subdir = input.get("memory_subdir", "default")
45 memory_id = input.get("memory_id")
46
47 if not memory_id:
48 return {"success": False, "error": "Memory ID is required for deletion"}
49
50 memory = await Memory.get_by_subdir(memory_subdir, preload_knowledge=False)
51
52 rem = await memory.delete_documents_by_ids([memory_id])
53
54 if len(rem) == 0:
55 return {
56 "success": False,
57 "error": f"Memory with ID '{memory_id}' not found",
58 }
59 else:
60 return {
61 "success": True,
62 "message": f"Memory {memory_id} deleted successfully",
63 }
64
65 except Exception as e:
66 return {"success": False, "error": f"Failed to delete memory: {str(e)}"}
67
68 async def _bulk_delete_memories(self, input: dict) -> dict:
69 """Delete multiple memories by IDs from the specified subdirectory."""
70 try:
71 memory_subdir = input.get("memory_subdir", "default")
72 memory_ids = input.get("memory_ids", [])
73
74 if not memory_ids:
75 return {
76 "success": False,
77 "error": "No memory IDs provided for bulk deletion",
78 }
79
80 if not isinstance(memory_ids, list):
81 return {
82 "success": False,
83 "error": "Memory IDs must be provided as a list",
84 }
85
86 # delete
87 memory = await Memory.get_by_subdir(memory_subdir, preload_knowledge=False)
88 rem = await memory.delete_documents_by_ids(memory_ids)
89
90 if len(rem) == len(memory_ids):
91 return {
92 "success": True,
93 "message": f"Successfully deleted {len(memory_ids)} memories",
94 }
95 elif len(rem) > 0:
96 return {
97 "success": True,
98 "message": f"Successfully deleted {len(rem)} memories. {len(memory_ids) - len(rem)} failed.",
99 }
100 else:
101 return {
102 "success": False,
103 "error": f"Failed to delete any memories.",
104 }
105
106 except Exception as e:
107 return {
108 "success": False,
109 "error": f"Failed to bulk delete memories: {str(e)}",
110 }
111
112 async def _get_current_memory_subdir(self, input: dict) -> dict:
113 """Get the current memory subdirectory from the active context."""
114 try:
115 # Try to get the context from the request
116 context_id = input.get("context_id", None)
117 if not context_id:
118 # Fallback to default if no context available
119 return {"success": True, "memory_subdir": "default"}
120
121 context = AgentContext.use(context_id)
122 if not context:
123 return {"success": True, "memory_subdir": "default"}
124
125 memory_subdir = get_context_memory_subdir(context)
126 return {"success": True, "memory_subdir": memory_subdir or "default"}
127
128 except Exception:
129 return {
130 "success": True, # Still success, just fallback to default
131 "memory_subdir": "default",
132 }
133
134 async def _get_memory_subdirs(self) -> dict:
135 """Get available memory subdirectories."""
136 try:
137 # Get subdirectories from memory folder
138 subdirs = get_existing_memory_subdirs()
139 return {"success": True, "subdirs": subdirs}
140 except Exception as e:
141 return {
142 "success": False,
143 "error": f"Failed to get memory subdirectories: {str(e)}",
144 "subdirs": ["default"],
145 }
146
147 async def _search_memories(self, input: dict) -> dict:
148 """Search memories in the specified subdirectory."""
149 try:
150 # Get search parameters
151 memory_subdir = input.get("memory_subdir", "default")
152 area_filter = input.get("area", "") # Filter by memory area
153 search_query = input.get("search", "") # Full-text search query
154 limit = input.get("limit", 100) # Number of results to return
155 threshold = input.get("threshold", 0.6) # Similarity threshold
156
157 memory = await Memory.get_by_subdir(memory_subdir, preload_knowledge=False)
158
159 memories = []
160
161 if search_query:
162 docs = await memory.search_similarity_threshold(
163 query=search_query,
164 limit=limit,
165 threshold=threshold,
166 filter=f"area == '{area_filter}'" if area_filter else "",
167 )
168 memories = docs
169 else:
170 # If no search query, get all memories from specified area(s)
171 all_docs = memory.db.get_all_docs()
172 for doc_id, doc in all_docs.items():
173 # Apply area filter if specified
174 if area_filter and doc.metadata.get("area", "") != area_filter:
175 continue
176 memories.append(doc)
177
178 # sort by timestamp
179 def get_sort_key(m):
180 timestamp = self._serialize_memory_timestamp(
181 m.metadata.get("timestamp", "0000-00-00 00:00:00")
182 )
183 return timestamp or "0000-00-00T00:00:00"
184
185 memories.sort(key=get_sort_key, reverse=True)
186
187 # Apply limit AFTER sorting to get the newest entries
188 if limit and len(memories) > limit:
189 memories = memories[:limit]
190
191 # Format memories for the dashboard
192 formatted_memories = [self._format_memory_for_dashboard(m) for m in memories]
193
194 # Get summary statistics
195 total_memories = len(formatted_memories)
196 knowledge_count = sum(
197 1 for m in formatted_memories if m["knowledge_source"]
198 )
199 conversation_count = total_memories - knowledge_count
200
201 # Get total count of all memories in database (unfiltered)
202 total_db_count = len(memory.db.get_all_docs())
203
204 return {
205 "success": True,
206 "memories": formatted_memories,
207 "total_count": total_memories,
208 "total_db_count": total_db_count,
209 "knowledge_count": knowledge_count,
210 "conversation_count": conversation_count,
211 "search_query": search_query,
212 "area_filter": area_filter,
213 "memory_subdir": memory_subdir,
214 }
215
216 except Exception as e:
217 return {"success": False, "error": str(e), "memories": [], "total_count": 0}
218
219 def _format_memory_for_dashboard(self, m: Document) -> dict:
220 """Format a memory document for the dashboard."""
221 metadata = dict(m.metadata)
222 similarity = metadata.get("_consolidation_similarity")
223 if isinstance(similarity, np.generic):
224 metadata["_consolidation_similarity"] = float(similarity)
225 timestamp = self._serialize_memory_timestamp(metadata.get("timestamp", "unknown"))
226 return {
227 "id": metadata.get("id", "unknown"),
228 "area": metadata.get("area", "unknown"),
229 "timestamp": timestamp,
230 # "content_preview": m.page_content[:200]
231 # + ("..." if len(m.page_content) > 200 else ""),
232 "content_full": m.page_content,
233 "knowledge_source": metadata.get("knowledge_source", False),
234 "source_file": metadata.get("source_file", ""),
235 "file_type": metadata.get("file_type", ""),
236 "consolidation_action": metadata.get("consolidation_action", ""),
237 "tags": metadata.get("tags", []),
238 "metadata": metadata, # Include full metadata for advanced users
239 }
240
241 def _serialize_memory_timestamp(self, value) -> str:
242 if not value or value == "unknown":
243 return "unknown"
244
245 if isinstance(value, str):
246 value = value.strip()
247 if not value or value == "unknown":
248 return "unknown"
249
250 localization = Localization.get()
251 if isinstance(value, str):
252 parsed = localization.localtime_str_to_utc_dt(value)
253 if parsed is None:
254 return value
255 return localization.utc_dt_to_localtime_str(parsed, timespec="seconds") or value
256
257 return localization.serialize_datetime(value) or str(value)
258
259 async def _update_memory(self, input: dict) -> dict:
260 try:
261 memory_subdir = input.get("memory_subdir")
262 original = input.get("original")
263 edited = input.get("edited")
264
265 if not memory_subdir or not original or not edited:
266 return {"success": False, "error": "Missing required parameters"}
267
268 doc = Document(
269 page_content=edited["content_full"],
270 metadata=edited["metadata"],
271 )
272
273 memory = await Memory.get_by_subdir(memory_subdir, preload_knowledge=False)
274 id = (await memory.update_documents([doc]))[0]
275 doc = memory.get_document_by_id(id)
276 formatted_doc = self._format_memory_for_dashboard(doc) if doc else None
277
278 return {"success": formatted_doc is not None, "memory": formatted_doc}
279 except Exception as e:
280 return {"success": False, "error": str(e), "memory": None}