agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/db/firestore/schemas.py

@@ -67,6 +67,7 @@ USER_MEMORY_COLLECTION_SCHEMA = [
     {"key": "agent_id"},
     {"key": "team_id"},
     {"key": "topics"},
+    {"key": "created_at"},
     {"key": "updated_at"},
     # Composite indexes for memory queries
     {"key": [("user_id", "ASCENDING"), ("agent_id", "ASCENDING")], "collection_group": False},
@@ -121,6 +122,44 @@ CULTURAL_KNOWLEDGE_COLLECTION_SCHEMA = [
     {"key": "updated_at"},
 ]
 
+TRACE_COLLECTION_SCHEMA = [
+    {"key": "trace_id", "unique": True},
+    {"key": "name"},
+    {"key": "status"},
+    {"key": "run_id"},
+    {"key": "session_id"},
+    {"key": "user_id"},
+    {"key": "agent_id"},
+    {"key": "team_id"},
+    {"key": "workflow_id"},
+    {"key": "start_time"},
+    {"key": "end_time"},
+    {"key": "created_at"},
+    # Composite indexes for common query patterns
+    {"key": [("session_id", "ASCENDING"), ("start_time", "DESCENDING")], "collection_group": False},
+    {"key": [("user_id", "ASCENDING"), ("start_time", "DESCENDING")], "collection_group": False},
+    {"key": [("agent_id", "ASCENDING"), ("start_time", "DESCENDING")], "collection_group": False},
+    {"key": [("team_id", "ASCENDING"), ("start_time", "DESCENDING")], "collection_group": False},
+    {"key": [("workflow_id", "ASCENDING"), ("start_time", "DESCENDING")], "collection_group": False},
+    {"key": [("run_id", "ASCENDING"), ("start_time", "DESCENDING")], "collection_group": False},
+    {"key": [("status", "ASCENDING"), ("start_time", "DESCENDING")], "collection_group": False},
+]
+
+SPAN_COLLECTION_SCHEMA = [
+    {"key": "span_id", "unique": True},
+    {"key": "trace_id"},
+    {"key": "parent_span_id"},
+    {"key": "name"},
+    {"key": "span_kind"},
+    {"key": "status_code"},
+    {"key": "start_time"},
+    {"key": "end_time"},
+    {"key": "created_at"},
+    # Composite indexes for common query patterns
+    {"key": [("trace_id", "ASCENDING"), ("start_time", "ASCENDING")], "collection_group": False},
+    {"key": [("parent_span_id", "ASCENDING"), ("start_time", "ASCENDING")], "collection_group": False},
+]
+
 
 def get_collection_indexes(collection_type: str) -> List[Dict[str, Any]]:
     """Get the index definitions for a specific collection type."""
@@ -131,6 +170,8 @@ def get_collection_indexes(collection_type: str) -> List[Dict[str, Any]]:
         "evals": EVAL_COLLECTION_SCHEMA,
         "knowledge": KNOWLEDGE_COLLECTION_SCHEMA,
         "culture": CULTURAL_KNOWLEDGE_COLLECTION_SCHEMA,
+        "traces": TRACE_COLLECTION_SCHEMA,
+        "spans": SPAN_COLLECTION_SCHEMA,
     }
 
     indexes = index_definitions.get(collection_type)

agno/db/firestore/utils.py

@@ -8,6 +8,7 @@ from uuid import uuid4
 
 from agno.db.firestore.schemas import get_collection_indexes
 from agno.db.schemas.culture import CulturalKnowledge
+from agno.db.utils import get_sort_value
 from agno.utils.log import log_debug, log_error, log_info, log_warning
 
 try:
