agno 2.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/db/redis/redis.py

@@ -1,8 +1,11 @@
 import time
 from datetime import date, datetime, timedelta, timezone
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 from uuid import uuid4
 
+if TYPE_CHECKING:
+    from agno.tracing.schemas import Span, Trace
+
 from agno.db.base import BaseDb, SessionType
 from agno.db.redis.utils import (
     apply_filters,
@@ -10,14 +13,17 @@ from agno.db.redis.utils import (
     apply_sorting,
     calculate_date_metrics,
     create_index_entries,
+    deserialize_cultural_knowledge_from_db,
     deserialize_data,
     fetch_all_sessions_data,
     generate_redis_key,
     get_all_keys_for_table,
     get_dates_to_calculate_metrics_for,
     remove_index_entries,
+    serialize_cultural_knowledge_for_db,
     serialize_data,
 )
+from agno.db.schemas.culture import CulturalKnowledge
 from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
 from agno.db.schemas.knowledge import KnowledgeRow
 from agno.db.schemas.memory import UserMemory
@@ -26,7 +32,7 @@ from agno.utils.log import log_debug, log_error, log_info
 from agno.utils.string import generate_id
 
 try:
-    from redis import Redis
+    from redis import Redis, RedisCluster
 except ImportError:
     raise ImportError("`redis` not installed. Please install it using `pip install redis`")
 
@@ -35,7 +41,7 @@ class RedisDb(BaseDb):
     def __init__(
         self,
         id: Optional[str] = None,
-        redis_client: Optional[Redis] = None,
+        redis_client: Optional[Union[Redis, RedisCluster]] = None,
         db_url: Optional[str] = None,
         db_prefix: str = "agno",
         expire: Optional[int] = None,
@@ -44,6 +50,9 @@ class RedisDb(BaseDb):
         metrics_table: Optional[str] = None,
         eval_table: Optional[str] = None,
         knowledge_table: Optional[str] = None,
+        culture_table: Optional[str] = None,
+        traces_table: Optional[str] = None,
+        spans_table: Optional[str] = None,
     ):
         """
         Interface for interacting with a Redis database.
@@ -53,6 +62,8 @@ class RedisDb(BaseDb):
             2. Use the db_url
             3. Raise an error if neither is provided
 
+        db_url only supports single-node Redis connections, if you need Redis Cluster support, provide a redis_client.
+
         Args:
             id (Optional[str]): The ID of the database.
             redis_client (Optional[Redis]): Redis client instance to use. If not provided a new client will be created.
@@ -64,6 +75,9 @@ class RedisDb(BaseDb):
             metrics_table (Optional[str]): Name of the table to store metrics
             eval_table (Optional[str]): Name of the table to store evaluation runs
             knowledge_table (Optional[str]): Name of the table to store knowledge documents
+            culture_table (Optional[str]): Name of the table to store cultural knowledge
+            traces_table (Optional[str]): Name of the table to store traces
+            spans_table (Optional[str]): Name of the table to store spans
 
         Raises:
             ValueError: If neither redis_client nor db_url is provided.
@@ -80,6 +94,9 @@ class RedisDb(BaseDb):
             metrics_table=metrics_table,
             eval_table=eval_table,
             knowledge_table=knowledge_table,
+            culture_table=culture_table,
+            traces_table=traces_table,
+            spans_table=spans_table,
         )
 
         self.db_prefix = db_prefix
@@ -94,6 +111,10 @@ class RedisDb(BaseDb):
 
     # -- DB methods --
 
+    def table_exists(self, table_name: str) -> bool:
+        """Redis implementation, always returns True."""
+        return True
+
     def _get_table_name(self, table_type: str) -> str:
         """Get the active table name for the given table type."""
         if table_type == "sessions":
@@ -111,6 +132,15 @@ class RedisDb(BaseDb):
         elif table_type == "knowledge":
             return self.knowledge_table_name
 
+        elif table_type == "culture":
+            return self.culture_table_name
+
+        elif table_type == "traces":
+            return self.trace_table_name
+
+        elif table_type == "spans":
+            return self.span_table_name
+
         else:
             raise ValueError(f"Unknown table type: {table_type}")
 
@@ -239,6 +269,14 @@ class RedisDb(BaseDb):
             log_error(f"Error getting all records for {table_type}: {e}")
             return []
 
+    def get_latest_schema_version(self):
+        """Get the latest version of the database schema."""
+        pass
+
+    def upsert_schema_version(self, version: str) -> None:
+        """Upsert the schema version into the database."""
+        pass
+
     # -- Session methods --
 
     def delete_session(self, session_id: str) -> bool:
@@ -318,8 +356,6 @@ class RedisDb(BaseDb):
             # Apply filters
             if user_id is not None and session.get("user_id") != user_id:
                 return None
-            if session_type is not None and session.get("session_type") != session_type:
-                return None
 
             if not deserialize:
                 return session
@@ -589,7 +625,7 @@ class RedisDb(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.
@@ -820,12 +856,14 @@ class RedisDb(BaseDb):
         self,
         limit: Optional[int] = None,
         page: Optional[int] = None,
+        user_id: Optional[str] = None,
     ) -> Tuple[List[Dict[str, Any]], int]:
         """Get user memory stats from Redis.
 
         Args:
             limit (Optional[int]): The maximum number of stats to return.
             page (Optional[int]): The page number to return.
+            user_id (Optional[str]): User ID for filtering.
 
         Returns:
             Tuple[List[Dict[str, Any]], int]: A tuple containing the list of stats and the total number of stats.
@@ -840,6 +878,9 @@ class RedisDb(BaseDb):
             user_stats = {}
             for memory in all_memories:
                 memory_user_id = memory.get("user_id")
+                # filter by user_id if provided
+                if user_id is not None and memory_user_id != user_id:
+                    continue
                 if memory_user_id is None:
                     continue
 
@@ -892,6 +933,9 @@ class RedisDb(BaseDb):
                 "memory_id": memory.memory_id,
                 "memory": memory.memory,
                 "topics": memory.topics,
+                "input": memory.input,
+                "feedback": memory.feedback,
+                "created_at": memory.created_at,
                 "updated_at": int(time.time()),
             }
 
