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

agno/db/postgres/postgres.py

@@ -307,7 +307,7 @@ class PostgresDb(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.
@@ -332,6 +332,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting sessions: {e}")
+            raise e
 
     def get_session(
         self,
@@ -390,7 +391,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from session table: {e}")
-            return None
+            raise e
 
     def get_sessions(
         self,
@@ -493,7 +494,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from session table: {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
@@ -560,7 +561,7 @@ class PostgresDb(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
@@ -700,7 +701,173 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception 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 insert or update multiple sessions.
+
+        Args:
+            sessions (List[Session]): The list of session data 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.
+        """
+        try:
+            if not sessions:
+                return []
+
+            table = self._get_table(table_type="sessions", create_table_if_not_found=True)
+            if table is None:
+                return []
+
+            # Group sessions by type for better handling
+            agent_sessions = [s for s in sessions if isinstance(s, AgentSession)]
+            team_sessions = [s for s in sessions if isinstance(s, TeamSession)]
+            workflow_sessions = [s for s in sessions if isinstance(s, WorkflowSession)]
+
+            results: List[Union[Session, Dict[str, Any]]] = []
+
+            # Bulk upsert agent sessions
+            if agent_sessions:
+                session_records = []
+                for agent_session in agent_sessions:
+                    session_dict = agent_session.to_dict()
+                    session_records.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"),
+                            "agent_data": session_dict.get("agent_data"),
+                            "session_data": session_dict.get("session_data"),
+                            "summary": session_dict.get("summary"),
+                            "metadata": session_dict.get("metadata"),
+                            "runs": session_dict.get("runs"),
+                            "created_at": session_dict.get("created_at"),
+                            "updated_at": int(time.time()),
+                        }
+                    )
+
+                with self.Session() as sess, sess.begin():
+                    stmt = postgresql.insert(table)
+                    update_columns = {
+                        col.name: stmt.excluded[col.name]
+                        for col in table.columns
+                        if col.name not in ["id", "session_id", "created_at"]
+                    }
+                    stmt = stmt.on_conflict_do_update(index_elements=["session_id"], set_=update_columns).returning(
+                        table
+                    )
+
+                    result = sess.execute(stmt, session_records)
+                    for row in result.fetchall():
+                        session_dict = dict(row._mapping)
+                        if deserialize:
+                            deserialized_agent_session = AgentSession.from_dict(session_dict)
+                            if deserialized_agent_session is None:
+                                continue
+                            results.append(deserialized_agent_session)
+                        else:
+                            results.append(session_dict)
+
+            # Bulk upsert team sessions
+            if team_sessions:
+                session_records = []
+                for team_session in team_sessions:
+                    session_dict = team_session.to_dict()
+                    session_records.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"),
+                            "team_data": session_dict.get("team_data"),
+                            "session_data": session_dict.get("session_data"),
+                            "summary": session_dict.get("summary"),
+                            "metadata": session_dict.get("metadata"),
+                            "runs": session_dict.get("runs"),
+                            "created_at": session_dict.get("created_at"),
+                            "updated_at": int(time.time()),
+                        }
+                    )
+
+                with self.Session() as sess, sess.begin():
+                    stmt = postgresql.insert(table)
+                    update_columns = {
+                        col.name: stmt.excluded[col.name]
+                        for col in table.columns
+                        if col.name not in ["id", "session_id", "created_at"]
+                    }
+                    stmt = stmt.on_conflict_do_update(index_elements=["session_id"], set_=update_columns).returning(
+                        table
+                    )
+
+                    result = sess.execute(stmt, session_records)
+                    for row in result.fetchall():
+                        session_dict = dict(row._mapping)
+                        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(session_dict)
+
+            # Bulk upsert workflow sessions
+            if workflow_sessions:
+                session_records = []
+                for workflow_session in workflow_sessions:
+                    session_dict = workflow_session.to_dict()
+                    session_records.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"),
+                            "workflow_data": session_dict.get("workflow_data"),
+                            "session_data": session_dict.get("session_data"),
+                            "summary": session_dict.get("summary"),
+                            "metadata": session_dict.get("metadata"),
+                            "runs": session_dict.get("runs"),
+                            "created_at": session_dict.get("created_at"),
+                            "updated_at": int(time.time()),
+                        }
+                    )
+
+                with self.Session() as sess, sess.begin():
+                    stmt = postgresql.insert(table)
+                    update_columns = {
+                        col.name: stmt.excluded[col.name]
+                        for col in table.columns
+                        if col.name not in ["id", "session_id", "created_at"]
+                    }
+                    stmt = stmt.on_conflict_do_update(index_elements=["session_id"], set_=update_columns).returning(
+                        table
+                    )
+
+                    result = sess.execute(stmt, session_records)
+                    for row in result.fetchall():
+                        session_dict = dict(row._mapping)
+                        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(session_dict)
+
+            return results
+
+        except Exception as e:
+            log_error(f"Exception bulk upserting sessions: {e}")
+            return []
 
     # -- Memory methods --
     def delete_user_memory(self, memory_id: str):
@@ -729,6 +896,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting user memory: {e}")
+            raise e
 
     def delete_user_memories(self, memory_ids: List[str]) -> None:
         """Delete user memories from the database.
@@ -755,6 +923,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting user memories: {e}")
+            raise e
 
     def get_all_memory_topics(self) -> List[str]:
         """Get all memory topics from the database.
@@ -814,7 +983,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from memory table: {e}")
-            return None
+            raise e
 
     def get_user_memories(
         self,
@@ -897,7 +1066,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from memory table: {e}")
-            return [] if deserialize else ([], 0)
+            raise e
 
     def clear_memories(self) -> None:
         """Delete all memories from the database.
@@ -914,7 +1083,8 @@ class PostgresDb(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
 
     def get_user_memory_stats(
         self, limit: Optional[int] = None, page: Optional[int] = None
@@ -981,7 +1151,7 @@ class PostgresDb(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
@@ -1043,7 +1213,80 @@ class PostgresDb(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 insert or update multiple memories in the database for improved performance.
+
+        Args:
+            memories (List[UserMemory]): The 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.
+        """
+        try:
+            if not memories:
+                return []
+
+            table = self._get_table(table_type="memories", create_table_if_not_found=True)
+            if table is None:
+                return []
+
+            # Prepare memory records for bulk insert
+            memory_records = []
+            current_time = int(time.time())
+
+            for memory in memories:
+                if memory.memory_id is None:
+                    memory.memory_id = str(uuid4())
+
+                memory_records.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": current_time,
+                    }
+                )
+
+            results: List[Union[UserMemory, Dict[str, Any]]] = []
+
+            with self.Session() as sess, sess.begin():
+                stmt = postgresql.insert(table)
+                update_columns = {
+                    col.name: stmt.excluded[col.name]
+                    for col in table.columns
+                    if col.name not in ["memory_id"]  # Don't update primary key
+                }
+                stmt = stmt.on_conflict_do_update(index_elements=["memory_id"], set_=update_columns).returning(table)
+
+                result = sess.execute(stmt, memory_records)
+                for row in result.fetchall():
+                    memory_dict = dict(row._mapping)
+                    if deserialize:
+                        deserialized_memory = UserMemory.from_dict(memory_dict)
+                        if deserialized_memory is None:
+                            continue
+                        results.append(deserialized_memory)
+                    else:
+                        results.append(memory_dict)
+
+            return results
+
+        except Exception as e:
+            log_error(f"Exception bulk upserting memories: {e}")
+            return []
 
     # -- Metrics methods --
     def _get_all_sessions_for_metrics_calculation(
@@ -1087,7 +1330,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception reading from sessions table: {e}")
-            return []
+            raise e
 
     def _get_metrics_calculation_starting_date(self, table: Table) -> Optional[date]:
         """Get the first date for which metrics calculation is needed:
@@ -1194,7 +1437,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception refreshing metrics: {e}")
-            return None
+            raise e
 
     def get_metrics(
         self,
@@ -1235,8 +1478,8 @@ class PostgresDb(BaseDb):
             return [row._mapping for row in result], latest_updated_at
 
         except Exception as e:
-            log_warning(f"Exception getting metrics: {e}")
-            return [], None
+            log_error(f"Exception getting metrics: {e}")
+            raise e
 
     # -- Knowledge methods --
     def delete_knowledge_content(self, id: str):
@@ -1245,17 +1488,18 @@ class PostgresDb(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)
 
         except Exception as e:
             log_error(f"Exception deleting knowledge content: {e}")
+            raise e
 
     def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
         """Get a knowledge row from the database.
@@ -1266,11 +1510,11 @@ class PostgresDb(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
-
         try:
+            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()
@@ -1281,7 +1525,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception getting knowledge content: {e}")
-            return None
+            raise e
 
     def get_knowledge_contents(
         self,
@@ -1305,11 +1549,11 @@ class PostgresDb(BaseDb):
         Raises:
             Exception: If an error occurs during retrieval.
         """
-        table = self._get_table(table_type="knowledge")
-        if table is None:
-            return [], 0
-
         try:
+            table = self._get_table(table_type="knowledge")
+            if table is None:
+                return [], 0
+
             with self.Session() as sess, sess.begin():
                 stmt = select(table)
 
@@ -1332,7 +1576,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Exception getting knowledge contents: {e}")
-            return [], 0
+            raise e
 
     def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
         """Upsert knowledge content in the database.
@@ -1410,7 +1654,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error upserting knowledge row: {e}")
-            return None
+            raise e
 
     # -- Eval methods --
     def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
@@ -1443,7 +1687,7 @@ class PostgresDb(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.
@@ -1467,6 +1711,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting eval run {eval_run_id}: {e}")
+            raise e
 
     def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
         """Delete multiple eval runs from the database.
@@ -1490,6 +1735,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
+            raise e
 
     def get_eval_run(
         self, eval_run_id: str, deserialize: Optional[bool] = True
@@ -1527,7 +1773,7 @@ class PostgresDb(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,
@@ -1622,7 +1868,7 @@ class PostgresDb(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
@@ -1658,7 +1904,7 @@ class PostgresDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error upserting eval run name {eval_run_id}: {e}")
-            return None
+            raise e
 
     # -- Migrations --