agno 1.8.1__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

agno/db/json/utils.py

@@ -0,0 +1,196 @@
+"""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.base import SessionType
+from agno.run.agent import RunOutput
+from agno.run.team import TeamRunOutput
+from agno.session.summary import SessionSummary
+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]]:
+    """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, [])
+        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.split(":")
+        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)]

agno/db/migrations/v1_to_v2.py

@@ -0,0 +1,162 @@
+"""Migration utility to migrate your Agno tables from v1 to v2"""
+
+from typing import Any, Dict, List, Optional, Union
+
+from sqlalchemy import text
+
+from agno.db.mysql.mysql import MySQLDb
+from agno.db.postgres.postgres import PostgresDb
+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
+
+
+def migrate(
+    db: Union[PostgresDb, MySQLDb, SqliteDb],
+    v1_db_schema: str,
+    agent_sessions_table_name: Optional[str] = None,
+    team_sessions_table_name: Optional[str] = None,
+    workflow_sessions_table_name: Optional[str] = None,
+    memories_table_name: Optional[str] = None,
+):
+    """Given a PostgresDb and table names, parse and migrate the tables' content to the corresponding v2 tables.
+
+    Args:
+        db: The database to migrate
+        v1_db_schema: The schema of the v1 tables
+        agent_sessions_table_name: The name of the agent sessions table. If not provided, the agent sessions table will not be migrated.
+        team_sessions_table_name: The name of the team sessions table. If not provided, the team sessions table will not be migrated.
+        workflow_sessions_table_name: The name of the workflow sessions table. If not provided, the workflow sessions table will not be migrated.
+        workflow_v2_sessions_table_name: The name of the workflow v2 sessions table. If not provided, the workflow v2 sessions table will not be migrated.
+        memories_table_name: The name of the memories table. If not provided, the memories table will not be migrated.
+    """
+    if agent_sessions_table_name:
+        db.migrate_table_from_v1_to_v2(
+            v1_db_schema=v1_db_schema,
+            v1_table_name=agent_sessions_table_name,
+            v1_table_type="agent_sessions",
+        )
+
+    if team_sessions_table_name:
+        db.migrate_table_from_v1_to_v2(
+            v1_db_schema=v1_db_schema,
+            v1_table_name=team_sessions_table_name,
+            v1_table_type="team_sessions",
+        )
+
+    if workflow_sessions_table_name:
+        db.migrate_table_from_v1_to_v2(
+            v1_db_schema=v1_db_schema,
+            v1_table_name=workflow_sessions_table_name,
+            v1_table_type="workflow_sessions",
+        )
+
+    if memories_table_name:
+        db.migrate_table_from_v1_to_v2(
+            v1_db_schema=v1_db_schema,
+            v1_table_name=memories_table_name,
+            v1_table_type="memories",
+        )
+
+
+def get_all_table_content(db, db_schema: str, table_name: str) -> list[dict[str, Any]]:
+    """Get all content from the given table"""
+    try:
+        with db.Session() as sess:
+            result = sess.execute(text(f"SELECT * FROM {db_schema}.{table_name}"))
+            return [row._asdict() for row in result]
+
+    except Exception as e:
+        log_error(f"Error getting all content from table {table_name}: {e}")
+        return []
+
+
+def parse_agent_sessions(v1_content: List[Dict[str, Any]]) -> List[AgentSession]:
+    """Parse v1 Agent sessions into v2 Agent sessions and Memories"""
+    sessions_v2 = []
+
+    for item in v1_content:
+        session = {
+            "agent_id": item.get("agent_id"),
+            "agent_data": item.get("agent_data"),
+            "session_id": item.get("session_id"),
+            "user_id": item.get("user_id"),
+            "session_data": item.get("session_data"),
+            "metadata": item.get("extra_data"),
+            "runs": item.get("memory", {}).get("runs"),
+            "created_at": item.get("created_at"),
+            "updated_at": item.get("updated_at"),
+        }
+        agent_session = AgentSession.from_dict(session)
+        if agent_session is not None:
+            sessions_v2.append(agent_session)
+
+    return sessions_v2
+
+
+def parse_team_sessions(v1_content: List[Dict[str, Any]]) -> List[TeamSession]:
+    """Parse v1 Team sessions into v2 Team sessions and Memories"""
+    sessions_v2 = []
+
+    for item in v1_content:
+        session = {
+            "team_id": item.get("team_id"),
+            "team_data": item.get("team_data"),
+            "session_id": item.get("session_id"),
+            "user_id": item.get("user_id"),
+            "session_data": item.get("session_data"),
+            "metadata": item.get("extra_data"),
+            "runs": item.get("memory", {}).get("runs"),
+            "created_at": item.get("created_at"),
+            "updated_at": item.get("updated_at"),
+        }
+        team_session = TeamSession.from_dict(session)
+        if team_session is not None:
+            sessions_v2.append(team_session)
+
+    return sessions_v2
+
+
+def parse_workflow_sessions(v1_content: List[Dict[str, Any]]) -> List[WorkflowSession]:
+    """Parse v1 Workflow sessions into v2 Workflow sessions"""
+    sessions_v2 = []
+
+    for item in v1_content:
+        session = {
+            "workflow_id": item.get("workflow_id"),
+            "workflow_data": item.get("workflow_data"),
+            "session_id": item.get("session_id"),
+            "user_id": item.get("user_id"),
+            "session_data": item.get("session_data"),
+            "metadata": item.get("extra_data"),
+            "created_at": item.get("created_at"),
+            "updated_at": item.get("updated_at"),
+            # Workflow v2 specific fields
+            "workflow_name": item.get("workflow_name"),
+            "runs": item.get("runs"),
+        }
+        workflow_session = WorkflowSession.from_dict(session)
+        if workflow_session is not None:
+            sessions_v2.append(workflow_session)
+
+    return sessions_v2
+
+
+def parse_memories(v1_content: List[Dict[str, Any]]) -> List[UserMemory]:
+    """Parse v1 Memories into v2 Memories"""
+    memories_v2 = []
+
+    for item in v1_content:
+        memory = {
+            "memory_id": item.get("memory_id"),
+            "memory": item.get("memory"),
+            "input": item.get("input"),
+            "updated_at": item.get("updated_at"),
+            "agent_id": item.get("agent_id"),
+            "team_id": item.get("team_id"),
+            "user_id": item.get("user_id"),
+        }
+        memories_v2.append(UserMemory.from_dict(memory))
+
+    return memories_v2

agno/db/mongo/__init__.py

@@ -0,0 +1,3 @@
+from agno.db.mongo.mongo import MongoDb
+
+__all__ = ["MongoDb"]