agno 2.1.4__py3-none-any.whl → 2.1.6__py3-none-any.whl

agno/db/singlestore/singlestore.py

@@ -797,7 +797,7 @@ class SingleStoreDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -841,6 +841,8 @@ class SingleStoreDb(BaseDb):
                     agent_data = []
                     for session in agent_sessions:
                         session_dict = session.to_dict()
+                        # Use preserved updated_at if flag is set, otherwise use current time
+                        updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                         agent_data.append(
                             {
                                 "session_id": session_dict.get("session_id"),
@@ -853,7 +855,7 @@ class SingleStoreDb(BaseDb):
                                 "summary": session_dict.get("summary"),
                                 "metadata": session_dict.get("metadata"),
                                 "created_at": session_dict.get("created_at"),
-                                "updated_at": session_dict.get("created_at"),
+                                "updated_at": updated_at,
                             }
                         )
 
@@ -867,7 +869,7 @@ class SingleStoreDb(BaseDb):
                             summary=stmt.inserted.summary,
                             metadata=stmt.inserted.metadata,
                             runs=stmt.inserted.runs,
-                            updated_at=int(time.time()),
+                            updated_at=stmt.inserted.updated_at,
                         )
                         sess.execute(stmt, agent_data)
 
@@ -890,6 +892,8 @@ class SingleStoreDb(BaseDb):
                     team_data = []
                     for session in team_sessions:
                         session_dict = session.to_dict()
+                        # Use preserved updated_at if flag is set, otherwise use current time
+                        updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                         team_data.append(
                             {
                                 "session_id": session_dict.get("session_id"),
@@ -902,7 +906,7 @@ class SingleStoreDb(BaseDb):
                                 "summary": session_dict.get("summary"),
                                 "metadata": session_dict.get("metadata"),
                                 "created_at": session_dict.get("created_at"),
-                                "updated_at": session_dict.get("created_at"),
+                                "updated_at": updated_at,
                             }
                         )
 
@@ -916,7 +920,7 @@ class SingleStoreDb(BaseDb):
                             summary=stmt.inserted.summary,
                             metadata=stmt.inserted.metadata,
                             runs=stmt.inserted.runs,
-                            updated_at=int(time.time()),
+                            updated_at=stmt.inserted.updated_at,
                         )
                         sess.execute(stmt, team_data)
 
@@ -939,6 +943,8 @@ class SingleStoreDb(BaseDb):
                     workflow_data = []
                     for session in workflow_sessions:
                         session_dict = session.to_dict()
+                        # Use preserved updated_at if flag is set, otherwise use current time
+                        updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
                         workflow_data.append(
                             {
                                 "session_id": session_dict.get("session_id"),
@@ -951,7 +957,7 @@ class SingleStoreDb(BaseDb):
                                 "summary": session_dict.get("summary"),
                                 "metadata": session_dict.get("metadata"),
                                 "created_at": session_dict.get("created_at"),
-                                "updated_at": session_dict.get("created_at"),
+                                "updated_at": updated_at,
                             }
                         )
 
@@ -965,7 +971,7 @@ class SingleStoreDb(BaseDb):
                             summary=stmt.inserted.summary,
                             metadata=stmt.inserted.metadata,
                             runs=stmt.inserted.runs,
-                            updated_at=int(time.time()),
+                            updated_at=stmt.inserted.updated_at,
                         )
                         sess.execute(stmt, workflow_data)
 
@@ -1336,7 +1342,7 @@ class SingleStoreDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.
@@ -1361,9 +1367,12 @@ class SingleStoreDb(BaseDb):
 
             # Prepare data for bulk insert
             memory_data = []
+            current_time = int(time.time())
             for memory in memories:
                 if memory.memory_id is None:
                     memory.memory_id = str(uuid4())
+                # Use preserved updated_at if flag is set, otherwise use current time
+                updated_at = memory.updated_at if preserve_updated_at else current_time
                 memory_data.append(
                     {
                         "memory_id": memory.memory_id,
@@ -1373,7 +1382,7 @@ class SingleStoreDb(BaseDb):
                         "agent_id": memory.agent_id,
                         "team_id": memory.team_id,
                         "topics": memory.topics,
-                        "updated_at": int(time.time()),
+                        "updated_at": updated_at,
                     }
                 )
 
