agno 2.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/db/json/utils.py

@@ -5,33 +5,10 @@ from datetime import date, datetime, timedelta, timezone
 from typing import Any, Dict, List, Optional
 from uuid import uuid4
 
-from agno.db.base import SessionType
-from agno.run.agent import RunOutput
-from agno.run.team import TeamRunOutput
-from agno.session.summary import SessionSummary
+from agno.db.schemas.culture import CulturalKnowledge
 from agno.utils.log import log_debug
 
 
-def hydrate_session(session: dict) -> dict:
-    """Convert nested dictionaries to their corresponding object types.
-
-    Args:
-        session (dict): The session dictionary to hydrate.
-
-    Returns:
-        dict: The hydrated session dictionary.
-    """
-    if session.get("summary") is not None:
-        session["summary"] = SessionSummary.from_dict(session["summary"])
-    if session.get("runs") is not None:
-        if session["session_type"] == SessionType.AGENT:
-            session["runs"] = [RunOutput.from_dict(run) for run in session["runs"]]
-        elif session["session_type"] == SessionType.TEAM:
-            session["runs"] = [TeamRunOutput.from_dict(run) for run in session["runs"]]
-
-    return session
-
-
 def apply_sorting(
     data: List[Dict[str, Any]], sort_by: Optional[str] = None, sort_order: Optional[str] = None
 ) -> List[Dict[str, Any]]:
@@ -101,7 +78,7 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
     all_user_ids = set()
 
     for session_type, sessions_count_key, runs_count_key in session_types:
-        sessions = sessions_data.get(session_type, [])
+        sessions = sessions_data.get(session_type, []) or []
         metrics[sessions_count_key] = len(sessions)
 
         for session in sessions:
@@ -122,7 +99,7 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
 
     model_metrics = []
     for model, count in model_counts.items():
-        model_id, model_provider = model.split(":")
+        model_id, model_provider = model.rsplit(":", 1)
         model_metrics.append({"model_id": model_id, "model_provider": model_provider, "count": count})
 
     metrics["users_count"] = len(all_user_ids)
@@ -194,3 +171,60 @@ def get_dates_to_calculate_metrics_for(starting_date: date) -> list[date]:
     if days_diff <= 0:
         return []
     return [starting_date + timedelta(days=x) for x in range(days_diff)]
+
+
+# -- Cultural Knowledge util methods --
+def serialize_cultural_knowledge_for_db(cultural_knowledge: CulturalKnowledge) -> Dict[str, Any]:
+    """Serialize a CulturalKnowledge object for database storage.
+
+    Converts the model's separate content, categories, and notes fields
+    into a single dict for the database content column.
+
+    Args:
+        cultural_knowledge (CulturalKnowledge): The cultural knowledge object to serialize.
+
+    Returns:
+        Dict[str, Any]: A dictionary with the content field as a dict containing content, categories, and notes.
+    """
+    content_dict: Dict[str, Any] = {}
+    if cultural_knowledge.content is not None:
+        content_dict["content"] = cultural_knowledge.content
+    if cultural_knowledge.categories is not None:
+        content_dict["categories"] = cultural_knowledge.categories
+    if cultural_knowledge.notes is not None:
+        content_dict["notes"] = cultural_knowledge.notes
+
+    return content_dict if content_dict else {}
+
+
+def deserialize_cultural_knowledge_from_db(db_row: Dict[str, Any]) -> CulturalKnowledge:
+    """Deserialize a database row to a CulturalKnowledge object.
+
+    The database stores content as a dict containing content, categories, and notes.
+    This method extracts those fields and converts them back to the model format.
+
+    Args:
+        db_row (Dict[str, Any]): The database row as a dictionary.
+
+    Returns:
+        CulturalKnowledge: The cultural knowledge object.
+    """
+    # Extract content, categories, and notes from the content field
+    content_json = db_row.get("content", {}) or {}
+
+    return CulturalKnowledge.from_dict(
+        {
+            "id": db_row.get("id"),
+            "name": db_row.get("name"),
+            "summary": db_row.get("summary"),
+            "content": content_json.get("content"),
+            "categories": content_json.get("categories"),
+            "notes": content_json.get("notes"),
+            "metadata": db_row.get("metadata"),
+            "input": db_row.get("input"),
+            "created_at": db_row.get("created_at"),
+            "updated_at": db_row.get("updated_at"),
+            "agent_id": db_row.get("agent_id"),
+            "team_id": db_row.get("team_id"),
+        }
+    )

