ai-token-tracker 0.1.0__py3-none-any.whl

ai_token_tracker/__init__.py

@@ -0,0 +1,33 @@
+"""Ai Token Tracker Python SDK."""
+
+from .classifier import DefaultLlmRequestClassifier
+from .envelope import map_tracked_call_to_envelope
+from .ingestion_client import AiTokenTrackerIngestionClient
+from .interception_scope import AiTokenTrackerInterceptionScope
+from .runtime import install_http_interception, uninstall_http_interception
+from .sdk_client import AiTokenTrackerSdkClient, CallScope
+from .types import (
+    AiTokenTrackerOptions,
+    IngestionEnvelope,
+    LlmRequestClassificationInput,
+    LlmRequestClassifier,
+    TrackResult,
+    TrackedLlmCall,
+)
+
+__all__ = [
+    "AiTokenTrackerIngestionClient",
+    "AiTokenTrackerOptions",
+    "AiTokenTrackerSdkClient",
+    "AiTokenTrackerInterceptionScope",
+    "CallScope",
+    "DefaultLlmRequestClassifier",
+    "IngestionEnvelope",
+    "LlmRequestClassificationInput",
+    "LlmRequestClassifier",
+    "TrackResult",
+    "TrackedLlmCall",
+    "install_http_interception",
+    "map_tracked_call_to_envelope",
+    "uninstall_http_interception",
+]

ai_token_tracker/classifier.py

@@ -0,0 +1,105 @@
+"""Request classification heuristics for LLM provider traffic."""
+
+from __future__ import annotations
+
+from urllib.parse import urlparse
+
+from .constants import (
+    PROVIDER_ANTHROPIC,
+    PROVIDER_GEMINI,
+    PROVIDER_OPENAI,
+    PROVIDER_UNIFY,
+    PROVIDER_UNKNOWN,
+)
+from .types import HeaderMap, LlmRequestClassificationInput, LlmRequestClassifier
+
+ANTHROPIC_HOST_INDICATORS = ("anthropic.com",)
+OPENAI_HOST_INDICATORS = ("openai.com", "openai.azure.com")
+GEMINI_HOST_INDICATORS = ("generativelanguage.googleapis.com", "aiplatform.googleapis.com")
+
+
+class DefaultLlmRequestClassifier(LlmRequestClassifier):
+    def try_classify(self, input_value: LlmRequestClassificationInput) -> str | None:
+        host = _parse_host(input_value.url)
+        if host is None:
+            return None
+
+        if _contains_any(host, ANTHROPIC_HOST_INDICATORS) or _has_header(input_value.headers, "anthropic-version"):
+            return PROVIDER_ANTHROPIC
+
+        if _contains_any(host, OPENAI_HOST_INDICATORS) or (
+            _has_authorization_bearer(input_value.headers) and _host_contains(host, "openai")
+        ):
+            return PROVIDER_OPENAI
+
+        if _contains_any(host, GEMINI_HOST_INDICATORS) or _has_header(input_value.headers, "x-goog-api-key"):
+            return PROVIDER_GEMINI
+
+        if _host_contains(host, "unify"):
+            return PROVIDER_UNIFY
+
+        if _looks_like_llm_path(input_value.url):
+            return PROVIDER_UNKNOWN
+
+        return None
+
+
+def _parse_host(url: str | None) -> str | None:
+    if not url:
+        return None
+
+    try:
+        parsed = urlparse(url)
+    except ValueError:
+        return None
+
+    return parsed.hostname
+
+
+def _contains_any(value: str, indicators: tuple[str, ...]) -> bool:
+    return any(_host_contains(value, indicator) for indicator in indicators)
+
+
+def _host_contains(host: str, substring: str) -> bool:
+    return substring.lower() in host.lower()
+
+
+def _has_header(headers: HeaderMap, header_name: str) -> bool:
+    target = header_name.lower()
+    return any(key.lower() == target for key in headers.keys())
+
+
+def _has_authorization_bearer(headers: HeaderMap) -> bool:
+    first_auth = _get_first_header_value(headers, "authorization")
+    if first_auth is None:
+        return False
+    return first_auth.lower().startswith("bearer ")
+
+
+def _get_first_header_value(headers: HeaderMap, header_name: str) -> str | None:
+    target = header_name.lower()
+    for key, values in headers.items():
+        if key.lower() != target:
+            continue
+        if len(values) == 0:
+            continue
+        return values[0]
+    return None
+
+
+def _looks_like_llm_path(url: str | None) -> bool:
+    if not url:
+        return False
+
+    try:
+        path = urlparse(url).path.lower()
+    except ValueError:
+        return False
+
+    return (
+        "/chat/completions" in path
+        or "/completions" in path
+        or "/responses" in path
+        or "/messages" in path
+        or ":generatecontent" in path
+    )

ai_token_tracker/constants.py

