agno 2.0.8__py3-none-any.whl → 2.0.9__py3-none-any.whl

agno/db/mongo/mongo.py

@@ -223,7 +223,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting session: {e}")
-            return False
+            raise e
 
     def delete_sessions(self, session_ids: List[str]) -> None:
         """Delete multiple sessions from the database.
@@ -241,6 +241,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting sessions: {e}")
+            raise e
 
     def get_session(
         self,
@@ -296,7 +297,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading session: {e}")
-            return None
+            raise e
 
     def get_sessions(
         self,
@@ -409,7 +410,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading sessions: {e}")
-            return [] if deserialize else ([], 0)
+            raise e
 
     def rename_session(
         self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
@@ -467,7 +468,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception renaming session: {e}")
-            return None
+            raise e
 
     def upsert_session(
         self, session: Session, deserialize: Optional[bool] = True
@@ -585,7 +586,145 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception upserting session: {e}")
-            return None
+            raise e
+
+    def upsert_sessions(
+        self, sessions: List[Session], deserialize: Optional[bool] = True
+    ) -> List[Union[Session, Dict[str, Any]]]:
+        """
+        Bulk upsert multiple sessions for improved performance on large datasets.
+
+        Args:
+            sessions (List[Session]): List of sessions to upsert.
+            deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
+
+        Returns:
+            List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
+
+        Raises:
+            Exception: If an error occurs during bulk upsert.
+        """
+        if not sessions:
+            return []
+
+        try:
+            collection = self._get_collection(table_type="sessions", create_collection_if_not_found=True)
+            if collection is None:
+                log_info("Sessions collection not available, falling back to individual upserts")
+                return [
+                    result
+                    for session in sessions
+                    if session is not None
+                    for result in [self.upsert_session(session, deserialize=deserialize)]
+                    if result is not None
+                ]
+
+            from pymongo import ReplaceOne
+
+            operations = []
+            results: List[Union[Session, Dict[str, Any]]] = []
+
+            for session in sessions:
+                if session is None:
+                    continue
+
+                serialized_session_dict = serialize_session_json_fields(session.to_dict())
+
+                if isinstance(session, AgentSession):
+                    record = {
+                        "session_id": serialized_session_dict.get("session_id"),
+                        "session_type": SessionType.AGENT.value,
+                        "agent_id": serialized_session_dict.get("agent_id"),
+                        "user_id": serialized_session_dict.get("user_id"),
+                        "runs": serialized_session_dict.get("runs"),
+                        "agent_data": serialized_session_dict.get("agent_data"),
+                        "session_data": serialized_session_dict.get("session_data"),
+                        "summary": serialized_session_dict.get("summary"),
+                        "metadata": serialized_session_dict.get("metadata"),
+                        "created_at": serialized_session_dict.get("created_at"),
+                        "updated_at": int(time.time()),
+                    }
+                elif isinstance(session, TeamSession):
+                    record = {
+                        "session_id": serialized_session_dict.get("session_id"),
+                        "session_type": SessionType.TEAM.value,
+                        "team_id": serialized_session_dict.get("team_id"),
+                        "user_id": serialized_session_dict.get("user_id"),
+                        "runs": serialized_session_dict.get("runs"),
+                        "team_data": serialized_session_dict.get("team_data"),
+                        "session_data": serialized_session_dict.get("session_data"),
+                        "summary": serialized_session_dict.get("summary"),
+                        "metadata": serialized_session_dict.get("metadata"),
+                        "created_at": serialized_session_dict.get("created_at"),
+                        "updated_at": int(time.time()),
+                    }
+                elif isinstance(session, WorkflowSession):
+                    record = {
+                        "session_id": serialized_session_dict.get("session_id"),
+                        "session_type": SessionType.WORKFLOW.value,
+                        "workflow_id": serialized_session_dict.get("workflow_id"),
+                        "user_id": serialized_session_dict.get("user_id"),
+                        "runs": serialized_session_dict.get("runs"),
+                        "workflow_data": serialized_session_dict.get("workflow_data"),
+                        "session_data": serialized_session_dict.get("session_data"),
+                        "summary": serialized_session_dict.get("summary"),
+                        "metadata": serialized_session_dict.get("metadata"),
+                        "created_at": serialized_session_dict.get("created_at"),
+                        "updated_at": int(time.time()),
+                    }
+                else:
+                    continue
+
+                operations.append(
+                    ReplaceOne(filter={"session_id": record["session_id"]}, replacement=record, upsert=True)
+                )
+
+            if operations:
+                # Execute bulk write
+                collection.bulk_write(operations)
+
+                # Fetch the results
+                session_ids = [session.session_id for session in sessions if session and session.session_id]
+                cursor = collection.find({"session_id": {"$in": session_ids}})
+
+                for doc in cursor:
+                    session_dict = deserialize_session_json_fields(doc)
+
+                    if deserialize:
+                        session_type = doc.get("session_type")
+                        if session_type == SessionType.AGENT.value:
+                            deserialized_agent_session = AgentSession.from_dict(session_dict)
+                            if deserialized_agent_session is None:
+                                continue
+                            results.append(deserialized_agent_session)
+
+                        elif session_type == SessionType.TEAM.value:
+                            deserialized_team_session = TeamSession.from_dict(session_dict)
+                            if deserialized_team_session is None:
+                                continue
+                            results.append(deserialized_team_session)
+
+                        elif session_type == SessionType.WORKFLOW.value:
+                            deserialized_workflow_session = WorkflowSession.from_dict(session_dict)
+                            if deserialized_workflow_session is None:
+                                continue
+                            results.append(deserialized_workflow_session)
+                    else:
+                        results.append(session_dict)
+
+            return results
+
+        except Exception as e:
+            log_error(f"Exception during bulk session upsert, falling back to individual upserts: {e}")
+
+            # Fallback to individual upserts
+            return [
+                result
+                for session in sessions
+                if session is not None
+                for result in [self.upsert_session(session, deserialize=deserialize)]
+                if result is not None
+            ]
 
     # -- Memory methods --
 
@@ -616,6 +755,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting memory: {e}")
+            raise e
 
     def delete_user_memories(self, memory_ids: List[str]) -> None:
         """Delete user memories from the database.
