agno 2.2.1__py3-none-any.whl → 2.2.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. agno/agent/agent.py +735 -574
  2. agno/culture/manager.py +22 -24
  3. agno/db/async_postgres/__init__.py +1 -1
  4. agno/db/dynamo/dynamo.py +0 -2
  5. agno/db/firestore/firestore.py +0 -2
  6. agno/db/gcs_json/gcs_json_db.py +0 -4
  7. agno/db/gcs_json/utils.py +0 -24
  8. agno/db/in_memory/in_memory_db.py +0 -3
  9. agno/db/json/json_db.py +4 -10
  10. agno/db/json/utils.py +0 -24
  11. agno/db/mongo/__init__.py +15 -1
  12. agno/db/mongo/async_mongo.py +1999 -0
  13. agno/db/mongo/mongo.py +0 -2
  14. agno/db/mysql/mysql.py +0 -3
  15. agno/db/postgres/__init__.py +1 -1
  16. agno/db/{async_postgres → postgres}/async_postgres.py +19 -22
  17. agno/db/postgres/postgres.py +7 -10
  18. agno/db/postgres/utils.py +106 -2
  19. agno/db/redis/redis.py +0 -2
  20. agno/db/singlestore/singlestore.py +0 -3
  21. agno/db/sqlite/__init__.py +2 -1
  22. agno/db/sqlite/async_sqlite.py +2269 -0
  23. agno/db/sqlite/sqlite.py +0 -2
  24. agno/db/sqlite/utils.py +96 -0
  25. agno/db/surrealdb/surrealdb.py +0 -6
  26. agno/knowledge/knowledge.py +3 -3
  27. agno/knowledge/reader/reader_factory.py +16 -0
  28. agno/knowledge/reader/tavily_reader.py +194 -0
  29. agno/memory/manager.py +28 -25
  30. agno/models/anthropic/claude.py +63 -6
  31. agno/models/base.py +251 -32
  32. agno/models/response.py +69 -0
  33. agno/os/router.py +7 -5
  34. agno/os/routers/memory/memory.py +2 -1
  35. agno/os/routers/memory/schemas.py +5 -2
  36. agno/os/schema.py +25 -20
  37. agno/os/utils.py +9 -2
  38. agno/run/agent.py +23 -30
  39. agno/run/base.py +17 -1
  40. agno/run/team.py +23 -29
  41. agno/run/workflow.py +17 -12
  42. agno/session/agent.py +3 -0
  43. agno/session/summary.py +4 -1
  44. agno/session/team.py +1 -1
  45. agno/team/team.py +599 -367
  46. agno/tools/dalle.py +2 -4
  47. agno/tools/eleven_labs.py +23 -25
  48. agno/tools/function.py +40 -0
  49. agno/tools/mcp/__init__.py +10 -0
  50. agno/tools/mcp/mcp.py +324 -0
  51. agno/tools/mcp/multi_mcp.py +347 -0
  52. agno/tools/mcp/params.py +24 -0
  53. agno/tools/slack.py +18 -3
  54. agno/tools/tavily.py +146 -0
  55. agno/utils/agent.py +366 -1
  56. agno/utils/mcp.py +92 -2
  57. agno/utils/media.py +166 -1
  58. agno/utils/print_response/workflow.py +17 -1
  59. agno/utils/team.py +89 -1
  60. agno/workflow/step.py +0 -1
  61. agno/workflow/types.py +10 -15
  62. {agno-2.2.1.dist-info → agno-2.2.3.dist-info}/METADATA +28 -25
  63. {agno-2.2.1.dist-info → agno-2.2.3.dist-info}/RECORD +66 -62
  64. agno/db/async_postgres/schemas.py +0 -139
  65. agno/db/async_postgres/utils.py +0 -347
  66. agno/tools/mcp.py +0 -679
  67. {agno-2.2.1.dist-info → agno-2.2.3.dist-info}/WHEEL +0 -0
  68. {agno-2.2.1.dist-info → agno-2.2.3.dist-info}/licenses/LICENSE +0 -0
  69. {agno-2.2.1.dist-info → agno-2.2.3.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,1999 @@
1
+ import time
2
+ from datetime import date, datetime, timedelta, timezone
3
+ from typing import Any, Dict, List, Optional, Tuple, Union
4
+ from uuid import uuid4
5
+
6
+ from agno.db.base import AsyncBaseDb, SessionType
7
+ from agno.db.mongo.utils import (
8
+ apply_pagination,
9
+ apply_sorting,
10
+ bulk_upsert_metrics,
11
+ calculate_date_metrics,
12
+ create_collection_indexes,
13
+ deserialize_cultural_knowledge_from_db,
14
+ fetch_all_sessions_data,
15
+ get_dates_to_calculate_metrics_for,
16
+ serialize_cultural_knowledge_for_db,
17
+ )
18
+ from agno.db.schemas.culture import CulturalKnowledge
19
+ from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
20
+ from agno.db.schemas.knowledge import KnowledgeRow
21
+ from agno.db.schemas.memory import UserMemory
22
+ from agno.db.utils import deserialize_session_json_fields
23
+ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
24
+ from agno.utils.log import log_debug, log_error, log_info
25
+ from agno.utils.string import generate_id
26
+
27
+ try:
28
+ import asyncio
29
+
30
+ from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection, AsyncIOMotorDatabase # type: ignore
31
+ except ImportError:
32
+ raise ImportError("`motor` not installed. Please install it using `pip install -U motor`")
33
+
34
+ try:
35
+ from pymongo import ReturnDocument
36
+ from pymongo.errors import OperationFailure
37
+ except ImportError:
38
+ raise ImportError("`pymongo` not installed. Please install it using `pip install -U pymongo`")
39
+
40
+
41
+ class AsyncMongoDb(AsyncBaseDb):
42
+ def __init__(
43
+ self,
44
+ db_client: Optional[AsyncIOMotorClient] = None,
45
+ db_name: Optional[str] = None,
46
+ db_url: Optional[str] = None,
47
+ session_collection: Optional[str] = None,
48
+ memory_collection: Optional[str] = None,
49
+ metrics_collection: Optional[str] = None,
50
+ eval_collection: Optional[str] = None,
51
+ knowledge_collection: Optional[str] = None,
52
+ culture_collection: Optional[str] = None,
53
+ id: Optional[str] = None,
54
+ ):
55
+ """
56
+ Async interface for interacting with a MongoDB database using Motor.
57
+
58
+ Args:
59
+ db_client (Optional[AsyncIOMotorClient]): The MongoDB async client to use.
60
+ db_name (Optional[str]): The name of the database to use.
61
+ db_url (Optional[str]): The database URL to connect to.
62
+ session_collection (Optional[str]): Name of the collection to store sessions.
63
+ memory_collection (Optional[str]): Name of the collection to store memories.
64
+ metrics_collection (Optional[str]): Name of the collection to store metrics.
65
+ eval_collection (Optional[str]): Name of the collection to store evaluation runs.
66
+ knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
67
+ culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
68
+ id (Optional[str]): ID of the database.
69
+
70
+ Raises:
71
+ ValueError: If neither db_url nor db_client is provided.
72
+ """
73
+ if id is None:
74
+ base_seed = db_url or str(db_client)
75
+ db_name_suffix = db_name if db_name is not None else "agno"
76
+ seed = f"{base_seed}#{db_name_suffix}"
77
+ id = generate_id(seed)
78
+
79
+ super().__init__(
80
+ id=id,
81
+ session_table=session_collection,
82
+ memory_table=memory_collection,
83
+ metrics_table=metrics_collection,
84
+ eval_table=eval_collection,
85
+ knowledge_table=knowledge_collection,
86
+ culture_table=culture_collection,
87
+ )
88
+
89
+ # Store configuration for lazy initialization
90
+ self._provided_client: Optional[AsyncIOMotorClient] = db_client
91
+ self.db_url: Optional[str] = db_url
92
+ self.db_name: str = db_name if db_name is not None else "agno"
93
+
94
+ if self._provided_client is None and self.db_url is None:
95
+ raise ValueError("One of db_url or db_client must be provided")
96
+
97
+ # Client and database will be lazily initialized per event loop
98
+ self._client: Optional[AsyncIOMotorClient] = None
99
+ self._database: Optional[AsyncIOMotorDatabase] = None
100
+ self._event_loop: Optional[asyncio.AbstractEventLoop] = None
101
+
102
+ def _ensure_client(self) -> AsyncIOMotorClient:
103
+ """
104
+ Ensure the Motor client is valid for the current event loop.
105
+
106
+ Motor's AsyncIOMotorClient is tied to the event loop it was created in.
107
+ If we detect a new event loop, we need to refresh the client.
108
+
109
+ Returns:
110
+ AsyncIOMotorClient: A valid client for the current event loop.
111
+ """
112
+ try:
113
+ current_loop = asyncio.get_running_loop()
114
+ except RuntimeError:
115
+ # No running loop, return existing client or create new one
116
+ if self._client is None:
117
+ if self._provided_client is not None:
118
+ self._client = self._provided_client
119
+ elif self.db_url is not None:
120
+ self._client = AsyncIOMotorClient(self.db_url)
121
+ log_debug("Created AsyncIOMotorClient outside event loop")
122
+ return self._client # type: ignore
123
+
124
+ # Check if we're in a different event loop
125
+ if self._event_loop is None or self._event_loop is not current_loop:
126
+ # New event loop detected, create new client
127
+ if self._provided_client is not None:
128
+ # User provided a client, use it but warn them
129
+ log_debug(
130
+ "New event loop detected. Using provided AsyncIOMotorClient, "
131
+ "which may cause issues if it was created in a different event loop."
132
+ )
133
+ self._client = self._provided_client
134
+ elif self.db_url is not None:
135
+ # Create a new client for this event loop
136
+ old_loop_id = id(self._event_loop) if self._event_loop else "None"
137
+ new_loop_id = id(current_loop)
138
+ log_debug(f"Event loop changed from {old_loop_id} to {new_loop_id}, creating new AsyncIOMotorClient")
139
+ self._client = AsyncIOMotorClient(self.db_url)
140
+
141
+ self._event_loop = current_loop
142
+ self._database = None # Reset database reference
143
+ # Clear collection caches when switching event loops
144
+ for attr in list(vars(self).keys()):
145
+ if attr.endswith("_collection"):
146
+ delattr(self, attr)
147
+
148
+ return self._client # type: ignore
149
+
150
+ @property
151
+ def db_client(self) -> AsyncIOMotorClient:
152
+ """Get the MongoDB client, ensuring it's valid for the current event loop."""
153
+ return self._ensure_client()
154
+
155
+ @property
156
+ def database(self) -> AsyncIOMotorDatabase:
157
+ """Get the MongoDB database, ensuring it's valid for the current event loop."""
158
+ try:
159
+ current_loop = asyncio.get_running_loop()
160
+ if self._database is None or self._event_loop != current_loop:
161
+ self._database = self.db_client[self.db_name]
162
+ except RuntimeError:
163
+ # No running loop - fallback to existing database or create new one
164
+ if self._database is None:
165
+ self._database = self.db_client[self.db_name]
166
+ return self._database
167
+
168
+ # -- DB methods --
169
+
170
+ def _should_reset_collection_cache(self) -> bool:
171
+ """Check if collection cache should be reset due to event loop change."""
172
+ try:
173
+ current_loop = asyncio.get_running_loop()
174
+ return self._event_loop is not current_loop
175
+ except RuntimeError:
176
+ return False
177
+
178
+ async def _get_collection(
179
+ self, table_type: str, create_collection_if_not_found: Optional[bool] = True
180
+ ) -> Optional[AsyncIOMotorCollection]:
181
+ """Get or create a collection based on table type.
182
+
183
+ Args:
184
+ table_type (str): The type of table to get or create.
185
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
186
+
187
+ Returns:
188
+ AsyncIOMotorCollection: The collection object.
189
+ """
190
+ # Ensure client is valid for current event loop before accessing collections
191
+ _ = self.db_client # This triggers _ensure_client()
192
+
193
+ # Check if collections need to be reset due to event loop change
194
+ reset_cache = self._should_reset_collection_cache()
195
+
196
+ if table_type == "sessions":
197
+ if reset_cache or not hasattr(self, "session_collection"):
198
+ if self.session_table_name is None:
199
+ raise ValueError("Session collection was not provided on initialization")
200
+ self.session_collection = await self._get_or_create_collection(
201
+ collection_name=self.session_table_name,
202
+ collection_type="sessions",
203
+ create_collection_if_not_found=create_collection_if_not_found,
204
+ )
205
+ return self.session_collection
206
+
207
+ if table_type == "memories":
208
+ if reset_cache or not hasattr(self, "memory_collection"):
209
+ if self.memory_table_name is None:
210
+ raise ValueError("Memory collection was not provided on initialization")
211
+ self.memory_collection = await self._get_or_create_collection(
212
+ collection_name=self.memory_table_name,
213
+ collection_type="memories",
214
+ create_collection_if_not_found=create_collection_if_not_found,
215
+ )
216
+ return self.memory_collection
217
+
218
+ if table_type == "metrics":
219
+ if reset_cache or not hasattr(self, "metrics_collection"):
220
+ if self.metrics_table_name is None:
221
+ raise ValueError("Metrics collection was not provided on initialization")
222
+ self.metrics_collection = await self._get_or_create_collection(
223
+ collection_name=self.metrics_table_name,
224
+ collection_type="metrics",
225
+ create_collection_if_not_found=create_collection_if_not_found,
226
+ )
227
+ return self.metrics_collection
228
+
229
+ if table_type == "evals":
230
+ if reset_cache or not hasattr(self, "eval_collection"):
231
+ if self.eval_table_name is None:
232
+ raise ValueError("Eval collection was not provided on initialization")
233
+ self.eval_collection = await self._get_or_create_collection(
234
+ collection_name=self.eval_table_name,
235
+ collection_type="evals",
236
+ create_collection_if_not_found=create_collection_if_not_found,
237
+ )
238
+ return self.eval_collection
239
+
240
+ if table_type == "knowledge":
241
+ if reset_cache or not hasattr(self, "knowledge_collection"):
242
+ if self.knowledge_table_name is None:
243
+ raise ValueError("Knowledge collection was not provided on initialization")
244
+ self.knowledge_collection = await self._get_or_create_collection(
245
+ collection_name=self.knowledge_table_name,
246
+ collection_type="knowledge",
247
+ create_collection_if_not_found=create_collection_if_not_found,
248
+ )
249
+ return self.knowledge_collection
250
+
251
+ if table_type == "culture":
252
+ if reset_cache or not hasattr(self, "culture_collection"):
253
+ if self.culture_table_name is None:
254
+ raise ValueError("Culture collection was not provided on initialization")
255
+ self.culture_collection = await self._get_or_create_collection(
256
+ collection_name=self.culture_table_name,
257
+ collection_type="culture",
258
+ create_collection_if_not_found=create_collection_if_not_found,
259
+ )
260
+ return self.culture_collection
261
+
262
+ raise ValueError(f"Unknown table type: {table_type}")
263
+
264
+ async def _get_or_create_collection(
265
+ self, collection_name: str, collection_type: str, create_collection_if_not_found: Optional[bool] = True
266
+ ) -> Optional[AsyncIOMotorCollection]:
267
+ """Get or create a collection with proper indexes.
268
+
269
+ Args:
270
+ collection_name (str): The name of the collection to get or create.
271
+ collection_type (str): The type of collection to get or create.
272
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
273
+
274
+ Returns:
275
+ Optional[AsyncIOMotorCollection]: The collection object.
276
+ """
277
+ try:
278
+ collection = self.database[collection_name]
279
+
280
+ if not hasattr(self, f"_{collection_name}_initialized"):
281
+ if not create_collection_if_not_found:
282
+ return None
283
+ # Note: Motor doesn't have sync create_index, so we use it as-is
284
+ # The indexes are created in the background
285
+ create_collection_indexes(collection, collection_type) # type: ignore
286
+ setattr(self, f"_{collection_name}_initialized", True)
287
+ log_debug(f"Initialized collection '{collection_name}'")
288
+ else:
289
+ log_debug(f"Collection '{collection_name}' already initialized")
290
+
291
+ return collection
292
+
293
+ except Exception as e:
294
+ log_error(f"Error getting collection {collection_name}: {e}")
295
+ raise
296
+
297
+ # -- Session methods --
298
+
299
+ async def delete_session(self, session_id: str) -> bool:
300
+ """Delete a session from the database.
301
+
302
+ Args:
303
+ session_id (str): The ID of the session to delete.
304
+
305
+ Returns:
306
+ bool: True if the session was deleted, False otherwise.
307
+
308
+ Raises:
309
+ Exception: If there is an error deleting the session.
310
+ """
311
+ try:
312
+ collection = await self._get_collection(table_type="sessions")
313
+ if collection is None:
314
+ return False
315
+
316
+ result = await collection.delete_one({"session_id": session_id})
317
+ if result.deleted_count == 0:
318
+ log_debug(f"No session found to delete with session_id: {session_id}")
319
+ return False
320
+ else:
321
+ log_debug(f"Successfully deleted session with session_id: {session_id}")
322
+ return True
323
+
324
+ except Exception as e:
325
+ log_error(f"Error deleting session: {e}")
326
+ raise e
327
+
328
+ async def delete_sessions(self, session_ids: List[str]) -> None:
329
+ """Delete multiple sessions from the database.
330
+
331
+ Args:
332
+ session_ids (List[str]): The IDs of the sessions to delete.
333
+ """
334
+ try:
335
+ collection = await self._get_collection(table_type="sessions")
336
+ if collection is None:
337
+ return
338
+
339
+ result = await collection.delete_many({"session_id": {"$in": session_ids}})
340
+ log_debug(f"Successfully deleted {result.deleted_count} sessions")
341
+
342
+ except Exception as e:
343
+ log_error(f"Error deleting sessions: {e}")
344
+ raise e
345
+
346
+ async def get_session(
347
+ self,
348
+ session_id: str,
349
+ session_type: SessionType,
350
+ user_id: Optional[str] = None,
351
+ deserialize: Optional[bool] = True,
352
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
353
+ """Read a session from the database.
354
+
355
+ Args:
356
+ session_id (str): The ID of the session to get.
357
+ session_type (SessionType): The type of session to get.
358
+ user_id (Optional[str]): The ID of the user to get the session for.
359
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
360
+
361
+ Returns:
362
+ Union[Session, Dict[str, Any], None]:
363
+ - When deserialize=True: Session object
364
+ - When deserialize=False: Session dictionary
365
+
366
+ Raises:
367
+ Exception: If there is an error reading the session.
368
+ """
369
+ try:
370
+ collection = await self._get_collection(table_type="sessions")
371
+ if collection is None:
372
+ return None
373
+
374
+ query = {"session_id": session_id}
375
+ if user_id is not None:
376
+ query["user_id"] = user_id
377
+ if session_type is not None:
378
+ query["session_type"] = session_type
379
+
380
+ result = await collection.find_one(query)
381
+ if result is None:
382
+ return None
383
+
384
+ session = deserialize_session_json_fields(result)
385
+ if not deserialize:
386
+ return session
387
+
388
+ if session_type == SessionType.AGENT:
389
+ return AgentSession.from_dict(session)
390
+ elif session_type == SessionType.TEAM:
391
+ return TeamSession.from_dict(session)
392
+ elif session_type == SessionType.WORKFLOW:
393
+ return WorkflowSession.from_dict(session)
394
+ else:
395
+ raise ValueError(f"Invalid session type: {session_type}")
396
+
397
+ except Exception as e:
398
+ log_error(f"Exception reading session: {e}")
399
+ raise e
400
+
401
+ async def get_sessions(
402
+ self,
403
+ session_type: Optional[SessionType] = None,
404
+ user_id: Optional[str] = None,
405
+ component_id: Optional[str] = None,
406
+ session_name: Optional[str] = None,
407
+ start_timestamp: Optional[int] = None,
408
+ end_timestamp: Optional[int] = None,
409
+ limit: Optional[int] = None,
410
+ page: Optional[int] = None,
411
+ sort_by: Optional[str] = None,
412
+ sort_order: Optional[str] = None,
413
+ deserialize: Optional[bool] = True,
414
+ ) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
415
+ """Get all sessions.
416
+
417
+ Args:
418
+ session_type (Optional[SessionType]): The type of session to get.
419
+ user_id (Optional[str]): The ID of the user to get the session for.
420
+ component_id (Optional[str]): The ID of the component to get the session for.
421
+ session_name (Optional[str]): The name of the session to filter by.
422
+ start_timestamp (Optional[int]): The start timestamp to filter sessions by.
423
+ end_timestamp (Optional[int]): The end timestamp to filter sessions by.
424
+ limit (Optional[int]): The limit of the sessions to get.
425
+ page (Optional[int]): The page number to get.
426
+ sort_by (Optional[str]): The field to sort the sessions by.
427
+ sort_order (Optional[str]): The order to sort the sessions by.
428
+ deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
429
+
430
+ Returns:
431
+ Union[List[AgentSession], List[TeamSession], List[WorkflowSession], Tuple[List[Dict[str, Any]], int]]:
432
+ - When deserialize=True: List of Session objects
433
+ - When deserialize=False: List of session dictionaries and the total count
434
+
435
+ Raises:
436
+ Exception: If there is an error reading the sessions.
437
+ """
438
+ try:
439
+ collection = await self._get_collection(table_type="sessions")
440
+ if collection is None:
441
+ return [] if deserialize else ([], 0)
442
+
443
+ # Filtering
444
+ query: Dict[str, Any] = {}
445
+ if user_id is not None:
446
+ query["user_id"] = user_id
447
+ if session_type is not None:
448
+ query["session_type"] = session_type
449
+ if component_id is not None:
450
+ if session_type == SessionType.AGENT:
451
+ query["agent_id"] = component_id
452
+ elif session_type == SessionType.TEAM:
453
+ query["team_id"] = component_id
454
+ elif session_type == SessionType.WORKFLOW:
455
+ query["workflow_id"] = component_id
456
+ if start_timestamp is not None:
457
+ query["created_at"] = {"$gte": start_timestamp}
458
+ if end_timestamp is not None:
459
+ if "created_at" in query:
460
+ query["created_at"]["$lte"] = end_timestamp
461
+ else:
462
+ query["created_at"] = {"$lte": end_timestamp}
463
+ if session_name is not None:
464
+ query["session_data.session_name"] = {"$regex": session_name, "$options": "i"}
465
+
466
+ # Get total count
467
+ total_count = await collection.count_documents(query)
468
+
469
+ cursor = collection.find(query)
470
+
471
+ # Sorting
472
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
473
+ if sort_criteria:
474
+ cursor = cursor.sort(sort_criteria)
475
+
476
+ # Pagination
477
+ query_args = apply_pagination({}, limit, page)
478
+ if query_args.get("skip"):
479
+ cursor = cursor.skip(query_args["skip"])
480
+ if query_args.get("limit"):
481
+ cursor = cursor.limit(query_args["limit"])
482
+
483
+ records = await cursor.to_list(length=None)
484
+ if records is None:
485
+ return [] if deserialize else ([], 0)
486
+ sessions_raw = [deserialize_session_json_fields(record) for record in records]
487
+
488
+ if not deserialize:
489
+ return sessions_raw, total_count
490
+
491
+ sessions: List[Union[AgentSession, TeamSession, WorkflowSession]] = []
492
+ for record in sessions_raw:
493
+ if session_type == SessionType.AGENT.value:
494
+ agent_session = AgentSession.from_dict(record)
495
+ if agent_session is not None:
496
+ sessions.append(agent_session)
497
+ elif session_type == SessionType.TEAM.value:
498
+ team_session = TeamSession.from_dict(record)
499
+ if team_session is not None:
500
+ sessions.append(team_session)
501
+ elif session_type == SessionType.WORKFLOW.value:
502
+ workflow_session = WorkflowSession.from_dict(record)
503
+ if workflow_session is not None:
504
+ sessions.append(workflow_session)
505
+
506
+ return sessions
507
+
508
+ except Exception as e:
509
+ log_error(f"Exception reading sessions: {e}")
510
+ raise e
511
+
512
+ async def rename_session(
513
+ self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
514
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
515
+ """Rename a session in the database.
516
+
517
+ Args:
518
+ session_id (str): The ID of the session to rename.
519
+ session_type (SessionType): The type of session to rename.
520
+ session_name (str): The new name of the session.
521
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
522
+
523
+ Returns:
524
+ Optional[Union[Session, Dict[str, Any]]]:
525
+ - When deserialize=True: Session object
526
+ - When deserialize=False: Session dictionary
527
+
528
+ Raises:
529
+ Exception: If there is an error renaming the session.
530
+ """
531
+ try:
532
+ collection = await self._get_collection(table_type="sessions")
533
+ if collection is None:
534
+ return None
535
+
536
+ try:
537
+ result = await collection.find_one_and_update(
538
+ {"session_id": session_id},
539
+ {"$set": {"session_data.session_name": session_name, "updated_at": int(time.time())}},
540
+ return_document=ReturnDocument.AFTER,
541
+ upsert=False,
542
+ )
543
+ except OperationFailure:
544
+ # If the update fails because session_data doesn't contain a session_name yet, we initialize session_data
545
+ result = await collection.find_one_and_update(
546
+ {"session_id": session_id},
547
+ {"$set": {"session_data": {"session_name": session_name}, "updated_at": int(time.time())}},
548
+ return_document=ReturnDocument.AFTER,
549
+ upsert=False,
550
+ )
551
+ if not result:
552
+ return None
553
+
554
+ deserialized_session = deserialize_session_json_fields(result)
555
+
556
+ if not deserialize:
557
+ return deserialized_session
558
+
559
+ if session_type == SessionType.AGENT.value:
560
+ return AgentSession.from_dict(deserialized_session)
561
+ elif session_type == SessionType.TEAM.value:
562
+ return TeamSession.from_dict(deserialized_session)
563
+ else:
564
+ return WorkflowSession.from_dict(deserialized_session)
565
+
566
+ except Exception as e:
567
+ log_error(f"Exception renaming session: {e}")
568
+ raise e
569
+
570
+ async def upsert_session(
571
+ self, session: Session, deserialize: Optional[bool] = True
572
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
573
+ """Insert or update a session in the database.
574
+
575
+ Args:
576
+ session (Session): The session to upsert.
577
+ deserialize (Optional[bool]): Whether to deserialize the session. Defaults to True.
578
+
579
+ Returns:
580
+ Optional[Union[Session, Dict[str, Any]]]: The upserted session.
581
+
582
+ Raises:
583
+ Exception: If there is an error upserting the session.
584
+ """
585
+ try:
586
+ collection = await self._get_collection(table_type="sessions", create_collection_if_not_found=True)
587
+ if collection is None:
588
+ return None
589
+
590
+ session_dict = session.to_dict()
591
+
592
+ if isinstance(session, AgentSession):
593
+ record = {
594
+ "session_id": session_dict.get("session_id"),
595
+ "session_type": SessionType.AGENT.value,
596
+ "agent_id": session_dict.get("agent_id"),
597
+ "user_id": session_dict.get("user_id"),
598
+ "runs": session_dict.get("runs"),
599
+ "agent_data": session_dict.get("agent_data"),
600
+ "session_data": session_dict.get("session_data"),
601
+ "summary": session_dict.get("summary"),
602
+ "metadata": session_dict.get("metadata"),
603
+ "created_at": session_dict.get("created_at"),
604
+ "updated_at": int(time.time()),
605
+ }
606
+
607
+ result = await collection.find_one_and_replace(
608
+ filter={"session_id": session_dict.get("session_id")},
609
+ replacement=record,
610
+ upsert=True,
611
+ return_document=ReturnDocument.AFTER,
612
+ )
613
+ if not result:
614
+ return None
615
+
616
+ session = result # type: ignore
617
+
618
+ if not deserialize:
619
+ return session
620
+
621
+ return AgentSession.from_dict(session) # type: ignore
622
+
623
+ elif isinstance(session, TeamSession):
624
+ record = {
625
+ "session_id": session_dict.get("session_id"),
626
+ "session_type": SessionType.TEAM.value,
627
+ "team_id": session_dict.get("team_id"),
628
+ "user_id": session_dict.get("user_id"),
629
+ "runs": session_dict.get("runs"),
630
+ "team_data": session_dict.get("team_data"),
631
+ "session_data": session_dict.get("session_data"),
632
+ "summary": session_dict.get("summary"),
633
+ "metadata": session_dict.get("metadata"),
634
+ "created_at": session_dict.get("created_at"),
635
+ "updated_at": int(time.time()),
636
+ }
637
+
638
+ result = await collection.find_one_and_replace(
639
+ filter={"session_id": session_dict.get("session_id")},
640
+ replacement=record,
641
+ upsert=True,
642
+ return_document=ReturnDocument.AFTER,
643
+ )
644
+ if not result:
645
+ return None
646
+
647
+ # MongoDB stores native objects, no deserialization needed for document fields
648
+ session = result # type: ignore
649
+
650
+ if not deserialize:
651
+ return session
652
+
653
+ return TeamSession.from_dict(session) # type: ignore
654
+
655
+ else:
656
+ record = {
657
+ "session_id": session_dict.get("session_id"),
658
+ "session_type": SessionType.WORKFLOW.value,
659
+ "workflow_id": session_dict.get("workflow_id"),
660
+ "user_id": session_dict.get("user_id"),
661
+ "runs": session_dict.get("runs"),
662
+ "workflow_data": session_dict.get("workflow_data"),
663
+ "session_data": session_dict.get("session_data"),
664
+ "summary": session_dict.get("summary"),
665
+ "metadata": session_dict.get("metadata"),
666
+ "created_at": session_dict.get("created_at"),
667
+ "updated_at": int(time.time()),
668
+ }
669
+
670
+ result = await collection.find_one_and_replace(
671
+ filter={"session_id": session_dict.get("session_id")},
672
+ replacement=record,
673
+ upsert=True,
674
+ return_document=ReturnDocument.AFTER,
675
+ )
676
+ if not result:
677
+ return None
678
+
679
+ session = result # type: ignore
680
+
681
+ if not deserialize:
682
+ return session
683
+
684
+ return WorkflowSession.from_dict(session) # type: ignore
685
+
686
+ except Exception as e:
687
+ log_error(f"Exception upserting session: {e}")
688
+ raise e
689
+
690
+ async def upsert_sessions(
691
+ self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
692
+ ) -> List[Union[Session, Dict[str, Any]]]:
693
+ """
694
+ Bulk upsert multiple sessions for improved performance on large datasets.
695
+
696
+ Args:
697
+ sessions (List[Session]): List of sessions to upsert.
698
+ deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
699
+ preserve_updated_at (bool): If True, preserve the updated_at from the session object.
700
+
701
+ Returns:
702
+ List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
703
+
704
+ Raises:
705
+ Exception: If an error occurs during bulk upsert.
706
+ """
707
+ if not sessions:
708
+ return []
709
+
710
+ try:
711
+ collection = await self._get_collection(table_type="sessions", create_collection_if_not_found=True)
712
+ if collection is None:
713
+ log_info("Sessions collection not available, falling back to individual upserts")
714
+ return [
715
+ result
716
+ for session in sessions
717
+ if session is not None
718
+ for result in [await self.upsert_session(session, deserialize=deserialize)]
719
+ if result is not None
720
+ ]
721
+
722
+ from pymongo import ReplaceOne
723
+
724
+ operations = []
725
+ results: List[Union[Session, Dict[str, Any]]] = []
726
+
727
+ for session in sessions:
728
+ if session is None:
729
+ continue
730
+
731
+ session_dict = session.to_dict()
732
+
733
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
734
+ updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
735
+
736
+ if isinstance(session, AgentSession):
737
+ record = {
738
+ "session_id": session_dict.get("session_id"),
739
+ "session_type": SessionType.AGENT.value,
740
+ "agent_id": session_dict.get("agent_id"),
741
+ "user_id": session_dict.get("user_id"),
742
+ "runs": session_dict.get("runs"),
743
+ "agent_data": session_dict.get("agent_data"),
744
+ "session_data": session_dict.get("session_data"),
745
+ "summary": session_dict.get("summary"),
746
+ "metadata": session_dict.get("metadata"),
747
+ "created_at": session_dict.get("created_at"),
748
+ "updated_at": updated_at,
749
+ }
750
+ elif isinstance(session, TeamSession):
751
+ record = {
752
+ "session_id": session_dict.get("session_id"),
753
+ "session_type": SessionType.TEAM.value,
754
+ "team_id": session_dict.get("team_id"),
755
+ "user_id": session_dict.get("user_id"),
756
+ "runs": session_dict.get("runs"),
757
+ "team_data": session_dict.get("team_data"),
758
+ "session_data": session_dict.get("session_data"),
759
+ "summary": session_dict.get("summary"),
760
+ "metadata": session_dict.get("metadata"),
761
+ "created_at": session_dict.get("created_at"),
762
+ "updated_at": updated_at,
763
+ }
764
+ elif isinstance(session, WorkflowSession):
765
+ record = {
766
+ "session_id": session_dict.get("session_id"),
767
+ "session_type": SessionType.WORKFLOW.value,
768
+ "workflow_id": session_dict.get("workflow_id"),
769
+ "user_id": session_dict.get("user_id"),
770
+ "runs": session_dict.get("runs"),
771
+ "workflow_data": session_dict.get("workflow_data"),
772
+ "session_data": session_dict.get("session_data"),
773
+ "summary": session_dict.get("summary"),
774
+ "metadata": session_dict.get("metadata"),
775
+ "created_at": session_dict.get("created_at"),
776
+ "updated_at": updated_at,
777
+ }
778
+ else:
779
+ continue
780
+
781
+ operations.append(
782
+ ReplaceOne(filter={"session_id": record["session_id"]}, replacement=record, upsert=True)
783
+ )
784
+
785
+ if operations:
786
+ # Execute bulk write
787
+ await collection.bulk_write(operations)
788
+
789
+ # Fetch the results
790
+ session_ids = [session.session_id for session in sessions if session and session.session_id]
791
+ cursor = collection.find({"session_id": {"$in": session_ids}})
792
+
793
+ async for doc in cursor:
794
+ session_dict = doc
795
+
796
+ if deserialize:
797
+ session_type = doc.get("session_type")
798
+ if session_type == SessionType.AGENT.value:
799
+ deserialized_agent_session = AgentSession.from_dict(session_dict)
800
+ if deserialized_agent_session is None:
801
+ continue
802
+ results.append(deserialized_agent_session)
803
+
804
+ elif session_type == SessionType.TEAM.value:
805
+ deserialized_team_session = TeamSession.from_dict(session_dict)
806
+ if deserialized_team_session is None:
807
+ continue
808
+ results.append(deserialized_team_session)
809
+
810
+ elif session_type == SessionType.WORKFLOW.value:
811
+ deserialized_workflow_session = WorkflowSession.from_dict(session_dict)
812
+ if deserialized_workflow_session is None:
813
+ continue
814
+ results.append(deserialized_workflow_session)
815
+ else:
816
+ results.append(session_dict)
817
+
818
+ return results
819
+
820
+ except Exception as e:
821
+ log_error(f"Exception during bulk session upsert, falling back to individual upserts: {e}")
822
+
823
+ # Fallback to individual upserts
824
+ return [
825
+ result
826
+ for session in sessions
827
+ if session is not None
828
+ for result in [await self.upsert_session(session, deserialize=deserialize)]
829
+ if result is not None
830
+ ]
831
+
832
+ # -- Memory methods --
833
+
834
+ async def delete_user_memory(self, memory_id: str, user_id: Optional[str] = None):
835
+ """Delete a user memory from the database.
836
+
837
+ Args:
838
+ memory_id (str): The ID of the memory to delete.
839
+ user_id (Optional[str]): The ID of the user to verify ownership. If provided, only delete if the memory belongs to this user.
840
+
841
+ Returns:
842
+ bool: True if the memory was deleted, False otherwise.
843
+
844
+ Raises:
845
+ Exception: If there is an error deleting the memory.
846
+ """
847
+ try:
848
+ collection = await self._get_collection(table_type="memories")
849
+ if collection is None:
850
+ return
851
+
852
+ query = {"memory_id": memory_id}
853
+ if user_id is not None:
854
+ query["user_id"] = user_id
855
+
856
+ result = await collection.delete_one(query)
857
+
858
+ success = result.deleted_count > 0
859
+ if success:
860
+ log_debug(f"Successfully deleted memory id: {memory_id}")
861
+ else:
862
+ log_debug(f"No memory found with id: {memory_id}")
863
+
864
+ except Exception as e:
865
+ log_error(f"Error deleting memory: {e}")
866
+ raise e
867
+
868
+ async def delete_user_memories(self, memory_ids: List[str], user_id: Optional[str] = None) -> None:
869
+ """Delete user memories from the database.
870
+
871
+ Args:
872
+ memory_ids (List[str]): The IDs of the memories to delete.
873
+ user_id (Optional[str]): The ID of the user to verify ownership. If provided, only delete memories that belong to this user.
874
+
875
+ Raises:
876
+ Exception: If there is an error deleting the memories.
877
+ """
878
+ try:
879
+ collection = await self._get_collection(table_type="memories")
880
+ if collection is None:
881
+ return
882
+
883
+ query: Dict[str, Any] = {"memory_id": {"$in": memory_ids}}
884
+ if user_id is not None:
885
+ query["user_id"] = user_id
886
+
887
+ result = await collection.delete_many(query)
888
+
889
+ if result.deleted_count == 0:
890
+ log_debug(f"No memories found with ids: {memory_ids}")
891
+
892
+ except Exception as e:
893
+ log_error(f"Error deleting memories: {e}")
894
+ raise e
895
+
896
+ async def get_all_memory_topics(self, user_id: Optional[str] = None) -> List[str]:
897
+ """Get all memory topics from the database.
898
+
899
+ Args:
900
+ user_id (Optional[str]): The ID of the user to filter by. Defaults to None.
901
+
902
+ Returns:
903
+ List[str]: The topics.
904
+
905
+ Raises:
906
+ Exception: If there is an error getting the topics.
907
+ """
908
+ try:
909
+ collection = await self._get_collection(table_type="memories")
910
+ if collection is None:
911
+ return []
912
+
913
+ query = {}
914
+ if user_id is not None:
915
+ query["user_id"] = user_id
916
+
917
+ topics = await collection.distinct("topics", query)
918
+ return [topic for topic in topics if topic]
919
+
920
+ except Exception as e:
921
+ log_error(f"Exception reading from collection: {e}")
922
+ raise e
923
+
924
+ async def get_user_memory(
925
+ self, memory_id: str, deserialize: Optional[bool] = True, user_id: Optional[str] = None
926
+ ) -> Optional[UserMemory]:
927
+ """Get a memory from the database.
928
+
929
+ Args:
930
+ memory_id (str): The ID of the memory to get.
931
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
932
+ user_id (Optional[str]): The ID of the user to verify ownership. If provided, only return the memory if it belongs to this user.
933
+
934
+ Returns:
935
+ Optional[UserMemory]:
936
+ - When deserialize=True: UserMemory object
937
+ - When deserialize=False: Memory dictionary
938
+
939
+ Raises:
940
+ Exception: If there is an error getting the memory.
941
+ """
942
+ try:
943
+ collection = await self._get_collection(table_type="memories")
944
+ if collection is None:
945
+ return None
946
+
947
+ query = {"memory_id": memory_id}
948
+ if user_id is not None:
949
+ query["user_id"] = user_id
950
+
951
+ result = await collection.find_one(query)
952
+ if result is None or not deserialize:
953
+ return result
954
+
955
+ # Remove MongoDB's _id field before creating UserMemory object
956
+ result_filtered = {k: v for k, v in result.items() if k != "_id"}
957
+ return UserMemory.from_dict(result_filtered)
958
+
959
+ except Exception as e:
960
+ log_error(f"Exception reading from collection: {e}")
961
+ raise e
962
+
963
+ async def get_user_memories(
964
+ self,
965
+ user_id: Optional[str] = None,
966
+ agent_id: Optional[str] = None,
967
+ team_id: Optional[str] = None,
968
+ topics: Optional[List[str]] = None,
969
+ search_content: Optional[str] = None,
970
+ limit: Optional[int] = None,
971
+ page: Optional[int] = None,
972
+ sort_by: Optional[str] = None,
973
+ sort_order: Optional[str] = None,
974
+ deserialize: Optional[bool] = True,
975
+ ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
976
+ """Get all memories from the database as UserMemory objects.
977
+
978
+ Args:
979
+ user_id (Optional[str]): The ID of the user to get the memories for.
980
+ agent_id (Optional[str]): The ID of the agent to get the memories for.
981
+ team_id (Optional[str]): The ID of the team to get the memories for.
982
+ topics (Optional[List[str]]): The topics to filter the memories by.
983
+ search_content (Optional[str]): The content to filter the memories by.
984
+ limit (Optional[int]): The limit of the memories to get.
985
+ page (Optional[int]): The page number to get.
986
+ sort_by (Optional[str]): The field to sort the memories by.
987
+ sort_order (Optional[str]): The order to sort the memories by.
988
+ deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
989
+
990
+ Returns:
991
+ Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
992
+ - When deserialize=True: List of UserMemory objects
993
+ - When deserialize=False: Tuple of (memory dictionaries, total count)
994
+
995
+ Raises:
996
+ Exception: If there is an error getting the memories.
997
+ """
998
+ try:
999
+ collection = await self._get_collection(table_type="memories")
1000
+ if collection is None:
1001
+ return [] if deserialize else ([], 0)
1002
+
1003
+ query: Dict[str, Any] = {}
1004
+ if user_id is not None:
1005
+ query["user_id"] = user_id
1006
+ if agent_id is not None:
1007
+ query["agent_id"] = agent_id
1008
+ if team_id is not None:
1009
+ query["team_id"] = team_id
1010
+ if topics is not None:
1011
+ query["topics"] = {"$in": topics}
1012
+ if search_content is not None:
1013
+ query["memory"] = {"$regex": search_content, "$options": "i"}
1014
+
1015
+ # Get total count
1016
+ total_count = await collection.count_documents(query)
1017
+
1018
+ # Apply sorting
1019
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1020
+
1021
+ # Apply pagination
1022
+ query_args = apply_pagination({}, limit, page)
1023
+
1024
+ cursor = collection.find(query)
1025
+ if sort_criteria:
1026
+ cursor = cursor.sort(sort_criteria)
1027
+ if query_args.get("skip"):
1028
+ cursor = cursor.skip(query_args["skip"])
1029
+ if query_args.get("limit"):
1030
+ cursor = cursor.limit(query_args["limit"])
1031
+
1032
+ records = await cursor.to_list(length=None)
1033
+ if not deserialize:
1034
+ return records, total_count
1035
+
1036
+ # Remove MongoDB's _id field before creating UserMemory objects
1037
+ return [UserMemory.from_dict({k: v for k, v in record.items() if k != "_id"}) for record in records]
1038
+
1039
+ except Exception as e:
1040
+ log_error(f"Exception reading from collection: {e}")
1041
+ raise e
1042
+
1043
+ async def get_user_memory_stats(
1044
+ self,
1045
+ limit: Optional[int] = None,
1046
+ page: Optional[int] = None,
1047
+ user_id: Optional[str] = None,
1048
+ ) -> Tuple[List[Dict[str, Any]], int]:
1049
+ """Get user memories stats.
1050
+
1051
+ Args:
1052
+ limit (Optional[int]): The limit of the memories to get.
1053
+ page (Optional[int]): The page number to get.
1054
+ user_id (Optional[str]): The ID of the user to filter by. Defaults to None.
1055
+
1056
+ Returns:
1057
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories stats and the total count.
1058
+
1059
+ Raises:
1060
+ Exception: If there is an error getting the memories stats.
1061
+ """
1062
+ try:
1063
+ collection = await self._get_collection(table_type="memories")
1064
+ if collection is None:
1065
+ return [], 0
1066
+
1067
+ match_stage: Dict[str, Any] = {"user_id": {"$ne": None}}
1068
+ if user_id is not None:
1069
+ match_stage["user_id"] = user_id
1070
+
1071
+ pipeline = [
1072
+ {"$match": match_stage},
1073
+ {
1074
+ "$group": {
1075
+ "_id": "$user_id",
1076
+ "total_memories": {"$sum": 1},
1077
+ "last_memory_updated_at": {"$max": "$updated_at"},
1078
+ }
1079
+ },
1080
+ {"$sort": {"last_memory_updated_at": -1}},
1081
+ ]
1082
+
1083
+ # Get total count
1084
+ count_pipeline = pipeline + [{"$count": "total"}]
1085
+ count_result = await collection.aggregate(count_pipeline).to_list(length=1)
1086
+ total_count = count_result[0]["total"] if count_result else 0
1087
+
1088
+ # Apply pagination
1089
+ if limit is not None:
1090
+ if page is not None:
1091
+ pipeline.append({"$skip": (page - 1) * limit}) # type: ignore
1092
+ pipeline.append({"$limit": limit}) # type: ignore
1093
+
1094
+ results = await collection.aggregate(pipeline).to_list(length=None)
1095
+
1096
+ formatted_results = [
1097
+ {
1098
+ "user_id": result["_id"],
1099
+ "total_memories": result["total_memories"],
1100
+ "last_memory_updated_at": result["last_memory_updated_at"],
1101
+ }
1102
+ for result in results
1103
+ ]
1104
+
1105
+ return formatted_results, total_count
1106
+
1107
+ except Exception as e:
1108
+ log_error(f"Exception getting user memory stats: {e}")
1109
+ raise e
1110
+
1111
+ async def upsert_user_memory(
1112
+ self, memory: UserMemory, deserialize: Optional[bool] = True
1113
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
1114
+ """Upsert a user memory in the database.
1115
+
1116
+ Args:
1117
+ memory (UserMemory): The memory to upsert.
1118
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
1119
+
1120
+ Returns:
1121
+ Optional[Union[UserMemory, Dict[str, Any]]]:
1122
+ - When deserialize=True: UserMemory object
1123
+ - When deserialize=False: Memory dictionary
1124
+
1125
+ Raises:
1126
+ Exception: If there is an error upserting the memory.
1127
+ """
1128
+ try:
1129
+ collection = await self._get_collection(table_type="memories", create_collection_if_not_found=True)
1130
+ if collection is None:
1131
+ return None
1132
+
1133
+ if memory.memory_id is None:
1134
+ memory.memory_id = str(uuid4())
1135
+
1136
+ update_doc = {
1137
+ "user_id": memory.user_id,
1138
+ "agent_id": memory.agent_id,
1139
+ "team_id": memory.team_id,
1140
+ "memory_id": memory.memory_id,
1141
+ "memory": memory.memory,
1142
+ "topics": memory.topics,
1143
+ "updated_at": int(time.time()),
1144
+ }
1145
+
1146
+ result = await collection.replace_one({"memory_id": memory.memory_id}, update_doc, upsert=True)
1147
+
1148
+ if result.upserted_id:
1149
+ update_doc["_id"] = result.upserted_id
1150
+
1151
+ if not deserialize:
1152
+ return update_doc
1153
+
1154
+ # Remove MongoDB's _id field before creating UserMemory object
1155
+ update_doc_filtered = {k: v for k, v in update_doc.items() if k != "_id"}
1156
+ return UserMemory.from_dict(update_doc_filtered)
1157
+
1158
+ except Exception as e:
1159
+ log_error(f"Exception upserting user memory: {e}")
1160
+ raise e
1161
+
1162
+ async def upsert_memories(
1163
+ self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
1164
+ ) -> List[Union[UserMemory, Dict[str, Any]]]:
1165
+ """
1166
+ Bulk upsert multiple user memories for improved performance on large datasets.
1167
+
1168
+ Args:
1169
+ memories (List[UserMemory]): List of memories to upsert.
1170
+ deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
1171
+ preserve_updated_at (bool): If True, preserve the updated_at from the memory object.
1172
+
1173
+ Returns:
1174
+ List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
1175
+
1176
+ Raises:
1177
+ Exception: If an error occurs during bulk upsert.
1178
+ """
1179
+ if not memories:
1180
+ return []
1181
+
1182
+ try:
1183
+ collection = await self._get_collection(table_type="memories", create_collection_if_not_found=True)
1184
+ if collection is None:
1185
+ log_info("Memories collection not available, falling back to individual upserts")
1186
+ return [
1187
+ result
1188
+ for memory in memories
1189
+ if memory is not None
1190
+ for result in [await self.upsert_user_memory(memory, deserialize=deserialize)]
1191
+ if result is not None
1192
+ ]
1193
+
1194
+ from pymongo import ReplaceOne
1195
+
1196
+ operations = []
1197
+ results: List[Union[UserMemory, Dict[str, Any]]] = []
1198
+
1199
+ current_time = int(time.time())
1200
+ for memory in memories:
1201
+ if memory is None:
1202
+ continue
1203
+
1204
+ if memory.memory_id is None:
1205
+ memory.memory_id = str(uuid4())
1206
+
1207
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
1208
+ updated_at = memory.updated_at if preserve_updated_at else current_time
1209
+
1210
+ record = {
1211
+ "user_id": memory.user_id,
1212
+ "agent_id": memory.agent_id,
1213
+ "team_id": memory.team_id,
1214
+ "memory_id": memory.memory_id,
1215
+ "memory": memory.memory,
1216
+ "topics": memory.topics,
1217
+ "updated_at": updated_at,
1218
+ }
1219
+
1220
+ operations.append(ReplaceOne(filter={"memory_id": memory.memory_id}, replacement=record, upsert=True))
1221
+
1222
+ if operations:
1223
+ # Execute bulk write
1224
+ await collection.bulk_write(operations)
1225
+
1226
+ # Fetch the results
1227
+ memory_ids = [memory.memory_id for memory in memories if memory and memory.memory_id]
1228
+ cursor = collection.find({"memory_id": {"$in": memory_ids}})
1229
+
1230
+ async for doc in cursor:
1231
+ if deserialize:
1232
+ # Remove MongoDB's _id field before creating UserMemory object
1233
+ doc_filtered = {k: v for k, v in doc.items() if k != "_id"}
1234
+ results.append(UserMemory.from_dict(doc_filtered))
1235
+ else:
1236
+ results.append(doc)
1237
+
1238
+ return results
1239
+
1240
+ except Exception as e:
1241
+ log_error(f"Exception during bulk memory upsert, falling back to individual upserts: {e}")
1242
+
1243
+ # Fallback to individual upserts
1244
+ return [
1245
+ result
1246
+ for memory in memories
1247
+ if memory is not None
1248
+ for result in [await self.upsert_user_memory(memory, deserialize=deserialize)]
1249
+ if result is not None
1250
+ ]
1251
+
1252
+ async def clear_memories(self) -> None:
1253
+ """Delete all memories from the database.
1254
+
1255
+ Raises:
1256
+ Exception: If an error occurs during deletion.
1257
+ """
1258
+ try:
1259
+ collection = await self._get_collection(table_type="memories")
1260
+ if collection is None:
1261
+ return
1262
+
1263
+ await collection.delete_many({})
1264
+
1265
+ except Exception as e:
1266
+ log_error(f"Exception deleting all memories: {e}")
1267
+ raise e
1268
+
1269
+ # -- Cultural Knowledge methods --
1270
+ async def clear_cultural_knowledge(self) -> None:
1271
+ """Delete all cultural knowledge from the database.
1272
+
1273
+ Raises:
1274
+ Exception: If an error occurs during deletion.
1275
+ """
1276
+ try:
1277
+ collection = await self._get_collection(table_type="culture")
1278
+ if collection is None:
1279
+ return
1280
+
1281
+ await collection.delete_many({})
1282
+
1283
+ except Exception as e:
1284
+ log_error(f"Exception deleting all cultural knowledge: {e}")
1285
+ raise e
1286
+
1287
+ async def delete_cultural_knowledge(self, id: str) -> None:
1288
+ """Delete cultural knowledge by ID.
1289
+
1290
+ Args:
1291
+ id (str): The ID of the cultural knowledge to delete.
1292
+
1293
+ Raises:
1294
+ Exception: If an error occurs during deletion.
1295
+ """
1296
+ try:
1297
+ collection = await self._get_collection(table_type="culture")
1298
+ if collection is None:
1299
+ return
1300
+
1301
+ await collection.delete_one({"id": id})
1302
+ log_debug(f"Deleted cultural knowledge with ID: {id}")
1303
+
1304
+ except Exception as e:
1305
+ log_error(f"Error deleting cultural knowledge: {e}")
1306
+ raise e
1307
+
1308
+ async def get_cultural_knowledge(
1309
+ self, id: str, deserialize: Optional[bool] = True
1310
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
1311
+ """Get cultural knowledge by ID.
1312
+
1313
+ Args:
1314
+ id (str): The ID of the cultural knowledge to retrieve.
1315
+ deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge object. Defaults to True.
1316
+
1317
+ Returns:
1318
+ Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The cultural knowledge if found, None otherwise.
1319
+
1320
+ Raises:
1321
+ Exception: If an error occurs during retrieval.
1322
+ """
1323
+ try:
1324
+ collection = await self._get_collection(table_type="culture")
1325
+ if collection is None:
1326
+ return None
1327
+
1328
+ result = await collection.find_one({"id": id})
1329
+ if result is None:
1330
+ return None
1331
+
1332
+ # Remove MongoDB's _id field
1333
+ result_filtered = {k: v for k, v in result.items() if k != "_id"}
1334
+
1335
+ if not deserialize:
1336
+ return result_filtered
1337
+
1338
+ return deserialize_cultural_knowledge_from_db(result_filtered)
1339
+
1340
+ except Exception as e:
1341
+ log_error(f"Error getting cultural knowledge: {e}")
1342
+ raise e
1343
+
1344
+ async def get_all_cultural_knowledge(
1345
+ self,
1346
+ agent_id: Optional[str] = None,
1347
+ team_id: Optional[str] = None,
1348
+ name: Optional[str] = None,
1349
+ limit: Optional[int] = None,
1350
+ page: Optional[int] = None,
1351
+ sort_by: Optional[str] = None,
1352
+ sort_order: Optional[str] = None,
1353
+ deserialize: Optional[bool] = True,
1354
+ ) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
1355
+ """Get all cultural knowledge with filtering and pagination.
1356
+
1357
+ Args:
1358
+ agent_id (Optional[str]): Filter by agent ID.
1359
+ team_id (Optional[str]): Filter by team ID.
1360
+ name (Optional[str]): Filter by name (case-insensitive partial match).
1361
+ limit (Optional[int]): Maximum number of results to return.
1362
+ page (Optional[int]): Page number for pagination.
1363
+ sort_by (Optional[str]): Field to sort by.
1364
+ sort_order (Optional[str]): Sort order ('asc' or 'desc').
1365
+ deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge objects. Defaults to True.
1366
+
1367
+ Returns:
1368
+ Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
1369
+ - When deserialize=True: List of CulturalKnowledge objects
1370
+ - When deserialize=False: Tuple with list of dictionaries and total count
1371
+
1372
+ Raises:
1373
+ Exception: If an error occurs during retrieval.
1374
+ """
1375
+ try:
1376
+ collection = await self._get_collection(table_type="culture")
1377
+ if collection is None:
1378
+ if not deserialize:
1379
+ return [], 0
1380
+ return []
1381
+
1382
+ # Build query
1383
+ query: Dict[str, Any] = {}
1384
+ if agent_id is not None:
1385
+ query["agent_id"] = agent_id
1386
+ if team_id is not None:
1387
+ query["team_id"] = team_id
1388
+ if name is not None:
1389
+ query["name"] = {"$regex": name, "$options": "i"}
1390
+
1391
+ # Get total count for pagination
1392
+ total_count = await collection.count_documents(query)
1393
+
1394
+ # Apply sorting
1395
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1396
+
1397
+ # Apply pagination
1398
+ query_args = apply_pagination({}, limit, page)
1399
+
1400
+ cursor = collection.find(query)
1401
+ if sort_criteria:
1402
+ cursor = cursor.sort(sort_criteria)
1403
+ if query_args.get("skip"):
1404
+ cursor = cursor.skip(query_args["skip"])
1405
+ if query_args.get("limit"):
1406
+ cursor = cursor.limit(query_args["limit"])
1407
+
1408
+ # Remove MongoDB's _id field from all results
1409
+ results_filtered = [{k: v for k, v in item.items() if k != "_id"} async for item in cursor]
1410
+
1411
+ if not deserialize:
1412
+ return results_filtered, total_count
1413
+
1414
+ return [deserialize_cultural_knowledge_from_db(item) for item in results_filtered]
1415
+
1416
+ except Exception as e:
1417
+ log_error(f"Error getting all cultural knowledge: {e}")
1418
+ raise e
1419
+
1420
+ async def upsert_cultural_knowledge(
1421
+ self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
1422
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
1423
+ """Upsert cultural knowledge in MongoDB.
1424
+
1425
+ Args:
1426
+ cultural_knowledge (CulturalKnowledge): The cultural knowledge to upsert.
1427
+ deserialize (Optional[bool]): Whether to deserialize the result. Defaults to True.
1428
+
1429
+ Returns:
1430
+ Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The upserted cultural knowledge.
1431
+
1432
+ Raises:
1433
+ Exception: If an error occurs during upsert.
1434
+ """
1435
+ try:
1436
+ collection = await self._get_collection(table_type="culture", create_collection_if_not_found=True)
1437
+ if collection is None:
1438
+ return None
1439
+
1440
+ # Serialize content, categories, and notes into a dict for DB storage
1441
+ content_dict = serialize_cultural_knowledge_for_db(cultural_knowledge)
1442
+
1443
+ # Create the document with serialized content
1444
+ update_doc = {
1445
+ "id": cultural_knowledge.id,
1446
+ "name": cultural_knowledge.name,
1447
+ "summary": cultural_knowledge.summary,
1448
+ "content": content_dict if content_dict else None,
1449
+ "metadata": cultural_knowledge.metadata,
1450
+ "input": cultural_knowledge.input,
1451
+ "created_at": cultural_knowledge.created_at,
1452
+ "updated_at": int(time.time()),
1453
+ "agent_id": cultural_knowledge.agent_id,
1454
+ "team_id": cultural_knowledge.team_id,
1455
+ }
1456
+
1457
+ result = await collection.replace_one({"id": cultural_knowledge.id}, update_doc, upsert=True)
1458
+
1459
+ if result.upserted_id:
1460
+ update_doc["_id"] = result.upserted_id
1461
+
1462
+ # Remove MongoDB's _id field
1463
+ doc_filtered = {k: v for k, v in update_doc.items() if k != "_id"}
1464
+
1465
+ if not deserialize:
1466
+ return doc_filtered
1467
+
1468
+ return deserialize_cultural_knowledge_from_db(doc_filtered)
1469
+
1470
+ except Exception as e:
1471
+ log_error(f"Error upserting cultural knowledge: {e}")
1472
+ raise e
1473
+
1474
+ # -- Metrics methods --
1475
+
1476
+ async def _get_all_sessions_for_metrics_calculation(
1477
+ self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
1478
+ ) -> List[Dict[str, Any]]:
1479
+ """Get all sessions of all types for metrics calculation."""
1480
+ try:
1481
+ collection = await self._get_collection(table_type="sessions")
1482
+ if collection is None:
1483
+ return []
1484
+
1485
+ query = {}
1486
+ if start_timestamp is not None:
1487
+ query["created_at"] = {"$gte": start_timestamp}
1488
+ if end_timestamp is not None:
1489
+ if "created_at" in query:
1490
+ query["created_at"]["$lte"] = end_timestamp
1491
+ else:
1492
+ query["created_at"] = {"$lte": end_timestamp}
1493
+
1494
+ projection = {
1495
+ "user_id": 1,
1496
+ "session_data": 1,
1497
+ "runs": 1,
1498
+ "created_at": 1,
1499
+ "session_type": 1,
1500
+ }
1501
+
1502
+ results = await collection.find(query, projection).to_list(length=None)
1503
+ return results
1504
+
1505
+ except Exception as e:
1506
+ log_error(f"Exception reading from sessions collection: {e}")
1507
+ return []
1508
+
1509
+ async def _get_metrics_calculation_starting_date(self, collection: AsyncIOMotorCollection) -> Optional[date]:
1510
+ """Get the first date for which metrics calculation is needed."""
1511
+ try:
1512
+ result = await collection.find_one({}, sort=[("date", -1)], limit=1)
1513
+
1514
+ if result is not None:
1515
+ result_date = datetime.strptime(result["date"], "%Y-%m-%d").date()
1516
+ if result.get("completed"):
1517
+ return result_date + timedelta(days=1)
1518
+ else:
1519
+ return result_date
1520
+
1521
+ # No metrics records. Return the date of the first recorded session.
1522
+ first_session_result = await self.get_sessions(
1523
+ sort_by="created_at", sort_order="asc", limit=1, deserialize=False
1524
+ )
1525
+ first_session_date = first_session_result[0][0]["created_at"] if first_session_result[0] else None # type: ignore
1526
+
1527
+ if first_session_date is None:
1528
+ return None
1529
+
1530
+ return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
1531
+
1532
+ except Exception as e:
1533
+ log_error(f"Exception getting metrics calculation starting date: {e}")
1534
+ return None
1535
+
1536
+ async def calculate_metrics(self) -> Optional[list[dict]]:
1537
+ """Calculate metrics for all dates without complete metrics."""
1538
+ try:
1539
+ collection = await self._get_collection(table_type="metrics", create_collection_if_not_found=True)
1540
+ if collection is None:
1541
+ return None
1542
+
1543
+ starting_date = await self._get_metrics_calculation_starting_date(collection)
1544
+ if starting_date is None:
1545
+ log_info("No session data found. Won't calculate metrics.")
1546
+ return None
1547
+
1548
+ dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
1549
+ if not dates_to_process:
1550
+ log_info("Metrics already calculated for all relevant dates.")
1551
+ return None
1552
+
1553
+ start_timestamp = int(
1554
+ datetime.combine(dates_to_process[0], datetime.min.time()).replace(tzinfo=timezone.utc).timestamp()
1555
+ )
1556
+ end_timestamp = int(
1557
+ datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time())
1558
+ .replace(tzinfo=timezone.utc)
1559
+ .timestamp()
1560
+ )
1561
+
1562
+ sessions = await self._get_all_sessions_for_metrics_calculation(
1563
+ start_timestamp=start_timestamp, end_timestamp=end_timestamp
1564
+ )
1565
+ all_sessions_data = fetch_all_sessions_data(
1566
+ sessions=sessions, dates_to_process=dates_to_process, start_timestamp=start_timestamp
1567
+ )
1568
+ if not all_sessions_data:
1569
+ log_info("No new session data found. Won't calculate metrics.")
1570
+ return None
1571
+
1572
+ results = []
1573
+ metrics_records = []
1574
+
1575
+ for date_to_process in dates_to_process:
1576
+ date_key = date_to_process.isoformat()
1577
+ sessions_for_date = all_sessions_data.get(date_key, {})
1578
+
1579
+ # Skip dates with no sessions
1580
+ if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
1581
+ continue
1582
+
1583
+ metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
1584
+ metrics_records.append(metrics_record)
1585
+
1586
+ if metrics_records:
1587
+ results = bulk_upsert_metrics(collection, metrics_records) # type: ignore
1588
+
1589
+ return results
1590
+
1591
+ except Exception as e:
1592
+ log_error(f"Error calculating metrics: {e}")
1593
+ raise e
1594
+
1595
+ async def get_metrics(
1596
+ self,
1597
+ starting_date: Optional[date] = None,
1598
+ ending_date: Optional[date] = None,
1599
+ ) -> Tuple[List[dict], Optional[int]]:
1600
+ """Get all metrics matching the given date range."""
1601
+ try:
1602
+ collection = await self._get_collection(table_type="metrics")
1603
+ if collection is None:
1604
+ return [], None
1605
+
1606
+ query = {}
1607
+ if starting_date:
1608
+ query["date"] = {"$gte": starting_date.isoformat()}
1609
+ if ending_date:
1610
+ if "date" in query:
1611
+ query["date"]["$lte"] = ending_date.isoformat()
1612
+ else:
1613
+ query["date"] = {"$lte": ending_date.isoformat()}
1614
+
1615
+ records = await collection.find(query).to_list(length=None)
1616
+ if not records:
1617
+ return [], None
1618
+
1619
+ # Get the latest updated_at
1620
+ latest_updated_at = max(record.get("updated_at", 0) for record in records)
1621
+
1622
+ return records, latest_updated_at
1623
+
1624
+ except Exception as e:
1625
+ log_error(f"Error getting metrics: {e}")
1626
+ raise e
1627
+
1628
+ # -- Knowledge methods --
1629
+
1630
+ async def delete_knowledge_content(self, id: str):
1631
+ """Delete a knowledge row from the database.
1632
+
1633
+ Args:
1634
+ id (str): The ID of the knowledge row to delete.
1635
+
1636
+ Raises:
1637
+ Exception: If an error occurs during deletion.
1638
+ """
1639
+ try:
1640
+ collection = await self._get_collection(table_type="knowledge")
1641
+ if collection is None:
1642
+ return
1643
+
1644
+ await collection.delete_one({"id": id})
1645
+
1646
+ log_debug(f"Deleted knowledge content with id '{id}'")
1647
+
1648
+ except Exception as e:
1649
+ log_error(f"Error deleting knowledge content: {e}")
1650
+ raise e
1651
+
1652
+ async def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
1653
+ """Get a knowledge row from the database.
1654
+
1655
+ Args:
1656
+ id (str): The ID of the knowledge row to get.
1657
+
1658
+ Returns:
1659
+ Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
1660
+
1661
+ Raises:
1662
+ Exception: If an error occurs during retrieval.
1663
+ """
1664
+ try:
1665
+ collection = await self._get_collection(table_type="knowledge")
1666
+ if collection is None:
1667
+ return None
1668
+
1669
+ result = await collection.find_one({"id": id})
1670
+ if result is None:
1671
+ return None
1672
+
1673
+ return KnowledgeRow.model_validate(result)
1674
+
1675
+ except Exception as e:
1676
+ log_error(f"Error getting knowledge content: {e}")
1677
+ raise e
1678
+
1679
+ async def get_knowledge_contents(
1680
+ self,
1681
+ limit: Optional[int] = None,
1682
+ page: Optional[int] = None,
1683
+ sort_by: Optional[str] = None,
1684
+ sort_order: Optional[str] = None,
1685
+ ) -> Tuple[List[KnowledgeRow], int]:
1686
+ """Get all knowledge contents from the database.
1687
+
1688
+ Args:
1689
+ limit (Optional[int]): The maximum number of knowledge contents to return.
1690
+ page (Optional[int]): The page number.
1691
+ sort_by (Optional[str]): The column to sort by.
1692
+ sort_order (Optional[str]): The order to sort by.
1693
+
1694
+ Returns:
1695
+ Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
1696
+
1697
+ Raises:
1698
+ Exception: If an error occurs during retrieval.
1699
+ """
1700
+ try:
1701
+ collection = await self._get_collection(table_type="knowledge")
1702
+ if collection is None:
1703
+ return [], 0
1704
+
1705
+ query: Dict[str, Any] = {}
1706
+
1707
+ # Get total count
1708
+ total_count = await collection.count_documents(query)
1709
+
1710
+ # Apply sorting
1711
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1712
+
1713
+ # Apply pagination
1714
+ query_args = apply_pagination({}, limit, page)
1715
+
1716
+ cursor = collection.find(query)
1717
+ if sort_criteria:
1718
+ cursor = cursor.sort(sort_criteria)
1719
+ if query_args.get("skip"):
1720
+ cursor = cursor.skip(query_args["skip"])
1721
+ if query_args.get("limit"):
1722
+ cursor = cursor.limit(query_args["limit"])
1723
+
1724
+ records = await cursor.to_list(length=None)
1725
+ knowledge_rows = [KnowledgeRow.model_validate(record) for record in records]
1726
+
1727
+ return knowledge_rows, total_count
1728
+
1729
+ except Exception as e:
1730
+ log_error(f"Error getting knowledge contents: {e}")
1731
+ raise e
1732
+
1733
+ async def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
1734
+ """Upsert knowledge content in the database.
1735
+
1736
+ Args:
1737
+ knowledge_row (KnowledgeRow): The knowledge row to upsert.
1738
+
1739
+ Returns:
1740
+ Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
1741
+
1742
+ Raises:
1743
+ Exception: If an error occurs during upsert.
1744
+ """
1745
+ try:
1746
+ collection = await self._get_collection(table_type="knowledge", create_collection_if_not_found=True)
1747
+ if collection is None:
1748
+ return None
1749
+
1750
+ update_doc = knowledge_row.model_dump()
1751
+ await collection.replace_one({"id": knowledge_row.id}, update_doc, upsert=True)
1752
+
1753
+ return knowledge_row
1754
+
1755
+ except Exception as e:
1756
+ log_error(f"Error upserting knowledge content: {e}")
1757
+ raise e
1758
+
1759
+ # -- Eval methods --
1760
+
1761
+ async def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
1762
+ """Create an EvalRunRecord in the database."""
1763
+ try:
1764
+ collection = await self._get_collection(table_type="evals", create_collection_if_not_found=True)
1765
+ if collection is None:
1766
+ return None
1767
+
1768
+ current_time = int(time.time())
1769
+ eval_dict = eval_run.model_dump()
1770
+ eval_dict["created_at"] = current_time
1771
+ eval_dict["updated_at"] = current_time
1772
+
1773
+ await collection.insert_one(eval_dict)
1774
+
1775
+ log_debug(f"Created eval run with id '{eval_run.run_id}'")
1776
+
1777
+ return eval_run
1778
+
1779
+ except Exception as e:
1780
+ log_error(f"Error creating eval run: {e}")
1781
+ raise e
1782
+
1783
+ async def delete_eval_run(self, eval_run_id: str) -> None:
1784
+ """Delete an eval run from the database."""
1785
+ try:
1786
+ collection = await self._get_collection(table_type="evals")
1787
+ if collection is None:
1788
+ return
1789
+
1790
+ result = await collection.delete_one({"run_id": eval_run_id})
1791
+
1792
+ if result.deleted_count == 0:
1793
+ log_debug(f"No eval run found with ID: {eval_run_id}")
1794
+ else:
1795
+ log_debug(f"Deleted eval run with ID: {eval_run_id}")
1796
+
1797
+ except Exception as e:
1798
+ log_error(f"Error deleting eval run {eval_run_id}: {e}")
1799
+ raise e
1800
+
1801
+ async def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
1802
+ """Delete multiple eval runs from the database."""
1803
+ try:
1804
+ collection = await self._get_collection(table_type="evals")
1805
+ if collection is None:
1806
+ return
1807
+
1808
+ result = await collection.delete_many({"run_id": {"$in": eval_run_ids}})
1809
+
1810
+ if result.deleted_count == 0:
1811
+ log_debug(f"No eval runs found with IDs: {eval_run_ids}")
1812
+ else:
1813
+ log_debug(f"Deleted {result.deleted_count} eval runs")
1814
+
1815
+ except Exception as e:
1816
+ log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
1817
+ raise e
1818
+
1819
+ async def get_eval_run_raw(self, eval_run_id: str) -> Optional[Dict[str, Any]]:
1820
+ """Get an eval run from the database as a raw dictionary."""
1821
+ try:
1822
+ collection = await self._get_collection(table_type="evals")
1823
+ if collection is None:
1824
+ return None
1825
+
1826
+ result = await collection.find_one({"run_id": eval_run_id})
1827
+ return result
1828
+
1829
+ except Exception as e:
1830
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1831
+ raise e
1832
+
1833
+ async def get_eval_run(
1834
+ self, eval_run_id: str, deserialize: Optional[bool] = True
1835
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1836
+ """Get an eval run from the database.
1837
+
1838
+ Args:
1839
+ eval_run_id (str): The ID of the eval run to get.
1840
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1841
+
1842
+ Returns:
1843
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1844
+ - When deserialize=True: EvalRunRecord object
1845
+ - When deserialize=False: EvalRun dictionary
1846
+
1847
+ Raises:
1848
+ Exception: If there is an error getting the eval run.
1849
+ """
1850
+ try:
1851
+ collection = await self._get_collection(table_type="evals")
1852
+ if collection is None:
1853
+ return None
1854
+
1855
+ eval_run_raw = await collection.find_one({"run_id": eval_run_id})
1856
+
1857
+ if not eval_run_raw:
1858
+ return None
1859
+
1860
+ if not deserialize:
1861
+ return eval_run_raw
1862
+
1863
+ return EvalRunRecord.model_validate(eval_run_raw)
1864
+
1865
+ except Exception as e:
1866
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1867
+ raise e
1868
+
1869
+ async def get_eval_runs(
1870
+ self,
1871
+ limit: Optional[int] = None,
1872
+ page: Optional[int] = None,
1873
+ sort_by: Optional[str] = None,
1874
+ sort_order: Optional[str] = None,
1875
+ agent_id: Optional[str] = None,
1876
+ team_id: Optional[str] = None,
1877
+ workflow_id: Optional[str] = None,
1878
+ model_id: Optional[str] = None,
1879
+ filter_type: Optional[EvalFilterType] = None,
1880
+ eval_type: Optional[List[EvalType]] = None,
1881
+ deserialize: Optional[bool] = True,
1882
+ ) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1883
+ """Get all eval runs from the database.
1884
+
1885
+ Args:
1886
+ limit (Optional[int]): The maximum number of eval runs to return.
1887
+ page (Optional[int]): The page number to return.
1888
+ sort_by (Optional[str]): The field to sort by.
1889
+ sort_order (Optional[str]): The order to sort by.
1890
+ agent_id (Optional[str]): The ID of the agent to filter by.
1891
+ team_id (Optional[str]): The ID of the team to filter by.
1892
+ workflow_id (Optional[str]): The ID of the workflow to filter by.
1893
+ model_id (Optional[str]): The ID of the model to filter by.
1894
+ eval_type (Optional[List[EvalType]]): The type of eval to filter by.
1895
+ filter_type (Optional[EvalFilterType]): The type of filter to apply.
1896
+ deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
1897
+
1898
+ Returns:
1899
+ Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1900
+ - When deserialize=True: List of EvalRunRecord objects
1901
+ - When deserialize=False: List of eval run dictionaries and the total count
1902
+
1903
+ Raises:
1904
+ Exception: If there is an error getting the eval runs.
1905
+ """
1906
+ try:
1907
+ collection = await self._get_collection(table_type="evals")
1908
+ if collection is None:
1909
+ return [] if deserialize else ([], 0)
1910
+
1911
+ query: Dict[str, Any] = {}
1912
+ if agent_id is not None:
1913
+ query["agent_id"] = agent_id
1914
+ if team_id is not None:
1915
+ query["team_id"] = team_id
1916
+ if workflow_id is not None:
1917
+ query["workflow_id"] = workflow_id
1918
+ if model_id is not None:
1919
+ query["model_id"] = model_id
1920
+ if eval_type is not None and len(eval_type) > 0:
1921
+ query["eval_type"] = {"$in": eval_type}
1922
+ if filter_type is not None:
1923
+ if filter_type == EvalFilterType.AGENT:
1924
+ query["agent_id"] = {"$ne": None}
1925
+ elif filter_type == EvalFilterType.TEAM:
1926
+ query["team_id"] = {"$ne": None}
1927
+ elif filter_type == EvalFilterType.WORKFLOW:
1928
+ query["workflow_id"] = {"$ne": None}
1929
+
1930
+ # Get total count
1931
+ total_count = await collection.count_documents(query)
1932
+
1933
+ # Apply default sorting by created_at desc if no sort parameters provided
1934
+ if sort_by is None:
1935
+ sort_criteria = [("created_at", -1)]
1936
+ else:
1937
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1938
+
1939
+ # Apply pagination
1940
+ query_args = apply_pagination({}, limit, page)
1941
+
1942
+ cursor = collection.find(query)
1943
+ if sort_criteria:
1944
+ cursor = cursor.sort(sort_criteria)
1945
+ if query_args.get("skip"):
1946
+ cursor = cursor.skip(query_args["skip"])
1947
+ if query_args.get("limit"):
1948
+ cursor = cursor.limit(query_args["limit"])
1949
+
1950
+ records = await cursor.to_list(length=None)
1951
+ if not records:
1952
+ return [] if deserialize else ([], 0)
1953
+
1954
+ if not deserialize:
1955
+ return records, total_count
1956
+
1957
+ return [EvalRunRecord.model_validate(row) for row in records]
1958
+
1959
+ except Exception as e:
1960
+ log_error(f"Exception getting eval runs: {e}")
1961
+ raise e
1962
+
1963
+ async def rename_eval_run(
1964
+ self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
1965
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1966
+ """Update the name of an eval run in the database.
1967
+
1968
+ Args:
1969
+ eval_run_id (str): The ID of the eval run to update.
1970
+ name (str): The new name of the eval run.
1971
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1972
+
1973
+ Returns:
1974
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1975
+ - When deserialize=True: EvalRunRecord object
1976
+ - When deserialize=False: EvalRun dictionary
1977
+
1978
+ Raises:
1979
+ Exception: If there is an error updating the eval run.
1980
+ """
1981
+ try:
1982
+ collection = await self._get_collection(table_type="evals")
1983
+ if collection is None:
1984
+ return None
1985
+
1986
+ result = await collection.find_one_and_update(
1987
+ {"run_id": eval_run_id}, {"$set": {"name": name, "updated_at": int(time.time())}}
1988
+ )
1989
+
1990
+ log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
1991
+
1992
+ if not result or not deserialize:
1993
+ return result
1994
+
1995
+ return EvalRunRecord.model_validate(result)
1996
+
1997
+ except Exception as e:
1998
+ log_error(f"Error updating eval run name {eval_run_id}: {e}")
1999
+ raise e