@@ -1389,7 +1398,7 @@ class SingleStoreDb(BaseDb):
                         user_id=stmt.inserted.user_id,
                         agent_id=stmt.inserted.agent_id,
                         team_id=stmt.inserted.team_id,
-                        updated_at=int(time.time()),
+                        updated_at=stmt.inserted.updated_at,
                     )
                     sess.execute(stmt, memory_data)
 

agno/db/sqlite/sqlite.py

@@ -36,9 +36,9 @@ except ImportError:
 class SqliteDb(BaseDb):
     def __init__(
         self,
+        db_file: Optional[str] = None,
         db_engine: Optional[Engine] = None,
         db_url: Optional[str] = None,
-        db_file: Optional[str] = None,
         session_table: Optional[str] = None,
         memory_table: Optional[str] = None,
         metrics_table: Optional[str] = None,
@@ -56,9 +56,9 @@ class SqliteDb(BaseDb):
             4. Create a new database in the current directory
 
         Args:
+            db_file (Optional[str]): The database file to connect to.
             db_engine (Optional[Engine]): The SQLAlchemy database engine to use.
             db_url (Optional[str]): The database URL to connect to.
-            db_file (Optional[str]): The database file to connect to.
             session_table (Optional[str]): Name of the table to store Agent, Team and Workflow sessions.
             memory_table (Optional[str]): Name of the table to store user memories.
             metrics_table (Optional[str]): Name of the table to store metrics.
@@ -664,7 +664,7 @@ class SqliteDb(BaseDb):
             raise e
 
     def upsert_sessions(
-        self, sessions: List[Session], deserialize: Optional[bool] = True
+        self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[Session, Dict[str, Any]]]:
         """
         Bulk upsert multiple sessions for improved performance on large datasets.
@@ -672,6 +672,7 @@ class SqliteDb(BaseDb):
         Args:
             sessions (List[Session]): List of sessions to upsert.
             deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
+            preserve_updated_at (bool): If True, preserve the updated_at from the session object.
 
         Returns:
             List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
@@ -715,6 +716,8 @@ class SqliteDb(BaseDb):
                     agent_data = []
                     for session in agent_sessions:
                         serialized_session = serialize_session_json_fields(session.to_dict())
+                        # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                        updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
                         agent_data.append(
                             {
                                 "session_id": serialized_session.get("session_id"),
@@ -727,7 +730,7 @@ class SqliteDb(BaseDb):
                                 "runs": serialized_session.get("runs"),
                                 "summary": serialized_session.get("summary"),
                                 "created_at": serialized_session.get("created_at"),
-                                "updated_at": serialized_session.get("created_at"),
+                                "updated_at": updated_at,
                             }
                         )
 
@@ -743,7 +746,7 @@ class SqliteDb(BaseDb):
                                 metadata=stmt.excluded.metadata,
                                 runs=stmt.excluded.runs,
                                 summary=stmt.excluded.summary,
-                                updated_at=int(time.time()),
+                                updated_at=stmt.excluded.updated_at,
                             ),
                         )
                         sess.execute(stmt, agent_data)
@@ -768,6 +771,8 @@ class SqliteDb(BaseDb):
                     team_data = []
                     for session in team_sessions:
                         serialized_session = serialize_session_json_fields(session.to_dict())
+                        # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                        updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
                         team_data.append(
                             {
                                 "session_id": serialized_session.get("session_id"),
@@ -777,7 +782,7 @@ class SqliteDb(BaseDb):
                                 "runs": serialized_session.get("runs"),
                                 "summary": serialized_session.get("summary"),
                                 "created_at": serialized_session.get("created_at"),
-                                "updated_at": serialized_session.get("created_at"),
+                                "updated_at": updated_at,
                                 "team_data": serialized_session.get("team_data"),
                                 "session_data": serialized_session.get("session_data"),
                                 "metadata": serialized_session.get("metadata"),
@@ -796,7 +801,7 @@ class SqliteDb(BaseDb):
                                 metadata=stmt.excluded.metadata,
                                 runs=stmt.excluded.runs,
                                 summary=stmt.excluded.summary,
-                                updated_at=int(time.time()),
+                                updated_at=stmt.excluded.updated_at,
                             ),
                         )
                         sess.execute(stmt, team_data)
@@ -821,6 +826,8 @@ class SqliteDb(BaseDb):
                     workflow_data = []
                     for session in workflow_sessions:
                         serialized_session = serialize_session_json_fields(session.to_dict())
+                        # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                        updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
                         workflow_data.append(
                             {
                                 "session_id": serialized_session.get("session_id"),
@@ -830,7 +837,7 @@ class SqliteDb(BaseDb):
                                 "runs": serialized_session.get("runs"),
                                 "summary": serialized_session.get("summary"),
                                 "created_at": serialized_session.get("created_at"),
-                                "updated_at": serialized_session.get("created_at"),
+                                "updated_at": updated_at,
                                 "workflow_data": serialized_session.get("workflow_data"),
                                 "session_data": serialized_session.get("session_data"),
                                 "metadata": serialized_session.get("metadata"),
@@ -849,7 +856,7 @@ class SqliteDb(BaseDb):
                                 metadata=stmt.excluded.metadata,
                                 runs=stmt.excluded.runs,
                                 summary=stmt.excluded.summary,
-                                updated_at=int(time.time()),
+                                updated_at=stmt.excluded.updated_at,
                             ),
                         )
                         sess.execute(stmt, workflow_data)
@@ -1224,7 +1231,7 @@ class SqliteDb(BaseDb):
             raise e
 
     def upsert_memories(
-        self, memories: List[UserMemory], deserialize: Optional[bool] = True
+        self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
     ) -> List[Union[UserMemory, Dict[str, Any]]]:
         """
         Bulk upsert multiple user memories for improved performance on large datasets.
@@ -1255,10 +1262,13 @@ class SqliteDb(BaseDb):
                 ]
             # Prepare bulk data
             bulk_data = []
