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

agno/db/surrealdb/surrealdb.py

@@ -1,6 +1,9 @@
 from datetime import date, datetime, timedelta, timezone
 from textwrap import dedent
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
+
+if TYPE_CHECKING:
+    from agno.tracing.schemas import Span, Trace
 
 from agno.db.base import BaseDb, SessionType
 from agno.db.postgres.utils import (
@@ -64,6 +67,8 @@ class SurrealDb(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,
         id: Optional[str] = None,
     ):
         """
@@ -71,6 +76,19 @@ class SurrealDb(BaseDb):
 
         Args:
             client: A blocking connection, either HTTP or WS
+            db_url: The URL of the SurrealDB database.
+            db_creds: The credentials for the SurrealDB database.
+            db_ns: The namespace for the SurrealDB database.
+            db_db: The database name for the SurrealDB database.
+            session_table: The name of the session table.
+            memory_table: The name of the memory table.
+            metrics_table: The name of the metrics table.
+            eval_table: The name of the eval table.
+            knowledge_table: The name of the knowledge table.
+            culture_table: The name of the culture table.
+            traces_table: The name of the traces table.
+            spans_table: The name of the spans table.
+            id: The ID of the database.
         """
         if id is None:
             base_seed = db_url
@@ -85,6 +103,8 @@ class SurrealDb(BaseDb):
             eval_table=eval_table,
             knowledge_table=knowledge_table,
             culture_table=culture_table,
+            traces_table=traces_table,
+            spans_table=spans_table,
         )
         self._client = client
         self._db_url = db_url
@@ -111,7 +131,9 @@ class SurrealDb(BaseDb):
             "knowledge": self.knowledge_table_name,
             "memories": self.memory_table_name,
             "sessions": self.session_table_name,
+            "spans": self.span_table_name,
             "teams": self._teams_table_name,
+            "traces": self.trace_table_name,
             "users": self._users_table_name,
             "workflows": self._workflows_table_name,
         }
@@ -159,6 +181,13 @@ class SurrealDb(BaseDb):
             table_name = self.eval_table_name
         elif table_type == "metrics":
             table_name = self.metrics_table_name
+        elif table_type == "traces":
+            table_name = self.trace_table_name
+        elif table_type == "spans":
+            # Ensure traces table exists before spans (for foreign key-like relationship)
+            if create_table_if_not_found:
+                self._get_table("traces", create_table_if_not_found=True)
+            table_name = self.span_table_name
         else:
             raise NotImplementedError(f"Unknown table type: {table_type}")
 
@@ -167,6 +196,14 @@ class SurrealDb(BaseDb):
 
         return table_name
 
