agno 2.3.3__py3-none-any.whl → 2.3.5__py3-none-any.whl

agno/db/mongo/async_mongo.py

@@ -4,6 +4,9 @@ from datetime import date, datetime, timedelta, timezone
 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 AsyncBaseDb, SessionType
 from agno.db.mongo.utils import (
     apply_pagination,
@@ -144,6 +147,8 @@ class AsyncMongoDb(AsyncBaseDb):
         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,
     ):
         """
@@ -165,6 +170,8 @@ class AsyncMongoDb(AsyncBaseDb):
             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:
@@ -185,6 +192,8 @@ class AsyncMongoDb(AsyncBaseDb):
             eval_table=eval_collection,
             knowledge_table=knowledge_collection,
             culture_table=culture_collection,
+            traces_table=traces_collection,
+            spans_table=spans_collection,
         )
 
         # Detect client type if provided
@@ -410,6 +419,28 @@ class AsyncMongoDb(AsyncBaseDb):
                 )
             return self.culture_collection
 
+        if table_type == "traces":
+            if reset_cache or not hasattr(self, "traces_collection"):
+                if self.trace_table_name is None:
+                    raise ValueError("Traces collection was not provided on initialization")
+                self.traces_collection = await 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":
+            if reset_cache or not hasattr(self, "spans_collection"):
+                if self.span_table_name is None:
+                    raise ValueError("Spans collection was not provided on initialization")
+                self.spans_collection = await 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}")
 
     async def _get_or_create_collection(
@@ -2158,3 +2189,494 @@ class AsyncMongoDb(AsyncBaseDb):
         except Exception as e:
             log_error(f"Error updating eval run name {eval_run_id}: {e}")
             raise e
+
+    # --- Traces ---
+    async 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 = await self._get_collection(table_type="traces", create_collection_if_not_found=True)
+            if collection is None:
+                return
+
+            # Check if trace already exists
+            existing = await collection.find_one({"trace_id": 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_values: Dict[str, Any] = {
+                    "end_time": trace.end_time.isoformat(),
+                    "duration_ms": recalculated_duration_ms,
+                    "status": trace.status,
+                    "name": trace.name if should_update_name else existing.get("name"),
+                }
+
+                # Update context fields ONLY if new value is not None (preserve non-null values)
+                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
+
+                log_debug(
+                    f"  Updating trace with context: run_id={update_values.get('run_id', 'unchanged')}, "
+                    f"session_id={update_values.get('session_id', 'unchanged')}, "
+                    f"user_id={update_values.get('user_id', 'unchanged')}, "
+                    f"agent_id={update_values.get('agent_id', 'unchanged')}, "
+                    f"team_id={update_values.get('team_id', 'unchanged')}, "
+                )
+
+                await collection.update_one({"trace_id": trace.trace_id}, {"$set": update_values})
+            else:
+                trace_dict = trace.to_dict()
+                trace_dict.pop("total_spans", None)
+                trace_dict.pop("error_count", None)
+                await collection.insert_one(trace_dict)
+
+        except Exception as e:
+            log_error(f"Error creating trace: {e}")
+            # Don't raise - tracing should not break the main application flow
+
+    async 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
+
+            collection = await self._get_collection(table_type="traces")
+            if collection is None:
+                return None
+
+            # Get spans collection for aggregation
+            spans_collection = await self._get_collection(table_type="spans")
+
+            query: Dict[str, Any] = {}
+            if trace_id:
+                query["trace_id"] = trace_id
+            elif run_id:
+                query["run_id"] = run_id
+            else:
+                log_debug("get_trace called without any filter parameters")
+                return None
+
+            # Find trace with sorting by most recent
+            result = await collection.find_one(query, sort=[("start_time", -1)])
+
+            if result:
+                # Calculate total_spans and error_count from spans collection
+                total_spans = 0
+                error_count = 0
+                if spans_collection is not None:
+                    total_spans = await spans_collection.count_documents({"trace_id": result["trace_id"]})
+                    error_count = await spans_collection.count_documents(
+                        {"trace_id": result["trace_id"], "status_code": "ERROR"}
+                    )
+
+                result["total_spans"] = total_spans
+                result["error_count"] = error_count
+                # Remove MongoDB's _id field
+                result.pop("_id", None)
+                return TraceSchema.from_dict(result)
+            return None
+
+        except Exception as e:
+            log_error(f"Error getting trace: {e}")
+            return None
+
+    async 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 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}"
+            )
+
+            collection = await self._get_collection(table_type="traces")
+            if collection is None:
+                log_debug("Traces collection not found")
+                return [], 0
+
+            # Get spans collection for aggregation
+            spans_collection = await self._get_collection(table_type="spans")
+
+            # Build query
+            query: Dict[str, Any] = {}
+            if run_id:
+                query["run_id"] = run_id
+            if session_id:
+                log_debug(f"Filtering by session_id={session_id}")
+                query["session_id"] = session_id
+            if user_id:
+                query["user_id"] = user_id
+            if agent_id:
+                query["agent_id"] = agent_id
+            if team_id:
+                query["team_id"] = team_id
+            if workflow_id:
+                query["workflow_id"] = workflow_id
+            if status:
+                query["status"] = status
+            if start_time:
+                query["start_time"] = {"$gte": start_time.isoformat()}
+            if end_time:
+                if "end_time" in query:
+                    query["end_time"]["$lte"] = end_time.isoformat()
+                else:
+                    query["end_time"] = {"$lte": end_time.isoformat()}
+
+            # Get total count
+            total_count = await collection.count_documents(query)
+            log_debug(f"Total matching traces: {total_count}")
+
+            # Apply pagination
+            skip = ((page or 1) - 1) * (limit or 20)
+            cursor = collection.find(query).sort("start_time", -1).skip(skip).limit(limit or 20)
+
+            results = await cursor.to_list(length=None)
+            log_debug(f"Returning page {page} with {len(results)} traces")
+
+            traces = []
+            for row in results:
+                # Calculate total_spans and error_count from spans collection
+                total_spans = 0
+                error_count = 0
+                if spans_collection is not None:
+                    total_spans = await spans_collection.count_documents({"trace_id": row["trace_id"]})
+                    error_count = await spans_collection.count_documents(
+                        {"trace_id": row["trace_id"], "status_code": "ERROR"}
+                    )
+
+                row["total_spans"] = total_spans
+                row["error_count"] = error_count
+                # Remove MongoDB's _id field
+                row.pop("_id", None)
+                traces.append(TraceSchema.from_dict(row))
+
+            return traces, total_count
+
+        except Exception as e:
+            log_error(f"Error getting traces: {e}")
+            return [], 0
+
+    async 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,
+                workflow_id, 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}"
+            )
+
+            collection = await self._get_collection(table_type="traces")
+            if collection is None:
+                log_debug("Traces collection not found")
+                return [], 0
+
+            # Build match stage
+            match_stage: Dict[str, Any] = {"session_id": {"$ne": None}}
+            if user_id:
+                match_stage["user_id"] = user_id
+            if agent_id:
+                match_stage["agent_id"] = agent_id
+            if team_id:
+                match_stage["team_id"] = team_id
+            if workflow_id:
+                match_stage["workflow_id"] = workflow_id
+            if start_time:
+                match_stage["created_at"] = {"$gte": start_time.isoformat()}
+            if end_time:
+                if "created_at" in match_stage:
+                    match_stage["created_at"]["$lte"] = end_time.isoformat()
+                else:
+                    match_stage["created_at"] = {"$lte": end_time.isoformat()}
+
+            # Build aggregation pipeline
+            pipeline: List[Dict[str, Any]] = [
+                {"$match": match_stage},
+                {
+                    "$group": {
+                        "_id": "$session_id",
+                        "user_id": {"$first": "$user_id"},
+                        "agent_id": {"$first": "$agent_id"},
+                        "team_id": {"$first": "$team_id"},
+                        "workflow_id": {"$first": "$workflow_id"},
+                        "total_traces": {"$sum": 1},
+                        "first_trace_at": {"$min": "$created_at"},
+                        "last_trace_at": {"$max": "$created_at"},
+                    }
+                },
+                {"$sort": {"last_trace_at": -1}},
+            ]
+
+            # Get total count
+            count_pipeline = pipeline + [{"$count": "total"}]
+            count_result = await collection.aggregate(count_pipeline).to_list(length=1)
+            total_count = count_result[0]["total"] if count_result else 0
+            log_debug(f"Total matching sessions: {total_count}")
+
+            # Apply pagination
+            skip = ((page or 1) - 1) * (limit or 20)
+            pipeline.append({"$skip": skip})
+            pipeline.append({"$limit": limit or 20})
+
+            results = await collection.aggregate(pipeline).to_list(length=None)
+            log_debug(f"Returning page {page} with {len(results)} session stats")
+
+            # Convert to list of dicts with datetime objects
+            stats_list = []
+            for row in results:
+                # Convert ISO strings to datetime objects
+                first_trace_at_str = row["first_trace_at"]
+                last_trace_at_str = row["last_trace_at"]
+
+                # Parse ISO format strings to datetime objects
+                first_trace_at = datetime.fromisoformat(first_trace_at_str.replace("Z", "+00:00"))
+                last_trace_at = datetime.fromisoformat(last_trace_at_str.replace("Z", "+00:00"))
+
+                stats_list.append(
+                    {
+                        "session_id": row["_id"],
+                        "user_id": row["user_id"],
+                        "agent_id": row["agent_id"],
+                        "team_id": row["team_id"],
+                        "workflow_id": row["workflow_id"],
+                        "total_traces": row["total_traces"],
+                        "first_trace_at": first_trace_at,
+                        "last_trace_at": last_trace_at,
+                    }
+                )
+
+            return stats_list, total_count
+
+        except Exception as e:
+            log_error(f"Error getting trace stats: {e}")
+            return [], 0
+
+    # --- Spans ---
+    async def create_span(self, span: "Span") -> None:
+        """Create a single span in the database.
+
+        Args:
+            span: The Span object to store.
+        """
+        try:
+            collection = await self._get_collection(table_type="spans", create_collection_if_not_found=True)
+            if collection is None:
+                return
+
+            await collection.insert_one(span.to_dict())
+
+        except Exception as e:
+            log_error(f"Error creating span: {e}")
+
+    async 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 = await self._get_collection(table_type="spans", create_collection_if_not_found=True)
+            if collection is None:
+                return
+
+            span_dicts = [span.to_dict() for span in spans]
+            await collection.insert_many(span_dicts)
+
+        except Exception as e:
+            log_error(f"Error creating spans batch: {e}")
+
+    async 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
+
+            collection = await self._get_collection(table_type="spans")
+            if collection is None:
+                return None
+
+            result = await collection.find_one({"span_id": span_id})
+            if result:
+                # Remove MongoDB's _id field
+                result.pop("_id", None)
+                return SpanSchema.from_dict(result)
+            return None
+
+        except Exception as e:
+            log_error(f"Error getting span: {e}")
+            return None
+
+    async 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
+
+            collection = await self._get_collection(table_type="spans")
+            if collection is None:
+                return []
+
+            # Build query
+            query: Dict[str, Any] = {}
+            if trace_id:
+                query["trace_id"] = trace_id
+            if parent_span_id:
+                query["parent_span_id"] = parent_span_id
+
+            cursor = collection.find(query).limit(limit or 1000)
+            results = await cursor.to_list(length=None)
+
+            spans = []
+            for row in results:
+                # Remove MongoDB's _id field
+                row.pop("_id", None)
+                spans.append(SpanSchema.from_dict(row))
+
+            return spans
+
+        except Exception as e:
+            log_error(f"Error getting spans: {e}")
+            return []