agno 2.2.13__py3-none-any.whl

agno/db/surrealdb/metrics.py

@@ -0,0 +1,292 @@
+from datetime import date, datetime, timedelta, timezone
+from textwrap import dedent
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from surrealdb import BlockingHttpSurrealConnection, BlockingWsSurrealConnection, RecordID
+
+from agno.db.base import SessionType
+from agno.db.surrealdb import utils
+from agno.db.surrealdb.models import desurrealize_session, surrealize_dates
+from agno.db.surrealdb.queries import WhereClause
+from agno.utils.log import log_error
+
+
+def get_all_sessions_for_metrics_calculation(
+    client: Union[BlockingWsSurrealConnection, BlockingHttpSurrealConnection],
+    table: str,
+    start_timestamp: Optional[datetime] = None,
+    end_timestamp: Optional[datetime] = None,
+) -> List[Dict[str, Any]]:
+    """
+    Get all sessions of all types (agent, team, workflow) as raw dictionaries.
+
+    Args:
+        start_timestamp (Optional[int]): The start timestamp to filter by. Defaults to None.
+        end_timestamp (Optional[int]): The end timestamp to filter by. Defaults to None.
+
+    Returns:
+        List[Dict[str, Any]]: List of session dictionaries with session_type field.
+
+    Raises:
+        Exception: If an error occurs during retrieval.
+    """
+    where = WhereClause()
+
+    # starting_date
+    if start_timestamp is not None:
+        where = where.and_("created_at", start_timestamp, ">=")
+
+    # ending_date
+    if end_timestamp is not None:
+        where = where.and_("created_at", end_timestamp, "<=")
+
+    where_clause, where_vars = where.build()
+
+    # Query
+    query = dedent(f"""
+        SELECT *
+        FROM {table}
+        {where_clause}
+    """)
+
+    results = utils.query(client, query, where_vars, dict)
+    return [desurrealize_session(x) for x in results]
+
+
+def get_metrics_calculation_starting_date(
+    client: Union[BlockingWsSurrealConnection, BlockingHttpSurrealConnection], table: str, get_sessions: Callable
+) -> Optional[date]:
+    """Get the first date for which metrics calculation is needed:
+
+    1. If there are metrics records, return the date of the first day without a complete metrics record.
+    2. If there are no metrics records, return the date of the first recorded session.
+    3. If there are no metrics records and no sessions records, return None.
+
+    Args:
+        table (Table): The table to get the starting date for.
+
+    Returns:
+        Optional[date]: The starting date for which metrics calculation is needed.
+    """
+    query = dedent(f"""
+        SELECT * FROM ONLY {table}
+        ORDER BY date DESC
+        LIMIT 1
+    """)
+    result = utils.query_one(client, query, {}, dict)
+    if result:
+        # 1. Return the date of the first day without a complete metrics record
+        result_date = result["date"]
+        assert isinstance(result_date, datetime)
+        result_date = result_date.date()
+
+        if result.get("completed"):
+            return result_date + timedelta(days=1)
+        else:
+            return result_date
+
+    # 2. No metrics records. Return the date of the first recorded session
+    first_session, _ = get_sessions(
+        session_type=SessionType.AGENT,  # this is ignored because of component_id=None and deserialize=False
+        sort_by="created_at",
+        sort_order="asc",
+        limit=1,
+        component_id=None,
+        deserialize=False,
+    )
+    assert isinstance(first_session, list)
+
+    first_session_date = first_session[0]["created_at"] if first_session else None
+
+    # 3. No metrics records and no sessions records. Return None
+    if first_session_date is None:
+        return None
+
+    # Handle different types for created_at
+    if isinstance(first_session_date, datetime):
+        return first_session_date.date()
+    elif isinstance(first_session_date, int):
+        # Assume it's a Unix timestamp
+        return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
+    elif isinstance(first_session_date, str):
+        # Try parsing as ISO format
+        return datetime.fromisoformat(first_session_date.replace("Z", "+00:00")).date()
+    else:
+        # If it's already a date object
+        if isinstance(first_session_date, date):
+            return first_session_date
+        raise ValueError(f"Unexpected type for created_at: {type(first_session_date)}")
+
+
+def bulk_upsert_metrics(
+    client: Union[BlockingWsSurrealConnection, BlockingHttpSurrealConnection],
+    table: str,
+    metrics_records: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    """Bulk upsert metrics into the database.
+
+    Args:
+        table (Table): The table to upsert into.
+        metrics_records (List[Dict[str, Any]]): The list of metrics records to upsert.
+
+    Returns:
+        list[dict]: The upserted metrics records.
+    """
+    if not metrics_records:
+        return []
+
+    metrics_records = [surrealize_dates(x) for x in metrics_records]
+
+    try:
+        results = []
+        from agno.utils.log import log_debug
+
+        for metric in metrics_records:
+            log_debug(f"Upserting metric: {metric}")  # Add this
+            result = utils.query_one(
+                client,
+                "UPSERT $record CONTENT $content",
+                {"record": RecordID(table, metric["id"]), "content": metric},
+                dict,
+            )
+            if result:
+                results.append(result)
+        return results
+
+    except Exception as e:
+        import traceback
+
+        log_error(traceback.format_exc())
+        log_error(f"Error upserting metrics: {e}")
+
+    return []
+
+
+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[Dict[str, Any]]): The sessions to process.
+        dates_to_process (list[date]): The dates to fetch session data for.
+        start_timestamp (int): The start timestamp (fallback if created_at is missing).
+
+    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:
+        created_at = session.get("created_at", start_timestamp)
+
+        # Handle different types for created_at
+        if isinstance(created_at, datetime):
+            session_date = created_at.date().isoformat()
+        elif isinstance(created_at, int):
+            session_date = datetime.fromtimestamp(created_at, tz=timezone.utc).date().isoformat()
+        elif isinstance(created_at, date):
+            session_date = created_at.isoformat()
+        else:
+            # Fallback to start_timestamp if type is unexpected
+            session_date = datetime.fromtimestamp(start_timestamp, tz=timezone.utc).date().isoformat()
+
+        if session_date in all_sessions_data:
+            session_type = session.get("session_type", "agent")  # Default to agent if missing
+            all_sessions_data[session_date][session_type].append(session)
+
+    return all_sessions_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 = datetime.now(timezone.utc)
+
+    return {
+        "id": date_to_process.isoformat(),  # Changed: Use date as ID (e.g., "2025-10-16")
+        "date": current_time.replace(hour=0, minute=0, second=0, microsecond=0),  # Date at midnight UTC
+        "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,
+    }