@@ -0,0 +1,32 @@
+"""Internal constants for Ai Token Tracker SDK."""
+
+from __future__ import annotations
+
+DEFAULT_API_BASE_URL = "http://localhost:8082"
+INTERNAL_AUTH_HEADER_NAME = "X-API-Key"
+INTERNAL_INGEST_PATH = "/ingest"
+DEFAULT_ENABLE_AUTO_INTERCEPTION = True
+DEFAULT_ENABLE_DIAGNOSTICS_FALLBACK = True
+
+HEADER_CONTENT_TYPE = "content-type"
+JSON_CONTENT_TYPE = "application/json"
+
+PROVIDER_ANTHROPIC = "anthropic"
+PROVIDER_OPENAI = "openai"
+PROVIDER_GEMINI = "gemini"
+PROVIDER_UNIFY = "unify"
+PROVIDER_UNKNOWN = "unknown"
+
+DEFAULT_SUCCESS_STATUS_CODE = 200
+DEFAULT_FAILURE_STATUS_CODE = 500
+
+ERR_TRACKING_FAILURE = "AiTokenTracker explicit ingestion failed."
+ERR_TRACKING_SERIALIZE_FAILURE = "AiTokenTracker explicit ingestion serialization failed."
+ERR_SCOPE_COMPLETE_FAILURE = "AiTokenTracker SDK scope completion failed."
+ERR_SCOPE_FAIL_FAILURE = "AiTokenTracker SDK scope failure tracking failed."
+ERR_SCOPE_SERIALIZE_REQUEST_FAILURE = "AiTokenTracker could not serialize SDK wrapper request payload."
+ERR_SCOPE_SERIALIZE_RESPONSE_FAILURE = "AiTokenTracker could not serialize SDK wrapper response payload."
+ERR_INTERCEPT_INGEST_FAILURE = "AiTokenTracker auto interception ingestion failed."
+ERR_INTERCEPT_CAPTURE_FAILURE = "AiTokenTracker failed to capture HTTP envelope for ingestion."
+ERR_FALLBACK_FAILURE = "AiTokenTracker diagnostics fallback event handling failed."
+ERR_NON_SUCCESS_STATUS = "AiTokenTracker ingest returned non-success status code: %s"

ai_token_tracker/envelope.py

@@ -0,0 +1,67 @@
+"""Envelope mapping helpers."""
+
+from __future__ import annotations
+
+from .types import HeaderMap, HttpMessageEnvelope, IngestionEnvelope, TrackedLlmCall
+
+_REDACTED_VALUE = "[REDACTED]"
+_SENSITIVE_HEADER_NAMES = {
+    "authorization",
+    "proxy-authorization",
+    "cookie",
+    "set-cookie",
+    "api-key",
+    "openai-organization",
+    "openai-project",
+    "anthropic-api-key",
+    "x-goog-api-key",
+}
+
+
+def map_tracked_call_to_envelope(llm_call: TrackedLlmCall) -> IngestionEnvelope:
+    return IngestionEnvelope(
+        providerHint=llm_call.provider_hint,
+        request=HttpMessageEnvelope(
+            method=llm_call.method,
+            url=llm_call.url,
+            statusCode=None,
+            headers=_redact_headers(_empty_headers_if_none(llm_call.request_headers)),
+            body=llm_call.request_body,
+        ),
+        response=HttpMessageEnvelope(
+            method=llm_call.method,
+            url=llm_call.url,
+            statusCode=llm_call.status_code,
+            headers=_redact_headers(_empty_headers_if_none(llm_call.response_headers)),
+            body=llm_call.response_body,
+        ),
+        customFilters=llm_call.custom_filters,
+    )
+
+
+def _empty_headers_if_none(value: HeaderMap | None) -> HeaderMap:
+    if value is None:
+        return {}
+    return value
+
+
+def _redact_headers(headers: HeaderMap) -> HeaderMap:
+    if not headers:
+        return {}
+
+    redacted: HeaderMap = {}
+    for key, values in headers.items():
+        redacted[key] = [_REDACTED_VALUE] if _is_sensitive_header(key) else list(values)
+
+    return redacted
+
+
+def _is_sensitive_header(header_name: str) -> bool:
+    normalized = header_name.lower()
+    if normalized == "x-api-key":
+        return False
+
+    if normalized.startswith("x-amz-"):
+        return True
+
+    return normalized in _SENSITIVE_HEADER_NAMES

ai_token_tracker/ingestion_client.py

