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

agno/db/json/utils.py

@@ -0,0 +1,230 @@
+"""Utility functions for the JSON database class."""
+
+import time
+from datetime import date, datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+from agno.db.schemas.culture import CulturalKnowledge
+from agno.utils.log import log_debug
+
+
+def apply_sorting(
+    data: List[Dict[str, Any]], sort_by: Optional[str] = None, sort_order: Optional[str] = None
+) -> List[Dict[str, Any]]:
+    """Apply sorting to the given data list.
+
+    Args:
+        data: The list of dictionaries to sort
+        sort_by: The field to sort by
+        sort_order: The sort order ('asc' or 'desc')
+
+    Returns:
+        The sorted list
+    """
+    if sort_by is None or not data:
+        return data
+
+    # Check if the sort field exists in the first item
+    if sort_by not in data[0]:
+        log_debug(f"Invalid sort field: '{sort_by}'. Will not apply any sorting.")
+        return data
+
+    try:
+        reverse_order = sort_order != "asc" if sort_order else True
+        return sorted(data, key=lambda x: x.get(sort_by, 0), reverse=reverse_order)
+    except Exception as e:
+        log_debug(f"Error sorting data by '{sort_by}': {e}")
+        return data
+
+
+def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
+    """Calculate metrics for the given single date.
+
+    Args:
+        date_to_process (date): The date to calculate metrics for.
+        sessions_data (dict): The sessions data to calculate metrics for.
+
+    Returns:
+        dict: The calculated metrics.
+    """
+    metrics = {
+        "users_count": 0,
+        "agent_sessions_count": 0,
+        "team_sessions_count": 0,
+        "workflow_sessions_count": 0,
+        "agent_runs_count": 0,
+        "team_runs_count": 0,
+        "workflow_runs_count": 0,
+    }
+    token_metrics = {
+        "input_tokens": 0,
+        "output_tokens": 0,
+        "total_tokens": 0,
+        "audio_total_tokens": 0,
+        "audio_input_tokens": 0,
+        "audio_output_tokens": 0,
+        "cache_read_tokens": 0,
+        "cache_write_tokens": 0,
+        "reasoning_tokens": 0,
+    }
+    model_counts: Dict[str, int] = {}
+
+    session_types = [
+        ("agent", "agent_sessions_count", "agent_runs_count"),
+        ("team", "team_sessions_count", "team_runs_count"),
+        ("workflow", "workflow_sessions_count", "workflow_runs_count"),
+    ]
+    all_user_ids = set()
+
+    for session_type, sessions_count_key, runs_count_key in session_types:
+        sessions = sessions_data.get(session_type, []) or []
+        metrics[sessions_count_key] = len(sessions)
+
+        for session in sessions:
+            if session.get("user_id"):
+                all_user_ids.add(session["user_id"])
+            metrics[runs_count_key] += len(session.get("runs", []))
+            if runs := session.get("runs", []):
+                for run in runs:
+                    if model_id := run.get("model"):
+                        model_provider = run.get("model_provider", "")
+                        model_counts[f"{model_id}:{model_provider}"] = (
+                            model_counts.get(f"{model_id}:{model_provider}", 0) + 1
+                        )
+
+            session_metrics = session.get("session_data", {}).get("session_metrics", {})
+            for field in token_metrics:
+                token_metrics[field] += session_metrics.get(field, 0)
+
+    model_metrics = []
+    for model, count in model_counts.items():
+        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)
+    current_time = int(time.time())
+
+    return {
+        "id": str(uuid4()),
+        "date": date_to_process.isoformat(),
+        "completed": date_to_process < datetime.now(timezone.utc).date(),
+        "token_metrics": token_metrics,
+        "model_metrics": model_metrics,
+        "created_at": current_time,
+        "updated_at": current_time,
+        "aggregation_period": "daily",
+        **metrics,
+    }
+
+
+def fetch_all_sessions_data(
+    sessions: List[Dict[str, Any]], dates_to_process: list[date], start_timestamp: int
+) -> Optional[dict]:
+    """Return all session data for the given dates, for all session types.
+
+    Args:
+        sessions: List of session dictionaries
+        dates_to_process (list[date]): The dates to fetch session data for.
+        start_timestamp (int): The starting timestamp for filtering
+
+    Returns:
+        dict: A dictionary with dates as keys and session data as values, for all session types.
+
+    Example:
+    {
+        "2000-01-01": {
+            "agent": [<session1>, <session2>, ...],
+            "team": [...],
+            "workflow": [...],
+        }
+    }
+    """
+    if not dates_to_process:
+        return None
+
+    all_sessions_data: Dict[str, Dict[str, List[Dict[str, Any]]]] = {
+        date_to_process.isoformat(): {"agent": [], "team": [], "workflow": []} for date_to_process in dates_to_process
+    }
+
+    for session in sessions:
+        session_date = (
+            datetime.fromtimestamp(session.get("created_at", start_timestamp), tz=timezone.utc).date().isoformat()
+        )
+        if session_date in all_sessions_data:
+            all_sessions_data[session_date][session["session_type"]].append(session)
+
+    return all_sessions_data
+
+
+def get_dates_to_calculate_metrics_for(starting_date: date) -> list[date]:
+    """Return the list of dates to calculate metrics for.
+
+    Args:
+        starting_date (date): The starting date to calculate metrics for.
+
+    Returns:
+        list[date]: The list of dates to calculate metrics for.
+    """
+    today = datetime.now(timezone.utc).date()
+    days_diff = (today - starting_date).days + 1
+    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