agno/db/surrealdb/models.py

@@ -0,0 +1,309 @@
+from dataclasses import asdict
+from datetime import date, datetime, timezone
+from textwrap import dedent
+from typing import Any, Dict, List, Literal, Optional, Sequence
+
+from surrealdb import RecordID
+
+from agno.db.base import SessionType
+from agno.db.schemas.culture import CulturalKnowledge
+from agno.db.schemas.evals import EvalRunRecord
+from agno.db.schemas.knowledge import KnowledgeRow
+from agno.db.schemas.memory import UserMemory
+from agno.session import Session
+from agno.session.agent import AgentSession
+from agno.session.team import TeamSession
+from agno.session.workflow import WorkflowSession
+
+TableType = Literal[
+    "agents",
+    "culture",
+    "evals",
+    "knowledge",
+    "memories",
+    "metrics",
+    "sessions",
+    "teams",
+    "users",
+    "workflows",
+]
+
+
+def deserialize_record_id(record: dict, agno_field: str, surreal_field: Optional[str] = None) -> dict:
+    if surreal_field is None:
+        surreal_field = agno_field
+    x = record.get(surreal_field)
+    if isinstance(x, RecordID):
+        record[agno_field] = x.id
+        if agno_field != surreal_field:
+            del record[surreal_field]
+    return record
+
+
+def surrealize_dates(record: dict) -> dict:
+    copy = record.copy()
+    for key, value in copy.items():
+        if isinstance(value, date):
+            copy[key] = datetime.combine(value, datetime.min.time()).replace(tzinfo=timezone.utc)
+        elif key in ["created_at", "updated_at"] and isinstance(value, (int, float)):
+            copy[key] = datetime.fromtimestamp(value).replace(tzinfo=timezone.utc)
+        elif key in ["created_at", "updated_at"] and isinstance(value, str):
+            # Handle ISO string format - convert back to datetime object for SurrealDB
+            try:
+                dt = datetime.fromisoformat(value)
+                if dt.tzinfo is None:
+                    dt = dt.replace(tzinfo=timezone.utc)
+                copy[key] = dt
+            except ValueError:
+                # If it's not a valid ISO format, leave it as is
+                pass
+        elif key in ["created_at", "updated_at"] and value is None:
+            # Set current time for None datetime fields
+            copy[key] = datetime.now(timezone.utc)
+        elif isinstance(value, datetime):
+            copy[key] = value.replace(tzinfo=timezone.utc)
+    return copy
+
+
+def desurrealize_dates(record: dict) -> dict:
+    copy = record.copy()
+    for key, value in copy.items():
+        if isinstance(value, datetime):
+            copy[key] = int(value.timestamp())
+    return copy
+
+
+def serialize_session(session: Session, table_names: dict[TableType, str]) -> dict:
+    _dict = session.to_dict()
+
+    if session.session_id is not None:
+        _dict["id"] = RecordID(table_names["sessions"], session.session_id)
+        del _dict["session_id"]
+
+    if isinstance(session, AgentSession):
+        _dict["agent"] = RecordID(table_names["agents"], session.agent_id)
+        del _dict["agent_id"]
+    elif isinstance(session, TeamSession):
+        _dict["team"] = RecordID(table_names["teams"], session.team_id)
+        del _dict["team_id"]
+    elif isinstance(session, WorkflowSession):
+        _dict["workflow"] = RecordID(table_names["workflows"], session.workflow_id)
+        del _dict["workflow_id"]
+
+    # surrealize dates
+    _dict = surrealize_dates(_dict)
+
+    return _dict
+
+
+def desurrealize_session(session_raw: dict, session_type: Optional[SessionType] = None) -> dict:
+    session_raw = deserialize_record_id(session_raw, "session_id", "id")
+    if session_type == SessionType.AGENT:
+        session_raw = deserialize_record_id(session_raw, "agent_id", "agent")
+    elif session_type == SessionType.TEAM:
+        session_raw = deserialize_record_id(session_raw, "team_id", "team")
+    elif session_type == SessionType.WORKFLOW:
+        session_raw = deserialize_record_id(session_raw, "workflow_id", "workflow")
+
+    session_raw = desurrealize_dates(session_raw)
+
+    if session_raw.get("agent_id"):
+        session_raw["session_type"] = SessionType.AGENT
+    elif session_raw.get("team_id"):
+        session_raw["session_type"] = SessionType.TEAM
+    elif session_raw.get("workflow_id"):
+        session_raw["session_type"] = SessionType.WORKFLOW
+
+    return session_raw
+
+
+def deserialize_session(session_type: SessionType, session_raw: dict) -> Optional[Session]:
+    session_raw = desurrealize_session(session_raw, session_type)
+
+    if session_type == SessionType.AGENT:
+        return AgentSession.from_dict(session_raw)
+    elif session_type == SessionType.TEAM:
+        return TeamSession.from_dict(session_raw)
+    elif session_type == SessionType.WORKFLOW:
+        return WorkflowSession.from_dict(session_raw)
+    else:
+        raise ValueError(f"Invalid session type: {session_type}")
+
+
+def deserialize_sessions(session_type: SessionType, sessions_raw: List[dict]) -> List[Session]:
+    return [x for x in [deserialize_session(session_type, x) for x in sessions_raw] if x is not None]
+
+
+def get_session_type(session: Session) -> SessionType:
+    if isinstance(session, AgentSession):
+        return SessionType.AGENT
+    elif isinstance(session, TeamSession):
+        return SessionType.TEAM
+    elif isinstance(session, WorkflowSession):
+        return SessionType.WORKFLOW
+    else:
+        raise ValueError(f"Invalid session instance: {type(session)}")
+
+
+def desurrealize_user_memory(memory_raw: dict) -> dict:
+    copy = memory_raw.copy()
+
+    copy = deserialize_record_id(copy, "memory_id", "id")
+    copy = deserialize_record_id(copy, "user_id", "user")
+    copy = deserialize_record_id(copy, "agent_id", "agent")
+    copy = deserialize_record_id(copy, "team_id", "team")
+    copy = deserialize_record_id(copy, "workflow_id", "workflow")
+
+    # TODO: is this ok? or should we cast datetimes to int? Like in desurrealize_session
+    # copy = desurrealize_dates(copy)
+    updated_at = copy.get("updated_at")
+    if not isinstance(updated_at, str):
+        copy["updated_at"] = str(updated_at)
+
+    return copy
+
+
+def deserialize_user_memory(memory_raw: dict) -> UserMemory:
+    return UserMemory.from_dict(desurrealize_user_memory(memory_raw))
+
+
+def deserialize_user_memories(memories_raw: Sequence[dict]) -> List[UserMemory]:
+    return [deserialize_user_memory(x) for x in memories_raw]
+
+
+def serialize_user_memory(memory: UserMemory, memory_table_name: str, user_table_name: str) -> dict:
+    dict_ = asdict(memory)
+    if memory.memory_id is not None:
+        dict_["id"] = RecordID(memory_table_name, memory.memory_id)
+        del dict_["memory_id"]
+    if memory.user_id is not None:
+        dict_["user"] = RecordID(user_table_name, memory.user_id)
+        del dict_["user_id"]
+
+    # surrealize dates
+    dict_ = surrealize_dates(dict_)
+
+    return dict_
+
+
+def deserialize_knowledge_row(knowledge_row_raw: dict) -> KnowledgeRow:
+    copy = knowledge_row_raw.copy()
+
+    copy = deserialize_record_id(copy, "id")
+    copy = desurrealize_dates(copy)
+
+    return KnowledgeRow.model_validate(copy)
+
+
+def serialize_knowledge_row(knowledge_row: KnowledgeRow, knowledge_table_name: str) -> dict:
+    dict_ = knowledge_row.model_dump()
+    if knowledge_row.id is not None:
+        dict_["id"] = RecordID(knowledge_table_name, knowledge_row.id)
+
+    # surrealize dates
+    dict_ = surrealize_dates(dict_)
+
+    return dict_
+
+
+def deserialize_cultural_knowledge(cultural_knowledge_raw: dict) -> CulturalKnowledge:
+    copy = cultural_knowledge_raw.copy()
+
+    copy = deserialize_record_id(copy, "id")
+    copy = desurrealize_dates(copy)
+
+    # Extract content, categories, and notes from the content field
+    content_json = copy.get("content", {}) or {}
+    if isinstance(content_json, dict):
+        copy["content"] = content_json.get("content")
+        copy["categories"] = content_json.get("categories")
+        copy["notes"] = content_json.get("notes")
+
+    return CulturalKnowledge.from_dict(copy)
+
+
+def serialize_cultural_knowledge(cultural_knowledge: CulturalKnowledge, culture_table_name: str) -> dict:
+    dict_ = asdict(cultural_knowledge)
+    if cultural_knowledge.id is not None:
+        dict_["id"] = RecordID(culture_table_name, cultural_knowledge.id)
+
+    # Serialize content, categories, and notes into a single content dict for DB storage
+    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
+
+    # Replace the separate fields with the combined content field
+    dict_["content"] = content_dict if content_dict else None
+    # Remove the now-redundant fields since they're in content
+    dict_.pop("categories", None)
+    dict_.pop("notes", None)
+
+    # surrealize dates
+    dict_ = surrealize_dates(dict_)
+
+    return dict_
+
+
+def desurrealize_eval_run_record(eval_run_record_raw: dict) -> dict:
+    copy = eval_run_record_raw.copy()
+
+    copy = deserialize_record_id(copy, "run_id", "id")
+    copy = deserialize_record_id(copy, "agent_id", "agent")
+    copy = deserialize_record_id(copy, "team_id", "team")
+    copy = deserialize_record_id(copy, "workflow_id", "workflow")
+
+    return copy
+
+
+def deserialize_eval_run_record(eval_run_record_raw: dict) -> EvalRunRecord:
+    return EvalRunRecord.model_validate(desurrealize_eval_run_record(eval_run_record_raw))
+
+
+def serialize_eval_run_record(eval_run_record: EvalRunRecord, table_names: dict[TableType, str]) -> dict:
+    dict_ = eval_run_record.model_dump()
+    if eval_run_record.run_id is not None:
+        dict_["id"] = RecordID(table_names["evals"], eval_run_record.run_id)
+        del dict_["run_id"]
+    if eval_run_record.agent_id is not None:
+        dict_["agent"] = RecordID(table_names["agents"], eval_run_record.agent_id)
+        del dict_["agent_id"]
+    if eval_run_record.team_id is not None:
+        dict_["team"] = RecordID(table_names["teams"], eval_run_record.team_id)
+        del dict_["team_id"]
+    if eval_run_record.workflow_id is not None:
+        dict_["workflow"] = RecordID(table_names["workflows"], eval_run_record.workflow_id)
+        del dict_["workflow_id"]
+    return dict_
+
+
+def get_schema(table_type: TableType, table_name: str) -> str:
+    define_table = f"DEFINE TABLE {table_name} SCHEMALESS;"
+    if table_type == "memories":
+        return dedent(f"""
+            {define_table}
+            DEFINE FIELD OVERWRITE updated_at ON {table_name} TYPE datetime VALUE time::now();
+            """)
+    elif table_type == "knowledge":
+        return dedent(f"""
+            {define_table}
+            DEFINE FIELD OVERWRITE created_at ON {table_name} TYPE datetime VALUE time::now();
+            DEFINE FIELD OVERWRITE updated_at ON {table_name} TYPE datetime VALUE time::now();
+            """)
+    elif table_type == "culture":
+        return dedent(f"""
+            {define_table}
+            DEFINE FIELD OVERWRITE created_at ON {table_name} TYPE datetime VALUE time::now();
+            DEFINE FIELD OVERWRITE updated_at ON {table_name} TYPE datetime VALUE time::now();
+            """)
+    elif table_type == "sessions":
+        return dedent(f"""
+            {define_table}
+            DEFINE FIELD OVERWRITE created_at ON {table_name} TYPE datetime VALUE time::now();
+            DEFINE FIELD OVERWRITE updated_at ON {table_name} TYPE datetime VALUE time::now();
+            """)
+    else:
+        return define_table

