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

agno/db/json/json_db.py

@@ -3,9 +3,12 @@ import os
 import time
 from datetime import date, datetime, timedelta, timezone
 from pathlib import Path
-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.json.utils import (
     apply_sorting,
@@ -34,6 +37,8 @@ class JsonDb(BaseDb):
         metrics_table: Optional[str] = None,
         eval_table: Optional[str] = None,
         knowledge_table: Optional[str] = None,
+        traces_table: Optional[str] = None,
+        spans_table: Optional[str] = None,
         id: Optional[str] = None,
     ):
         """
@@ -47,6 +52,8 @@ class JsonDb(BaseDb):
             metrics_table (Optional[str]): Name of the JSON file to store metrics.
             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.
+            traces_table (Optional[str]): Name of the JSON file to store run traces.
+            spans_table (Optional[str]): Name of the JSON file to store span events.
             id (Optional[str]): ID of the database.
         """
         if id is None:
@@ -61,6 +68,8 @@ class JsonDb(BaseDb):
             metrics_table=metrics_table,
             eval_table=eval_table,
             knowledge_table=knowledge_table,
+            traces_table=traces_table,
+            spans_table=spans_table,
         )
 
         # Create the directory where the JSON files will be stored, if it doesn't exist
@@ -1338,3 +1347,431 @@ class JsonDb(BaseDb):
         except Exception as e:
             log_error(f"Error upserting cultural knowledge: {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:
+            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, 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.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.
+        """
+        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).
+        """
+        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:
+                session_id = t.get("session_id")
+                if not 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 session_id not in session_stats:
+                    session_stats[session_id] = {
+                        "session_id": 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[session_id]["total_traces"] += 1
+                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
+
+            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 []