@@ -638,6 +778,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting memories: {e}")
+            raise e
 
     def get_all_memory_topics(self) -> List[str]:
         """Get all memory topics from the database.
@@ -658,7 +799,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from collection: {e}")
-            return []
+            raise e
 
     def get_user_memory(self, memory_id: str, deserialize: Optional[bool] = True) -> Optional[UserMemory]:
         """Get a memory from the database.
@@ -690,7 +831,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from collection: {e}")
-            return None
+            raise e
 
     def get_user_memories(
         self,
@@ -769,7 +910,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from collection: {e}")
-            return []
+            raise e
 
     def get_user_memory_stats(
         self,
@@ -831,7 +972,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception getting user memory stats: {e}")
-            return [], 0
+            raise e
 
     def upsert_user_memory(
         self, memory: UserMemory, deserialize: Optional[bool] = True
@@ -882,7 +1023,92 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception upserting user memory: {e}")
-            return None
+            raise e
+
+    def upsert_memories(
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+    ) -> List[Union[UserMemory, Dict[str, Any]]]:
+        """
+        Bulk upsert multiple user memories for improved performance on large datasets.
+
+        Args:
+            memories (List[UserMemory]): List of memories to upsert.
+            deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
+
+        Returns:
+            List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
+
+        Raises:
+            Exception: If an error occurs during bulk upsert.
+        """
+        if not memories:
+            return []
+
+        try:
+            collection = self._get_collection(table_type="memories", create_collection_if_not_found=True)
+            if collection is None:
+                log_info("Memories collection not available, falling back to individual upserts")
+                return [
+                    result
+                    for memory in memories
+                    if memory is not None
+                    for result in [self.upsert_user_memory(memory, deserialize=deserialize)]
+                    if result is not None
+                ]
+
+            from pymongo import ReplaceOne
+
+            operations = []
+            results: List[Union[UserMemory, Dict[str, Any]]] = []
+
+            for memory in memories:
+                if memory is None:
+                    continue
+
+                if memory.memory_id is None:
+                    memory.memory_id = str(uuid4())
+
+                record = {
+                    "user_id": memory.user_id,
+                    "agent_id": memory.agent_id,
+                    "team_id": memory.team_id,
+                    "memory_id": memory.memory_id,
+                    "memory": memory.memory,
+                    "topics": memory.topics,
+                    "updated_at": int(time.time()),
+                }
+
+                operations.append(ReplaceOne(filter={"memory_id": memory.memory_id}, replacement=record, upsert=True))
+
+            if operations:
+                # Execute bulk write
+                collection.bulk_write(operations)
+
+                # Fetch the results
+                memory_ids = [memory.memory_id for memory in memories if memory and memory.memory_id]
+                cursor = collection.find({"memory_id": {"$in": memory_ids}})
+
+                for doc in cursor:
+                    if deserialize:
+                        # Remove MongoDB's _id field before creating UserMemory object
+                        doc_filtered = {k: v for k, v in doc.items() if k != "_id"}
+                        results.append(UserMemory.from_dict(doc_filtered))
+                    else:
+                        results.append(doc)
+
+            return results
+
+        except Exception as e:
+            log_error(f"Exception during bulk memory upsert, falling back to individual upserts: {e}")
+
+            # Fallback to individual upserts
+            return [
+                result
+                for memory in memories
+                if memory is not None
+                for result in [self.upsert_user_memory(memory, deserialize=deserialize)]
+                if result is not None
+            ]
 
     def clear_memories(self) -> None:
         """Delete all memories from the database.
@@ -898,9 +1124,8 @@ class MongoDb(BaseDb):
             collection.delete_many({})
 
         except Exception as e:
-            from agno.utils.log import log_warning
-
-            log_warning(f"Exception deleting all memories: {e}")
+            log_error(f"Exception deleting all memories: {e}")
+            raise e
 
     # -- Metrics methods --
 
@@ -1052,7 +1277,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error getting metrics: {e}")
-            return [], None
+            raise e
 
     # -- Knowledge methods --
 
@@ -1076,7 +1301,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting knowledge content: {e}")
-            raise
+            raise e
 
     def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
         """Get a knowledge row from the database.
@@ -1103,7 +1328,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error getting knowledge content: {e}")
-            return None
+            raise e
 
     def get_knowledge_contents(
         self,
@@ -1158,7 +1383,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error getting knowledge contents: {e}")
-            return [], 0
+            raise e
 
     def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
         """Upsert knowledge content in the database.
@@ -1184,7 +1409,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error upserting knowledge content: {e}")
-            return None
+            raise e
 
     # -- Eval methods --
 
@@ -1208,7 +1433,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error creating eval run: {e}")
-            return None
+            raise e
 
     def delete_eval_run(self, eval_run_id: str) -> None:
         """Delete an eval run from the database."""
@@ -1226,7 +1451,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting eval run {eval_run_id}: {e}")
-            raise
+            raise e
 
     def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
         """Delete multiple eval runs from the database."""
@@ -1244,7 +1469,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
-            raise
+            raise e
 
     def get_eval_run_raw(self, eval_run_id: str) -> Optional[Dict[str, Any]]:
         """Get an eval run from the database as a raw dictionary."""
@@ -1258,7 +1483,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception getting eval run {eval_run_id}: {e}")
-            return None
+            raise e
 
     def get_eval_run(self, eval_run_id: str, deserialize: Optional[bool] = True) -> Optional[EvalRunRecord]:
         """Get an eval run from the database.
@@ -1292,7 +1517,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception getting eval run {eval_run_id}: {e}")
-            return None
+            raise e
 
     def get_eval_runs(
         self,
@@ -1386,8 +1611,8 @@ class MongoDb(BaseDb):
             return [EvalRunRecord.model_validate(row) for row in records]
 
         except Exception as e:
-            log_debug(f"Exception getting eval runs: {e}")
-            return [] if deserialize else ([], 0)
+            log_error(f"Exception getting eval runs: {e}")
+            raise e
 
     def rename_eval_run(
         self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
@@ -1425,7 +1650,7 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error updating eval run name {eval_run_id}: {e}")
-            raise
+            raise e
 
     def migrate_table_from_v1_to_v2(self, v1_db_schema: str, v1_table_name: str, v1_table_type: str):
         """Migrate all content in the given collection to the right v2 collection"""