agno 2.2.13__py3-none-any.whl

agno/db/mongo/schemas.py

@@ -0,0 +1,87 @@
+"""MongoDB collection schemas and related utilities"""
+
+from typing import Any, Dict, List
+
+SESSION_COLLECTION_SCHEMA = [
+    {"key": "session_id", "unique": True},
+    {"key": "user_id"},
+    {"key": "session_type"},
+    {"key": "agent_id"},
+    {"key": "team_id"},
+    {"key": "workflow_id"},
+    {"key": "created_at"},
+    {"key": "updated_at"},
+]
+
+MEMORY_COLLECTION_SCHEMA = [
+    {"key": "memory_id", "unique": True},
+    {"key": "user_id"},
+    {"key": "agent_id"},
+    {"key": "team_id"},
+    {"key": "topics"},
+    {"key": "updated_at"},
+]
+
+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"},
+    {"key": [("date", 1), ("aggregation_period", 1)], "unique": True},
+]
+
+CULTURAL_KNOWLEDGE_COLLECTION_SCHEMA = [
+    {"key": "id", "unique": True},
+    {"key": "name"},
+    {"key": "agent_id"},
+    {"key": "team_id"},
+    {"key": "created_at"},
+    {"key": "updated_at"},
+]
+
+
+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": MEMORY_COLLECTION_SCHEMA,
+        "metrics": METRICS_COLLECTION_SCHEMA,
+        "evals": EVAL_COLLECTION_SCHEMA,
+        "knowledge": KNOWLEDGE_COLLECTION_SCHEMA,
+        "culture": CULTURAL_KNOWLEDGE_COLLECTION_SCHEMA,
+    }
+
+    indexes = index_definitions.get(collection_type)
+    if not indexes:
+        raise ValueError(f"Unknown collection type: {collection_type}")
+
+    return indexes  # type: ignore[return-value]

agno/db/mongo/utils.py

@@ -0,0 +1,259 @@
+"""Utility functions for the MongoDB 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.mongo.schemas import get_collection_indexes
+from agno.db.schemas.culture import CulturalKnowledge
+from agno.utils.log import log_error, log_warning
+
+try:
+    from pymongo.collection import Collection
+except ImportError:
+    raise ImportError("`pymongo` not installed. Please install it using `pip install pymongo`")
+
+
+# -- DB util methods --
+def create_collection_indexes(collection: Collection, collection_type: str) -> None:
+    """Create all required indexes for a collection"""
+    try:
+        indexes = get_collection_indexes(collection_type)
+        for index_spec in indexes:
+            key = index_spec["key"]
+            unique = index_spec.get("unique", False)
+
+            if isinstance(key, list):
+                collection.create_index(key, unique=unique)
+            else:
+                collection.create_index([(key, 1)], unique=unique)
+
+    except Exception as e:
+        log_warning(f"Error creating indexes for {collection_type} collection: {e}")
+
+
+def apply_sorting(
+    query_args: Dict[str, Any], sort_by: Optional[str] = None, sort_order: Optional[str] = None
+) -> List[tuple]:
+    """Apply sorting to MongoDB query."""
+    if sort_by is None:
+        return []
+
+    sort_direction = 1 if sort_order == "asc" else -1
+    return [(sort_by, sort_direction)]
+
+
+def apply_pagination(
+    query_args: Dict[str, Any], limit: Optional[int] = None, page: Optional[int] = None
+) -> Dict[str, Any]:
+    """Apply pagination to MongoDB query."""
+    if limit is not None:
+        query_args["limit"] = limit
+        if page is not None:
+            query_args["skip"] = (page - 1) * limit
+    return query_args
+
+
+# -- 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, []) or []
+        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", [])
+            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.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,
+        "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: Collection, metrics_records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Bulk upsert metrics into the database.
+
+    Args:
+        collection (Collection): The collection 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 = []
+    for record in metrics_records:
+        record["date"] = record["date"].isoformat() if isinstance(record["date"], date) else record["date"]
+        try:
+            collection.replace_one(
+                {
+                    "date": record["date"],
+                    "aggregation_period": record["aggregation_period"],
+                },
+                record,
+                upsert=True,
+            )
+
+            results.append(record)
+
+        except Exception as e:
+            log_error(f"Error upserting metrics record: {e}")
+            continue
+
+    return results
+
+
+# -- 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.
+    MongoDB stores as BSON which natively supports nested documents.
+
+    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/mysql/__init__.py

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