@@ -912,7 +956,7 @@ class RedisDb(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.
@@ -1475,3 +1519,623 @@ class RedisDb(BaseDb):
         except Exception as e:
             log_error(f"Error updating eval run name {eval_run_id}: {e}")
             raise
+
+    # -- Cultural Knowledge methods --
+    def clear_cultural_knowledge(self) -> None:
+        """Delete all cultural knowledge from the database.
+
+        Raises:
+            Exception: If an error occurs during deletion.
+        """
+        try:
+            keys = get_all_keys_for_table(redis_client=self.redis_client, prefix=self.db_prefix, table_type="culture")
+
+            if keys:
+                self.redis_client.delete(*keys)
+
+        except Exception as e:
+            log_error(f"Exception deleting all cultural knowledge: {e}")
+            raise e
+
+    def delete_cultural_knowledge(self, id: str) -> None:
+        """Delete cultural knowledge by ID.
+
+        Args:
+            id (str): The ID of the cultural knowledge to delete.
+
+        Raises:
+            Exception: If an error occurs during deletion.
+        """
+        try:
+            if self._delete_record("culture", id, index_fields=["name", "agent_id", "team_id"]):
+                log_debug(f"Successfully deleted cultural knowledge id: {id}")
+            else:
+                log_debug(f"No cultural knowledge found with id: {id}")
+
+        except Exception as e:
+            log_error(f"Error deleting cultural knowledge: {e}")
+            raise e
+
+    def get_cultural_knowledge(
+        self, id: str, deserialize: Optional[bool] = True
+    ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
+        """Get cultural knowledge by ID.
+
+        Args:
+            id (str): The ID of the cultural knowledge to retrieve.
+            deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge object. Defaults to True.
+
+        Returns:
+            Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The cultural knowledge if found, None otherwise.
+
+        Raises:
+            Exception: If an error occurs during retrieval.
+        """
+        try:
+            cultural_knowledge = self._get_record("culture", id)
+
+            if cultural_knowledge is None:
+                return None
+
+            if not deserialize:
+                return cultural_knowledge
+
+            return deserialize_cultural_knowledge_from_db(cultural_knowledge)
+
+        except Exception as e:
+            log_error(f"Error getting cultural knowledge: {e}")
+            raise e
+
+    def get_all_cultural_knowledge(
+        self,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        name: Optional[str] = None,
+        limit: Optional[int] = None,
+        page: Optional[int] = None,
+        sort_by: Optional[str] = None,
+        sort_order: Optional[str] = None,
+        deserialize: Optional[bool] = True,
+    ) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
+        """Get all cultural knowledge with filtering and pagination.
+
+        Args:
+            agent_id (Optional[str]): Filter by agent ID.
+            team_id (Optional[str]): Filter by team ID.
+            name (Optional[str]): Filter by name (case-insensitive partial match).
+            limit (Optional[int]): Maximum number of results to return.
+            page (Optional[int]): Page number for pagination.
+            sort_by (Optional[str]): Field to sort by.
+            sort_order (Optional[str]): Sort order ('asc' or 'desc').
+            deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge objects. Defaults to True.
+
+        Returns:
+            Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
+                - When deserialize=True: List of CulturalKnowledge objects
+                - When deserialize=False: Tuple with list of dictionaries and total count
+
+        Raises:
+            Exception: If an error occurs during retrieval.
+        """
+        try:
+            all_cultural_knowledge = self._get_all_records("culture")
+
+            # Apply filters
+            filtered_items = []
+            for item in all_cultural_knowledge:
+                if agent_id is not None and item.get("agent_id") != agent_id:
+                    continue
+                if team_id is not None and item.get("team_id") != team_id:
+                    continue
+                if name is not None and name.lower() not in item.get("name", "").lower():
+                    continue
+
+                filtered_items.append(item)
+
+            sorted_items = apply_sorting(records=filtered_items, sort_by=sort_by, sort_order=sort_order)
+            paginated_items = apply_pagination(records=sorted_items, limit=limit, page=page)
+
+            if not deserialize:
+                return paginated_items, len(filtered_items)
+
+            return [deserialize_cultural_knowledge_from_db(item) for item in paginated_items]
+
+        except Exception as e:
+            log_error(f"Error getting all cultural knowledge: {e}")
+            raise e
+
+    def upsert_cultural_knowledge(
+        self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
+    ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
+        """Upsert cultural knowledge in Redis.
+
+        Args:
+            cultural_knowledge (CulturalKnowledge): The cultural knowledge to upsert.
+            deserialize (Optional[bool]): Whether to deserialize the result. Defaults to True.
+
+        Returns:
+            Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The upserted cultural knowledge.
+
+        Raises:
+            Exception: If an error occurs during upsert.
+        """
+        try:
+            # Serialize content, categories, and notes into a dict for DB storage
+            content_dict = serialize_cultural_knowledge_for_db(cultural_knowledge)
+            item_id = cultural_knowledge.id or str(uuid4())
+
+            # Create the item dict with serialized content
+            data = {
+                "id": item_id,
+                "name": cultural_knowledge.name,
+                "summary": cultural_knowledge.summary,
+                "content": content_dict if content_dict else None,
+                "metadata": cultural_knowledge.metadata,
+                "input": cultural_knowledge.input,
+                "created_at": cultural_knowledge.created_at,
+                "updated_at": int(time.time()),
+                "agent_id": cultural_knowledge.agent_id,
+                "team_id": cultural_knowledge.team_id,
+            }
+
+            success = self._store_record("culture", item_id, data, index_fields=["name", "agent_id", "team_id"])
+
+            if not success:
+                return None
+
+            if not deserialize:
+                return data
+
+            return deserialize_cultural_knowledge_from_db(data)
+
+        except Exception as e:
+            log_error(f"Error upserting cultural knowledge: {e}")
+            raise e
+
+    # --- Traces ---
+    def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
+
+        Args:
+            trace: The Trace object to store (one per trace_id).
+        """
+        try:
+            # Check if trace already exists
+            existing = self._get_record("traces", trace.trace_id)
+
+            if existing:
+                # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
+                def get_component_level(
+                    workflow_id: Optional[str], team_id: Optional[str], agent_id: Optional[str], name: str
+                ) -> int:
+                    # Check if name indicates a root span
+                    is_root_name = ".run" in name or ".arun" in name
+
+                    if not is_root_name:
+                        return 0  # Child span (not a root)
+                    elif workflow_id:
+                        return 3  # Workflow root
+                    elif team_id:
+                        return 2  # Team root
+                    elif agent_id:
+                        return 1  # Agent root
+                    else:
+                        return 0  # Unknown
+
+                existing_level = get_component_level(
+                    existing.get("workflow_id"),
+                    existing.get("team_id"),
+                    existing.get("agent_id"),
+                    existing.get("name", ""),
+                )
+                new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
+
+                # Only update name if new trace is from a higher or equal level
+                should_update_name = new_level > existing_level
+
+                # Parse existing start_time to calculate correct duration
+                existing_start_time_str = existing.get("start_time")
+                if isinstance(existing_start_time_str, str):
+                    existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
+                else:
+                    existing_start_time = trace.start_time
+
+                recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
+
+                # Update existing record
+                existing["end_time"] = trace.end_time.isoformat()
+                existing["duration_ms"] = recalculated_duration_ms
+                existing["status"] = trace.status
+                if should_update_name:
+                    existing["name"] = trace.name
+
+                # Update context fields ONLY if new value is not None (preserve non-null values)
+                if trace.run_id is not None:
+                    existing["run_id"] = trace.run_id
+                if trace.session_id is not None:
+                    existing["session_id"] = trace.session_id
+                if trace.user_id is not None:
+                    existing["user_id"] = trace.user_id
+                if trace.agent_id is not None:
+                    existing["agent_id"] = trace.agent_id
+                if trace.team_id is not None:
+                    existing["team_id"] = trace.team_id
+                if trace.workflow_id is not None:
+                    existing["workflow_id"] = trace.workflow_id
+
+                log_debug(
+                    f"  Updating trace with context: run_id={existing.get('run_id', 'unchanged')}, "
+                    f"session_id={existing.get('session_id', 'unchanged')}, "
+                    f"user_id={existing.get('user_id', 'unchanged')}, "
+                    f"agent_id={existing.get('agent_id', 'unchanged')}, "
+                    f"team_id={existing.get('team_id', 'unchanged')}, "
+                )
+
+                self._store_record(
+                    "traces",
+                    trace.trace_id,
+                    existing,
+                    index_fields=["run_id", "session_id", "user_id", "agent_id", "team_id", "workflow_id", "status"],
+                )
+            else:
+                trace_dict = trace.to_dict()
+                trace_dict.pop("total_spans", None)
+                trace_dict.pop("error_count", None)
+                self._store_record(
+                    "traces",
+                    trace.trace_id,
+                    trace_dict,
+                    index_fields=["run_id", "session_id", "user_id", "agent_id", "team_id", "workflow_id", "status"],
+                )
+
+        except Exception as e:
+            log_error(f"Error creating trace: {e}")
+            # Don't raise - tracing should not break the main application flow
+
+    def get_trace(
+        self,
+        trace_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+    ):
+        """Get a single trace by trace_id or other filters.
+
+        Args:
+            trace_id: The unique trace identifier.
+            run_id: Filter by run ID (returns first match).
+
+        Returns:
+            Optional[Trace]: The trace if found, None otherwise.
+
+        Note:
+            If multiple filters are provided, trace_id takes precedence.
+            For other filters, the most recent trace is returned.
+        """
+        try:
+            from agno.tracing.schemas import Trace as TraceSchema
+
+            if trace_id:
+                result = self._get_record("traces", trace_id)
+                if result:
+                    # Calculate total_spans and error_count
+                    all_spans = self._get_all_records("spans")
+                    trace_spans = [s for s in all_spans if s.get("trace_id") == trace_id]
+                    result["total_spans"] = len(trace_spans)
+                    result["error_count"] = len([s for s in trace_spans if s.get("status_code") == "ERROR"])
+                    return TraceSchema.from_dict(result)
+                return None
+
+            elif run_id:
+                all_traces = self._get_all_records("traces")
+                matching = [t for t in all_traces if t.get("run_id") == run_id]
+                if matching:
+                    # Sort by start_time descending and get most recent
+                    matching.sort(key=lambda x: x.get("start_time", ""), reverse=True)
+                    result = matching[0]
+                    # Calculate total_spans and error_count
+                    all_spans = self._get_all_records("spans")
+                    trace_spans = [s for s in all_spans if s.get("trace_id") == result.get("trace_id")]
+                    result["total_spans"] = len(trace_spans)
+                    result["error_count"] = len([s for s in trace_spans if s.get("status_code") == "ERROR"])
+                    return TraceSchema.from_dict(result)
+                return None
+
+            else:
+                log_debug("get_trace called without any filter parameters")
+                return None
+
+        except Exception as e:
+            log_error(f"Error getting trace: {e}")
+            return None
+
+    def get_traces(
+        self,
+        run_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        workflow_id: Optional[str] = None,
+        status: Optional[str] = None,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        limit: Optional[int] = 20,
+        page: Optional[int] = 1,
+    ) -> tuple[List, int]:
+        """Get traces matching the provided filters.
+
+        Args:
+            run_id: Filter by run ID.
+            session_id: Filter by session ID.
+            user_id: Filter by user ID.
+            agent_id: Filter by agent ID.
+            team_id: Filter by team ID.
+            workflow_id: Filter by workflow ID.
+            status: Filter by status (OK, ERROR, UNSET).
+            start_time: Filter traces starting after this datetime.
+            end_time: Filter traces ending before this datetime.
+            limit: Maximum number of traces to return per page.
+            page: Page number (1-indexed).
+
+        Returns:
+            tuple[List[Trace], int]: Tuple of (list of matching traces, total count).
+        """
+        try:
+            from agno.tracing.schemas import Trace as TraceSchema
+
+            log_debug(
+                f"get_traces called with filters: run_id={run_id}, session_id={session_id}, "
+                f"user_id={user_id}, agent_id={agent_id}, page={page}, limit={limit}"
+            )
+
+            all_traces = self._get_all_records("traces")
+            all_spans = self._get_all_records("spans")
+
+            # Apply filters
+            filtered_traces = []
+            for trace in all_traces:
+                if run_id and trace.get("run_id") != run_id:
+                    continue
+                if session_id and trace.get("session_id") != session_id:
+                    continue
+                if user_id and trace.get("user_id") != user_id:
+                    continue
+                if agent_id and trace.get("agent_id") != agent_id:
+                    continue
+                if team_id and trace.get("team_id") != team_id:
+                    continue
+                if workflow_id and trace.get("workflow_id") != workflow_id:
+                    continue
+                if status and trace.get("status") != status:
+                    continue
+                if start_time:
+                    trace_start = trace.get("start_time", "")
+                    if trace_start and trace_start < start_time.isoformat():
+                        continue
+                if end_time:
+                    trace_end = trace.get("end_time", "")
+                    if trace_end and trace_end > end_time.isoformat():
+                        continue
+
+                filtered_traces.append(trace)
+
+            total_count = len(filtered_traces)
+
+            # Sort by start_time descending
+            filtered_traces.sort(key=lambda x: x.get("start_time", ""), reverse=True)
+
+            # Apply pagination
+            paginated_traces = apply_pagination(records=filtered_traces, limit=limit, page=page)
+
+            traces = []
+            for row in paginated_traces:
+                # Calculate total_spans and error_count
+                trace_spans = [s for s in all_spans if s.get("trace_id") == row.get("trace_id")]
+                row["total_spans"] = len(trace_spans)
+                row["error_count"] = len([s for s in trace_spans if s.get("status_code") == "ERROR"])
+                traces.append(TraceSchema.from_dict(row))
+
+            return traces, total_count
+
+        except Exception as e:
+            log_error(f"Error getting traces: {e}")
+            return [], 0
+
+    def get_trace_stats(
+        self,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        workflow_id: Optional[str] = None,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        limit: Optional[int] = 20,
+        page: Optional[int] = 1,
+    ) -> tuple[List[Dict[str, Any]], int]:
+        """Get trace statistics grouped by session.
+
+        Args:
+            user_id: Filter by user ID.
+            agent_id: Filter by agent ID.
+            team_id: Filter by team ID.
+            workflow_id: Filter by workflow ID.
+            start_time: Filter sessions with traces created after this datetime.
+            end_time: Filter sessions with traces created before this datetime.
+            limit: Maximum number of sessions to return per page.
+            page: Page number (1-indexed).
+
+        Returns:
+            tuple[List[Dict], int]: Tuple of (list of session stats dicts, total count).
+                Each dict contains: session_id, user_id, agent_id, team_id, total_traces,
+                first_trace_at, last_trace_at.
+        """
+        try:
+            log_debug(
+                f"get_trace_stats called with filters: user_id={user_id}, agent_id={agent_id}, "
+                f"workflow_id={workflow_id}, team_id={team_id}, "
+                f"start_time={start_time}, end_time={end_time}, page={page}, limit={limit}"
+            )
+
+            all_traces = self._get_all_records("traces")
+
+            # Filter traces and group by session_id
+            session_stats: Dict[str, Dict[str, Any]] = {}
+            for trace in all_traces:
+                trace_session_id = trace.get("session_id")
+                if not trace_session_id:
+                    continue
+
+                # Apply filters
+                if user_id and trace.get("user_id") != user_id:
+                    continue
+                if agent_id and trace.get("agent_id") != agent_id:
+                    continue
+                if team_id and trace.get("team_id") != team_id:
+                    continue
+                if workflow_id and trace.get("workflow_id") != workflow_id:
+                    continue
+
+                created_at = trace.get("created_at", "")
+                if start_time and created_at < start_time.isoformat():
+                    continue
+                if end_time and created_at > end_time.isoformat():
+                    continue
+
+                if trace_session_id not in session_stats:
+                    session_stats[trace_session_id] = {
+                        "session_id": trace_session_id,
+                        "user_id": trace.get("user_id"),
+                        "agent_id": trace.get("agent_id"),
+                        "team_id": trace.get("team_id"),
+                        "workflow_id": trace.get("workflow_id"),
+                        "total_traces": 0,
+                        "first_trace_at": created_at,
+                        "last_trace_at": created_at,
+                    }
+
+                session_stats[trace_session_id]["total_traces"] += 1
+                if created_at < session_stats[trace_session_id]["first_trace_at"]:
+                    session_stats[trace_session_id]["first_trace_at"] = created_at
+                if created_at > session_stats[trace_session_id]["last_trace_at"]:
+                    session_stats[trace_session_id]["last_trace_at"] = created_at
+
+            # Convert to list and sort by last_trace_at descending
+            stats_list = list(session_stats.values())
+            stats_list.sort(key=lambda x: x.get("last_trace_at", ""), reverse=True)
+
+            total_count = len(stats_list)
+
+            # Apply pagination
+            paginated_stats = apply_pagination(records=stats_list, limit=limit, page=page)
+
+            # Convert ISO strings to datetime objects
+            for stat in paginated_stats:
+                first_trace_at_str = stat["first_trace_at"]
+                last_trace_at_str = stat["last_trace_at"]
+                stat["first_trace_at"] = datetime.fromisoformat(first_trace_at_str.replace("Z", "+00:00"))
+                stat["last_trace_at"] = datetime.fromisoformat(last_trace_at_str.replace("Z", "+00:00"))
+
+            return paginated_stats, total_count
+
+        except Exception as e:
+            log_error(f"Error getting trace stats: {e}")
+            return [], 0
+
+    # --- Spans ---
+    def create_span(self, span: "Span") -> None:
+        """Create a single span in the database.
+
+        Args:
+            span: The Span object to store.
+        """
+        try:
+            self._store_record(
+                "spans",
+                span.span_id,
+                span.to_dict(),
+                index_fields=["trace_id", "parent_span_id"],
+            )
+
+        except Exception as e:
+            log_error(f"Error creating span: {e}")
+
+    def create_spans(self, spans: List) -> None:
+        """Create multiple spans in the database as a batch.
+
+        Args:
+            spans: List of Span objects to store.
+        """
+        if not spans:
+            return
+
+        try:
+            for span in spans:
+                self._store_record(
+                    "spans",
+                    span.span_id,
+                    span.to_dict(),
+                    index_fields=["trace_id", "parent_span_id"],
+                )
+
+        except Exception as e:
+            log_error(f"Error creating spans batch: {e}")
+
+    def get_span(self, span_id: str):
+        """Get a single span by its span_id.
+
+        Args:
+            span_id: The unique span identifier.
+
+        Returns:
+            Optional[Span]: The span if found, None otherwise.
+        """
+        try:
+            from agno.tracing.schemas import Span as SpanSchema
+
+            result = self._get_record("spans", span_id)
+            if result:
+                return SpanSchema.from_dict(result)
+            return None
+
+        except Exception as e:
+            log_error(f"Error getting span: {e}")
+            return None
+
+    def get_spans(
+        self,
+        trace_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+        limit: Optional[int] = 1000,
+    ) -> List:
+        """Get spans matching the provided filters.
+
+        Args:
+            trace_id: Filter by trace ID.
+            parent_span_id: Filter by parent span ID.
+            limit: Maximum number of spans to return.
+
+        Returns:
+            List[Span]: List of matching spans.
+        """
+        try:
+            from agno.tracing.schemas import Span as SpanSchema
+
+            all_spans = self._get_all_records("spans")
+
+            # Apply filters
+            filtered_spans = []
+            for span in all_spans:
+                if trace_id and span.get("trace_id") != trace_id:
+                    continue
+                if parent_span_id and span.get("parent_span_id") != parent_span_id:
+                    continue
+                filtered_spans.append(span)
+
+            # Apply limit
+            if limit:
+                filtered_spans = filtered_spans[:limit]
+
+            return [SpanSchema.from_dict(s) for s in filtered_spans]
+
+        except Exception as e:
+            log_error(f"Error getting spans: {e}")
+            return []