ai-token-tracker 0.1.0

Python SDK for capturing LLM traffic and ingesting envelopes into Ai Token Tracker

Ai Token Tracker Python SDK

ai-token-tracker captures outbound LLM traffic and posts ingest envelopes to Ai Token Tracker.

Feature parity targets:

1. Automatic HTTP interception (requests, httpx sync/async)
2. SDK wrapper integration (begin_llm_call scope)
3. Diagnostics fallback (metadata-only safety net)

Requirements

• Python 3.10+

Installation

pip install ai-token-tracker

Optional extras:

pip install ai-token-tracker[requests]
pip install ai-token-tracker[httpx]
pip install ai-token-tracker[all]

Defaults

• API base URL: http://localhost:8082 (internal, hardcoded)
• enable_auto_interception=True
• enable_diagnostics_fallback=True
• Ingest auth header (internal): X-API-Key
• Ingest path (internal): /ingest

Quick Start (Plain Script)

from ai_token_tracker import (
    AiTokenTrackerIngestionClient,
    AiTokenTrackerOptions,
    AiTokenTrackerSdkClient,
)

ingestion = AiTokenTrackerIngestionClient(
    AiTokenTrackerOptions(auth_token="atk_your_key")
)

sdk = AiTokenTrackerSdkClient(ingestion)

scope = sdk.begin_llm_call(
    provider_hint="openai",
    method="POST",
    url="https://api.openai.com/v1/responses",
    request={"model": "gpt-4.1-mini", "input": "hello"},
    custom_filters={"workflow": "quickstart"},
)

try:
    provider_response = {"id": "resp_123", "output": [{"type": "output_text", "text": "hi"}]}
    scope.complete(provider_response, status_code=200)
except Exception as ex:
    scope.fail(ex, status_code=500)
    raise
finally:
    ingestion.close()

Automatic HTTP Interception

Interception is explicit install + option-gated:

import requests
from ai_token_tracker import (
    AiTokenTrackerIngestionClient,
    AiTokenTrackerOptions,
    install_http_interception,
)

client = AiTokenTrackerIngestionClient(
    AiTokenTrackerOptions(
        auth_token="atk_your_key",
        enable_auto_interception=True,
    )
)

install_http_interception(client)

# Classified request is captured and ingested automatically.
requests.post(
    "https://api.openai.com/v1/responses",
    headers={"Authorization": "Bearer sk_test"},
    json={"model": "gpt-4.1-mini", "input": "hello"},
    timeout=10,
)

client.close()  # unpatches interceptors + shuts down background worker

Scoped Custom Filters For Automatic Interception

Use interception scope to attach custom filters without wrapper integration:

import requests
from ai_token_tracker import (
    AiTokenTrackerIngestionClient,
    AiTokenTrackerInterceptionScope,
    AiTokenTrackerOptions,
    install_http_interception,
)

client = AiTokenTrackerIngestionClient(
    AiTokenTrackerOptions(auth_token="atk_your_key", enable_auto_interception=True)
)
install_http_interception(client)

scope = AiTokenTrackerInterceptionScope.create({"workflow": "social_post_generation"})
scope.add_custom_filters({"job_id": "job-42", "customer_id": "cust-19"})

scope.open()
try:
    requests.post(
        "https://api.openai.com/v1/responses",
        headers={"Authorization": "Bearer sk_test"},
        json={"model": "gpt-4.1-mini", "input": "hello"},
        timeout=10,
    )
finally:
    scope.close()

Always close scope in finally to avoid leaking filters to later calls.

SDK Wrapper Example

from ai_token_tracker import AiTokenTrackerIngestionClient, AiTokenTrackerOptions, AiTokenTrackerSdkClient

ingestion = AiTokenTrackerIngestionClient(AiTokenTrackerOptions(auth_token="atk_your_key"))
tracker = AiTokenTrackerSdkClient(ingestion)