agno/db/migrations/manager.py

@@ -0,0 +1,199 @@
+import importlib
+from typing import Optional, Union
+
+from packaging import version as packaging_version
+from packaging.version import Version
+
+from agno.db.base import AsyncBaseDb, BaseDb
+from agno.utils.log import log_error, log_info, log_warning
+
+
+class MigrationManager:
+    """Manager class to handle database migrations"""
+
+    available_versions: list[tuple[str, Version]] = [
+        ("v2_0_0", packaging_version.parse("2.0.0")),
+        ("v2_3_0", packaging_version.parse("2.3.0")),
+    ]
+
+    def __init__(self, db: Union[AsyncBaseDb, BaseDb]):
+        self.db = db
+
+    @property
+    def latest_schema_version(self) -> Version:
+        return self.available_versions[-1][1]
+
+    async def up(self, target_version: Optional[str] = None, table_type: Optional[str] = None, force: bool = False):
+        """Handle executing an up migration.
+
+        Args:
+            target_version: The version to migrate to, e.g. "v3.0.0". If not provided, the latest available version will be used.
+            table_type: The type of table to migrate. If not provided, all table types will be considered.
+        """
+
+        # If not target version is provided, use the latest available version
+        if not target_version:
+            _target_version = self.latest_schema_version
+            log_info(
+                f"No target version provided. Will migrate to the latest available version: {str(_target_version)}"
+            )
+        else:
+            _target_version = packaging_version.parse(target_version)
+
+        # Select tables to migrate
+        if table_type:
+            if table_type not in ["memory", "session", "metrics", "eval", "knowledge", "culture"]:
+                log_warning(
+                    f"Invalid table type: {table_type}. Use one of: memory, session, metrics, eval, knowledge, culture"
+                )
+                return
+            tables = [(table_type, getattr(self.db, f"{table_type}_table_name"))]
+        else:
+            tables = [
+                ("memories", self.db.memory_table_name),
+                ("sessions", self.db.session_table_name),
+                ("metrics", self.db.metrics_table_name),
+                ("evals", self.db.eval_table_name),
+                ("knowledge", self.db.knowledge_table_name),
+                ("culture", self.db.culture_table_name),
+            ]
+
+        # Handle migrations for each table separately (extend in future if needed):
+        for table_type, table_name in tables:
+            if isinstance(self.db, AsyncBaseDb):
+                current_version = packaging_version.parse(await self.db.get_latest_schema_version(table_name))
+            else:
+                current_version = packaging_version.parse(self.db.get_latest_schema_version(table_name))
+
+            if current_version is None:
+                log_warning(f"Skipping up migration: No version found for table {table_name}.")
+                continue
+
+            # If the target version is less or equal to the current version, no migrations needed
+            if _target_version <= current_version and not force:
+                log_warning(
+                    f"Skipping up migration: the version of table '{table_name}' ({current_version}) is less or equal to the target version ({_target_version})."
+                )
+                continue
+
+            log_info(
+                f"Starting database migration for table {table_name}. Current version: {current_version}. Target version: {_target_version}."
+            )
+
+            # Find files after the current version
+            latest_version = None
+            migration_executed = False
+            for version, normalised_version in self.available_versions:
+                if normalised_version > current_version:
+                    if target_version and normalised_version > _target_version:
+                        break
+
+                    log_info(f"Applying migration {normalised_version} on {table_name}")
+                    migration_executed = await self._up_migration(version, table_type, table_name)
+                    if migration_executed:
+                        log_info(f"Successfully applied migration {normalised_version} on table {table_name}")
+                    else:
+                        log_info(f"Skipping application of migration {normalised_version} on table {table_name}")
+
+                    latest_version = normalised_version.public
+
+            if migration_executed and latest_version:
+                log_info(f"Storing version {latest_version} in database for table {table_name}")
+                if isinstance(self.db, AsyncBaseDb):
+                    await self.db.upsert_schema_version(table_name, latest_version)
+                else:
+                    self.db.upsert_schema_version(table_name, latest_version)
+                log_info(f"Successfully stored version {latest_version} in database for table {table_name}")
+            log_info("----------------------------------------------------------")
+
+    async def _up_migration(self, version: str, table_type: str, table_name: str) -> bool:
+        """Run the database-specific logic to handle an up migration.
+
+        Args:
+            version: The version to migrate to, e.g. "v3.0.0"
+        """
+        migration_module = importlib.import_module(f"agno.db.migrations.versions.{version}")
+
+        try:
+            if isinstance(self.db, AsyncBaseDb):
+                return await migration_module.async_up(self.db, table_type, table_name)
+            else:
+                return migration_module.up(self.db, table_type, table_name)
+        except Exception as e:
+            log_error(f"Error running migration to version {version}: {e}")
+            raise
+
+    async def down(self, target_version: str, table_type: Optional[str] = None, force: bool = False):
+        """Handle executing a down migration.
+
+        Args:
+            target_version: The version to migrate to. e.g. "v2.3.0"
+            table_type: The type of table to migrate. If not provided, all table types will be considered.
+        """
+        _target_version = packaging_version.parse(target_version)
+
+        # Select tables to migrate
+        if table_type:
+            if table_type not in ["memory", "session", "metrics", "eval", "knowledge", "culture"]:
+                log_warning(
+                    f"Invalid table type: {table_type}. Use one of: memory, session, metrics, eval, knowledge, culture"
+                )
+                return
+            tables = [(table_type, getattr(self.db, f"{table_type}_table_name"))]
+        else:
+            tables = [
+                ("memories", self.db.memory_table_name),
+                ("sessions", self.db.session_table_name),
+                ("metrics", self.db.metrics_table_name),
+                ("evals", self.db.eval_table_name),
+                ("knowledge", self.db.knowledge_table_name),
+                ("culture", self.db.culture_table_name),
+            ]
+
+        for table_type, table_name in tables:
+            if isinstance(self.db, AsyncBaseDb):
+                current_version = packaging_version.parse(await self.db.get_latest_schema_version(table_name))
+            else:
+                current_version = packaging_version.parse(self.db.get_latest_schema_version(table_name))
+
+            if _target_version >= current_version and not force:
+                log_warning(
+                    f"Skipping down migration: the version of table '{table_name}' ({current_version}) is less or equal to the target version ({_target_version})."
+                )
+                continue
+
+            migration_executed = False
+            # Run down migration for all versions between target and current (include down of current version)
+            # Apply down migrations in reverse order to ensure dependencies are met
+            for version, normalised_version in reversed(self.available_versions):
+                if normalised_version > _target_version:
+                    log_info(f"Reverting migration {normalised_version} on table {table_name}")
+                    migration_executed = await self._down_migration(version, table_type, table_name)
+                    if migration_executed:
+                        log_info(f"Successfully reverted migration {normalised_version} on table {table_name}")
+                    else:
+                        log_info(f"Skipping revert of migration {normalised_version} on table {table_name}")
+
+            if migration_executed:
+                log_info(f"Storing version {_target_version} in database for table {table_name}")
+                if isinstance(self.db, AsyncBaseDb):
+                    await self.db.upsert_schema_version(table_name, _target_version.public)
+                else:
+                    self.db.upsert_schema_version(table_name, _target_version.public)
+                log_info(f"Successfully stored version {_target_version} in database for table {table_name}")
+
+    async def _down_migration(self, version: str, table_type: str, table_name: str) -> bool:
+        """Run the database-specific logic to handle a down migration.
+
+        Args:
+            version: The version to migrate to, e.g. "v3.0.0"
+        """
+        migration_module = importlib.import_module(f"agno.db.migrations.versions.{version}")
+        try:
+            if isinstance(self.db, AsyncBaseDb):
+                return await migration_module.async_down(self.db, table_type, table_name)
+            else:
+                return migration_module.down(self.db, table_type, table_name)
+        except Exception as e:
+            log_error(f"Error running migration to version {version}: {e}")
+            raise

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