agno 2.1.4__py3-none-any.whl → 2.1.6__py3-none-any.whl

agno/db/firestore/firestore.py

@@ -558,7 +558,7 @@ class FirestoreDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -939,7 +939,7 @@ class FirestoreDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.

agno/db/firestore/utils.py

@@ -164,11 +164,13 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
             if session.get("user_id"):
                 all_user_ids.add(session["user_id"])
             runs = session.get("runs", []) or []
-            metrics[runs_count_key] += len(runs)
 
-            if runs := session.get("runs", []):
+            if runs:
                 if isinstance(runs, str):
                     runs = json.loads(runs)
+
+                metrics[runs_count_key] += len(runs)
+
                 for run in runs:
                     if model_id := run.get("model"):
                         model_provider = run.get("model_provider", "")

agno/db/gcs_json/gcs_json_db.py

@@ -412,7 +412,7 @@ class GcsJsonDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -704,7 +704,7 @@ class GcsJsonDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.

agno/db/in_memory/in_memory_db.py

@@ -309,7 +309,7 @@ class InMemoryDb(BaseDb):
         return False
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -590,7 +590,7 @@ class InMemoryDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.

agno/db/json/json_db.py

@@ -398,7 +398,7 @@ class JsonDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -703,7 +703,7 @@ class JsonDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.

agno/db/migrations/v1_to_v2.py

@@ -1,15 +1,13 @@
 """Migration utility to migrate your Agno tables from v1 to v2"""
 
+import gc
 import json
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, cast
 
 from sqlalchemy import text
 
-from agno.db.mongo.mongo import MongoDb
-from agno.db.mysql.mysql import MySQLDb
-from agno.db.postgres.postgres import PostgresDb
+from agno.db.base import BaseDb
 from agno.db.schemas.memory import UserMemory
-from agno.db.sqlite.sqlite import SqliteDb
 from agno.session import AgentSession, TeamSession, WorkflowSession
 from agno.utils.log import log_error, log_info, log_warning
 