scope = tracker.begin_llm_call(
    provider_hint="anthropic",
    method="POST",
    url="https://api.anthropic.com/v1/messages",
    request={"model": "claude-sonnet-4-6", "max_tokens": 256, "messages": [{"role": "user", "content": "hello"}]},
    custom_filters={"tenant": "acme", "job": "summarization"},
)

scope.complete({"id": "msg_123"}, status_code=200)
ingestion.close()

Diagnostics Fallback Setup

Diagnostics fallback is best-effort metadata capture for uncovered traffic.

from ai_token_tracker import AiTokenTrackerIngestionClient, AiTokenTrackerOptions, install_http_interception

client = AiTokenTrackerIngestionClient(
    AiTokenTrackerOptions(
        auth_token="atk_your_key",
        enable_diagnostics_fallback=True,
    )
)
install_http_interception(client)

Fallback behavior:

• captures method, url, status, headers
• does not guarantee request/response bodies
• intended for visibility and coverage-gap detection, not full fidelity auditing

FastAPI Example

from contextlib import asynccontextmanager
from fastapi import FastAPI
from ai_token_tracker import (
    AiTokenTrackerIngestionClient,
    AiTokenTrackerOptions,
    AiTokenTrackerSdkClient,
    install_http_interception,
)

ingestion = AiTokenTrackerIngestionClient(
    AiTokenTrackerOptions(auth_token="atk_your_key", enable_auto_interception=True)
)
tracker = AiTokenTrackerSdkClient(ingestion)

@asynccontextmanager
async def lifespan(app: FastAPI):
    install_http_interception(ingestion)
    yield
    ingestion.close()

app = FastAPI(lifespan=lifespan)

@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}

Background Worker / Job Pipeline Example

from ai_token_tracker import AiTokenTrackerIngestionClient, AiTokenTrackerOptions, AiTokenTrackerSdkClient

ingestion = AiTokenTrackerIngestionClient(AiTokenTrackerOptions(auth_token="atk_your_key"))
tracker = AiTokenTrackerSdkClient(ingestion)

def run_job(job_id: str, prompt: str) -> None:
    scope = tracker.begin_llm_call(
        provider_hint="openai",
        method="POST",
        url="https://api.openai.com/v1/responses",
        request={"model": "gpt-4.1-mini", "input": prompt},
        custom_filters={"job_id": job_id, "pipeline": "nightly"},
    )
    try:
        response = {"id": f"resp_{job_id}", "output": [{"type": "output_text", "text": "done"}]}
        scope.complete(response)
    except Exception as ex:
        scope.fail(ex)
        raise

run_job("42", "summarize daily usage")
ingestion.close()

Custom Filters Guidance

custom_filters is key/value metadata attached to each envelope:

• keep keys stable (tenant, workflow, job_id, environment)
• keep values low-cardinality when possible
• avoid putting secrets or full prompt bodies in filters

Verification Checklist

1. Use SDK with valid auth_token (ingest URL is internal and fixed by SDK).
2. Create client with valid auth_token.
3. For interception, set option flag and call install_http_interception(client).
4. Send known LLM request (/v1/responses or /chat/completions).
5. Confirm ingest API receives POST /ingest with X-API-Key and expected envelope JSON.
6. Confirm app flow still works if ingest endpoint is unavailable (tracking paths are non-throwing).

Error Handling + Logging

• track(...) and scoped complete(...) / fail(...) are non-throwing for ingest failures.
• Warnings logged for ingestion failures and non-2xx ingest status.
• Debug logs used for interception/fallback internal capture issues.

Public API Summary

• AiTokenTrackerOptions(auth_token, enable_auto_interception=True, enable_diagnostics_fallback=True)
• AiTokenTrackerIngestionClient.track(...) -> TrackResult
• AiTokenTrackerIngestionClient.track_async(...) -> TrackResult
• AiTokenTrackerIngestionClient.close()
• AiTokenTrackerSdkClient.begin_llm_call(...) -> CallScope
• CallScope.complete(...), CallScope.fail(...)
• CallScope.complete_async(...), CallScope.fail_async(...)
• install_http_interception(client), uninstall_http_interception(client)