xtrace-sdk 0.0.3__tar.gz

xtrace_sdk-0.0.3/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: xtrace-sdk
+Version: 0.0.3
+Summary: XTrace Python SDK
+Requires-Python: >=3.9
+Requires-Dist: requests>=2.31

xtrace_sdk-0.0.3/pyproject.toml

@@ -0,0 +1,15 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "xtrace-sdk"
+version = "0.0.3"
+description = "XTrace Python SDK"
+requires-python = ">=3.9"
+dependencies = [
+  "requests>=2.31",
+]
+
+[tool.setuptools]
+packages = ["xtrace_sdk"]

xtrace_sdk-0.0.3/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

xtrace_sdk-0.0.3/xtrace_sdk/__init__.py

@@ -0,0 +1,4 @@
+from .client import XTraceClient
+from .openai import observe_openai
+
+__all__ = ["XTraceClient", "observe_openai"]

xtrace_sdk-0.0.3/xtrace_sdk/client.py

@@ -0,0 +1,316 @@
+import atexit
+import json
+import os
+import queue
+import threading
+import time
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional, Union
+
+import requests
+
+
+@dataclass
+class XTraceConfig:
+    base_url: str
+    api_key: str
+    default_project_id: str
+    queue_max_size: int = 10_000
+    batch_max_size: int = 100
+    flush_interval_s: float = 0.5
+    request_timeout_s: float = 2.0
+    max_retries: int = 3
+
+
+class XTraceClient:
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        default_project_id: Optional[str] = None,
+        *,
+        queue_max_size: int = 10_000,
+        batch_max_size: int = 100,
+        flush_interval_s: float = 0.5,
+        request_timeout_s: float = 2.0,
+        max_retries: int = 3,
+    ) -> None:
+        base_url = base_url or os.environ.get("XTRACE_BASE_URL") or "http://127.0.0.1:8742"
+        api_key = api_key or os.environ.get("XTRACE_API_KEY") or os.environ.get("XTRACE_BEARER_TOKEN")
+        if not api_key:
+            raise ValueError("missing api_key (env XTRACE_API_KEY or XTRACE_BEARER_TOKEN)")
+
+        default_project_id = default_project_id or os.environ.get("XTRACE_PROJECT_ID") or "default"
+
+        self._cfg = XTraceConfig(
+            base_url=base_url.rstrip("/"),
+            api_key=api_key,
+            default_project_id=default_project_id,
+            queue_max_size=queue_max_size,
+            batch_max_size=batch_max_size,
+            flush_interval_s=flush_interval_s,
+            request_timeout_s=request_timeout_s,
+            max_retries=max_retries,
+        )
+
+        self._q: "queue.Queue[Dict[str, Any]]" = queue.Queue(maxsize=self._cfg.queue_max_size)
+        self._stop = threading.Event()
+        self._worker = threading.Thread(target=self._run, name="xtrace-ingest", daemon=True)
+
+        self.dropped_events = 0
+        self.sent_batches = 0
+        self.failed_batches = 0
+
+        self._worker.start()
+        atexit.register(self.shutdown, timeout_s=1.0)
+
+    @staticmethod
+    def new_id() -> str:
+        return str(uuid.uuid4())
+
+    def enqueue_batch(self, payload: Dict[str, Any]) -> None:
+        try:
+            self._q.put_nowait(payload)
+        except queue.Full:
+            self.dropped_events += 1
+
+    def flush(self, timeout_s: float = 5.0) -> None:
+        deadline = time.time() + timeout_s
+        while time.time() < deadline:
+            if self._q.empty():
+                return
+            time.sleep(0.05)
+
+    def shutdown(self, timeout_s: float = 1.0) -> None:
+        if self._stop.is_set():
+            return
+        self._stop.set()
+        self.flush(timeout_s=timeout_s)
+
+    def _run(self) -> None:
+        session = requests.Session()
+        url = f"{self._cfg.base_url}/v1/l/batch"
+        headers = {
+            "Authorization": f"Bearer {self._cfg.api_key}",
+            "Content-Type": "application/json",
+        }
+
+        buf: List[Dict[str, Any]] = []
+        last_flush = time.time()
+
+        while not self._stop.is_set():
+            timeout = max(0.0, self._cfg.flush_interval_s - (time.time() - last_flush))
+            try:
+                item = self._q.get(timeout=timeout)
+                buf.append(item)
+            except queue.Empty:
+                pass
+
+            if not buf:
+                continue
+
+            if len(buf) < self._cfg.batch_max_size and (time.time() - last_flush) < self._cfg.flush_interval_s:
+                continue
+
+            payloads = self._split_by_trace(buf)
+            buf.clear()
+            last_flush = time.time()
+
+            for payload in payloads:
+                ok = self._post_with_retry(session, url, headers, payload)
+                if ok:
+                    self.sent_batches += 1
+                else:
+                    self.failed_batches += 1
+
+        # best-effort drain
+        drain: List[Dict[str, Any]] = []
+        while True:
+            try:
+                drain.append(self._q.get_nowait())
+            except queue.Empty:
+                break
+        if drain:
+            for payload in self._split_by_trace(drain):
+                ok = self._post_with_retry(session, url, headers, payload)
+                if ok:
+                    self.sent_batches += 1
+                else:
+                    self.failed_batches += 1
+
+    def _split_by_trace(self, items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        grouped: Dict[str, Dict[str, Any]] = {}
+
+        for it in items:
+            trace = it.get("trace")
+            observations = it.get("observations") or []
+
+            trace_id = None
+            if isinstance(trace, dict):
+                trace_id = trace.get("id")
+            if trace_id is None and observations:
+                first = observations[0]
+                if isinstance(first, dict):
+                    trace_id = first.get("traceId")
+
+            if trace_id is None:
+                continue
+
+            if trace_id not in grouped:
+                grouped[trace_id] = {"trace": None, "observations": []}
+
+            if grouped[trace_id]["trace"] is None and trace is not None:
+                grouped[trace_id]["trace"] = trace
+
+            grouped[trace_id]["observations"].extend(observations)
+
+        return list(grouped.values())
+
+    def _post_with_retry(
+        self,
+        session: requests.Session,
+        url: str,
+        headers: Dict[str, str],
+        payload: Dict[str, Any],
+    ) -> bool:
+        delay = 0.2
+        for attempt in range(self._cfg.max_retries + 1):
+            try:
+                resp = session.post(url, headers=headers, json=payload, timeout=self._cfg.request_timeout_s)
+                if 200 <= resp.status_code < 300:
+                    return True
+                if resp.status_code == 429 or 500 <= resp.status_code < 600:
+                    time.sleep(delay)
+                    delay = min(delay * 2.0, 5.0)
+                    continue
+                return False
+            except requests.RequestException:
+                if attempt >= self._cfg.max_retries:
+                    return False
+                time.sleep(delay)
+                delay = min(delay * 2.0, 5.0)
+        return False
+
+    def push_metrics(self, metrics: List[Dict[str, Any]]) -> Dict[str, Any]:
+        url = f"{self._cfg.base_url}/v1/metrics/batch"
+        headers = {
+            "Authorization": f"Bearer {self._cfg.api_key}",
+            "Content-Type": "application/json",
+        }
+        payload = {"metrics": [self._normalize_metric_point(m) for m in metrics]}
+
+        delay = 0.2
+        with requests.Session() as session:
+            for attempt in range(self._cfg.max_retries + 1):
+                try:
+                    resp = session.post(
+                        url,
+                        headers=headers,
+                        json=payload,
+                        timeout=self._cfg.request_timeout_s,
+                    )
+                    if 200 <= resp.status_code < 300:
+                        return resp.json()
+                    if resp.status_code == 429 or 500 <= resp.status_code < 600:
+                        if attempt >= self._cfg.max_retries:
+                            break
+                        time.sleep(delay)
+                        delay = min(delay * 2.0, 5.0)
+                        continue
+                    resp.raise_for_status()
+                except requests.RequestException:
+                    if attempt >= self._cfg.max_retries:
+                        break
+                    time.sleep(delay)
+                    delay = min(delay * 2.0, 5.0)
+
+        raise RuntimeError("push_metrics failed")
+
+    def metrics_names(self) -> List[str]:
+        url = f"{self._cfg.base_url}/api/public/metrics/names"
+        headers = {
+            "Authorization": f"Bearer {self._cfg.api_key}",
+            "Content-Type": "application/json",
+        }
+
+        with requests.Session() as session:
+            resp = session.get(url, headers=headers, timeout=self._cfg.request_timeout_s)
+            resp.raise_for_status()
+            data = resp.json()
+
+        names = data.get("data") if isinstance(data, dict) else None
+        if not isinstance(names, list):
+            raise RuntimeError("invalid metrics_names response")
+        return names
+
+    def metrics_query(
+        self,
+        *,
+        name: str,
+        from_ts: Optional[Union[str, datetime]] = None,
+        to_ts: Optional[Union[str, datetime]] = None,
+        labels: Optional[Dict[str, str]] = None,
+        step: Optional[str] = None,
+        agg: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        url = f"{self._cfg.base_url}/api/public/metrics/query"
+        headers = {
+            "Authorization": f"Bearer {self._cfg.api_key}",
+            "Content-Type": "application/json",
+        }
+
+        params: Dict[str, str] = {"name": name}
+        if from_ts is not None:
+            params["from"] = self._normalize_ts(from_ts)
+        if to_ts is not None:
+            params["to"] = self._normalize_ts(to_ts)
+        if labels is not None:
+            params["labels"] = json.dumps(labels, separators=(",", ":"), ensure_ascii=False)
+        if step is not None:
+            params["step"] = step
+        if agg is not None:
+            params["agg"] = agg
+
+        with requests.Session() as session:
+            resp = session.get(url, headers=headers, params=params, timeout=self._cfg.request_timeout_s)
+            resp.raise_for_status()
+            data = resp.json()
+
+        if not isinstance(data, dict) or "data" not in data:
+            raise RuntimeError("invalid metrics_query response")
+        return data
+
+    @staticmethod
+    def _normalize_ts(ts: Union[str, datetime]) -> str:
+        if isinstance(ts, str):
+            return ts
+        if ts.tzinfo is None:
+            ts = ts.replace(tzinfo=timezone.utc)
+        return ts.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
+
+    def _normalize_metric_point(self, m: Dict[str, Any]) -> Dict[str, Any]:
+        name = m.get("name")
+        value = m.get("value")
+        timestamp = m.get("timestamp")
+
+        if not isinstance(name, str) or not name:
+            raise ValueError("metric.name must be non-empty string")
+        if not isinstance(value, (int, float)):
+            raise ValueError("metric.value must be number")
+        if timestamp is None:
+            raise ValueError("metric.timestamp is required")
+
+        labels = m.get("labels")
+        if labels is None:
+            labels = {}
+        if not isinstance(labels, dict):
+            raise ValueError("metric.labels must be dict")
+
+        return {
+            "name": name,
+            "labels": {str(k): str(v) for k, v in labels.items()},
+            "value": float(value),
+            "timestamp": self._normalize_ts(timestamp),
+        }

xtrace_sdk-0.0.3/xtrace_sdk/openai.py

@@ -0,0 +1,364 @@
+import time
+from typing import Any, Dict, Iterable, Optional
+
+from .client import XTraceClient
+
+
+def observe_openai(
+    client: Any,
+    *,
+    xtrace: XTraceClient,
+    name: Optional[str] = None,
+    user_id: Optional[str] = None,
+    session_id: Optional[str] = None,
+    tags: Optional[list[str]] = None,
+    metadata: Optional[dict] = None,
+    project_id: Optional[str] = None,
+) -> Any:
+    return _OpenAIClientWrapper(
+        client,
+        xtrace=xtrace,
+        name=name,
+        user_id=user_id,
+        session_id=session_id,
+        tags=tags or [],
+        metadata=metadata or {},
+        project_id=project_id,
+    )
+
+
+class _OpenAIClientWrapper:
+    def __init__(
+        self,
+        client: Any,
+        *,
+        xtrace: XTraceClient,
+        name: Optional[str],
+        user_id: Optional[str],
+        session_id: Optional[str],
+        tags: list[str],
+        metadata: dict,
+        project_id: Optional[str],
+    ) -> None:
+        self._client = client
+        self._xtrace = xtrace
+        self._name = name
+        self._user_id = user_id
+        self._session_id = session_id
+        self._tags = tags
+        self._metadata = metadata
+        self._project_id = project_id
+
+    @property
+    def chat(self) -> Any:
+        return _ChatWrapper(self)
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._client, item)
+
+
+class _ChatWrapper:
+    def __init__(self, root: _OpenAIClientWrapper) -> None:
+        self._root = root
+
+    @property
+    def completions(self) -> Any:
+        return _ChatCompletionsWrapper(self._root)
+
+
+class _ChatCompletionsWrapper:
+    def __init__(self, root: _OpenAIClientWrapper) -> None:
+        self._root = root
+
+    def create(self, *args: Any, **kwargs: Any) -> Any:
+        trace_id = self._root._xtrace.new_id()
+        obs_id = self._root._xtrace.new_id()
+
+        start = time.time()
+        stream = bool(kwargs.get("stream", False))
+
+        result = self._root._client.chat.completions.create(*args, **kwargs)
+
+        if stream:
+            return _StreamWrapper(
+                result,
+                root=self._root,
+                trace_id=trace_id,
+                obs_id=obs_id,
+                start=start,
+                kwargs=kwargs,
+            )
+
+        latency = time.time() - start
+        _record_non_stream(
+            root=self._root,
+            trace_id=trace_id,
+            obs_id=obs_id,
+            start=start,
+            latency=latency,
+            response=result,
+            kwargs=kwargs,
+        )
+        return result
+
+
+class _StreamWrapper:
+    def __init__(
+        self,
+        inner: Iterable[Any],
+        *,
+        root: _OpenAIClientWrapper,
+        trace_id: str,
+        obs_id: str,
+        start: float,
+        kwargs: Dict[str, Any],
+    ) -> None:
+        self._inner = iter(inner)
+        self._root = root
+        self._trace_id = trace_id
+        self._obs_id = obs_id
+        self._start = start
+        self._kwargs = kwargs
+        self._ttfb: Optional[float] = None
+        self._output_parts: list[str] = []
+        self._usage: Optional[Dict[str, Any]] = None
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            chunk = next(self._inner)
+        except StopIteration:
+            self._finalize()
+            raise
+        else:
+            delta = _extract_stream_delta_text(chunk)
+            if delta:
+                if self._ttfb is None:
+                    self._ttfb = time.time() - self._start
+                self._output_parts.append(delta)
+
+            usage = _extract_stream_usage(chunk)
+            if usage is not None:
+                self._usage = usage
+            return chunk
+
+    def close(self) -> None:
+        self._finalize()
+
+    def _finalize(self) -> None:
+        latency = time.time() - self._start
+        output_text = "".join(self._output_parts) if self._output_parts else None
+
+        payload = _build_payload(
+            root=self._root,
+            trace_id=self._trace_id,
+            obs_id=self._obs_id,
+            start=self._start,
+            latency=latency,
+            ttfb=self._ttfb,
+            output_text=output_text,
+            usage=self._usage,
+            kwargs=self._kwargs,
+        )
+        self._root._xtrace.enqueue_batch(payload)
+
+
+def _record_non_stream(
+    *,
+    root: _OpenAIClientWrapper,
+    trace_id: str,
+    obs_id: str,
+    start: float,
+    latency: float,
+    response: Any,
+    kwargs: Dict[str, Any],
+) -> None:
+    usage = _extract_usage(response)
+    output_text = _extract_output_text(response)
+
+    payload = _build_payload(
+        root=root,
+        trace_id=trace_id,
+        obs_id=obs_id,
+        start=start,
+        latency=latency,
+        ttfb=None,
+        output_text=output_text,
+        usage=usage,
+        kwargs=kwargs,
+    )
+    root._xtrace.enqueue_batch(payload)
+
+
+def _build_payload(
+    *,
+    root: _OpenAIClientWrapper,
+    trace_id: str,
+    obs_id: str,
+    start: float,
+    latency: float,
+    ttfb: Optional[float],
+    output_text: Optional[str],
+    usage: Optional[Dict[str, Any]],
+    kwargs: Dict[str, Any],
+) -> Dict[str, Any]:
+    ts_iso = _to_iso(start)
+    end_iso = _to_iso(start + latency)
+    completion_iso = _to_iso(start + ttfb) if ttfb is not None else None
+
+    trace = {
+        "id": trace_id,
+        "timestamp": ts_iso,
+        "name": kwargs.get("name") or root._name,
+        "userId": root._user_id,
+        "sessionId": root._session_id,
+        "tags": root._tags,
+        "metadata": root._metadata or None,
+        "projectId": root._project_id or root._xtrace._cfg.default_project_id,
+        "latency": latency,
+        "totalCost": None,
+    }
+
+    messages = kwargs.get("messages")
+    model = kwargs.get("model")
+
+    obs = {
+        "id": obs_id,
+        "traceId": trace_id,
+        "type": "GENERATION",
+        "name": "chat",
+        "startTime": ts_iso,
+        "endTime": end_iso,
+        "completionStartTime": completion_iso,
+        "model": model,
+        "modelParameters": None,
+        "input": messages,
+        "output": output_text,
+        "usage": usage,
+        "level": "DEFAULT",
+        "statusMessage": None,
+        "parentObservationId": None,
+        "promptId": None,
+        "promptName": None,
+        "promptVersion": None,
+        "modelId": None,
+        "inputPrice": None,
+        "outputPrice": None,
+        "totalPrice": None,
+        "calculatedInputCost": None,
+        "calculatedOutputCost": None,
+        "calculatedTotalCost": None,
+        "latency": latency,
+        "timeToFirstToken": ttfb,
+        "completionTokens": _safe_get_int(usage, ["output"]),
+        "promptTokens": _safe_get_int(usage, ["input"]),
+        "totalTokens": _safe_get_int(usage, ["total"]),
+        "unit": _safe_get_str(usage, ["unit"]),
+        "metadata": kwargs.get("metadata"),
+        "projectId": root._project_id or root._xtrace._cfg.default_project_id,
+    }
+
+    return {"trace": trace, "observations": [obs]}
+
+
+def _to_iso(ts: float) -> str:
+    return time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(ts)) + f".{int((ts % 1) * 1_000_000):06d}Z"
+
+
+def _extract_output_text(resp: Any) -> Optional[str]:
+    try:
+        choice0 = resp.choices[0]
+        msg = getattr(choice0, "message", None)
+        if msg is not None:
+            return getattr(msg, "content", None)
+    except Exception:
+        return None
+
+
+def _extract_stream_usage(chunk: Any) -> Optional[Dict[str, Any]]:
+    # Best-effort: depends on provider and whether stream_options.include_usage is enabled.
+    try:
+        u = getattr(chunk, "usage", None)
+        if u is not None:
+            prompt = getattr(u, "prompt_tokens", None)
+            completion = getattr(u, "completion_tokens", None)
+            total = getattr(u, "total_tokens", None)
+            if prompt is None and completion is None and total is None:
+                return None
+            return {
+                "input": int(prompt or 0),
+                "output": int(completion or 0),
+                "total": int(total or ((prompt or 0) + (completion or 0))),
+                "unit": "TOKENS",
+            }
+
+        # allow custom injected dict
+        u2 = getattr(chunk, "xtrace_usage", None)
+        if isinstance(u2, dict):
+            return u2
+        return None
+    except Exception:
+        return None
+    return None
+
+
+def _extract_usage(resp: Any) -> Optional[Dict[str, Any]]:
+    try:
+        u = getattr(resp, "usage", None)
+        if u is None:
+            return None
+        prompt = getattr(u, "prompt_tokens", None)
+        completion = getattr(u, "completion_tokens", None)
+        total = getattr(u, "total_tokens", None)
+        if prompt is None and completion is None and total is None:
+            return None
+        return {
+            "input": int(prompt or 0),
+            "output": int(completion or 0),
+            "total": int(total or ((prompt or 0) + (completion or 0))),
+            "unit": "TOKENS",
+        }
+    except Exception:
+        return None
+
+
+def _extract_stream_delta_text(chunk: Any) -> Optional[str]:
+    try:
+        choices = getattr(chunk, "choices", None)
+        if not choices:
+            return None
+        delta = getattr(choices[0], "delta", None)
+        if delta is None:
+            return None
+        return getattr(delta, "content", None)
+    except Exception:
+        return None
+
+
+def _safe_get_int(d: Optional[Dict[str, Any]], path: list[str]) -> Optional[int]:
+    if not d:
+        return None
+    cur: Any = d
+    for p in path:
+        if not isinstance(cur, dict) or p not in cur:
+            return None
+        cur = cur[p]
+    try:
+        return int(cur)
+    except Exception:
+        return None
+
+
+def _safe_get_str(d: Optional[Dict[str, Any]], path: list[str]) -> Optional[str]:
+    if not d:
+        return None
+    cur: Any = d
+    for p in path:
+        if not isinstance(cur, dict) or p not in cur:
+            return None
+        cur = cur[p]
+    if cur is None:
+        return None
+    return str(cur)

xtrace_sdk-0.0.3/xtrace_sdk.egg-info/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: xtrace-sdk
+Version: 0.0.3
+Summary: XTrace Python SDK
+Requires-Python: >=3.9
+Requires-Dist: requests>=2.31

xtrace_sdk-0.0.3/xtrace_sdk.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+pyproject.toml
+xtrace_sdk/__init__.py
+xtrace_sdk/client.py
+xtrace_sdk/openai.py
+xtrace_sdk.egg-info/PKG-INFO
+xtrace_sdk.egg-info/SOURCES.txt
+xtrace_sdk.egg-info/dependency_links.txt
+xtrace_sdk.egg-info/requires.txt
+xtrace_sdk.egg-info/top_level.txt

xtrace_sdk-0.0.3/xtrace_sdk.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

xtrace_sdk-0.0.3/xtrace_sdk.egg-info/requires.txt

@@ -0,0 +1 @@
+requests>=2.31

xtrace_sdk-0.0.3/xtrace_sdk.egg-info/top_level.txt

@@ -0,0 +1 @@
+xtrace_sdk