@@ -0,0 +1,146 @@
+"""Ingestion client for Ai Token Tracker envelopes."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict
+from typing import Callable, Sequence
+from urllib.error import HTTPError, URLError
+from urllib.request import Request, urlopen
+
+from .classifier import DefaultLlmRequestClassifier
+from .constants import (
+    ERR_NON_SUCCESS_STATUS,
+    ERR_TRACKING_FAILURE,
+    ERR_TRACKING_SERIALIZE_FAILURE,
+    HEADER_CONTENT_TYPE,
+    INTERNAL_AUTH_HEADER_NAME,
+    JSON_CONTENT_TYPE,
+)
+from .envelope import map_tracked_call_to_envelope
+from .interception import (
+    _InterceptionSubscriber,
+    install_interception,
+    mark_ingest_call_end,
+    mark_ingest_call_start,
+)
+from .options import resolve_options, to_ingest_url
+from .types import (
+    AiTokenTrackerOptions,
+    LlmRequestClassificationInput,
+    LlmRequestClassifier,
+    LoggerLike,
+    ResolvedAiTokenTrackerOptions,
+    TrackResult,
+    TrackedLlmCall,
+)
+
+
+class AiTokenTrackerIngestionClient:
+    def __init__(
+        self,
+        options: AiTokenTrackerOptions,
+        *,
+        logger: LoggerLike | None = None,
+        request_classifiers: Sequence[LlmRequestClassifier] | None = None,
+    ) -> None:
+        self._options: ResolvedAiTokenTrackerOptions = resolve_options(options)
+        self._logger: LoggerLike = logger if logger is not None else logging.getLogger("ai_token_tracker")
+        self._classifier: LlmRequestClassifier = _CompositeRequestClassifier(
+            list(request_classifiers or []) + [DefaultLlmRequestClassifier()]
+        )
+        self._dispose_interception: Callable[[], None] | None = None
+        self._lock = threading.RLock()
+        self._executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="ai-token-tracker")
+
+    def track(self, call: TrackedLlmCall) -> TrackResult:
+        envelope = map_tracked_call_to_envelope(call)
+
+        try:
+            payload = json.dumps(asdict(envelope), separators=(",", ":"))
+        except Exception as ex:
+            self._logger.warning(ERR_TRACKING_SERIALIZE_FAILURE, exc_info=ex)
+            return TrackResult(success=False, status_code=None, error=ex)
+
+        request = Request(
+            to_ingest_url(self._options),
+            data=payload.encode("utf-8"),
+            method="POST",
+            headers={
+                HEADER_CONTENT_TYPE: JSON_CONTENT_TYPE,
+                INTERNAL_AUTH_HEADER_NAME: self._options.auth_token,
+            },
+        )
+
+        mark_ingest_call_start()
+        try:
+            with urlopen(request, timeout=10) as response:  # noqa: S310
+                status_code = response.getcode()
+                if status_code < 200 or status_code > 299:
+                    self._logger.warning(ERR_NON_SUCCESS_STATUS, status_code)
+                    return TrackResult(success=False, status_code=status_code, error=None)
+                return TrackResult(success=True, status_code=status_code, error=None)
+        except HTTPError as ex:
+            self._logger.warning(ERR_NON_SUCCESS_STATUS, ex.code)
+            return TrackResult(success=False, status_code=ex.code, error=None)
+        except (URLError, TimeoutError, OSError) as ex:
+            self._logger.warning(ERR_TRACKING_FAILURE, exc_info=ex)
+            return TrackResult(success=False, status_code=None, error=ex)
+        except Exception as ex:
+            self._logger.warning(ERR_TRACKING_FAILURE, exc_info=ex)
+            return TrackResult(success=False, status_code=None, error=ex)
+        finally:
+            mark_ingest_call_end()
+
+    async def track_async(self, call: TrackedLlmCall) -> TrackResult:
+        return await asyncio.to_thread(self.track, call)
+
+    def install_http_interception(self) -> None:
+        if not self._options.enable_auto_interception and not self._options.enable_diagnostics_fallback:
+            return
+
+        with self._lock:
+            if self._dispose_interception is not None:
+                return
+
+            subscriber = _InterceptionSubscriber(
+                enable_auto_interception=self._options.enable_auto_interception,
+                enable_diagnostics_fallback=self._options.enable_diagnostics_fallback,
+                classifier=self._classifier,
+                on_tracked_call=self._submit_background_track,
+                logger=self._logger,
+            )
+            self._dispose_interception = install_interception(subscriber)
+
+    def uninstall_http_interception(self) -> None:
+        with self._lock:
+            if self._dispose_interception is None:
+                return
+            self._dispose_interception()
+            self._dispose_interception = None
+
+    def close(self) -> None:
+        self.uninstall_http_interception()
+        self._executor.shutdown(wait=False, cancel_futures=True)
+
+    def _submit_background_track(self, call: TrackedLlmCall) -> None:
+        try:
+            self._executor.submit(self.track, call)
+        except Exception as ex:
+            self._logger.debug("AiTokenTracker background ingestion submission failed.", exc_info=ex)
+
+
+class _CompositeRequestClassifier(LlmRequestClassifier):
+    def __init__(self, classifiers: list[LlmRequestClassifier]) -> None:
+        self._classifiers = classifiers
+
+    def try_classify(self, input_value: LlmRequestClassificationInput) -> str | None:
+        for classifier in self._classifiers:
+            provider_hint = classifier.try_classify(input_value)
+            if provider_hint is not None:
+                return provider_hint
+        return None