+    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
+
     def _query(
         self,
         query: str,
@@ -1351,3 +1388,568 @@ class SurrealDb(BaseDb):
         if not raw or not deserialize:
             return raw
         return deserialize_eval_run_record(raw)
+
+    # --- 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:
+            table = self._get_table("traces", create_table_if_not_found=True)
+            record = RecordID(table, trace.trace_id)
+
+            # Check if trace exists
+            existing = self._query_one("SELECT * FROM ONLY $record", {"record": record}, dict)
+
+            if existing:
+                # 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 = existing.get("start_time")
+                if isinstance(existing_start_time, datetime):
+                    recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
+                else:
+                    recalculated_duration_ms = trace.duration_ms
+
+                # Build update query
+                update_fields = [
+                    "end_time = $end_time",
+                    "duration_ms = $duration_ms",
+                    "status = $status",
+                ]
+                update_vars: Dict[str, Any] = {
+                    "record": record,
+                    "end_time": trace.end_time,
+                    "duration_ms": recalculated_duration_ms,
+                    "status": trace.status,
+                }
+
+                if should_update_name:
+                    update_fields.append("name = $name")
+                    update_vars["name"] = trace.name
+
+                # Update context fields only if new value is not None
+                if trace.run_id is not None:
+                    update_fields.append("run_id = $run_id")
+                    update_vars["run_id"] = trace.run_id
+                if trace.session_id is not None:
+                    update_fields.append("session_id = $session_id")
+                    update_vars["session_id"] = trace.session_id
+                if trace.user_id is not None:
+                    update_fields.append("user_id = $user_id")
+                    update_vars["user_id"] = trace.user_id
+                if trace.agent_id is not None:
+                    update_fields.append("agent_id = $agent_id")
+                    update_vars["agent_id"] = trace.agent_id
+                if trace.team_id is not None:
+                    update_fields.append("team_id = $team_id")
+                    update_vars["team_id"] = trace.team_id
+                if trace.workflow_id is not None:
+                    update_fields.append("workflow_id = $workflow_id")
+                    update_vars["workflow_id"] = trace.workflow_id
+
+                update_query = f"UPDATE ONLY $record SET {', '.join(update_fields)}"
+                self._query_one(update_query, update_vars, dict)
+            else:
+                # Create new trace
+                trace_dict = trace.to_dict()
+                trace_dict.pop("total_spans", None)
+                trace_dict.pop("error_count", None)
+
+                # Convert datetime fields
+                if isinstance(trace_dict.get("start_time"), str):
+                    trace_dict["start_time"] = datetime.fromisoformat(trace_dict["start_time"].replace("Z", "+00:00"))
+                if isinstance(trace_dict.get("end_time"), str):
+                    trace_dict["end_time"] = datetime.fromisoformat(trace_dict["end_time"].replace("Z", "+00:00"))
+                if isinstance(trace_dict.get("created_at"), str):
+                    trace_dict["created_at"] = datetime.fromisoformat(trace_dict["created_at"].replace("Z", "+00:00"))
+
+                self._query_one(
+                    "CREATE ONLY $record CONTENT $content",
+                    {"record": record, "content": trace_dict},
+                    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:
+            table = self._get_table("traces", create_table_if_not_found=False)
+            spans_table = self._get_table("spans", create_table_if_not_found=False)
+
+            if trace_id:
+                record = RecordID(table, trace_id)
+                trace_data = self._query_one("SELECT * FROM ONLY $record", {"record": record}, dict)
+            elif run_id:
+                query = dedent(f"""
+                    SELECT * FROM {table}
+                    WHERE run_id = $run_id
+                    ORDER BY start_time DESC
+                    LIMIT 1
+                """)
+                trace_data = self._query_one(query, {"run_id": run_id}, dict)
+            else:
+                log_debug("get_trace called without any filter parameters")
+                return None
+
+            if not trace_data:
+                return None
+
+            # Calculate total_spans and error_count
+            id_obj = trace_data.get("id")
+            trace_id_val = trace_data.get("trace_id") or (id_obj.id if id_obj is not None else None)
+            if trace_id_val:
+                count_query = f"SELECT count() as total FROM {spans_table} WHERE trace_id = $trace_id GROUP ALL"
+                count_result = self._query_one(count_query, {"trace_id": trace_id_val}, dict)
+                trace_data["total_spans"] = count_result.get("total", 0) if count_result else 0
+
+                error_query = f"SELECT count() as total FROM {spans_table} WHERE trace_id = $trace_id AND status_code = 'ERROR' GROUP ALL"
+                error_result = self._query_one(error_query, {"trace_id": trace_id_val}, dict)
+                trace_data["error_count"] = error_result.get("total", 0) if error_result else 0
+
+            # Deserialize
+            return self._deserialize_trace(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:
+            table = self._get_table("traces", create_table_if_not_found=False)
+            spans_table = self._get_table("spans", create_table_if_not_found=False)
+
+            # Build where clause
+            where = WhereClause()
+            if run_id:
+                where.and_("run_id", run_id)
+            if session_id:
+                where.and_("session_id", session_id)
+            if user_id:
+                where.and_("user_id", user_id)
+            if agent_id:
+                where.and_("agent_id", agent_id)
+            if team_id:
+                where.and_("team_id", team_id)
+            if workflow_id:
+                where.and_("workflow_id", workflow_id)
+            if status:
+                where.and_("status", status)
+            if start_time:
+                where.and_("start_time", start_time, ">=")
+            if end_time:
+                where.and_("end_time", end_time, "<=")
+
+            where_clause, where_vars = where.build()
+
+            # Total count
+            total_count = self._count(table, where_clause, where_vars)
+
+            # Query with pagination
+            order_limit_start_clause = order_limit_start("start_time", "DESC", limit, page)
+            query = dedent(f"""
+                SELECT * FROM {table}
+                {where_clause}
+                {order_limit_start_clause}
+            """)
+            traces_raw = self._query(query, where_vars, dict)
+
+            # Add total_spans and error_count to each trace
+            result_traces = []
+            for trace_data in traces_raw:
+                id_obj = trace_data.get("id")
+                trace_id_val = trace_data.get("trace_id") or (id_obj.id if id_obj is not None else None)
+                if trace_id_val:
+                    count_query = f"SELECT count() as total FROM {spans_table} WHERE trace_id = $trace_id GROUP ALL"
+                    count_result = self._query_one(count_query, {"trace_id": trace_id_val}, dict)
+                    trace_data["total_spans"] = count_result.get("total", 0) if count_result else 0
+
+                    error_query = f"SELECT count() as total FROM {spans_table} WHERE trace_id = $trace_id AND status_code = 'ERROR' GROUP ALL"
+                    error_result = self._query_one(error_query, {"trace_id": trace_id_val}, dict)
+                    trace_data["error_count"] = error_result.get("total", 0) if error_result else 0
+
+                result_traces.append(self._deserialize_trace(trace_data))
+
+            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:
+            table = self._get_table("traces", create_table_if_not_found=False)
+
+            # Build where clause
+            where = WhereClause()
+            where.and_("!!session_id", True, "=")  # Ensure session_id is not null
+            if user_id:
+                where.and_("user_id", user_id)
+            if agent_id:
+                where.and_("agent_id", agent_id)
+            if team_id:
+                where.and_("team_id", team_id)
+            if workflow_id:
+                where.and_("workflow_id", workflow_id)
+            if start_time:
+                where.and_("created_at", start_time, ">=")
+            if end_time:
+                where.and_("created_at", end_time, "<=")
+
+            where_clause, where_vars = where.build()
+
+            # Get total count of unique sessions
+            count_query = dedent(f"""
+                SELECT count() as total FROM (
+                    SELECT session_id FROM {table}
+                    {where_clause}
+                    GROUP BY session_id
+                ) GROUP ALL
+            """)
+            count_result = self._query_one(count_query, where_vars, dict)
+            total_count = count_result.get("total", 0) if count_result else 0
+
+            # Query with aggregation
+            order_limit_start_clause = order_limit_start("last_trace_at", "DESC", limit, page)
+            query = dedent(f"""
+                SELECT
+                    session_id,
+                    user_id,
+                    agent_id,
+                    team_id,
+                    workflow_id,
+                    count() AS total_traces,
+                    time::min(created_at) AS first_trace_at,
+                    time::max(created_at) AS last_trace_at
+                FROM {table}
+                {where_clause}
+                GROUP BY session_id, user_id, agent_id, team_id, workflow_id
+                {order_limit_start_clause}
+            """)
+            results = self._query(query, where_vars, dict)
+
+            # Convert datetime objects
+            stats_list = []
+            for row in results:
+                stat = dict(row)
+                if isinstance(stat.get("first_trace_at"), datetime):
+                    pass  # Keep as datetime
+                if isinstance(stat.get("last_trace_at"), datetime):
+                    pass  # Keep as datetime
+                stats_list.append(stat)
+
+            return stats_list, total_count
+
+        except Exception as e:
+            log_error(f"Error getting trace stats: {e}")
+            return [], 0
+
+    def _deserialize_trace(self, trace_data: dict) -> "Trace":
+        """Helper to deserialize a trace record from SurrealDB."""
+        from agno.tracing.schemas import Trace
+
+        # Handle RecordID for id field
+        if isinstance(trace_data.get("id"), RecordID):
+            if "trace_id" not in trace_data or not trace_data["trace_id"]:
+                trace_data["trace_id"] = trace_data["id"].id
+            del trace_data["id"]
+
+        # Convert datetime to ISO string for Trace.from_dict
+        for field in ["start_time", "end_time", "created_at"]:
+            if isinstance(trace_data.get(field), datetime):
+                trace_data[field] = trace_data[field].isoformat()
+
+        return Trace.from_dict(trace_data)
+
+    # --- Spans ---
+    def create_span(self, span: "Span") -> None:
+        """Create a single span in the database.
+
+        Args:
+            span: The Span object to store.
+        """
+        try:
+            table = self._get_table("spans", create_table_if_not_found=True)
+            record = RecordID(table, span.span_id)
+
+            span_dict = span.to_dict()
+
+            # Convert datetime fields
+            if isinstance(span_dict.get("start_time"), str):
+                span_dict["start_time"] = datetime.fromisoformat(span_dict["start_time"].replace("Z", "+00:00"))
+            if isinstance(span_dict.get("end_time"), str):
+                span_dict["end_time"] = datetime.fromisoformat(span_dict["end_time"].replace("Z", "+00:00"))
+            if isinstance(span_dict.get("created_at"), str):
+                span_dict["created_at"] = datetime.fromisoformat(span_dict["created_at"].replace("Z", "+00:00"))
+
+            self._query_one(
+                "CREATE ONLY $record CONTENT $content",
+                {"record": record, "content": span_dict},
+                dict,
+            )
+
+        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:
+            table = self._get_table("spans", create_table_if_not_found=True)
+
+            for span in spans:
+                record = RecordID(table, span.span_id)
+                span_dict = span.to_dict()
+
+                # Convert datetime fields
+                if isinstance(span_dict.get("start_time"), str):
+                    span_dict["start_time"] = datetime.fromisoformat(span_dict["start_time"].replace("Z", "+00:00"))
+                if isinstance(span_dict.get("end_time"), str):
+                    span_dict["end_time"] = datetime.fromisoformat(span_dict["end_time"].replace("Z", "+00:00"))
+                if isinstance(span_dict.get("created_at"), str):
+                    span_dict["created_at"] = datetime.fromisoformat(span_dict["created_at"].replace("Z", "+00:00"))
+
+                self._query_one(
+                    "CREATE ONLY $record CONTENT $content",
+                    {"record": record, "content": span_dict},
+                    dict,
+                )
+
+        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:
+            table = self._get_table("spans", create_table_if_not_found=False)
+            record = RecordID(table, span_id)
+
+            span_data = self._query_one("SELECT * FROM ONLY $record", {"record": record}, dict)
+            if not span_data:
+                return None
+
+            return self._deserialize_span(span_data)
+
+        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:
+            table = self._get_table("spans", create_table_if_not_found=False)
+
+            # Build where clause
+            where = WhereClause()
+            if trace_id:
+                where.and_("trace_id", trace_id)
+            if parent_span_id:
+                where.and_("parent_span_id", parent_span_id)
+
+            where_clause, where_vars = where.build()
+
+            # Query
+            limit_clause = f"LIMIT {limit}" if limit else ""
+            query = dedent(f"""
+                SELECT * FROM {table}
+                {where_clause}
+                ORDER BY start_time ASC
+                {limit_clause}
+            """)
+            spans_raw = self._query(query, where_vars, dict)
+
+            return [self._deserialize_span(s) for s in spans_raw]
+
+        except Exception as e:
+            log_error(f"Error getting spans: {e}")
+            return []
+
+    def _deserialize_span(self, span_data: dict) -> "Span":
+        """Helper to deserialize a span record from SurrealDB."""
+        from agno.tracing.schemas import Span
+
+        # Handle RecordID for id field
+        if isinstance(span_data.get("id"), RecordID):
+            if "span_id" not in span_data or not span_data["span_id"]:
+                span_data["span_id"] = span_data["id"].id
+            del span_data["id"]
+
+        # Convert datetime to ISO string for Span.from_dict
+        for field in ["start_time", "end_time", "created_at"]:
+            if isinstance(span_data.get(field), datetime):
+                span_data[field] = span_data[field].isoformat()
+
+        return Span.from_dict(span_data)
+
+    # -- 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 SurrealDb")
+
+    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 SurrealDb")
+
+    def delete_learning(self, id: str) -> bool:
+        raise NotImplementedError("Learning methods not yet implemented for SurrealDb")
+
+    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 SurrealDb")