@@ -315,7 +313,7 @@ def convert_v1_fields_to_v2(data: Dict[str, Any]) -> Dict[str, Any]:
 
 
 def migrate(
-    db: Union[PostgresDb, MySQLDb, SqliteDb, MongoDb],
+    db: BaseDb,
     v1_db_schema: str,
     agent_sessions_table_name: Optional[str] = None,
     team_sessions_table_name: Optional[str] = None,
@@ -372,7 +370,7 @@ def migrate(
 
 
 def migrate_table_in_batches(
-    db: Union[PostgresDb, MySQLDb, SqliteDb, MongoDb],
+    db: BaseDb,
     v1_db_schema: str,
     v1_table_name: str,
     v1_table_type: str,
@@ -410,7 +408,7 @@ def migrate_table_in_batches(
                 if hasattr(db, "Session"):
                     db.Session.remove()  # type: ignore
 
-                db.upsert_sessions(sessions)  # type: ignore
+                db.upsert_sessions(sessions, preserve_updated_at=True)  # type: ignore
                 total_migrated += len(sessions)
                 log_info(f"Bulk upserted {len(sessions)} sessions in batch {batch_count}")
 
@@ -420,21 +418,35 @@ def migrate_table_in_batches(
                 if hasattr(db, "Session"):
                     db.Session.remove()  # type: ignore
 
-                db.upsert_memories(memories)
+                db.upsert_memories(memories, preserve_updated_at=True)
                 total_migrated += len(memories)
                 log_info(f"Bulk upserted {len(memories)} memories in batch {batch_count}")
 
         log_info(f"Completed batch {batch_count}: migrated {batch_size_actual} records")
 
+        # Explicit cleanup to free memory before next batch
+        del batch_content
+        if v1_table_type in ["agent_sessions", "team_sessions", "workflow_sessions"]:
+            del sessions
+        elif v1_table_type == "memories":
+            del memories
+
+        # Force garbage collection to return memory to OS
+        # This is necessary because Python's memory allocator retains memory after large operations
+        # See: https://github.com/sqlalchemy/sqlalchemy/issues/4616
+        gc.collect()
+
     log_info(f"✅ Migration completed for table {v1_table_name}: {total_migrated} total records migrated")
 
 
-def get_table_content_in_batches(
-    db: Union[PostgresDb, MySQLDb, SqliteDb, MongoDb], db_schema: str, table_name: str, batch_size: int = 5000
-):
+def get_table_content_in_batches(db: BaseDb, db_schema: str, table_name: str, batch_size: int = 5000):
     """Get table content in batches to avoid memory issues with large tables"""
     try:
-        if isinstance(db, MongoDb):
+        if type(db).__name__ == "MongoDb":
+            from agno.db.mongo.mongo import MongoDb
+
+            db = cast(MongoDb, db)
+
             # MongoDB implementation with cursor and batching
             collection = db.database[table_name]
             cursor = collection.find({}).batch_size(batch_size)
@@ -455,6 +467,24 @@ def get_table_content_in_batches(
                 yield batch
         else:
             # SQL database implementations (PostgresDb, MySQLDb, SqliteDb)
+            if type(db).__name__ == "PostgresDb":
+                from agno.db.postgres.postgres import PostgresDb
+
+                db = cast(PostgresDb, db)
+
+            elif type(db).__name__ == "MySQLDb":
+                from agno.db.mysql.mysql import MySQLDb
+
+                db = cast(MySQLDb, db)
+
+            elif type(db).__name__ == "SqliteDb":
+                from agno.db.sqlite.sqlite import SqliteDb
+
+                db = cast(SqliteDb, db)
+
+            else:
+                raise ValueError(f"Invalid database type: {type(db).__name__}")
+
             offset = 0
             while True:
                 # Create a new session for each batch to avoid transaction conflicts

agno/db/mongo/mongo.py

@@ -588,7 +588,7 @@ class MongoDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -596,6 +596,7 @@ class MongoDb(BaseDb):
         Args:
             sessions (List[Session]): List of sessions to upsert.
             deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
+            preserve_updated_at (bool): If True, preserve the updated_at from the session object.
 
         Returns:
             List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
@@ -629,6 +630,9 @@ class MongoDb(BaseDb):
 
                 session_dict = session.to_dict()
 
+                # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
+
                 if isinstance(session, AgentSession):
                     record = {
                         "session_id": session_dict.get("session_id"),
@@ -641,7 +645,7 @@ class MongoDb(BaseDb):
                         "summary": session_dict.get("summary"),
                         "metadata": session_dict.get("metadata"),
                         "created_at": session_dict.get("created_at"),
-                        "updated_at": int(time.time()),
+                        "updated_at": updated_at,
                     }
                 elif isinstance(session, TeamSession):
                     record = {
@@ -655,7 +659,7 @@ class MongoDb(BaseDb):
                         "summary": session_dict.get("summary"),
                         "metadata": session_dict.get("metadata"),
                         "created_at": session_dict.get("created_at"),
-                        "updated_at": int(time.time()),
+                        "updated_at": updated_at,
                     }
                 elif isinstance(session, WorkflowSession):
                     record = {
@@ -669,7 +673,7 @@ class MongoDb(BaseDb):
                         "summary": session_dict.get("summary"),
                         "metadata": session_dict.get("metadata"),
                         "created_at": session_dict.get("created_at"),
-                        "updated_at": int(time.time()),
+                        "updated_at": updated_at,
                     }
                 else:
                     continue
@@ -1044,7 +1048,7 @@ class MongoDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.
@@ -1079,6 +1083,7 @@ class MongoDb(BaseDb):
             operations = []
             results: List[Union[UserMemory, Dict[str, Any]]] = []
 
+            current_time = int(time.time())
             for memory in memories:
                 if memory is None:
                     continue
@@ -1086,6 +1091,9 @@ class MongoDb(BaseDb):
                 if memory.memory_id is None:
                     memory.memory_id = str(uuid4())
 
+                # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                updated_at = memory.updated_at if preserve_updated_at else current_time
+
                 record = {
                     "user_id": memory.user_id,
                     "agent_id": memory.agent_id,
@@ -1093,7 +1101,7 @@ class MongoDb(BaseDb):
                     "memory_id": memory.memory_id,
                     "memory": memory.memory,
                     "topics": memory.topics,
-                    "updated_at": int(time.time()),
+                    "updated_at": updated_at,
                 }
 
                 operations.append(ReplaceOne(filter={"memory_id": memory.memory_id}, replacement=record, upsert=True))

agno/db/mongo/utils.py

@@ -16,8 +16,6 @@ except ImportError:
 
 
 # -- DB util methods --
-
-
 def create_collection_indexes(collection: Collection, collection_type: str) -> None:
     """Create all required indexes for a collection"""
     try:
@@ -58,8 +56,6 @@ def apply_pagination(
 
 
 # -- Metrics util methods --
-
-
 def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
     """Calculate metrics for the given single date."""
     metrics = {

agno/db/mysql/mysql.py

@@ -706,7 +706,7 @@ class MySQLDb(BaseDb):
             return None
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -714,6 +714,7 @@ class MySQLDb(BaseDb):
         Args:
             sessions (List[Session]): List of sessions to upsert.
             deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
+            preserve_updated_at (bool): If True, preserve the updated_at from the session object.
 
         Returns:
             List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
@@ -758,6 +759,8 @@ class MySQLDb(BaseDb):
                     agent_data = []
                     for session in agent_sessions:
                         session_dict = session.to_dict()
+                        # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                        updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                         agent_data.append(
                             {
                                 "session_id": session_dict.get("session_id"),
@@ -770,7 +773,7 @@ class MySQLDb(BaseDb):
                                 "summary": session_dict.get("summary"),
                                 "metadata": session_dict.get("metadata"),
                                 "created_at": session_dict.get("created_at"),
-                                "updated_at": session_dict.get("created_at"),
+                                "updated_at": updated_at,
                             }
                         )
 
@@ -784,7 +787,7 @@ class MySQLDb(BaseDb):
                             summary=stmt.inserted.summary,
                             metadata=stmt.inserted.metadata,
                             runs=stmt.inserted.runs,
-                            updated_at=int(time.time()),
+                            updated_at=stmt.inserted.updated_at,
                         )
                         sess.execute(stmt, agent_data)
 
@@ -808,6 +811,8 @@ class MySQLDb(BaseDb):
                     team_data = []
                     for session in team_sessions:
                         session_dict = session.to_dict()
+                        # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                        updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                         team_data.append(
                             {
                                 "session_id": session_dict.get("session_id"),
@@ -820,7 +825,7 @@ class MySQLDb(BaseDb):
                                 "summary": session_dict.get("summary"),
                                 "metadata": session_dict.get("metadata"),
                                 "created_at": session_dict.get("created_at"),
-                                "updated_at": session_dict.get("created_at"),
+                                "updated_at": updated_at,
                             }
                         )
 
@@ -834,7 +839,7 @@ class MySQLDb(BaseDb):
                             summary=stmt.inserted.summary,
                             metadata=stmt.inserted.metadata,
                             runs=stmt.inserted.runs,
-                            updated_at=int(time.time()),
+                            updated_at=stmt.inserted.updated_at,
                         )
                         sess.execute(stmt, team_data)
 
@@ -858,6 +863,8 @@ class MySQLDb(BaseDb):
                     workflow_data = []
                     for session in workflow_sessions:
                         session_dict = session.to_dict()
+                        # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                        updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                         workflow_data.append(
                             {
                                 "session_id": session_dict.get("session_id"),
@@ -870,7 +877,7 @@ class MySQLDb(BaseDb):
                                 "summary": session_dict.get("summary"),
                                 "metadata": session_dict.get("metadata"),
                                 "created_at": session_dict.get("created_at"),
-                                "updated_at": session_dict.get("created_at"),
+                                "updated_at": updated_at,
                             }
                         )
 
@@ -884,7 +891,7 @@ class MySQLDb(BaseDb):
                             summary=stmt.inserted.summary,
                             metadata=stmt.inserted.metadata,
                             runs=stmt.inserted.runs,
-                            updated_at=int(time.time()),
+                            updated_at=stmt.inserted.updated_at,
                         )
                         sess.execute(stmt, workflow_data)
 
@@ -1281,7 +1288,7 @@ class MySQLDb(BaseDb):
             return None
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.
@@ -1313,10 +1320,13 @@ class MySQLDb(BaseDb):
 
             # Prepare bulk data
             bulk_data = []
+            current_time = int(time.time())
             for memory in memories:
                 if memory.memory_id is None:
                     memory.memory_id = str(uuid4())
 
+                # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                updated_at = memory.updated_at if preserve_updated_at else current_time
                 bulk_data.append(
                     {
                         "memory_id": memory.memory_id,
@@ -1326,7 +1336,7 @@ class MySQLDb(BaseDb):
                         "agent_id": memory.agent_id,
                         "team_id": memory.team_id,
                         "topics": memory.topics,
-                        "updated_at": int(time.time()),
+                        "updated_at": updated_at,
                     }
                 )
 
@@ -1341,7 +1351,7 @@ class MySQLDb(BaseDb):
                     input=stmt.inserted.input,
                     agent_id=stmt.inserted.agent_id,
                     team_id=stmt.inserted.team_id,
-                    updated_at=int(time.time()),
+                    updated_at=stmt.inserted.updated_at,
                 )
                 sess.execute(stmt, bulk_data)
 
@@ -1654,9 +1664,9 @@ class MySQLDb(BaseDb):
                     if page is not None:
                         stmt = stmt.offset((page - 1) * limit)
 
-                    result = sess.execute(stmt).fetchall()
-                    if not result:
-                        return [], 0
+                result = sess.execute(stmt).fetchall()
+                if not result:
+                    return [], 0
 
                 return [KnowledgeRow.model_validate(record._mapping) for record in result], total_count
 

agno/db/postgres/postgres.py

@@ -704,7 +704,7 @@ class PostgresDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk insert or update multiple sessions.
@@ -712,6 +712,7 @@ class PostgresDb(BaseDb):
         Args:
             sessions (List[Session]): The list of session data to upsert.
             deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
+            preserve_updated_at (bool): If True, preserve the updated_at from the session object.
 
         Returns:
             List[Union[Session, Dict[str, Any]]]: List of upserted sessions
@@ -739,6 +740,8 @@ class PostgresDb(BaseDb):
                 session_records = []
                 for agent_session in agent_sessions:
                     session_dict = agent_session.to_dict()
+                    # Use preserved updated_at if flag is set (even if None), otherwise use current time
+                    updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                     session_records.append(
                         {
                             "session_id": session_dict.get("session_id"),
@@ -751,7 +754,7 @@ class PostgresDb(BaseDb):
                             "metadata": session_dict.get("metadata"),
                             "runs": session_dict.get("runs"),
                             "created_at": session_dict.get("created_at"),
-                            "updated_at": int(time.time()),
+                            "updated_at": updated_at,
                         }
                     )
 
@@ -782,6 +785,8 @@ class PostgresDb(BaseDb):
                 session_records = []
                 for team_session in team_sessions:
                     session_dict = team_session.to_dict()
+                    # Use preserved updated_at if flag is set (even if None), otherwise use current time
+                    updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                     session_records.append(
                         {
                             "session_id": session_dict.get("session_id"),
@@ -794,7 +799,7 @@ class PostgresDb(BaseDb):
                             "metadata": session_dict.get("metadata"),
                             "runs": session_dict.get("runs"),
                             "created_at": session_dict.get("created_at"),
-                            "updated_at": int(time.time()),
+                            "updated_at": updated_at,
                         }
                     )
 
@@ -825,6 +830,8 @@ class PostgresDb(BaseDb):
                 session_records = []
                 for workflow_session in workflow_sessions:
                     session_dict = workflow_session.to_dict()
+                    # Use preserved updated_at if flag is set (even if None), otherwise use current time
+                    updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                     session_records.append(
                         {
                             "session_id": session_dict.get("session_id"),
@@ -837,7 +844,7 @@ class PostgresDb(BaseDb):
                             "metadata": session_dict.get("metadata"),
                             "runs": session_dict.get("runs"),
                             "created_at": session_dict.get("created_at"),
-                            "updated_at": int(time.time()),
+                            "updated_at": updated_at,
                         }
                     )
 
@@ -1234,7 +1241,7 @@ class PostgresDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk insert or update multiple memories in the database for improved performance.
@@ -1242,6 +1249,8 @@ class PostgresDb(BaseDb):
         Args:
             memories (List[UserMemory]): The list of memories to upsert.
             deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
+            preserve_updated_at (bool): If True, preserve the updated_at from the memory object.
+                                       If False (default), set updated_at to current time.
 
         Returns:
             List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories
@@ -1265,6 +1274,8 @@ class PostgresDb(BaseDb):
                 if memory.memory_id is None:
                     memory.memory_id = str(uuid4())
 
+                # Use preserved updated_at if flag is set (even if None), otherwise use current time
+                updated_at = memory.updated_at if preserve_updated_at else current_time
                 memory_records.append(
                     {
                         "memory_id": memory.memory_id,
@@ -1274,7 +1285,7 @@ class PostgresDb(BaseDb):
                         "agent_id": memory.agent_id,
                         "team_id": memory.team_id,
                         "topics": memory.topics,
-                        "updated_at": current_time,
+                        "updated_at": updated_at,
                     }
                 )
 

agno/db/redis/redis.py

@@ -589,7 +589,7 @@ class RedisDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -912,7 +912,7 @@ class RedisDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.