agno/db/surrealdb/queries.py

@@ -0,0 +1,71 @@
+from textwrap import dedent
+from typing import Any, Final, Literal, Optional
+
+OPERATOR = Literal["=", "!=", "<=", ">=", "~", "IN", "CONTAINSANY"]
+
+COUNT_QUERY: Final[str] = dedent("""
+    RETURN (
+        SELECT count(id) AS count
+        {group_fields}
+        FROM {table}
+        {where_clause}
+        {group_clause}
+    )[0] OR {{count: 0}}
+""")
+
+
+class WhereClause:
+    def __init__(self):
+        self._conditions = []
+        self._params = {}
+        self._param_count = 0
+
+    def _add_filter(self, field: str, operator: str, value: Any):
+        param_name = f"p{self._param_count}"
+        self._params[param_name] = value
+        self._param_count += 1
+
+        condition = f"{field} {operator} ${param_name}"
+        if not self._conditions:
+            self._conditions.append(condition)
+        else:
+            self._conditions.append("AND")
+            self._conditions.append(condition)
+        return self
+
+    def and_(self, field: str, value: Any, operator: OPERATOR = "="):
+        return self._add_filter(field, operator, value)
+
+    def build(self) -> tuple[str, dict[str, Any]]:
+        if not self._conditions:
+            return "", {}
+        return "WHERE " + " ".join(self._conditions), self._params
+
+
+def order_limit_start(
+    sort_by: Optional[str] = None,
+    sort_order: Optional[str] = None,
+    limit: Optional[int] = None,
+    page: Optional[int] = None,
+) -> str:
+    if sort_order is not None:
+        if "desc" in sort_order.lower():
+            sort_order = "DESC"
+        else:
+            sort_order = "ASC"
+
+    order_clause = f"ORDER BY {sort_by} {sort_order or ''}" if sort_by is not None else ""
+
+    if limit is not None:
+        limit_clause = f"LIMIT {limit}"
+        if page is not None:
+            offset = (page - 1) * limit
+            start_clause = f"START {offset}"
+        else:
+            start_clause = ""
+    else:
+        limit_clause = ""
+        start_clause = ""
+
+    clauses = [order_clause, limit_clause, start_clause]
+    return " ".join(clause for clause in clauses if clause)