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

agno/db/singlestore/singlestore.py

@@ -393,7 +393,7 @@ class SingleStoreDb(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 all given sessions from the database.
@@ -418,6 +418,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting sessions: {e}")
+            raise e
 
     def get_session(
         self,
@@ -476,7 +477,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from session table: {e}")
-            return None
+            raise e
 
     def get_sessions(
         self,
@@ -580,8 +581,8 @@ class SingleStoreDb(BaseDb):
                 raise ValueError(f"Invalid session type: {session_type}")
 
         except Exception as e:
-            log_debug(f"Exception reading from session table: {e}")
-            return []
+            log_error(f"Exception reading from session table: {e}")
+            raise e
 
     def rename_session(
         self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
@@ -643,7 +644,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error renaming session: {e}")
-            return None
+            raise e
 
     def upsert_session(self, session: Session, deserialize: Optional[bool] = True) -> Optional[Session]:
         """
@@ -793,7 +794,200 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error upserting into sessions table: {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:
+            table = self._get_table(table_type="sessions", create_table_if_not_found=True)
+            if table is None:
+                return []
+
+            # Group sessions by type for batch processing
+            agent_sessions = []
+            team_sessions = []
+            workflow_sessions = []
+
+            for session in sessions:
+                if isinstance(session, AgentSession):
+                    agent_sessions.append(session)
+                elif isinstance(session, TeamSession):
+                    team_sessions.append(session)
+                elif isinstance(session, WorkflowSession):
+                    workflow_sessions.append(session)
+
+            results: List[Union[Session, Dict[str, Any]]] = []
+
+            with self.Session() as sess, sess.begin():
+                # Bulk upsert agent sessions
+                if agent_sessions:
+                    agent_data = []
+                    for session in agent_sessions:
+                        session_dict = session.to_dict()
+                        agent_data.append(
+                            {
+                                "session_id": session_dict.get("session_id"),
+                                "session_type": SessionType.AGENT.value,
+                                "agent_id": session_dict.get("agent_id"),
+                                "user_id": session_dict.get("user_id"),
+                                "runs": session_dict.get("runs"),
+                                "agent_data": session_dict.get("agent_data"),
+                                "session_data": session_dict.get("session_data"),
+                                "summary": session_dict.get("summary"),
+                                "metadata": session_dict.get("metadata"),
+                                "created_at": session_dict.get("created_at"),
+                                "updated_at": session_dict.get("created_at"),
+                            }
+                        )
+
+                    if agent_data:
+                        stmt = mysql.insert(table)
+                        stmt = stmt.on_duplicate_key_update(
+                            agent_id=stmt.inserted.agent_id,
+                            user_id=stmt.inserted.user_id,
+                            agent_data=stmt.inserted.agent_data,
+                            session_data=stmt.inserted.session_data,
+                            summary=stmt.inserted.summary,
+                            metadata=stmt.inserted.metadata,
+                            runs=stmt.inserted.runs,
+                            updated_at=int(time.time()),
+                        )
+                        sess.execute(stmt, agent_data)
+
+                        # Fetch the results for agent sessions
+                        agent_ids = [session.session_id for session in agent_sessions]
+                        select_stmt = select(table).where(table.c.session_id.in_(agent_ids))
+                        result = sess.execute(select_stmt).fetchall()
+
+                        for row in result:
+                            if deserialize:
+                                deserialized_session = AgentSession.from_dict(session_dict)
+                                if deserialized_session is None:
+                                    continue
+                                results.append(deserialized_session)
+                            else:
+                                results.append(dict(row._mapping))
+
+                # Bulk upsert team sessions
+                if team_sessions:
+                    team_data = []
+                    for session in team_sessions:
+                        session_dict = session.to_dict()
+                        team_data.append(
+                            {
+                                "session_id": session_dict.get("session_id"),
+                                "session_type": SessionType.TEAM.value,
+                                "team_id": session_dict.get("team_id"),
+                                "user_id": session_dict.get("user_id"),
+                                "runs": session_dict.get("runs"),
+                                "team_data": session_dict.get("team_data"),
+                                "session_data": session_dict.get("session_data"),
+                                "summary": session_dict.get("summary"),
+                                "metadata": session_dict.get("metadata"),
+                                "created_at": session_dict.get("created_at"),
+                                "updated_at": session_dict.get("created_at"),
+                            }
+                        )
+
+                    if team_data:
+                        stmt = mysql.insert(table)
+                        stmt = stmt.on_duplicate_key_update(
+                            team_id=stmt.inserted.team_id,
+                            user_id=stmt.inserted.user_id,
+                            team_data=stmt.inserted.team_data,
+                            session_data=stmt.inserted.session_data,
+                            summary=stmt.inserted.summary,
+                            metadata=stmt.inserted.metadata,
+                            runs=stmt.inserted.runs,
+                            updated_at=int(time.time()),
+                        )
+                        sess.execute(stmt, team_data)
+
+                        # Fetch the results for team sessions
+                        team_ids = [session.session_id for session in team_sessions]
+                        select_stmt = select(table).where(table.c.session_id.in_(team_ids))
+                        result = sess.execute(select_stmt).fetchall()
+
+                        for row in result:
+                            if deserialize:
+                                deserialized_team_session = TeamSession.from_dict(session_dict)
+                                if deserialized_team_session is None:
+                                    continue
+                                results.append(deserialized_team_session)
+                            else:
+                                results.append(dict(row._mapping))
+
+                # Bulk upsert workflow sessions
+                if workflow_sessions:
+                    workflow_data = []
+                    for session in workflow_sessions:
+                        session_dict = session.to_dict()
+                        workflow_data.append(
+                            {
+                                "session_id": session_dict.get("session_id"),
+                                "session_type": SessionType.WORKFLOW.value,
+                                "workflow_id": session_dict.get("workflow_id"),
+                                "user_id": session_dict.get("user_id"),
+                                "runs": session_dict.get("runs"),
+                                "workflow_data": session_dict.get("workflow_data"),
+                                "session_data": session_dict.get("session_data"),
+                                "summary": session_dict.get("summary"),
+                                "metadata": session_dict.get("metadata"),
+                                "created_at": session_dict.get("created_at"),
+                                "updated_at": session_dict.get("created_at"),
+                            }
+                        )
+
+                    if workflow_data:
+                        stmt = mysql.insert(table)
+                        stmt = stmt.on_duplicate_key_update(
+                            workflow_id=stmt.inserted.workflow_id,
+                            user_id=stmt.inserted.user_id,
+                            workflow_data=stmt.inserted.workflow_data,
+                            session_data=stmt.inserted.session_data,
+                            summary=stmt.inserted.summary,
+                            metadata=stmt.inserted.metadata,
+                            runs=stmt.inserted.runs,
+                            updated_at=int(time.time()),
+                        )
+                        sess.execute(stmt, workflow_data)
+
+                        # Fetch the results for workflow sessions
+                        workflow_ids = [session.session_id for session in workflow_sessions]
+                        select_stmt = select(table).where(table.c.session_id.in_(workflow_ids))
+                        result = sess.execute(select_stmt).fetchall()
+
+                        for row in result:
+                            if deserialize:
+                                deserialized_workflow_session = WorkflowSession.from_dict(session_dict)
+                                if deserialized_workflow_session is None:
+                                    continue
+                                results.append(deserialized_workflow_session)
+                            else:
+                                results.append(dict(row._mapping))
+
+            return results
+
+        except Exception as e:
+            log_error(f"Exception during bulk session upsert: {e}")
+            return []
 
     # -- Memory methods --
     def delete_user_memory(self, memory_id: str):
@@ -825,6 +1019,7 @@ class SingleStoreDb(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.
@@ -848,6 +1043,7 @@ class SingleStoreDb(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.
@@ -875,7 +1071,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from memory table: {e}")
-            return []
+            raise e
 
     def get_user_memory(self, memory_id: str, deserialize: Optional[bool] = True) -> Optional[UserMemory]:
         """Get a memory from the database.
@@ -911,7 +1107,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from memory table: {e}")
-            return None
+            raise e
 
     def get_user_memories(
         self,
@@ -995,7 +1191,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from memory table: {e}")
-            return []
+            raise e
 
     def get_user_memory_stats(
         self, limit: Optional[int] = None, page: Optional[int] = None
@@ -1062,7 +1258,7 @@ class SingleStoreDb(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
@@ -1126,7 +1322,83 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error 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:
+            table = self._get_table(table_type="memories", create_table_if_not_found=True)
+            if table is None:
+                return []
+
+            # Prepare data for bulk insert
+            memory_data = []
+            for memory in memories:
+                if memory.memory_id is None:
+                    memory.memory_id = str(uuid4())
+                memory_data.append(
+                    {
+                        "memory_id": memory.memory_id,
+                        "memory": memory.memory,
+                        "input": memory.input,
+                        "user_id": memory.user_id,
+                        "agent_id": memory.agent_id,
+                        "team_id": memory.team_id,
+                        "topics": memory.topics,
+                        "updated_at": int(time.time()),
+                    }
+                )
+
+            results: List[Union[UserMemory, Dict[str, Any]]] = []
+
+            with self.Session() as sess, sess.begin():
+                if memory_data:
+                    stmt = mysql.insert(table)
+                    stmt = stmt.on_duplicate_key_update(
+                        memory=stmt.inserted.memory,
+                        topics=stmt.inserted.topics,
+                        input=stmt.inserted.input,
+                        user_id=stmt.inserted.user_id,
+                        agent_id=stmt.inserted.agent_id,
+                        team_id=stmt.inserted.team_id,
+                        updated_at=int(time.time()),
+                    )
+                    sess.execute(stmt, memory_data)
+
+                    # Fetch the results
+                    memory_ids = [memory.memory_id for memory in memories if memory.memory_id]
+                    select_stmt = select(table).where(table.c.memory_id.in_(memory_ids))
+                    result = sess.execute(select_stmt).fetchall()
+
+                    for row in result:
+                        memory_raw = dict(row._mapping)
+                        if deserialize:
+                            results.append(UserMemory.from_dict(memory_raw))
+                        else:
+                            results.append(memory_raw)
+
+            return results
+
+        except Exception as e:
+            log_error(f"Exception during bulk memory upsert: {e}")
+            return []
 
     def clear_memories(self) -> None:
         """Delete all memories from the database.
@@ -1143,7 +1415,8 @@ class SingleStoreDb(BaseDb):
                 sess.execute(table.delete())
 
         except Exception as e:
-            log_warning(f"Exception deleting all memories: {e}")
+            log_error(f"Exception deleting all memories: {e}")
+            raise e
 
     # -- Metrics methods --
     def _get_all_sessions_for_metrics_calculation(
@@ -1285,7 +1558,7 @@ class SingleStoreDb(BaseDb):
             return metrics_records
 
         except Exception as e:
-            log_error(f"Error refreshing metrics: {e}")
+            log_error(f"Error calculating metrics: {e}")
             raise e
 
     def get_metrics(
@@ -1328,7 +1601,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error getting metrics: {e}")
-            return [], None
+            raise e
 
     # -- Knowledge methods --
 
@@ -1338,15 +1611,19 @@ class SingleStoreDb(BaseDb):
         Args:
             id (str): The ID of the knowledge row to delete.
         """
-        table = self._get_table(table_type="knowledge")
-        if table is None:
-            return
+        try:
+            table = self._get_table(table_type="knowledge")
+            if table is None:
+                return
 
-        with self.Session() as sess, sess.begin():
-            stmt = table.delete().where(table.c.id == id)
-            sess.execute(stmt)
+            with self.Session() as sess, sess.begin():
+                stmt = table.delete().where(table.c.id == id)
+                sess.execute(stmt)
 
-        log_debug(f"Deleted knowledge content with id '{id}'")
+            log_debug(f"Deleted knowledge content with id '{id}'")
+        except Exception as e:
+            log_error(f"Error deleting knowledge content: {e}")
+            raise e
 
     def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
         """Get a knowledge row from the database.
@@ -1357,16 +1634,20 @@ class SingleStoreDb(BaseDb):
         Returns:
             Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
         """
-        table = self._get_table(table_type="knowledge")
-        if table is None:
-            return None
-
-        with self.Session() as sess, sess.begin():
-            stmt = select(table).where(table.c.id == id)
-            result = sess.execute(stmt).fetchone()
-            if result is None:
+        try:
+            table = self._get_table(table_type="knowledge")
+            if table is None:
                 return None
-            return KnowledgeRow.model_validate(result._mapping)
+
+            with self.Session() as sess, sess.begin():
+                stmt = select(table).where(table.c.id == id)
+                result = sess.execute(stmt).fetchone()
+                if result is None:
+                    return None
+                return KnowledgeRow.model_validate(result._mapping)
+        except Exception as e:
+            log_error(f"Error getting knowledge content: {e}")
+            raise e
 
     def get_knowledge_contents(
         self,
@@ -1419,7 +1700,7 @@ class SingleStoreDb(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.
@@ -1464,7 +1745,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error upserting knowledge row: {e}")
-            return None
+            raise e
 
     # -- Eval methods --
 
@@ -1498,7 +1779,7 @@ class SingleStoreDb(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.
@@ -1521,7 +1802,7 @@ class SingleStoreDb(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.
@@ -1544,7 +1825,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
-            raise
+            raise e
 
     def get_eval_run(
         self, eval_run_id: str, deserialize: Optional[bool] = True
@@ -1582,7 +1863,7 @@ class SingleStoreDb(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,
@@ -1677,7 +1958,7 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception getting eval runs: {e}")
-            return [] if deserialize else ([], 0)
+            raise e
 
     def rename_eval_run(
         self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
@@ -1716,4 +1997,4 @@ class SingleStoreDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error renaming eval run {eval_run_id}: {e}")
-            raise
+            raise e