@@ -126,19 +127,33 @@ def apply_pagination(query, limit: Optional[int] = None, page: Optional[int] = N
 def apply_sorting_to_records(
     records: List[Dict[str, Any]], sort_by: Optional[str] = None, sort_order: Optional[str] = None
 ) -> List[Dict[str, Any]]:
-    """Apply sorting to in-memory records (for cases where Firestore query sorting isn't possible)."""
-    if sort_by is None:
-        return records
+    """Apply sorting to in-memory records (for cases where Firestore query sorting isn't possible).
+
+    Args:
+        records: The list of dictionaries to sort
+        sort_by: The field to sort by
+        sort_order: The sort order ('asc' or 'desc')
 
-    def get_sort_key(record):
-        value = record.get(sort_by, 0)
-        if value is None:
-            return 0 if records and isinstance(records[0].get(sort_by, 0), (int, float)) else ""
-        return value
+    Returns:
+        The sorted list
+
+    Note:
+        If sorting by "updated_at", will fallback to "created_at" in case of None.
+    """
+    if sort_by is None or not records:
+        return records
 
     try:
-        is_reverse = sort_order == "desc"
-        return sorted(records, key=get_sort_key, reverse=is_reverse)
+        is_descending = sort_order == "desc"
+
+        # Sort using the helper function that handles updated_at -> created_at fallback
+        sorted_records = sorted(
+            records,
+            key=lambda x: (get_sort_value(x, sort_by) is None, get_sort_value(x, sort_by)),
+            reverse=is_descending,
+        )
+
+        return sorted_records
     except Exception as e:
         log_warning(f"Error sorting Firestore records: {e}")
         return records

agno/db/gcs_json/gcs_json_db.py

@@ -1,9 +1,12 @@
 import json
 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.gcs_json.utils import (
     apply_sorting,
@@ -38,6 +41,8 @@ class GcsJsonDb(BaseDb):
         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,
         project: Optional[str] = None,
         credentials: Optional[Any] = None,
         id: Optional[str] = None,
@@ -54,6 +59,8 @@ class GcsJsonDb(BaseDb):
             eval_table (Optional[str]): Name of the JSON file to store evaluation runs.
             knowledge_table (Optional[str]): Name of the JSON file to store knowledge content.
             culture_table (Optional[str]): Name of the JSON file to store cultural knowledge.
+            traces_table (Optional[str]): Name of the JSON file to store traces.
+            spans_table (Optional[str]): Name of the JSON file to store spans.
             project (Optional[str]): GCP project ID. If None, uses default project.
             location (Optional[str]): GCS bucket location. If None, uses default location.
             credentials (Optional[Any]): GCP credentials. If None, uses default credentials.
@@ -72,6 +79,8 @@ class GcsJsonDb(BaseDb):
             eval_table=eval_table,
             knowledge_table=knowledge_table,
             culture_table=culture_table,
+            traces_table=traces_table,
+            spans_table=spans_table,
         )
 
         self.bucket_name = bucket_name
@@ -142,6 +151,14 @@ class GcsJsonDb(BaseDb):
             log_error(f"Error writing to the {blob_name} JSON file in GCS: {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:
@@ -627,13 +644,14 @@ class GcsJsonDb(BaseDb):
             raise e
 
     def get_user_memory_stats(
-        self, limit: Optional[int] = None, page: Optional[int] = None
+        self, limit: Optional[int] = None, page: Optional[int] = None, user_id: Optional[str] = None
     ) -> Tuple[List[Dict[str, Any]], int]:
         """Get user memory statistics.
 
         Args:
             limit (Optional[int]): Maximum number of results to return.
             page (Optional[int]): Page number for pagination.
+            user_id (Optional[str]): User ID for filtering.
 
         Returns:
             Tuple[List[Dict[str, Any]], int]: List of user memory statistics and total count.
@@ -644,7 +662,9 @@ class GcsJsonDb(BaseDb):
 
             for memory in 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:
                     if memory_user_id not in user_stats:
                         user_stats[memory_user_id] = {
@@ -1333,3 +1353,486 @@ class GcsJsonDb(BaseDb):
         except Exception as e:
             log_warning(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:
+            traces = self._read_json_file(self.trace_table_name, create_table_if_not_found=True)
+
+            # Check if trace exists
+            existing_idx = None
+            for i, existing in enumerate(traces):
+                if existing.get("trace_id") == trace.trace_id:
+                    existing_idx = i
+                    break
+
+            if existing_idx is not None:
+                existing = traces[existing_idx]
+
+                # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
+                def get_component_level(workflow_id: Any, team_id: Any, agent_id: Any, name: str) -> int:
+                    is_root_name = ".run" in name or ".arun" in name
+                    if not is_root_name:
+                        return 0
+                    elif workflow_id:
+                        return 3
+                    elif team_id:
+                        return 2
+                    elif agent_id:
+                        return 1
+                    else:
+                        return 0
+
+                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)
+                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 trace
+                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
+                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
+
+                traces[existing_idx] = existing
+            else:
+                # Add new trace
+                trace_dict = trace.to_dict()
+                trace_dict.pop("total_spans", None)
+                trace_dict.pop("error_count", None)
+                traces.append(trace_dict)
+
+            self._write_json_file(self.trace_table_name, traces)
+
+        except Exception as e:
+            log_error(f"Error creating trace: {e}")
+
+    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
+
+            traces = self._read_json_file(self.trace_table_name, create_table_if_not_found=False)
+            if not traces:
+                return None
+
+            # Get spans for calculating total_spans and error_count
+            spans = self._read_json_file(self.span_table_name, create_table_if_not_found=False)
+
+            # Filter traces
+            filtered = []
+            for t in traces:
+                if trace_id and t.get("trace_id") == trace_id:
+                    filtered.append(t)
+                    break
+                elif run_id and t.get("run_id") == run_id:
+                    filtered.append(t)
+
+            if not filtered:
+                return None
+
+            # Sort by start_time desc and get first
+            filtered.sort(key=lambda x: x.get("start_time", ""), reverse=True)
+            trace_data = filtered[0]
+
+            # Calculate total_spans and error_count
+            trace_spans = [s for s in spans if s.get("trace_id") == trace_data.get("trace_id")]
+            trace_data["total_spans"] = len(trace_spans)
+            trace_data["error_count"] = sum(1 for s in trace_spans if s.get("status_code") == "ERROR")
+
+            return Trace.from_dict(trace_data)
+
+        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 with pagination.
+
+        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
+
+            traces = self._read_json_file(self.trace_table_name, create_table_if_not_found=False)
+            if not traces:
+                return [], 0
+
+            # Get spans for calculating total_spans and error_count
+            spans = self._read_json_file(self.span_table_name, create_table_if_not_found=False)
+
+            # Apply filters
+            filtered = []
+            for t in traces:
+                if run_id and t.get("run_id") != run_id:
+                    continue
+                if session_id and t.get("session_id") != session_id:
+                    continue
+                if user_id and t.get("user_id") != user_id:
+                    continue
+                if agent_id and t.get("agent_id") != agent_id:
+                    continue
+                if team_id and t.get("team_id") != team_id:
+                    continue
+                if workflow_id and t.get("workflow_id") != workflow_id:
+                    continue
+                if status and t.get("status") != status:
+                    continue
+                if start_time:
+                    trace_start = t.get("start_time", "")
+                    if trace_start < start_time.isoformat():
+                        continue
+                if end_time:
+                    trace_end = t.get("end_time", "")
+                    if trace_end > end_time.isoformat():
+                        continue
+                filtered.append(t)
+
+            total_count = len(filtered)
+
+            # Sort by start_time desc
+            filtered.sort(key=lambda x: x.get("start_time", ""), reverse=True)
+
+            # Apply pagination
+            if limit and page:
+                start_idx = (page - 1) * limit
+                filtered = filtered[start_idx : start_idx + limit]
+
+            # Add total_spans and error_count to each trace
+            result_traces = []
+            for t in filtered:
+                trace_spans = [s for s in spans if s.get("trace_id") == t.get("trace_id")]
+                t["total_spans"] = len(trace_spans)
+                t["error_count"] = sum(1 for s in trace_spans if s.get("status_code") == "ERROR")
+                result_traces.append(Trace.from_dict(t))
+
+            return result_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, workflow_id, total_traces,
+                first_trace_at, last_trace_at.
+        """
+        try:
+            traces = self._read_json_file(self.trace_table_name, create_table_if_not_found=False)
+            if not traces:
+                return [], 0
+
+            # Group by session_id
+            session_stats: Dict[str, Dict[str, Any]] = {}
+
+            for t in traces:
+                trace_session_id = t.get("session_id")
+                if not trace_session_id:
+                    continue
+
+                # Apply filters
+                if user_id and t.get("user_id") != user_id:
+                    continue
+                if agent_id and t.get("agent_id") != agent_id:
+                    continue
+                if team_id and t.get("team_id") != team_id:
+                    continue
+                if workflow_id and t.get("workflow_id") != workflow_id:
+                    continue
+
+                created_at = t.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": t.get("user_id"),
+                        "agent_id": t.get("agent_id"),
+                        "team_id": t.get("team_id"),
+                        "workflow_id": t.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 and session_stats[trace_session_id]["first_trace_at"]:
+                    if created_at < session_stats[trace_session_id]["first_trace_at"]:
+                        session_stats[trace_session_id]["first_trace_at"] = created_at
+                if created_at and session_stats[trace_session_id]["last_trace_at"]:
+                    if created_at > session_stats[trace_session_id]["last_trace_at"]:
+                        session_stats[trace_session_id]["last_trace_at"] = created_at
+
+            stats_list = list(session_stats.values())
+            total_count = len(stats_list)
+
+            # Sort by last_trace_at desc
+            stats_list.sort(key=lambda x: x.get("last_trace_at", ""), reverse=True)
+
+            # Apply pagination
+            if limit and page:
+                start_idx = (page - 1) * limit
+                stats_list = stats_list[start_idx : start_idx + limit]
+
+            # Convert ISO strings to datetime objects
+            for stat in stats_list:
+                first_at = stat.get("first_trace_at", "")
+                last_at = stat.get("last_trace_at", "")
+                if first_at:
+                    stat["first_trace_at"] = datetime.fromisoformat(first_at.replace("Z", "+00:00"))
+                if last_at:
+                    stat["last_trace_at"] = datetime.fromisoformat(last_at.replace("Z", "+00:00"))
+
+            return stats_list, 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:
+            spans = self._read_json_file(self.span_table_name, create_table_if_not_found=True)
+            spans.append(span.to_dict())
+            self._write_json_file(self.span_table_name, spans)
+
+        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:
+            existing_spans = self._read_json_file(self.span_table_name, create_table_if_not_found=True)
+            for span in spans:
+                existing_spans.append(span.to_dict())
+            self._write_json_file(self.span_table_name, existing_spans)
+
+        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
+
+            spans = self._read_json_file(self.span_table_name, create_table_if_not_found=False)
+
+            for s in spans:
+                if s.get("span_id") == span_id:
+                    return Span.from_dict(s)
+
+            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
+
+            spans = self._read_json_file(self.span_table_name, create_table_if_not_found=False)
+            if not spans:
+                return []
+
+            # Apply filters
+            filtered = []
+            for s in spans:
+                if trace_id and s.get("trace_id") != trace_id:
+                    continue
+                if parent_span_id and s.get("parent_span_id") != parent_span_id:
+                    continue
+                filtered.append(s)
+
+            # Apply limit
+            if limit:
+                filtered = filtered[:limit]
+
+            return [Span.from_dict(s) for s in filtered]
+
+        except Exception as e:
+            log_error(f"Error getting spans: {e}")
+            return []
+
+    # -- Learning methods (stubs) --
+    def get_learning(
+        self,
+        learning_type: str,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        entity_id: Optional[str] = None,
+        entity_type: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        raise NotImplementedError("Learning methods not yet implemented for GcsJsonDb")
+
+    def upsert_learning(
+        self,
+        id: str,
+        learning_type: str,
+        content: Dict[str, Any],
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        entity_id: Optional[str] = None,
+        entity_type: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        raise NotImplementedError("Learning methods not yet implemented for GcsJsonDb")
+
+    def delete_learning(self, id: str) -> bool:
+        raise NotImplementedError("Learning methods not yet implemented for GcsJsonDb")
+
+    def get_learnings(
+        self,
+        learning_type: Optional[str] = None,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        entity_id: Optional[str] = None,
+        entity_type: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> List[Dict[str, Any]]:
+        raise NotImplementedError("Learning methods not yet implemented for GcsJsonDb")