+            current_time = int(time.time())
             for memory in memories:
                 if memory.memory_id is None:
                     memory.memory_id = str(uuid4())
 
+                # Use preserved updated_at if flag is set and value exists, otherwise use current time
+                updated_at = memory.updated_at if preserve_updated_at else current_time
                 bulk_data.append(
                     {
                         "user_id": memory.user_id,
@@ -1267,7 +1277,7 @@ class SqliteDb(BaseDb):
                         "memory_id": memory.memory_id,
                         "memory": memory.memory,
                         "topics": memory.topics,
-                        "updated_at": int(time.time()),
+                        "updated_at": updated_at,
                     }
                 )
 
@@ -1284,7 +1294,7 @@ class SqliteDb(BaseDb):
                         input=stmt.excluded.input,
                         agent_id=stmt.excluded.agent_id,
                         team_id=stmt.excluded.team_id,
-                        updated_at=int(time.time()),
+                        updated_at=stmt.excluded.updated_at,
                     ),
                 )
                 sess.execute(stmt, bulk_data)

agno/db/sqlite/utils.py

@@ -179,9 +179,11 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
         for session in sessions:
             if session.get("user_id"):
                 all_user_ids.add(session["user_id"])
-            metrics[runs_count_key] += len(session.get("runs", []))
+
+            # Parse runs from JSON string
             if runs := session.get("runs", []):
-                runs = json.loads(runs)
+                runs = json.loads(runs) if isinstance(runs, str) else runs
+                metrics[runs_count_key] += len(runs)
                 for run in runs:
                     if model_id := run.get("model"):
                         model_provider = run.get("model_provider", "")
@@ -189,7 +191,10 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
                             model_counts.get(f"{model_id}:{model_provider}", 0) + 1
                         )
 
-            session_data = json.loads(session.get("session_data", {}))
+            # Parse session_data from JSON string
+            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)

agno/db/surrealdb/__init__.py

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

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,
+    }