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

agno/db/firestore/schemas.py

@@ -0,0 +1,130 @@
+"""Firestore collection schemas and related utilities"""
+
+from typing import Any, Dict, List
+
+SESSION_COLLECTION_SCHEMA = [
+    {"key": "session_id"},
+    {"key": "user_id"},
+    {"key": "session_type"},
+    {"key": "agent_id"},
+    {"key": "team_id"},
+    {"key": "workflow_id"},
+    {"key": "created_at"},
+    {"key": "updated_at"},
+    {"key": "session_data.session_name"},
+    # Composite indexes for get_sessions queries with sorting
+    # These match the actual query patterns: filters + created_at ordering
+    {"key": [("session_type", "ASCENDING"), ("created_at", "DESCENDING")], "collection_group": False},
+    {
+        "key": [("session_type", "ASCENDING"), ("agent_id", "ASCENDING"), ("created_at", "DESCENDING")],
+        "collection_group": False,
+    },
+    {
+        "key": [("session_type", "ASCENDING"), ("team_id", "ASCENDING"), ("created_at", "DESCENDING")],
+        "collection_group": False,
+    },
+    {
+        "key": [("session_type", "ASCENDING"), ("workflow_id", "ASCENDING"), ("created_at", "DESCENDING")],
+        "collection_group": False,
+    },
+    # For user-specific queries with sorting
+    {
+        "key": [("user_id", "ASCENDING"), ("session_type", "ASCENDING"), ("created_at", "DESCENDING")],
+        "collection_group": False,
+    },
+    {
+        "key": [
+            ("user_id", "ASCENDING"),
+            ("session_type", "ASCENDING"),
+            ("agent_id", "ASCENDING"),
+            ("created_at", "DESCENDING"),
+        ],
+        "collection_group": False,
+    },
+    {
+        "key": [
+            ("user_id", "ASCENDING"),
+            ("session_type", "ASCENDING"),
+            ("team_id", "ASCENDING"),
+            ("created_at", "DESCENDING"),
+        ],
+        "collection_group": False,
+    },
+    {
+        "key": [
+            ("user_id", "ASCENDING"),
+            ("session_type", "ASCENDING"),
+            ("workflow_id", "ASCENDING"),
+            ("created_at", "DESCENDING"),
+        ],
+        "collection_group": False,
+    },
+]
+
+USER_MEMORY_COLLECTION_SCHEMA = [
+    {"key": "memory_id", "unique": True},
+    {"key": "user_id"},
+    {"key": "agent_id"},
+    {"key": "team_id"},
+    {"key": "topics"},
+    {"key": "updated_at"},
+    # Composite indexes for memory queries
+    {"key": [("user_id", "ASCENDING"), ("agent_id", "ASCENDING")], "collection_group": False},
+    {"key": [("user_id", "ASCENDING"), ("team_id", "ASCENDING")], "collection_group": False},
+    {"key": [("user_id", "ASCENDING"), ("workflow_id", "ASCENDING")], "collection_group": False},
+]
+
+EVAL_COLLECTION_SCHEMA = [
+    {"key": "run_id", "unique": True},
+    {"key": "eval_type"},
+    {"key": "eval_input"},
+    {"key": "agent_id"},
+    {"key": "team_id"},
+    {"key": "workflow_id"},
+    {"key": "model_id"},
+    {"key": "created_at"},
+    {"key": "updated_at"},
+]
+
+KNOWLEDGE_COLLECTION_SCHEMA = [
+    {"key": "id", "unique": True},
+    {"key": "name"},
+    {"key": "description"},
+    {"key": "type"},
+    {"key": "status"},
+    {"key": "status_message"},
+    {"key": "metadata"},
+    {"key": "size"},
+    {"key": "linked_to"},
+    {"key": "access_count"},
+    {"key": "created_at"},
+    {"key": "updated_at"},
+    {"key": "external_id"},
+]
+
+METRICS_COLLECTION_SCHEMA = [
+    {"key": "id", "unique": True},
+    {"key": "date"},
+    {"key": "aggregation_period"},
+    {"key": "created_at"},
+    {"key": "updated_at"},
+    # Composite index for metrics uniqueness (same as MongoDB)
+    {"key": [("date", "ASCENDING"), ("aggregation_period", "ASCENDING")], "collection_group": False, "unique": True},
+]
+
+
+def get_collection_indexes(collection_type: str) -> List[Dict[str, Any]]:
+    """Get the index definitions for a specific collection type."""
+    index_definitions = {
+        "sessions": SESSION_COLLECTION_SCHEMA,
+        "memories": USER_MEMORY_COLLECTION_SCHEMA,
+        "metrics": METRICS_COLLECTION_SCHEMA,
+        "evals": EVAL_COLLECTION_SCHEMA,
+        "knowledge": KNOWLEDGE_COLLECTION_SCHEMA,
+    }
+
+    indexes = index_definitions.get(collection_type)
+    if not indexes:
+        raise ValueError(f"Unknown collection type: {collection_type}")
+
+    return indexes  # type: ignore

agno/db/firestore/utils.py

