agno 2.3.4__py3-none-any.whl → 2.3.6__py3-none-any.whl

agno/db/firestore/firestore.py

@@ -1,8 +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.firestore.utils import (
     apply_pagination,
@@ -45,6 +49,8 @@ class FirestoreDb(BaseDb):
         eval_collection: Optional[str] = None,
         knowledge_collection: Optional[str] = None,
         culture_collection: Optional[str] = None,
+        traces_collection: Optional[str] = None,
+        spans_collection: Optional[str] = None,
         id: Optional[str] = None,
     ):
         """
@@ -59,6 +65,8 @@ class FirestoreDb(BaseDb):
             eval_collection (Optional[str]): Name of the collection to store evaluation runs.
             knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
             culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
+            traces_collection (Optional[str]): Name of the collection to store traces.
+            spans_collection (Optional[str]): Name of the collection to store spans.
             id (Optional[str]): ID of the database.
 
         Raises:
@@ -76,6 +84,8 @@ class FirestoreDb(BaseDb):
             eval_table=eval_collection,
             knowledge_table=knowledge_collection,
             culture_table=culture_collection,
+            traces_table=traces_collection,
+            spans_table=spans_collection,
         )
 
         _client: Optional[Client] = db_client
@@ -176,6 +186,30 @@ class FirestoreDb(BaseDb):
                 )
             return self.culture_collection
 
+        if table_type == "traces":
+            if not hasattr(self, "traces_collection"):
+                if self.trace_table_name is None:
+                    raise ValueError("Traces collection was not provided on initialization")
+                self.traces_collection = self._get_or_create_collection(
+                    collection_name=self.trace_table_name,
+                    collection_type="traces",
+                    create_collection_if_not_found=create_collection_if_not_found,
+                )
+            return self.traces_collection
+
+        if table_type == "spans":
+            # Ensure traces collection exists first (spans reference traces)
+            self._get_collection("traces", create_collection_if_not_found=create_collection_if_not_found)
+            if not hasattr(self, "spans_collection"):
+                if self.span_table_name is None:
+                    raise ValueError("Spans collection was not provided on initialization")
+                self.spans_collection = self._get_or_create_collection(
+                    collection_name=self.span_table_name,
+                    collection_type="spans",
+                    create_collection_if_not_found=create_collection_if_not_found,
+                )
+            return self.spans_collection
+
         raise ValueError(f"Unknown table type: {table_type}")
 
     def _get_or_create_collection(
@@ -1805,3 +1839,541 @@ class FirestoreDb(BaseDb):
         except Exception as e:
             log_error(f"Error updating eval run name {eval_run_id}: {e}")
             raise e
+
+    # --- Traces ---
+    def create_trace(self, trace: "Trace") -> None:
+        """Create a single trace record in the database.
+
+        Args:
+            trace: The Trace object to store (one per trace_id).
+        """
+        try:
+            collection_ref = self._get_collection(table_type="traces", create_collection_if_not_found=True)
+            if collection_ref is None:
+                return
+
+            # Check if trace already exists
+            docs = collection_ref.where(filter=FieldFilter("trace_id", "==", trace.trace_id)).limit(1).stream()
+            existing_doc = None
+            existing_data = None
+            for doc in docs:
+                existing_doc = doc
+                existing_data = doc.to_dict()
+                break
+
+            if existing_data and existing_doc is not None:
+                # Update existing trace
+                def get_component_level(workflow_id, team_id, agent_id, name):
+                    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_data.get("workflow_id"),
+                    existing_data.get("team_id"),
+                    existing_data.get("agent_id"),
+                    existing_data.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_data.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_values: Dict[str, Any] = {
+                    "end_time": trace.end_time.isoformat(),
+                    "duration_ms": recalculated_duration_ms,
+                    "status": trace.status,
+                }
+
+                if should_update_name:
+                    update_values["name"] = trace.name
+
+                # Update context fields only if new value is not None
+                if trace.run_id is not None:
+                    update_values["run_id"] = trace.run_id
+                if trace.session_id is not None:
+                    update_values["session_id"] = trace.session_id
+                if trace.user_id is not None:
+                    update_values["user_id"] = trace.user_id
+                if trace.agent_id is not None:
+                    update_values["agent_id"] = trace.agent_id
+                if trace.team_id is not None:
+                    update_values["team_id"] = trace.team_id
+                if trace.workflow_id is not None:
+                    update_values["workflow_id"] = trace.workflow_id
+
+                existing_doc.reference.update(update_values)
+            else:
+                # Create new trace with initialized counters
+                trace_dict = trace.to_dict()
+                trace_dict["total_spans"] = 0
+                trace_dict["error_count"] = 0
+                collection_ref.add(trace_dict)
+
+        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
+
+            collection_ref = self._get_collection(table_type="traces")
+            if collection_ref is None:
+                return None
+
+            if trace_id:
+                docs = collection_ref.where(filter=FieldFilter("trace_id", "==", trace_id)).limit(1).stream()
+            elif run_id:
+                from google.cloud.firestore import Query
+
+                docs = (
+                    collection_ref.where(filter=FieldFilter("run_id", "==", run_id))
+                    .order_by("start_time", direction=Query.DESCENDING)
+                    .limit(1)
+                    .stream()
+                )
+            else:
+                log_debug("get_trace called without any filter parameters")
+                return None
+
+            for doc in docs:
+                trace_data = doc.to_dict()
+                # Use stored values (default to 0 if not present)
+                trace_data.setdefault("total_spans", 0)
+                trace_data.setdefault("error_count", 0)
+                return Trace.from_dict(trace_data)
+
+            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
+
+            collection_ref = self._get_collection(table_type="traces")
+            if collection_ref is None:
+                return [], 0
+
+            query = collection_ref
+
+            # Apply filters
+            if run_id:
+                query = query.where(filter=FieldFilter("run_id", "==", run_id))
+            if session_id:
+                query = query.where(filter=FieldFilter("session_id", "==", session_id))
+            if user_id:
+                query = query.where(filter=FieldFilter("user_id", "==", user_id))
+            if agent_id:
+                query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
+            if team_id:
+                query = query.where(filter=FieldFilter("team_id", "==", team_id))
+            if workflow_id:
+                query = query.where(filter=FieldFilter("workflow_id", "==", workflow_id))
+            if status:
+                query = query.where(filter=FieldFilter("status", "==", status))
+            if start_time:
+                query = query.where(filter=FieldFilter("start_time", ">=", start_time.isoformat()))
+            if end_time:
+                query = query.where(filter=FieldFilter("end_time", "<=", end_time.isoformat()))
+
+            # Get all matching documents
+            docs = query.stream()
+            all_records = [doc.to_dict() for doc in docs]
+
+            # Sort by start_time descending
+            all_records.sort(key=lambda x: x.get("start_time", ""), reverse=True)
+
+            # Get total count
+            total_count = len(all_records)
+
+            # Apply pagination
+            if limit and page:
+                offset = (page - 1) * limit
+                paginated_records = all_records[offset : offset + limit]
+            elif limit:
+                paginated_records = all_records[:limit]
+            else:
+                paginated_records = all_records
+
+            # Convert to Trace objects with stored span counts
+            traces = []
+            for trace_data in paginated_records:
+                trace_data.setdefault("total_spans", 0)
+                trace_data.setdefault("error_count", 0)
+                traces.append(Trace.from_dict(trace_data))
+
+            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, workflow_id, total_traces,
+                first_trace_at, last_trace_at.
+        """
+        try:
+            collection_ref = self._get_collection(table_type="traces")
+            if collection_ref is None:
+                return [], 0
+
+            query = collection_ref
+
+            # Apply filters
+            if user_id:
+                query = query.where(filter=FieldFilter("user_id", "==", user_id))
+            if agent_id:
+                query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
+            if team_id:
+                query = query.where(filter=FieldFilter("team_id", "==", team_id))
+            if workflow_id:
+                query = query.where(filter=FieldFilter("workflow_id", "==", workflow_id))
+            if start_time:
+                query = query.where(filter=FieldFilter("created_at", ">=", start_time.isoformat()))
+            if end_time:
+                query = query.where(filter=FieldFilter("created_at", "<=", end_time.isoformat()))
+
+            # Get all matching documents
+            docs = query.stream()
+
+            # Aggregate by session_id
+            session_stats: Dict[str, Dict[str, Any]] = {}
+            for doc in docs:
+                trace_data = doc.to_dict()
+                session_id = trace_data.get("session_id")
+                if not session_id:
+                    continue
+
+                if session_id not in session_stats:
+                    session_stats[session_id] = {
+                        "session_id": session_id,
+                        "user_id": trace_data.get("user_id"),
+                        "agent_id": trace_data.get("agent_id"),
+                        "team_id": trace_data.get("team_id"),
+                        "workflow_id": trace_data.get("workflow_id"),
+                        "total_traces": 0,
+                        "first_trace_at": trace_data.get("created_at"),
+                        "last_trace_at": trace_data.get("created_at"),
+                    }
+
+                session_stats[session_id]["total_traces"] += 1
+
+                created_at = trace_data.get("created_at")
+                if (
+                    created_at
+                    and session_stats[session_id]["first_trace_at"]
+                    and session_stats[session_id]["last_trace_at"]
+                ):
+                    if created_at < session_stats[session_id]["first_trace_at"]:
+                        session_stats[session_id]["first_trace_at"] = created_at
+                    if created_at > session_stats[session_id]["last_trace_at"]:
+                        session_stats[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)
+
+            # Convert datetime strings to datetime objects
+            for stat in stats_list:
+                first_trace_at = stat["first_trace_at"]
+                last_trace_at = stat["last_trace_at"]
+                if isinstance(first_trace_at, str):
+                    stat["first_trace_at"] = datetime.fromisoformat(first_trace_at.replace("Z", "+00:00"))
+                if isinstance(last_trace_at, str):
+                    stat["last_trace_at"] = datetime.fromisoformat(last_trace_at.replace("Z", "+00:00"))
+
+            # Get total count
+            total_count = len(stats_list)
+
+            # Apply pagination
+            if limit and page:
+                offset = (page - 1) * limit
+                paginated_stats = stats_list[offset : offset + limit]
+            elif limit:
+                paginated_stats = stats_list[:limit]
+            else:
+                paginated_stats = stats_list
+
+            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:
+            collection_ref = self._get_collection(table_type="spans", create_collection_if_not_found=True)
+            if collection_ref is None:
+                return
+
+            span_dict = span.to_dict()
+            # Serialize attributes as JSON string
+            if "attributes" in span_dict and isinstance(span_dict["attributes"], dict):
+                span_dict["attributes"] = json.dumps(span_dict["attributes"])
+
+            collection_ref.add(span_dict)
+
+            # Increment total_spans and error_count on trace
+            traces_collection = self._get_collection(table_type="traces")
+            if traces_collection:
+                try:
+                    docs = (
+                        traces_collection.where(filter=FieldFilter("trace_id", "==", span.trace_id)).limit(1).stream()
+                    )
+                    for doc in docs:
+                        trace_data = doc.to_dict()
+                        current_total = trace_data.get("total_spans", 0)
+                        current_errors = trace_data.get("error_count", 0)
+
+                        update_values = {"total_spans": current_total + 1}
+                        if span.status_code == "ERROR":
+                            update_values["error_count"] = current_errors + 1
+
+                        doc.reference.update(update_values)
+                        break
+                except Exception as update_error:
+                    log_debug(f"Could not update trace span counts: {update_error}")
+
+        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:
+            collection_ref = self._get_collection(table_type="spans", create_collection_if_not_found=True)
+            if collection_ref is None:
+                return
+
+            # Firestore batch has a limit of 500 operations
+            batch = self.db_client.batch()
+            batch_count = 0
+
+            for span in spans:
+                span_dict = span.to_dict()
+                # Serialize attributes as JSON string
+                if "attributes" in span_dict and isinstance(span_dict["attributes"], dict):
+                    span_dict["attributes"] = json.dumps(span_dict["attributes"])
+
+                doc_ref = collection_ref.document()
+                batch.set(doc_ref, span_dict)
+                batch_count += 1
+
+                # Commit batch if reaching limit
+                if batch_count >= 500:
+                    batch.commit()
+                    batch = self.db_client.batch()
+                    batch_count = 0
+
+            # Commit remaining operations
+            if batch_count > 0:
+                batch.commit()
+
+            # Update trace with total_spans and error_count
+            trace_id = spans[0].trace_id
+            spans_count = len(spans)
+            error_count = sum(1 for s in spans if s.status_code == "ERROR")
+
+            traces_collection = self._get_collection(table_type="traces")
+            if traces_collection:
+                try:
+                    docs = traces_collection.where(filter=FieldFilter("trace_id", "==", trace_id)).limit(1).stream()
+                    for doc in docs:
+                        trace_data = doc.to_dict()
+                        current_total = trace_data.get("total_spans", 0)
+                        current_errors = trace_data.get("error_count", 0)
+
+                        doc.reference.update(
+                            {
+                                "total_spans": current_total + spans_count,
+                                "error_count": current_errors + error_count,
+                            }
+                        )
+                        break
+                except Exception as update_error:
+                    log_debug(f"Could not update trace span counts: {update_error}")
+
+        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
+
+            collection_ref = self._get_collection(table_type="spans")
+            if collection_ref is None:
+                return None
+
+            docs = collection_ref.where(filter=FieldFilter("span_id", "==", span_id)).limit(1).stream()
+
+            for doc in docs:
+                span_data = doc.to_dict()
+                # Deserialize attributes from JSON string
+                if "attributes" in span_data and isinstance(span_data["attributes"], str):
+                    span_data["attributes"] = json.loads(span_data["attributes"])
+                return Span.from_dict(span_data)
+
+            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
+
+            collection_ref = self._get_collection(table_type="spans")
+            if collection_ref is None:
+                return []
+
+            query = collection_ref
+
+            if trace_id:
+                query = query.where(filter=FieldFilter("trace_id", "==", trace_id))
+            if parent_span_id:
+                query = query.where(filter=FieldFilter("parent_span_id", "==", parent_span_id))
+
+            if limit:
+                query = query.limit(limit)
+
+            docs = query.stream()
+
+            spans = []
+            for doc in docs:
+                span_data = doc.to_dict()
+                # Deserialize attributes from JSON string
+                if "attributes" in span_data and isinstance(span_data["attributes"], str):
+                    span_data["attributes"] = json.loads(span_data["attributes"])
+                spans.append(Span.from_dict(span_data))
+
+            return spans
+
+        except Exception as e:
+            log_error(f"Error getting spans: {e}")
+            return []

agno/db/firestore/schemas.py

@@ -122,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."""
@@ -132,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)