simplai-sdk 0.1.0__py3-none-any.whl

traces/agents/client.py

@@ -0,0 +1,697 @@
+from __future__ import annotations
+
+import json
+import time
+from typing import Any, Dict, List, Optional
+
+import httpx
+
+from constants import (
+    DEFAULT_BASE_URL,
+    METRICS_V1_AGENT_PATH,
+    METRICS_V1_TOOL_PATH,
+    TRACE_V1_DETAILS_PATH,
+    TRACE_V1_FETCH_PATH,
+    TRACE_V2_AGGREGATE_OUTPUT_PATH,
+    TRACE_V2_DETAILS_PATH,
+    TRACE_V2_DOWNLOAD_PATH,
+    TRACE_V2_FETCH_PATH,
+    TRACE_V2_TREE_PATH,
+)
+from .models import (
+    FetchTraceFilters,
+    MetricsRequestDto,
+    PageableResponse,
+    RagTrace,
+    TraceError,
+    TraceFilters,
+    TraceNode,
+)
+
+
+class TraceClient:
+    """Low-level HTTP client for the Simplai tracer API.
+
+    This class is reusable and manages underlying HTTP clients for efficiency.
+
+    Args:
+        api_key: PIM-SID key used for authenticating with the Simplai edge service.
+        base_url: Base URL of the Simplai edge service.
+        timeout: Default request timeout in seconds.
+        max_retries: Number of retries for transient HTTP errors.
+        backoff_factor: Base factor (in seconds) used for exponential backoff.
+        user_id: Optional user ID for trace operations.
+        tenant_id: Optional tenant ID (defaults to "1").
+        project_id: Optional project ID.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        *,
+        base_url: str = DEFAULT_BASE_URL,
+        timeout: float = 30.0,
+        max_retries: int = 3,
+        backoff_factor: float = 0.5,
+        user_id: Optional[str] = None,
+        tenant_id: str = "1",
+        project_id: Optional[int] = None,
+        seller_id: Optional[str] = None,
+        client_id: Optional[str] = None,
+        seller_profile_id: Optional[str] = None,
+    ) -> None:
+        self.api_key = api_key
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self.backoff_factor = backoff_factor
+        self.user_id = user_id
+        self.tenant_id = tenant_id
+        self.project_id = project_id
+        self.seller_id = seller_id
+        self.client_id = client_id
+        self.seller_profile_id = seller_profile_id
+
+        self._sync_client: Optional[httpx.Client] = None
+        self._async_client: Optional[httpx.AsyncClient] = None
+
+    # ------------------------------------------------------------------
+    # Internal HTTP helpers
+    # ------------------------------------------------------------------
+
+    def _get_sync_client(self) -> httpx.Client:
+        if self._sync_client is None:
+            self._sync_client = httpx.Client(timeout=self.timeout)
+        return self._sync_client
+
+    def _get_async_client(self) -> httpx.AsyncClient:
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(timeout=self.timeout)
+        return self._async_client
+
+    def _headers(self) -> Dict[str, str]:
+        headers = {
+            "PIM-SID": self.api_key,
+            "Content-Type": "application/json",
+        }
+        if self.tenant_id:
+            headers["X-TENANT-ID"] = str(self.tenant_id)
+        if self.user_id:
+            headers["X-USER-ID"] = str(self.user_id)
+        if self.project_id:
+            headers["X-PROJECT-ID"] = str(self.project_id)
+        if self.seller_id:
+            headers["X-SELLER-ID"] = str(self.seller_id)
+        if self.client_id:
+            headers["X-CLIENT-ID"] = str(self.client_id)
+        if self.seller_profile_id:
+            headers["X-SELLER-PROFILE-ID"] = str(self.seller_profile_id)
+        return headers
+
+    def _request_with_retries_sync(
+        self, method: str, url: str, **kwargs: Any
+    ) -> httpx.Response:
+        """Make HTTP request with retry logic."""
+        client = self._get_sync_client()
+        last_exception = None
+
+        for attempt in range(self.max_retries + 1):
+            try:
+                response = client.request(method, url, **kwargs)
+                # Don't retry on 4xx errors (except 429)
+                if response.status_code < 500 or response.status_code == 429:
+                    return response
+                last_exception = None
+            except (httpx.NetworkError, httpx.TimeoutException) as e:
+                last_exception = e
+                if attempt < self.max_retries:
+                    wait_time = self.backoff_factor * (2 ** attempt)
+                    time.sleep(wait_time)
+                else:
+                    raise TraceError(f"Request failed after {self.max_retries + 1} attempts: {e}")
+
+        if last_exception:
+            raise TraceError(f"Request failed: {last_exception}")
+
+        # If we get here, it's a 5xx error that we should retry
+        raise TraceError(f"Request failed with status {response.status_code}")
+
+    # ------------------------------------------------------------------
+    # Request-level traces (RAG traces) - V1 API
+    # ------------------------------------------------------------------
+
+    def fetch_traces(
+        self,
+        filters: Optional[FetchTraceFilters] = None,
+        *,
+        page: int = 0,
+        size: int = 20,
+        sort: str = "id",
+        direction: str = "DESC",
+    ) -> PageableResponse:
+        """Fetch paginated request-level traces (RAG traces).
+
+        Args:
+            filters: Optional filters for querying traces.
+            page: Page number (0-indexed).
+            size: Page size.
+            sort: Sort field.
+            direction: Sort direction (ASC or DESC).
+
+        Returns:
+            PageableResponse containing paginated RagTrace objects.
+        """
+        url = self.base_url + TRACE_V1_FETCH_PATH
+        headers = self._headers()
+
+        # Build query parameters
+        params: Dict[str, Any] = {
+            "page": page,
+            "size": size,
+            "sort": f"{sort},{direction}",
+        }
+
+        if filters:
+            filters.tenant_id = filters.tenant_id or int(self.tenant_id)
+            params.update(filters.to_query_params())
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to fetch traces: {response.status_code} - {response.text}")
+
+        data = response.json()
+        return self._parse_pageable_response(data, RagTrace)
+
+    def get_trace_details(
+        self, *, trace_id: Optional[str] = None, request_id: Optional[str] = None
+    ) -> RagTrace:
+        """Get details of a specific request-level trace.
+
+        Args:
+            trace_id: Trace ID (optional, but either trace_id or request_id required).
+            request_id: Request ID (optional, but either trace_id or request_id required).
+
+        Returns:
+            RagTrace object with full details.
+        """
+        if not trace_id and not request_id:
+            raise ValueError("Either trace_id or request_id must be provided")
+
+        url = self.base_url + TRACE_V1_DETAILS_PATH
+        headers = self._headers()
+
+        params: Dict[str, Any] = {}
+        if trace_id:
+            params["traceId"] = trace_id
+        if request_id:
+            params["requestId"] = request_id
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to get trace details: {response.status_code} - {response.text}")
+
+        data = response.json()
+        if "data" in data:
+            trace_data = data["data"]
+        else:
+            trace_data = data
+
+        return self._parse_rag_trace(trace_data)
+
+    # ------------------------------------------------------------------
+    # Tree-based traces - V2 API
+    # ------------------------------------------------------------------
+
+    def fetch_trace_nodes(
+        self,
+        filters: Optional[TraceFilters] = None,
+        *,
+        page: int = 0,
+        size: int = 20,
+        sort: str = "id",
+        direction: str = "DESC",
+    ) -> PageableResponse:
+        """Fetch paginated tree-based trace nodes (root traces).
+
+        Args:
+            filters: Optional filters for querying traces.
+            page: Page number (0-indexed).
+            size: Page size.
+            sort: Sort field.
+            direction: Sort direction (ASC or DESC).
+
+        Returns:
+            PageableResponse containing paginated TraceNode objects.
+        """
+        url = self.base_url + TRACE_V2_FETCH_PATH
+        headers = self._headers()
+
+        params: Dict[str, Any] = {
+            "page": page,
+            "size": size,
+            "sort": f"{sort},{direction}",
+        }
+
+        if filters:
+            filters.tenant_id = filters.tenant_id or int(self.tenant_id)
+            filters.project_id = filters.project_id or (int(self.project_id) if self.project_id else None)
+            params.update(filters.to_query_params())
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to fetch trace nodes: {response.status_code} - {response.text}")
+
+        data = response.json()
+        return self._parse_pageable_response(data, TraceNode)
+
+    def get_aggregate_output(
+        self,
+        filters: Optional[TraceFilters] = None,
+        *,
+        page: int = 0,
+        size: int = 20,
+        sort: str = "id",
+        direction: str = "DESC",
+    ) -> Dict[str, Any]:
+        """Get aggregate output traces (summary view).
+
+        Args:
+            filters: Optional filters for querying traces.
+            page: Page number (0-indexed).
+            size: Page size.
+            sort: Sort field.
+            direction: Sort direction (ASC or DESC).
+
+        Returns:
+            Dictionary containing aggregate trace data.
+        """
+        url = self.base_url + TRACE_V2_AGGREGATE_OUTPUT_PATH
+        headers = self._headers()
+
+        params: Dict[str, Any] = {
+            "page": page,
+            "size": size,
+            "sort": f"{sort},{direction}",
+        }
+
+        if filters:
+            filters.tenant_id = filters.tenant_id or int(self.tenant_id)
+            filters.project_id = filters.project_id or (int(self.project_id) if self.project_id else None)
+            params.update(filters.to_query_params())
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to get aggregate output: {response.status_code} - {response.text}")
+
+        data = response.json()
+        if "data" in data:
+            return data["data"]
+        return data
+
+    def get_trace_tree(
+        self, tree_id: str, node_id: str, *, max_depth: Optional[int] = None, raw_response: bool = False
+    ) -> Union[TraceNode, Dict[str, Any]]:
+        """Get a trace sub-tree starting from a specific node.
+
+        Args:
+            tree_id: Tree ID of the trace.
+            node_id: Node ID to start the tree from.
+            max_depth: Maximum depth to traverse (optional).
+            raw_response: If True, return raw JSON response instead of parsed TraceNode.
+
+        Returns:
+            TraceNode representing the root of the sub-tree, or raw JSON dict if raw_response=True.
+        """
+        url = self.base_url + TRACE_V2_TREE_PATH
+        headers = self._headers()
+
+        params: Dict[str, Any] = {
+            "treeId": tree_id,
+            "nodeId": node_id,
+        }
+        if max_depth:
+            params["maxDepth"] = max_depth
+
+        # Include empty data payload like the working cURL
+        data = {"query": "", "variables": {}}
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params, json=data)
+
+        if response.status_code != 200:
+            error_detail = response.text
+            try:
+                error_json = response.json()
+                if isinstance(error_json, dict):
+                    error_detail = json.dumps(error_json, indent=2)
+            except:
+                pass
+            raise TraceError(
+                f"Failed to get trace tree: {response.status_code} - {error_detail}\n"
+                f"URL: {url}\n"
+                f"Params: treeId={tree_id}, nodeId={node_id}"
+            )
+
+        data = response.json()
+
+        if raw_response:
+            return data
+
+        if "result" in data:
+            node_data = data["result"]
+        elif "data" in data:
+            node_data = data["data"]
+        else:
+            node_data = data
+
+        return self._parse_trace_node(node_data)
+
+    def get_trace_node_details(self, tree_id: str, node_id: str) -> TraceNode:
+        """Get details of a specific trace node.
+
+        Args:
+            tree_id: Tree ID of the trace.
+            node_id: Node ID to get details for.
+
+        Returns:
+            TraceNode object with full details.
+        """
+        url = self.base_url + TRACE_V2_DETAILS_PATH
+        headers = self._headers()
+
+        params = {
+            "treeId": tree_id,
+            "nodeId": node_id,
+        }
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to get trace node details: {response.status_code} - {response.text}")
+
+        data = response.json()
+
+        if raw_response:
+            return data
+
+        if "result" in data:
+            node_data = data["result"]
+        elif "data" in data:
+            node_data = data["data"]
+        else:
+            node_data = data
+
+        return self._parse_trace_node(node_data)
+
+    def download_traces(
+        self, filters: Optional[TraceFilters] = None
+    ) -> bytes:
+        """Download traces as CSV.
+
+        Args:
+            filters: Optional filters for querying traces.
+
+        Returns:
+            Bytes containing CSV file content.
+        """
+        url = self.base_url + TRACE_V2_DOWNLOAD_PATH
+        headers = self._headers()
+
+        params: Dict[str, Any] = {}
+        if filters:
+            filters.tenant_id = filters.tenant_id or int(self.tenant_id)
+            filters.project_id = filters.project_id or (int(self.project_id) if self.project_id else None)
+            params.update(filters.to_query_params())
+
+        # POST request with filters in body
+        body = {}
+        if filters:
+            body = {k: v for k, v in filters.to_query_params().items() if v}
+
+        response = self._request_with_retries_sync("POST", url, headers=headers, json=body, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to download traces: {response.status_code} - {response.text}")
+
+        return response.content
+
+    # ------------------------------------------------------------------
+    # Metrics API
+    # ------------------------------------------------------------------
+
+    def get_agent_metrics(
+        self, metrics_request: Optional[MetricsRequestDto] = None
+    ) -> Dict[str, Any]:
+        """Get agent-level metrics from traces.
+
+        Args:
+            metrics_request: Optional metrics request with filters.
+
+        Returns:
+            Dictionary containing aggregated agent metrics.
+        """
+        url = self.base_url + METRICS_V1_AGENT_PATH
+        headers = self._headers()
+
+        params: Dict[str, Any] = {}
+        if metrics_request:
+            metrics_request.tenant_id = metrics_request.tenant_id or int(self.tenant_id)
+            metrics_request.project_id = metrics_request.project_id or (int(self.project_id) if self.project_id else None)
+            params.update(metrics_request.to_query_params())
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to get agent metrics: {response.status_code} - {response.text}")
+
+        data = response.json()
+        if "data" in data:
+            return data["data"]
+        return data
+
+    def get_tool_metrics(
+        self, metrics_request: Optional[MetricsRequestDto] = None
+    ) -> Dict[str, Any]:
+        """Get tool-level metrics from traces.
+
+        Args:
+            metrics_request: Optional metrics request with filters.
+
+        Returns:
+            Dictionary containing aggregated tool metrics.
+        """
+        url = self.base_url + METRICS_V1_TOOL_PATH
+        headers = self._headers()
+
+        params: Dict[str, Any] = {}
+        if metrics_request:
+            metrics_request.tenant_id = metrics_request.tenant_id or int(self.tenant_id)
+            metrics_request.project_id = metrics_request.project_id or (int(self.project_id) if self.project_id else None)
+            params.update(metrics_request.to_query_params())
+
+        response = self._request_with_retries_sync("GET", url, headers=headers, params=params)
+
+        if response.status_code != 200:
+            raise TraceError(f"Failed to get tool metrics: {response.status_code} - {response.text}")
+
+        data = response.json()
+        if "data" in data:
+            return data["data"]
+        return data
+
+    # ------------------------------------------------------------------
+    # Response parsing helpers
+    # ------------------------------------------------------------------
+
+    def _parse_pageable_response(self, data: Dict[str, Any], item_class: type) -> PageableResponse:
+        """Parse paginated response from API."""
+        if "data" in data:
+            page_data = data["data"]
+        else:
+            page_data = data
+
+        # Extract pagination info
+        content = page_data.get("content", [])
+        total_elements = page_data.get("totalElements", len(content))
+        total_pages = page_data.get("totalPages", 1)
+        page_number = page_data.get("number", 0)
+        page_size = page_data.get("size", len(content))
+
+        # Parse items
+        parsed_items = []
+        for item in content:
+            if item_class == RagTrace:
+                parsed_items.append(self._parse_rag_trace(item))
+            elif item_class == TraceNode:
+                parsed_items.append(self._parse_trace_node(item))
+            else:
+                parsed_items.append(item)
+
+        return PageableResponse(
+            success=data.get("success", True),
+            status=data.get("status", "OK"),
+            data=parsed_items,
+            total_elements=total_elements,
+            total_pages=total_pages,
+            page_number=page_number,
+            page_size=page_size,
+        )
+
+    def _parse_rag_trace(self, data: Dict[str, Any]) -> RagTrace:
+        """Parse RagTrace from API response."""
+        # Handle datetime strings
+        from datetime import datetime
+
+        created_at = None
+        updated_at = None
+        if "created_at" in data and data["created_at"]:
+            try:
+                created_at = datetime.fromisoformat(data["created_at"].replace("Z", "+00:00"))
+            except:
+                pass
+
+        if "updated_at" in data and data["updated_at"]:
+            try:
+                updated_at = datetime.fromisoformat(data["updated_at"].replace("Z", "+00:00"))
+            except:
+                pass
+
+        # Parse eval_metrics if present
+        eval_metrics = None
+        if "eval_metrics" in data and data["eval_metrics"]:
+            from .models import EvalMetrics
+            eval_metrics = [EvalMetrics(**metric) if isinstance(metric, dict) else metric for metric in data["eval_metrics"]]
+
+        return RagTrace(
+            id=data.get("id"),
+            request_id=data.get("request_id"),
+            conversation_id=data.get("conversation_id"),
+            user_id=data.get("user_id"),
+            username=data.get("username"),
+            tenant_id=data.get("tenant_id"),
+            workflow_id=data.get("workflow_id"),
+            workflow_name=data.get("workflow_name"),
+            agent_id=data.get("agent_id"),
+            agent_name=data.get("agent_name"),
+            kb_id=data.get("kb_id"),
+            kb_name=data.get("kb_name"),
+            model_id=data.get("model_id"),
+            model_name=data.get("model_name"),
+            model_input_tokens=data.get("model_input_tokens"),
+            model_output_tokens=data.get("model_output_tokens"),
+            kb_query_tokens=data.get("kb_query_tokens"),
+            context_tokens=data.get("context_tokens"),
+            query=data.get("query"),
+            prompt=data.get("prompt"),
+            context=data.get("context"),
+            response=data.get("response"),
+            response_time=data.get("response_time"),
+            model_response_time=data.get("model_response_time"),
+            kb_response_time=data.get("kb_response_time"),
+            user_feedback=data.get("user_feedback"),
+            source=data.get("source", []),
+            created_at=created_at,
+            updated_at=updated_at,
+            evaluation_started_by=data.get("evaluation_started_by"),
+            evaluation_started_by_name=data.get("evaluation_started_by_name"),
+            eval_metrics=eval_metrics,
+            status=data.get("status"),
+            ground_truth=data.get("ground_truth"),
+            payload=data,
+        )
+
+    def _parse_trace_node(self, data: Dict[str, Any]) -> TraceNode:
+        """Parse TraceNode from API response."""
+        from datetime import datetime
+        from .models import TraceData, ChildInfo
+
+        # Handle datetime strings
+        created_at = None
+        execution_time = None
+        if "created_at" in data and data["created_at"]:
+            try:
+                created_at = datetime.fromisoformat(data["created_at"].replace("Z", "+00:00"))
+            except:
+                pass
+
+        if "execution_time" in data and data["execution_time"]:
+            try:
+                execution_time = datetime.fromisoformat(data["execution_time"].replace("Z", "+00:00"))
+            except:
+                pass
+
+        # Parse nested data (only pass known TraceData keys; API may return extra fields e.g. user_message)
+        trace_data = None
+        if "data" in data and data["data"]:
+            d = data["data"]
+            if isinstance(d, dict):
+                known_data_keys = {"input", "output", "error", "metadata"}
+                filtered = {k: d[k] for k in known_data_keys if k in d}
+                try:
+                    trace_data = TraceData(**filtered)
+                except (TypeError, ValueError):
+                    trace_data = None  # keep raw in payload
+            else:
+                trace_data = d
+
+        # Parse children recursively
+        children = None
+        if "children" in data and data["children"]:
+            children = [self._parse_trace_node(child) for child in data["children"]]
+
+        # Parse descendants (only pass known ChildInfo keys)
+        descendants = None
+        if "descendants" in data and data["descendants"]:
+            result_descendants = []
+            for d in data["descendants"]:
+                if isinstance(d, dict):
+                    known_child_keys = {"node_id", "name", "entity_type"}
+                    filtered_child = {k: d[k] for k in known_child_keys if k in d}
+                    try:
+                        result_descendants.append(ChildInfo(**filtered_child))
+                    except (TypeError, ValueError):
+                        result_descendants.append(d)
+                else:
+                    result_descendants.append(d)
+            descendants = result_descendants
+
+        return TraceNode(
+            id=data.get("id"),
+            tenant_id=data.get("tenant_id"),
+            project_id=data.get("project_id"),
+            user_id=data.get("user_id"),
+            source=data.get("source"),
+            is_bulk_run_parent=data.get("is_bulk_run_parent"),
+            node_id=data.get("node_id"),
+            name=data.get("name"),
+            entity_type=data.get("entity_type"),
+            entity_id=data.get("entity_id"),
+            tree_id=data.get("tree_id"),
+            parent_id=data.get("parent_id"),
+            parent_span_id=data.get("parent_span_id"),
+            children=children,
+            descendants=descendants,
+            path=data.get("path"),
+            data=trace_data,
+            has_children=data.get("has_children"),
+            created_at=created_at,
+            execution_time=execution_time,
+            profile_id=data.get("profile_id"),
+            profile=data.get("profile"),
+            is_root=data.get("is_root"),
+            sub_entity_id=data.get("sub_entity_id"),
+            completion_tokens=data.get("completion_tokens"),
+            prompt_tokens=data.get("prompt_tokens"),
+            total_tokens=data.get("total_tokens"),
+            latency=data.get("latency"),
+            context_window_utilization=data.get("context_window_utilization"),
+            status=data.get("status"),
+            origin=data.get("origin"),
+            app_id=data.get("app_id"),
+            child_count=data.get("child_count"),
+            payload=data,
+        )
+