@@ -0,0 +1,278 @@
+"""Utility functions for the Firestore database class."""
+
+import json
+import time
+from datetime import date, datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+from agno.db.firestore.schemas import get_collection_indexes
+from agno.utils.log import log_debug, log_error, log_info, log_warning
+
+try:
+    from google.cloud.firestore import Client  # type: ignore[import-untyped]
+    from google.cloud.firestore_admin_v1 import FirestoreAdminClient, Index  # type: ignore[import-untyped]
+except ImportError:
+    raise ImportError(
+        "`google-cloud-firestore` not installed. Please install it using `pip install google-cloud-firestore`"
+    )
+
+
+# -- DB util methods --
+
+
+def create_collection_indexes(client: Client, collection_name: str, collection_type: str) -> None:
+    """Create all required indexes for a collection including composite indexes.
+
+    This function automatically creates both single-field and composite indexes.
+    """
+    try:
+        indexes = get_collection_indexes(collection_type)
+        composite_indexes = []
+
+        # Get all composite indexes
+        for idx in indexes:
+            if isinstance(idx["key"], list):
+                composite_indexes.append(idx)
+
+        # Create composite indexes programmatically
+        if composite_indexes:
+            _create_composite_indexes(client, collection_name, composite_indexes)
+            log_debug(f"Collection '{collection_name}' initialized")
+
+    except Exception as e:
+        log_warning(f"Error processing indexes for {collection_type} collection: {e}")
+
+
+def _create_composite_indexes(client: Client, collection_name: str, composite_indexes: List[Dict[str, Any]]) -> None:
+    """Create composite indexes using Firestore Admin API."""
+    try:
+        project_id = client.project
+        if not project_id:
+            log_warning("Cannot create composite indexes: project_id not available from client")
+            return
+
+        admin_client = FirestoreAdminClient()
+
+        created_count = 0
+        for idx_spec in composite_indexes:
+            try:
+                # Build index fields
+                fields = []
+                for field_name, direction in idx_spec["key"]:
+                    field_direction = (
+                        Index.IndexField.Order.ASCENDING
+                        if direction == "ASCENDING"
+                        else Index.IndexField.Order.DESCENDING
+                    )
+                    fields.append(Index.IndexField(field_path=field_name, order=field_direction))
+
+                # Create index definition
+                index = Index(
+                    query_scope=Index.QueryScope.COLLECTION
+                    if not idx_spec.get("collection_group", True)
+                    else Index.QueryScope.COLLECTION_GROUP,
+                    fields=fields,
+                )
+
+                # Note: Firestore doesn't support unique constraints on composite indexes
+                if idx_spec.get("unique", False):
+                    log_debug(
+                        f"Unique constraint on composite index ignored for {collection_name} - not supported by Firestore"
+                    )
+
+                # Create the index
+                parent_path = f"projects/{project_id}/databases/(default)/collectionGroups/{collection_name}"
+                admin_client.create_index(parent=parent_path, index=index)
+                created_count += 1
+
+            except Exception as e:
+                if "already exists" in str(e).lower():
+                    continue
+                else:
+                    log_error(f"Error creating composite index: {e}")
+
+    except Exception as e:
+        log_warning(f"Error initializing Firestore Admin client for composite indexes: {e}")
+        log_info("Fallback: You can create composite indexes manually via Firebase Console or gcloud CLI")
+
+
+def apply_sorting(query, sort_by: Optional[str] = None, sort_order: Optional[str] = None):
+    """Apply sorting to Firestore query."""
+    if sort_by is None:
+        return query
+
+    from google.cloud.firestore import Query
+
+    if sort_order == "asc":
+        return query.order_by(sort_by, direction=Query.ASCENDING)
+    else:
+        return query.order_by(sort_by, direction=Query.DESCENDING)
+
+
+def apply_pagination(query, limit: Optional[int] = None, page: Optional[int] = None):
+    """Apply pagination to Firestore query."""
+    if limit is not None:
+        query = query.limit(limit)
+        if page is not None and page > 1:
+            # Note: Firestore pagination typically uses cursor-based pagination
+            # For offset-based pagination, we'd need to skip documents
+            offset = (page - 1) * limit
+            query = query.offset(offset)
+    return query
+
+
+# -- Metrics util methods --
+
+
+def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
+    """Calculate metrics for the given single date."""
+    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"])
+            runs = session.get("runs", []) or []
+            metrics[runs_count_key] += len(runs)
+
+            if runs := session.get("runs", []):
+                if isinstance(runs, str):
+                    runs = json.loads(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_data = session.get("session_data", {})
+            if isinstance(session_data, str):
+                session_data = json.loads(session_data)
+            session_metrics = 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,
+        "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."""
+    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."""
+    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)]
+
+
+def bulk_upsert_metrics(collection_ref, metrics_records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Bulk upsert metrics into the database.
+
+    Args:
+        collection_ref: The Firestore collection reference to upsert the metrics into.
+        metrics_records (List[Dict[str, Any]]): The list of metrics records to upsert.
+
+    Returns:
+        The list of upserted metrics records.
+    """
+    if not metrics_records:
+        return []
+
+    results = []
+    batch = collection_ref._client.batch()
+
+    for i, record in enumerate(metrics_records):
+        record["date"] = record["date"].isoformat() if isinstance(record["date"], date) else record["date"]
+        try:
+            # Create a unique document ID based on date and aggregation period
+            doc_id = f"{record['date']}_{record['aggregation_period']}"
+            doc_ref = collection_ref.document(doc_id)
+            batch.set(doc_ref, record, merge=True)
+            results.append(record)
+
+            # Firestore batch limit is 500 operations
+            if (i + 1) % 500 == 0:
+                batch.commit()
+                batch = collection_ref._client.batch()
+
+        except Exception as e:
+            log_error(f"Error preparing metrics record for batch: {e}")
+            continue
+
+    # Commit remaining operations
+    if len(metrics_records) % 500 != 0:
+        try:
+            batch.commit()
+        except Exception as e:
+            log_error(f"Error committing metrics batch: {e}")
+
+    return results

agno/db/gcs_json/__init__.py

@@ -0,0 +1,3 @@
+from agno.db.gcs_json.gcs_json_db import GcsJsonDb
+
+__all__ = ["GcsJsonDb"]