From 00f0d6167c94e6ed524de4af7e6a3f90ed09b528 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 7 Oct 2026 07:15:58 +0000 Subject: [PATCH 1/2] feat(collectors): add sessionless LiteLLM usage collector Summary: - Add collectors.extraction with shared helpers for timestamps, latency, TTFT, tokens, cache, cost, and status/error fields. - Add UsageRequestNormalizer that maps LiteLLM spend-log rows (fixture and native SpendLogs shapes) into usage_requests without a session, skips rows lacking request_id/call_id with diagnostics, and preserves optional benchmark join fields in redacted metadata. - Add ProxyKeyAttributionResolver that resolves spend-log key references to proxy_keys by litellm_key_id, ephemeral virtual-key mapping, or alias; raw key references are never persisted. - Add LiteLLMUsageCollector that fetches /spend/logs for all keys (summarize=false), filters to the window, links known benchmark sessions, and writes idempotently via SQLUsageRequestRepository. - Add UsageReconciliationReport with mapped/partial/skipped, attribution, session-link, duplicate, and missing-field counts. - Document the collector flow and metadata redaction boundary. Rationale: - COE-373: usage ingestion must work for all keys, not only benchmark sessions, while keeping the session-bound benchmark collector unchanged. - ADR-002 forbids persisting raw virtual key values; unattributed keys are reported by alias or sha256 fingerprint. Tests: - pytest tests/unit/test_usage_normalizer.py tests/unit/test_collectors.py - make lint format-check type-check test (701 passed) Co-Authored-By: bot_apk --- docs/data-model-and-observability.md | 14 +- docs/security-and-operations.md | 4 + src/collectors/__init__.py | 15 + src/collectors/extraction.py | 295 ++++++++++++ src/collectors/key_attribution.py | 122 +++++ src/collectors/litellm_collector.py | 251 +++++++++- src/collectors/normalize_requests.py | 526 +++++++++++++++++++++ tests/unit/test_usage_normalizer.py | 660 +++++++++++++++++++++++++++ 8 files changed, 1884 insertions(+), 3 deletions(-) create mode 100644 src/collectors/extraction.py create mode 100644 src/collectors/key_attribution.py create mode 100644 tests/unit/test_usage_normalizer.py diff --git a/docs/data-model-and-observability.md b/docs/data-model-and-observability.md index e6958cf..eccee2b 100644 --- a/docs/data-model-and-observability.md +++ b/docs/data-model-and-observability.md @@ -236,6 +236,18 @@ Use LiteLLM spend logs or structured logs to capture: Collectors must resolve the LiteLLM virtual key ID to the stable `proxy_key_id` (and denormalized `key_alias`) from the `proxy_keys` registry through the configured collector mapping mechanism, then drop the raw key ID before writing `usage_requests` rows. If no registry entry exists, the collector stores `proxy_key_id = null` and `key_alias = null`, and logs a warning. +#### Sessionless usage collector + +`collectors.LiteLLMUsageCollector` collects all-key usage without a benchmark session: + +1. Fetch `GET /spend/logs?start_date=&end_date=&summarize=false` (LiteLLM filters by calendar date; rows are then filtered to the half-open window `[start_time, end_time)` by `startTime`). +2. Normalize each row with `collectors.UsageRequestNormalizer` into a `usage_requests` row. Field extraction (timestamps, latency, TTFT, tokens, cache, cost, status/errors) is shared with benchmark request normalization through `collectors.extraction`. +3. Resolve key attribution with `collectors.ProxyKeyAttributionResolver`: registry `litellm_key_id` match on the spend-log key reference, then an optional in-memory virtual-key-to-alias mapping, then the LiteLLM-reported alias (`api_key_alias` / `metadata.user_api_key_alias`). Matches set `proxy_key_id`, `key_alias`, and denormalize `owner`/`team`/`customer` into `request_metadata`. Unmatched rows keep `proxy_key_id = null` and `key_alias = null`, store the non-secret reported alias as `request_metadata.litellm_key_alias`, and are counted as unattributed. +4. Preserve optional benchmark join fields in `request_metadata` (`benchmark_session_id`, `experiment_id`, `variant_id`, `task_card_id`, `harness_profile`, `trace_id`). Sources, in priority order: `metadata.spend_logs_metadata`, `metadata.requester_metadata`, `metadata`, `request_tags` (`key:value` / `key=value`). Accepted keys include the credential-service tags (`benchmark_*`) and short forms (`session_id`, `experiment`, `variant`, `task_card`, `harness`). The `benchmark_session_id` FK column is set only when the value is a UUID of a known session; LiteLLM's own top-level `session_id` is stored as `litellm_session_id` and never treated as a benchmark session. +5. Write idempotently through `SQLUsageRequestRepository.create_many` (duplicate `litellm_call_id` values are skipped). `dry_run=True` normalizes and reports without writing. + +Rows without `request_id` or `call_id` are skipped with reason `missing_stable_request_id`. Rows missing best-effort fields are accepted as **partial**. The `UsageReconciliationReport` (`to_dict()` / `to_markdown()`) reports mapped, partial, skipped, out-of-window, attributed, unattributed, session-linked, written, and duplicate counts; missing-field counts keyed by LiteLLM source field name (`endTime`, `spend`, `api_key_alias`, `proxy_key_id`, ...); unattributed keys by alias or `sha256:` fingerprint; categorized fetch errors (`auth_error`, `http_error`, `connection_error`, `invalid_response`, `repository_error`); and up to 100 per-row diagnostics. + ### Prometheus Use Prometheus for: @@ -423,7 +435,7 @@ The following table maps LiteLLM `/spend/logs` fields to the canonical `usage_re | `requested_model` | — (metadata) | **Stable** | Model alias sent by the client; useful for audit | | `provider` | `provider_id` | **Stable** | Short provider slug (e.g. `openai`, `fireworks`) | | `custom_llm_provider` | `provider_route` | **Best-effort** | Full provider string; may differ from `provider` | -| `spend` | `cost_usd` (planned) | **Best-effort** | Cost in USD. May be `0.0` for failed requests or when provider pricing is not configured | +| `spend` | `cost_usd` | **Best-effort** | Cost in USD. May be `0.0` for failed requests or when provider pricing is not configured | | `total_tokens` | `input_tokens + output_tokens` | **Stable** | Sum of prompt + completion tokens | | `prompt_tokens` | `input_tokens` | **Stable** | Input side token count | | `completion_tokens` | `output_tokens` | **Stable** | Output side token count | diff --git a/docs/security-and-operations.md b/docs/security-and-operations.md index b39c9bc..c91d1fe 100644 --- a/docs/security-and-operations.md +++ b/docs/security-and-operations.md @@ -64,6 +64,10 @@ Supported secret patterns: - AWS credentials - Environment variables with secret-like names (PASSWORD, SECRET, KEY, TOKEN, CREDENTIAL) +### Usage collection metadata + +The sessionless usage collector (`collectors.LiteLLMUsageCollector`) never copies arbitrary LiteLLM metadata into `usage_requests.request_metadata`. Only an allowlist of scalar join and routing fields is kept (benchmark join fields, `trace_id`, `stream`, `end_user`, `team_id`, `team_alias`, `litellm_session_id`, `request_tags`, retry/fallback counts), each passed through `RedactionFilter` and truncated to 255 characters. Prompt/response content (`messages`, `response`, `proxy_server_request`) and raw key references (`api_key`, `metadata.user_api_key`) are never persisted; unattributed keys are reported only by alias or a truncated `sha256:` fingerprint. Error messages are redacted before storage. + ### Content Capture Configuration Content capture is controlled by `ContentCaptureConfig`: diff --git a/src/collectors/__init__.py b/src/collectors/__init__.py index b5110c6..07d3ced 100644 --- a/src/collectors/__init__.py +++ b/src/collectors/__init__.py @@ -1,9 +1,12 @@ """Collectors: LiteLLM collection, Prometheus collection, normalization, and rollup jobs.""" +from collectors.key_attribution import KeyAttribution, ProxyKeyAttributionResolver from collectors.litellm_collector import ( CollectionDiagnostics, IngestWatermark, LiteLLMCollector, + LiteLLMUsageCollector, + UsageCollectionResult, ) from collectors.metric_catalog import MetricCatalog from collectors.normalization import NormalizationJob @@ -12,6 +15,10 @@ RequestNormalizer, RequestNormalizerJob, UnmappedRowDiagnostics, + UsageNormalizationResult, + UsageReconciliationReport, + UsageRequestNormalizer, + UsageRowDiagnostics, ) from collectors.prometheus_collector import PrometheusCollector from collectors.rollup_job import RollupJob @@ -19,13 +26,21 @@ __all__ = [ "CollectionDiagnostics", "IngestWatermark", + "KeyAttribution", "LiteLLMCollector", + "LiteLLMUsageCollector", "MetricCatalog", "NormalizationJob", "PrometheusCollector", + "ProxyKeyAttributionResolver", "ReconciliationReport", "RequestNormalizer", "RequestNormalizerJob", "RollupJob", "UnmappedRowDiagnostics", + "UsageCollectionResult", + "UsageNormalizationResult", + "UsageReconciliationReport", + "UsageRequestNormalizer", + "UsageRowDiagnostics", ] diff --git a/src/collectors/extraction.py b/src/collectors/extraction.py new file mode 100644 index 0000000..bc56bb3 --- /dev/null +++ b/src/collectors/extraction.py @@ -0,0 +1,295 @@ +"""Shared field-extraction helpers for LiteLLM spend-log records. + +These helpers are pure functions over raw LiteLLM records so that benchmark +request normalization and sessionless usage normalization interpret latency, +TTFT, tokens, cache, cost, and error fields identically. + +Unit conventions follow docs/data-model-and-observability.md: +``latency``/``total_latency``/``ttft``/``time_to_first_token`` are seconds, +``request_duration_ms`` is milliseconds, and all outputs are milliseconds. +""" + +import json +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any + +STARTED_AT_FIELDS = ("startTime", "start_time", "timestamp", "created_at") +FINISHED_AT_FIELDS = ("endTime", "end_time") +COMPLETION_START_FIELDS = ("completion_start_time", "completionStartTime") +LATENCY_SECONDS_FIELDS = ("latency", "total_latency") +TTFT_SECONDS_FIELDS = ("ttft", "time_to_first_token") +SUCCESS_STATUSES = frozenset({"success", "succeeded", "ok"}) +FAILURE_STATUSES = frozenset({"failure", "failed", "error"}) + + +@dataclass(frozen=True) +class TokenCounts: + """Token counts extracted from a raw record.""" + + input_tokens: int | None + output_tokens: int | None + total_tokens: int | None + + +@dataclass(frozen=True) +class CacheFields: + """Cache counters extracted from a raw record.""" + + cache_hit: bool | None + cached_input_tokens: int | None + cache_write_tokens: int | None + + +@dataclass(frozen=True) +class ErrorFields: + """Status and error details extracted from a raw record.""" + + status: str + error_code: str | None + error_message: str | None + + @property + def is_error(self) -> bool: + """Whether the record represents a failed request.""" + return self.status in FAILURE_STATUSES + + +def first_present(data: dict[str, Any], *keys: str) -> Any: + """Return the first value that is present and not None or empty string.""" + for key in keys: + value = data.get(key) + if value is not None and value != "": + return value + return None + + +def to_int(value: Any) -> int | None: + """Coerce a value to int, returning None when not numeric.""" + if value is None or isinstance(value, bool): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def to_float(value: Any) -> float | None: + """Coerce a value to float, returning None when not numeric.""" + if value is None or isinstance(value, bool): + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def parse_bool(value: Any) -> bool | None: + """Parse LiteLLM boolean values, which may be serialized as strings.""" + if value is None: + return None + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"true", "1", "yes"}: + return True + if lowered in {"false", "0", "no", "none", ""}: + return False + return None + + +def parse_timestamp(value: Any) -> datetime | None: + """Parse ISO-8601 strings, epoch numbers, or datetimes into aware UTC datetimes.""" + if value is None or value == "": + return None + parsed: datetime | None = None + if isinstance(value, datetime): + parsed = value + elif isinstance(value, (int, float)) and not isinstance(value, bool): + try: + parsed = datetime.fromtimestamp(value, tz=UTC) + except (OverflowError, OSError, ValueError): + return None + elif isinstance(value, str): + try: + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + if parsed is None: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=UTC) + return parsed.astimezone(UTC) + + +def parse_json_object(value: Any) -> dict[str, Any]: + """Return a dict from a dict or JSON-encoded string, else an empty dict.""" + if isinstance(value, dict): + return value + if isinstance(value, str) and value.strip().startswith("{"): + try: + decoded = json.loads(value) + except ValueError: + return {} + return decoded if isinstance(decoded, dict) else {} + return {} + + +def parse_json_list(value: Any) -> list[Any]: + """Return a list from a list or JSON-encoded string, else an empty list.""" + if isinstance(value, list): + return value + if isinstance(value, str) and value.strip().startswith("["): + try: + decoded = json.loads(value) + except ValueError: + return [] + return decoded if isinstance(decoded, list) else [] + return [] + + +def extract_record_metadata(raw: dict[str, Any]) -> dict[str, Any]: + """Return the record's ``metadata`` object (LiteLLM may JSON-encode it).""" + return parse_json_object(raw.get("metadata")) + + +def extract_started_at(raw: dict[str, Any]) -> datetime | None: + """Extract the request start timestamp.""" + return parse_timestamp(first_present(raw, *STARTED_AT_FIELDS)) + + +def extract_finished_at(raw: dict[str, Any]) -> datetime | None: + """Extract the request end timestamp.""" + return parse_timestamp(first_present(raw, *FINISHED_AT_FIELDS)) + + +def _elapsed_ms(start: datetime | None, end: datetime | None) -> float | None: + if start is None or end is None or end < start: + return None + return round((end - start).total_seconds() * 1000, 3) + + +def extract_latency_ms(raw: dict[str, Any]) -> float | None: + """Extract total latency in milliseconds. + + Prefers explicit seconds fields, then ``request_duration_ms``, then the + ``endTime - startTime`` difference. + """ + seconds = to_float(first_present(raw, *LATENCY_SECONDS_FIELDS)) + if seconds is not None: + return round(seconds * 1000, 3) + duration_ms = to_float(raw.get("request_duration_ms")) + if duration_ms is not None: + return duration_ms + return _elapsed_ms(extract_started_at(raw), extract_finished_at(raw)) + + +def extract_ttft_ms(raw: dict[str, Any]) -> float | None: + """Extract time-to-first-token in milliseconds. + + Uses ``ttft``/``time_to_first_token`` (seconds) when present; otherwise + derives ``round((completion_start_time - startTime) * 1000)``. + """ + seconds = to_float(first_present(raw, *TTFT_SECONDS_FIELDS)) + if seconds is not None: + return round(seconds * 1000, 3) + started_at = extract_started_at(raw) + completion_start = parse_timestamp(first_present(raw, *COMPLETION_START_FIELDS)) + elapsed = _elapsed_ms(started_at, completion_start) + return float(round(elapsed)) if elapsed is not None else None + + +def extract_token_counts(raw: dict[str, Any]) -> TokenCounts: + """Extract input/output/total token counts from top-level or ``usage`` fields.""" + usage = parse_json_object(raw.get("usage")) + sources = (raw, usage) + + def _first_int(*keys: str) -> int | None: + for source in sources: + value = to_int(first_present(source, *keys)) + if value is not None: + return value + return None + + input_tokens = _first_int("prompt_tokens", "input_tokens") + output_tokens = _first_int("completion_tokens", "output_tokens") + total_tokens = _first_int("total_tokens") + if total_tokens is None and input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + return TokenCounts(input_tokens, output_tokens, total_tokens) + + +def extract_cache_fields(raw: dict[str, Any]) -> CacheFields: + """Extract cache hit flag and cache token counters. + + Falls back to ``metadata.additional_usage_values`` (LiteLLM provider usage) + and ``usage.prompt_tokens_details.cached_tokens``. + """ + metadata = extract_record_metadata(raw) + additional = parse_json_object(metadata.get("additional_usage_values")) + usage = parse_json_object(raw.get("usage")) + prompt_details = parse_json_object(usage.get("prompt_tokens_details")) + + cached_input_tokens = to_int(raw.get("cached_input_tokens")) + if cached_input_tokens is None: + cached_input_tokens = to_int( + first_present(additional, "cache_read_input_tokens", "cached_tokens") + ) + if cached_input_tokens is None: + cached_input_tokens = to_int(prompt_details.get("cached_tokens")) + + cache_write_tokens = to_int(raw.get("cache_write_tokens")) + if cache_write_tokens is None: + cache_write_tokens = to_int(additional.get("cache_creation_input_tokens")) + + cache_hit = parse_bool(first_present(raw, "cache_hit", "cached")) + return CacheFields(cache_hit, cached_input_tokens, cache_write_tokens) + + +def extract_cost_usd(raw: dict[str, Any]) -> float | None: + """Extract request spend in USD.""" + return to_float(first_present(raw, "spend", "cost", "response_cost")) + + +def _error_message(value: Any) -> str | None: + if value is None or value == "" or value is False: + return None + if isinstance(value, dict): + message = first_present(value, "error_message", "message", "error") + return str(message) if message is not None else None + if value is True: + return None + return str(value) + + +def extract_error_fields(raw: dict[str, Any]) -> ErrorFields: + """Extract normalized status, error code, and error message. + + Reads top-level ``status``/``error``/``error_code`` and falls back to + LiteLLM's ``metadata.status`` and ``metadata.error_information``. + """ + metadata = extract_record_metadata(raw) + error_info = parse_json_object(metadata.get("error_information")) + raw_error = raw.get("error") + + error_message = _error_message(raw_error) or _error_message(error_info) + if error_message is None: + response = parse_json_object(raw.get("response")) + error_message = _error_message(response.get("error")) + + error_code_value = first_present(raw, "error_code") + if error_code_value is None and isinstance(raw_error, dict): + error_code_value = first_present(raw_error, "code", "status_code") + if error_code_value is None: + error_code_value = first_present(error_info, "error_code", "status_code") + error_code = str(error_code_value) if error_code_value is not None else None + + status_value = first_present(raw, "status") or first_present(metadata, "status") + status = str(status_value).strip().lower() if status_value is not None else "" + if not status: + has_error = error_message is not None or error_code is not None or raw_error is True + status = "failure" if has_error else "success" + return ErrorFields(status=status, error_code=error_code, error_message=error_message) diff --git a/src/collectors/key_attribution.py b/src/collectors/key_attribution.py new file mode 100644 index 0000000..3a147c4 --- /dev/null +++ b/src/collectors/key_attribution.py @@ -0,0 +1,122 @@ +"""Resolve LiteLLM spend-log key references to the ``proxy_keys`` registry. + +Per ADR-002, the LiteLLM virtual key reference carried on spend logs is used +only for lookup and is never persisted. Unresolved keys are reported by a +non-reversible fingerprint so operators can detect unattributed traffic. +""" + +import hashlib +import logging +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Protocol +from uuid import UUID + +from benchmark_core.db.models import ProxyKey as ProxyKeyORM + +logger = logging.getLogger(__name__) + + +class ProxyKeyLookup(Protocol): + """Subset of the proxy key repository used for attribution.""" + + async def get_by_litellm_key_id(self, litellm_key_id: str) -> ProxyKeyORM | None: ... + + async def get_by_alias(self, key_alias: str) -> ProxyKeyORM | None: ... + + +@dataclass(frozen=True) +class KeyAttribution: + """Registry attribution for one usage record.""" + + proxy_key_id: UUID + key_alias: str + litellm_key_id: str | None + owner: str | None + team: str | None + customer: str | None + matched_by: str + + +def key_fingerprint(key_reference: str) -> str: + """Return a short, non-reversible fingerprint for a key reference.""" + digest = hashlib.sha256(key_reference.encode("utf-8")).hexdigest() + return f"sha256:{digest[:12]}" + + +class ProxyKeyAttributionResolver: + """Resolve spend-log key references to registry entries, with caching. + + Lookup order: registry ``litellm_key_id`` match on the spend-log key + reference, then the operator-supplied ephemeral ``virtual_key_aliases`` + mapping (ADR-002 default mechanism), then the alias reported by LiteLLM. + """ + + def __init__( + self, + repository: ProxyKeyLookup, + virtual_key_aliases: Mapping[str, str] | None = None, + ) -> None: + self._repository = repository + self._virtual_key_aliases = dict(virtual_key_aliases or {}) + self._by_key_ref: dict[str, ProxyKeyORM | None] = {} + self._by_alias: dict[str, ProxyKeyORM | None] = {} + + async def resolve( + self, + key_reference: str | None, + reported_alias: str | None, + ) -> KeyAttribution | None: + """Resolve attribution for a record. + + Args: + key_reference: LiteLLM key reference from the spend log (hashed + ``api_key`` or ``metadata.user_api_key``). Lookup only. + reported_alias: ``api_key_alias`` reported by LiteLLM, if any. + + Returns: + Attribution when a registry entry matches, otherwise None. + """ + if key_reference: + orm = await self._lookup_key_ref(key_reference) + if orm is not None: + return self._to_attribution(orm, "litellm_key_id") + mapped_alias = self._virtual_key_aliases.get(key_reference) + if mapped_alias: + orm = await self._lookup_alias(mapped_alias) + if orm is not None: + return self._to_attribution(orm, "virtual_key_mapping") + if reported_alias: + orm = await self._lookup_alias(reported_alias) + if orm is not None: + return self._to_attribution(orm, "key_alias") + logger.warning( + "Usage record key not found in proxy_keys registry (key=%s, alias=%s)", + key_fingerprint(key_reference) if key_reference else None, + reported_alias, + ) + return None + + async def _lookup_key_ref(self, key_reference: str) -> ProxyKeyORM | None: + if key_reference not in self._by_key_ref: + self._by_key_ref[key_reference] = await self._repository.get_by_litellm_key_id( + key_reference + ) + return self._by_key_ref[key_reference] + + async def _lookup_alias(self, alias: str) -> ProxyKeyORM | None: + if alias not in self._by_alias: + self._by_alias[alias] = await self._repository.get_by_alias(alias) + return self._by_alias[alias] + + @staticmethod + def _to_attribution(orm: ProxyKeyORM, matched_by: str) -> KeyAttribution: + return KeyAttribution( + proxy_key_id=orm.id, + key_alias=orm.key_alias, + litellm_key_id=orm.litellm_key_id, + owner=orm.owner, + team=orm.team, + customer=orm.customer, + matched_by=matched_by, + ) diff --git a/src/collectors/litellm_collector.py b/src/collectors/litellm_collector.py index fdc26fb..be6e53f 100644 --- a/src/collectors/litellm_collector.py +++ b/src/collectors/litellm_collector.py @@ -2,14 +2,23 @@ import contextlib from dataclasses import dataclass, field -from datetime import UTC, datetime -from typing import Any +from datetime import UTC, datetime, timedelta +from typing import Any, Protocol from uuid import UUID import httpx +from benchmark_core.db.models import UsageRequest as UsageRequestORM from benchmark_core.models import Request from benchmark_core.repositories import RequestRepository +from benchmark_core.security import get_redaction_filter +from collectors.key_attribution import ProxyKeyAttributionResolver, key_fingerprint +from collectors.normalize_requests import ( + UsageNormalizationResult, + UsageReconciliationReport, + UsageRequestNormalizer, + apply_key_attribution, +) @dataclass @@ -354,3 +363,241 @@ def normalize_request( cache_hit=cache_hit, metadata=metadata, ) + + +# ============================================================================= +# Sessionless all-key usage collection (usage_requests) +# ============================================================================= + + +class UsageRequestWriter(Protocol): + """Subset of the usage request repository used by the collector.""" + + async def create_many( + self, requests: list[UsageRequestORM] + ) -> tuple[list[UsageRequestORM], int]: ... + + +class BenchmarkSessionLookup(Protocol): + """Subset of the session repository used to validate join FKs.""" + + async def exists(self, id: UUID) -> bool: ... + + +@dataclass +class UsageCollectionResult: + """Result of a sessionless usage collection run.""" + + usage_requests: list[UsageRequestORM] + report: UsageReconciliationReport + + +def _as_utc(value: datetime) -> datetime: + return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC) + + +class LiteLLMUsageCollector: + """Collect all-key LiteLLM spend logs into ``usage_requests``. + + Unlike :class:`LiteLLMCollector`, no benchmark session is required. Rows are + fetched for a time window, normalized, attributed to the ``proxy_keys`` + registry when possible, linked to benchmark sessions when metadata carries + a known session UUID, and written idempotently by ``litellm_call_id``. + """ + + SPEND_LOGS_PATH = "/spend/logs" + MAX_ERROR_DETAIL_LENGTH = 200 + + def __init__( + self, + base_url: str, + api_key: str, + repository: UsageRequestWriter | None = None, + key_resolver: ProxyKeyAttributionResolver | None = None, + session_lookup: BenchmarkSessionLookup | None = None, + normalizer: UsageRequestNormalizer | None = None, + transport: httpx.AsyncBaseTransport | None = None, + timeout: float = 60.0, + ) -> None: + self._base_url = base_url.rstrip("/") + self._api_key = api_key + self._repository = repository + self._key_resolver = key_resolver + self._session_lookup = session_lookup + self._normalizer = normalizer or UsageRequestNormalizer() + self._transport = transport + self._timeout = timeout + self._redaction = get_redaction_filter() + self._session_exists: dict[UUID, bool] = {} + + async def collect( + self, + start_time: datetime, + end_time: datetime, + dry_run: bool = False, + ) -> UsageCollectionResult: + """Fetch, normalize, and (unless dry run) persist usage for a window. + + Args: + start_time: Inclusive window start (naive values are treated as UTC). + end_time: Exclusive window end (naive values are treated as UTC). + dry_run: When True, nothing is written to the database. + + Returns: + Normalized rows and the reconciliation report. Fetch failures are + reported in ``report.errors`` with a category instead of raising. + """ + start, end = _as_utc(start_time), _as_utc(end_time) + if start >= end: + raise ValueError("start_time must be earlier than end_time") + report = UsageReconciliationReport(dry_run=dry_run) + raw_records = await self.fetch_spend_logs(start, end, report) + if raw_records is None: + return UsageCollectionResult(usage_requests=[], report=report) + usage_requests = await self.process_records( + raw_records, report, start_time=start, end_time=end, dry_run=dry_run + ) + return UsageCollectionResult(usage_requests=usage_requests, report=report) + + async def fetch_spend_logs( + self, + start_time: datetime, + end_time: datetime, + report: UsageReconciliationReport, + ) -> list[Any] | None: + """Fetch individual spend-log rows for all keys in a window. + + LiteLLM filters ``/spend/logs`` by calendar date, so the request spans + every date touched by the window; rows are filtered precisely later. + Returns None and records a categorized error on failure. + """ + params = { + "start_date": start_time.date().isoformat(), + "end_date": (end_time.date() + timedelta(days=1)).isoformat(), + "summarize": "false", + } + headers = {"Authorization": f"Bearer {self._api_key}"} + try: + async with httpx.AsyncClient( + timeout=self._timeout, transport=self._transport + ) as client: + response = await client.get( + f"{self._base_url}{self.SPEND_LOGS_PATH}", headers=headers, params=params + ) + response.raise_for_status() + payload = response.json() + except httpx.HTTPStatusError as exc: + status = exc.response.status_code + category = "auth_error" if status in (401, 403) else "http_error" + detail = self._redaction.redact_string(exc.response.text)[ + : self.MAX_ERROR_DETAIL_LENGTH + ] + report.add_error(category, f"GET {self.SPEND_LOGS_PATH} returned {status}: {detail}") + return None + except httpx.RequestError as exc: + report.add_error( + "connection_error", + f"GET {self.SPEND_LOGS_PATH} failed: {type(exc).__name__}", + ) + return None + except ValueError: + report.add_error("invalid_response", f"{self.SPEND_LOGS_PATH} returned non-JSON body") + return None + + if isinstance(payload, list): + return payload + if isinstance(payload, dict): + for key in ("data", "logs"): + if isinstance(payload.get(key), list): + records: list[Any] = payload[key] + return records + report.add_error( + "invalid_response", + f"{self.SPEND_LOGS_PATH} returned {type(payload).__name__}, expected a list of rows", + ) + return None + + async def process_records( + self, + raw_records: list[Any], + report: UsageReconciliationReport, + start_time: datetime | None = None, + end_time: datetime | None = None, + dry_run: bool = False, + ) -> list[UsageRequestORM]: + """Normalize, attribute, and persist already-fetched raw records.""" + start = _as_utc(start_time) if start_time is not None else None + end = _as_utc(end_time) if end_time is not None else None + accepted: list[UsageRequestORM] = [] + seen_ids: set[str] = set() + + for index, raw in enumerate(raw_records): + result = self._normalizer.normalize(raw, row_index=index) + usage = result.usage_request + if usage is not None and self._outside_window(usage.started_at, start, end): + report.out_of_window_count += 1 + continue + if usage is not None: + await self._attribute(result, report) + await self._link_session(result, report) + if usage.litellm_call_id in seen_ids: + report.duplicate_count += 1 + result.diagnostics.notes.append("duplicate litellm_call_id within batch") + continue + seen_ids.add(usage.litellm_call_id) + accepted.append(usage) + report.record_row(result.diagnostics) + + if dry_run or self._repository is None or not accepted: + return accepted + try: + created, skipped = await self._repository.create_many(accepted) + except Exception as exc: + report.add_error("repository_error", f"{type(exc).__name__}: {exc}") + raise + report.written_count = len(created) + report.duplicate_count += skipped + return created + + @staticmethod + def _outside_window( + started_at: datetime | None, start: datetime | None, end: datetime | None + ) -> bool: + if started_at is None: + return False + return (start is not None and started_at < start) or (end is not None and started_at >= end) + + async def _attribute( + self, result: UsageNormalizationResult, report: UsageReconciliationReport + ) -> None: + if self._key_resolver is None: + return + attribution = await self._key_resolver.resolve( + result.key_reference, result.reported_key_alias + ) + apply_key_attribution(result, attribution) + if attribution is not None: + report.attributed_count += 1 + else: + label = result.reported_key_alias or ( + key_fingerprint(result.key_reference) if result.key_reference else "" + ) + report.record_unattributed(label) + + async def _link_session( + self, result: UsageNormalizationResult, report: UsageReconciliationReport + ) -> None: + usage = result.usage_request + candidate = result.benchmark_session_candidate + if usage is None or candidate is None: + return + if self._session_lookup is not None: + if candidate not in self._session_exists: + self._session_exists[candidate] = await self._session_lookup.exists(candidate) + if not self._session_exists[candidate]: + result.diagnostics.notes.append( + "benchmark_session_id not found in sessions; preserved in request_metadata only" + ) + return + usage.benchmark_session_id = candidate + report.session_linked_count += 1 diff --git a/src/collectors/normalize_requests.py b/src/collectors/normalize_requests.py index 5d7a2cb..70fb041 100644 --- a/src/collectors/normalize_requests.py +++ b/src/collectors/normalize_requests.py @@ -4,13 +4,32 @@ and generates reconciliation reports for unmapped rows. """ +import uuid from dataclasses import dataclass, field from datetime import UTC, datetime from typing import Any from uuid import UUID from benchmark_core.db.models import Request as RequestORM +from benchmark_core.db.models import UsageRequest as UsageRequestORM from benchmark_core.repositories.request_repository import SQLRequestRepository +from benchmark_core.security import RedactionFilter +from collectors.extraction import ( + extract_cache_fields, + extract_cost_usd, + extract_error_fields, + extract_finished_at, + extract_latency_ms, + extract_record_metadata, + extract_started_at, + extract_token_counts, + extract_ttft_ms, + first_present, + parse_bool, + parse_json_list, + parse_json_object, +) +from collectors.key_attribution import KeyAttribution @dataclass @@ -546,3 +565,510 @@ async def run_with_validation( # For now, delegate to run() - session validation can be added # when session repository integration is needed return await self.run(raw_requests) + + +# ============================================================================= +# Sessionless usage normalization (usage_requests) +# ============================================================================= + +USAGE_OUTCOME_MAPPED = "mapped" +USAGE_OUTCOME_PARTIAL = "partial" +USAGE_OUTCOME_SKIPPED = "skipped" + +SKIP_MISSING_STABLE_ID = "missing_stable_request_id" +SKIP_INVALID_RECORD = "invalid_record_type" + +# Canonical join field -> accepted source keys, in priority order. The +# ``benchmark_*`` names are the tags written by the session credential service. +USAGE_SESSION_JOIN_FIELDS: dict[str, tuple[str, ...]] = { + "benchmark_session_id": ("benchmark_session_id", "session_id"), + "experiment_id": ("benchmark_experiment_id", "experiment_id", "experiment"), + "variant_id": ("benchmark_variant_id", "variant_id", "variant"), + "task_card_id": ("benchmark_task_card_id", "task_card_id", "task_card"), + "harness_profile": ("benchmark_harness_profile", "harness_profile", "harness"), +} +USAGE_TRACE_FIELDS = ("trace_id", "span_id", "parent_span_id") +MAX_METADATA_STRING_LENGTH = 255 +MAX_REQUEST_TAGS = 50 +MAX_ERROR_MESSAGE_LENGTH = 2000 +MAX_USAGE_ROW_DIAGNOSTICS = 100 + + +@dataclass +class UsageRowDiagnostics: + """Per-row normalization diagnostics for usage collection. + + ``missing_fields`` uses LiteLLM source field names so operators can map + gaps directly to proxy configuration or the spend-log field inventory. + """ + + row_index: int | None = None + litellm_call_id: str | None = None + outcome: str = USAGE_OUTCOME_MAPPED + reason: str = "" + missing_fields: list[str] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + raw_keys: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + """Convert diagnostics to a JSON-serializable dictionary.""" + return { + "row_index": self.row_index, + "litellm_call_id": self.litellm_call_id, + "outcome": self.outcome, + "reason": self.reason, + "missing_fields": list(self.missing_fields), + "notes": list(self.notes), + "raw_keys": list(self.raw_keys), + } + + +@dataclass +class UsageNormalizationResult: + """Outcome of normalizing one raw spend-log record. + + ``key_reference`` is the LiteLLM key reference used only for registry + lookup; it is never written to ``usage_requests``. + """ + + usage_request: UsageRequestORM | None + diagnostics: UsageRowDiagnostics + key_reference: str | None = field(default=None, repr=False) + reported_key_alias: str | None = None + benchmark_session_candidate: UUID | None = None + + +@dataclass +class UsageReconciliationReport: + """Reconciliation report for all-key usage collection.""" + + dry_run: bool = False + total_rows: int = 0 + mapped_count: int = 0 + partial_count: int = 0 + skipped_count: int = 0 + out_of_window_count: int = 0 + attributed_count: int = 0 + unattributed_count: int = 0 + session_linked_count: int = 0 + written_count: int = 0 + duplicate_count: int = 0 + missing_field_counts: dict[str, int] = field(default_factory=dict) + skip_reason_counts: dict[str, int] = field(default_factory=dict) + unattributed_keys: dict[str, int] = field(default_factory=dict) + errors: list[dict[str, str]] = field(default_factory=list) + rows: list[UsageRowDiagnostics] = field(default_factory=list) + + @property + def accepted_count(self) -> int: + """Rows accepted for persistence (fully or partially mapped).""" + return self.mapped_count + self.partial_count + + def record_row(self, diagnostics: UsageRowDiagnostics) -> None: + """Record the final outcome of one in-window row.""" + self.total_rows += 1 + if diagnostics.outcome == USAGE_OUTCOME_SKIPPED: + self.skipped_count += 1 + reason = diagnostics.reason or "unknown" + self.skip_reason_counts[reason] = self.skip_reason_counts.get(reason, 0) + 1 + elif diagnostics.outcome == USAGE_OUTCOME_PARTIAL: + self.partial_count += 1 + else: + self.mapped_count += 1 + for name in diagnostics.missing_fields: + self.missing_field_counts[name] = self.missing_field_counts.get(name, 0) + 1 + if diagnostics.outcome != USAGE_OUTCOME_MAPPED and len(self.rows) < ( + MAX_USAGE_ROW_DIAGNOSTICS + ): + self.rows.append(diagnostics) + + def record_unattributed(self, key_label: str) -> None: + """Record a row whose key could not be resolved in ``proxy_keys``.""" + self.unattributed_count += 1 + self.unattributed_keys[key_label] = self.unattributed_keys.get(key_label, 0) + 1 + + def add_error(self, category: str, message: str) -> None: + """Record a collection-level error (fetch, response, repository).""" + self.errors.append({"category": category, "message": message}) + + def to_dict(self) -> dict[str, Any]: + """Generate a JSON-serializable report.""" + return { + "summary": { + "dry_run": self.dry_run, + "total_rows": self.total_rows, + "accepted_count": self.accepted_count, + "mapped_count": self.mapped_count, + "partial_count": self.partial_count, + "skipped_count": self.skipped_count, + "out_of_window_count": self.out_of_window_count, + "attributed_count": self.attributed_count, + "unattributed_count": self.unattributed_count, + "session_linked_count": self.session_linked_count, + "written_count": self.written_count, + "duplicate_count": self.duplicate_count, + }, + "missing_field_counts": dict(self.missing_field_counts), + "skip_reason_counts": dict(self.skip_reason_counts), + "unattributed_keys": dict(self.unattributed_keys), + "errors": list(self.errors), + "rows": [row.to_dict() for row in self.rows], + } + + def to_markdown(self) -> str: + """Generate a markdown report.""" + lines = [ + "# Usage Collection Reconciliation Report", + "", + "## Summary", + "", + f"- **Dry run**: {self.dry_run}", + f"- **Total rows**: {self.total_rows}", + f"- **Mapped**: {self.mapped_count}", + f"- **Partially mapped**: {self.partial_count}", + f"- **Skipped**: {self.skipped_count}", + f"- **Outside window**: {self.out_of_window_count}", + f"- **Attributed to proxy_keys**: {self.attributed_count}", + f"- **Unattributed**: {self.unattributed_count}", + f"- **Linked to benchmark session**: {self.session_linked_count}", + f"- **Written**: {self.written_count}", + f"- **Duplicates skipped**: {self.duplicate_count}", + "", + ] + sections: list[tuple[str, str, dict[str, int]]] = [ + ("Missing Source Fields", "Field", self.missing_field_counts), + ("Skip Reasons", "Reason", self.skip_reason_counts), + ("Unattributed Keys", "Key", self.unattributed_keys), + ] + for title, column, counts in sections: + if not counts: + continue + lines.extend([f"## {title}", "", f"| {column} | Count |", "|---|---|"]) + for name, count in sorted(counts.items(), key=lambda item: (-item[1], item[0])): + lines.append(f"| {name} | {count} |") + lines.append("") + if self.errors: + lines.extend(["## Errors", ""]) + lines.extend(f"- `{error['category']}`: {error['message']}" for error in self.errors) + lines.append("") + if self.rows: + lines.extend(["## Rows Needing Attention (first 10)", ""]) + for row in self.rows[:10]: + label = row.litellm_call_id or f"row {row.row_index}" + lines.append(f"### {label} ({row.outcome})") + lines.append("") + if row.reason: + lines.append(f"- **Reason**: {row.reason}") + if row.missing_fields: + lines.append(f"- **Missing fields**: {', '.join(row.missing_fields)}") + for note in row.notes: + lines.append(f"- **Note**: {note}") + lines.append("") + return "\n".join(lines) + + +def apply_key_attribution( + result: UsageNormalizationResult, + attribution: KeyAttribution | None, +) -> None: + """Apply registry attribution to a normalized usage row (ADR-002). + + On a registry match the stable ``proxy_key_id`` and registry alias/ID are + written and owner/team/customer are denormalized into metadata. Without a + match, ``proxy_key_id`` and ``key_alias`` stay null; the LiteLLM-reported + alias (non-secret) is kept in metadata for retroactive matching. + """ + usage = result.usage_request + if usage is None: + return + metadata = dict(usage.request_metadata or {}) + if attribution is not None: + usage.proxy_key_id = attribution.proxy_key_id + usage.key_alias = attribution.key_alias + usage.litellm_key_id = attribution.litellm_key_id + metadata["key_attribution"] = attribution.matched_by + for name in ("owner", "team", "customer"): + value = getattr(attribution, name) + if value is not None: + metadata[name] = value + else: + usage.proxy_key_id = None + usage.key_alias = None + usage.litellm_key_id = None + metadata["key_attribution"] = "unresolved" + if result.reported_key_alias: + metadata["litellm_key_alias"] = result.reported_key_alias + if "proxy_key_id" not in result.diagnostics.missing_fields: + result.diagnostics.missing_fields.append("proxy_key_id") + result.diagnostics.notes.append( + "key not found in proxy_keys registry; register the key alias to attribute usage" + ) + if result.diagnostics.outcome == USAGE_OUTCOME_MAPPED: + result.diagnostics.outcome = USAGE_OUTCOME_PARTIAL + usage.request_metadata = metadata + + +class UsageRequestNormalizer: + """Normalize raw LiteLLM spend-log records into ``usage_requests`` rows. + + No benchmark session is required. Optional session join fields are + preserved when present. Only allowlisted, redacted metadata is stored; + prompt/response content and raw key references are never copied. + """ + + def __init__(self, redaction_filter: RedactionFilter | None = None) -> None: + self._redaction = redaction_filter or RedactionFilter() + + def normalize( + self, + raw_data: Any, + row_index: int | None = None, + ) -> UsageNormalizationResult: + """Normalize one raw spend-log record. + + Args: + raw_data: Raw record from LiteLLM ``/spend/logs``. + row_index: Position of the record in the fetched batch. + + Returns: + Result with the ORM row (None when skipped) and diagnostics. + """ + if not isinstance(raw_data, dict): + return UsageNormalizationResult( + usage_request=None, + diagnostics=UsageRowDiagnostics( + row_index=row_index, + outcome=USAGE_OUTCOME_SKIPPED, + reason=SKIP_INVALID_RECORD, + notes=[f"expected JSON object, got {type(raw_data).__name__}"], + ), + ) + + diagnostics = UsageRowDiagnostics(row_index=row_index, raw_keys=sorted(raw_data.keys())) + litellm_call_id = first_present(raw_data, "request_id", "call_id", "litellm_call_id") + if litellm_call_id is None: + diagnostics.outcome = USAGE_OUTCOME_SKIPPED + diagnostics.reason = SKIP_MISSING_STABLE_ID + diagnostics.missing_fields = ["request_id", "call_id"] + return UsageNormalizationResult(usage_request=None, diagnostics=diagnostics) + diagnostics.litellm_call_id = str(litellm_call_id) + + record_metadata = extract_record_metadata(raw_data) + started_at = extract_started_at(raw_data) + finished_at = extract_finished_at(raw_data) + tokens = extract_token_counts(raw_data) + cache = extract_cache_fields(raw_data) + errors = extract_error_fields(raw_data) + cost_usd = extract_cost_usd(raw_data) + latency_ms = extract_latency_ms(raw_data) + ttft_ms = extract_ttft_ms(raw_data) + stream = parse_bool(raw_data.get("stream")) + + resolved_model = first_present(raw_data, "model", "model_id") + requested_model = first_present(raw_data, "requested_model", "model_group") + provider_route = first_present(raw_data, "custom_llm_provider") + provider = first_present(raw_data, "provider") or provider_route + if provider is None and isinstance(resolved_model, str) and "/" in resolved_model: + provider = resolved_model.split("/", 1)[0] + + key_reference = first_present(raw_data, "api_key") or first_present( + record_metadata, "user_api_key_hash", "user_api_key" + ) + reported_alias = first_present(raw_data, "api_key_alias", "key_alias") or first_present( + record_metadata, "user_api_key_alias" + ) + + join_fields, session_candidate = self._extract_join_fields( + raw_data, record_metadata, diagnostics + ) + request_metadata = self._build_metadata(raw_data, record_metadata, join_fields, stream) + + error_message = errors.error_message + if error_message is not None: + error_message = self._redaction.redact_string(error_message)[:MAX_ERROR_MESSAGE_LENGTH] + + cache_hit = cache.cache_hit + if cache_hit is None and cache.cached_input_tokens: + cache_hit = True + + request_id = first_present(raw_data, "request_id") + usage = UsageRequestORM( + id=uuid.uuid4(), + litellm_call_id=str(litellm_call_id), + request_id=str(request_id) if request_id is not None else None, + key_alias=str(reported_alias) if reported_alias is not None else None, + litellm_key_id=None, + proxy_key_id=None, + benchmark_session_id=None, + provider=str(provider) if provider is not None else None, + provider_route=str(provider_route) if provider_route is not None else None, + requested_model=str(requested_model) if requested_model is not None else None, + resolved_model=str(resolved_model) if resolved_model is not None else None, + route=self._optional_str(first_present(raw_data, "call_type", "route")), + started_at=started_at, + finished_at=finished_at, + latency_ms=latency_ms, + ttft_ms=ttft_ms, + input_tokens=tokens.input_tokens, + output_tokens=tokens.output_tokens, + cached_input_tokens=cache.cached_input_tokens, + cache_write_tokens=cache.cache_write_tokens, + cost_usd=cost_usd, + status=errors.status, + error_code=errors.error_code, + error_message=error_message, + cache_hit=cache_hit, + request_metadata=request_metadata, + created_at=datetime.now(UTC), + ) + + missing = diagnostics.missing_fields + if started_at is None: + missing.append("startTime") + if finished_at is None: + missing.append("endTime") + if resolved_model is None: + missing.append("model") + if provider is None: + missing.append("provider") + if tokens.input_tokens is None: + missing.append("prompt_tokens") + if tokens.output_tokens is None: + missing.append("completion_tokens") + if cost_usd is None: + missing.append("spend") + if latency_ms is None: + missing.append("latency") + if stream and ttft_ms is None: + missing.append("ttft") + if errors.is_error and errors.error_code is None: + missing.append("error_code") + if key_reference is None and reported_alias is None: + missing.append("api_key") + elif reported_alias is None: + missing.append("api_key_alias") + diagnostics.outcome = USAGE_OUTCOME_PARTIAL if missing else USAGE_OUTCOME_MAPPED + + return UsageNormalizationResult( + usage_request=usage, + diagnostics=diagnostics, + key_reference=str(key_reference) if key_reference is not None else None, + reported_key_alias=str(reported_alias) if reported_alias is not None else None, + benchmark_session_candidate=session_candidate, + ) + + @staticmethod + def _optional_str(value: Any) -> str | None: + return str(value) if value is not None else None + + def _metadata_sources( + self, raw_data: dict[str, Any], record_metadata: dict[str, Any] + ) -> list[dict[str, Any]]: + sources = [ + parse_json_object(record_metadata.get("spend_logs_metadata")), + parse_json_object(record_metadata.get("requester_metadata")), + record_metadata, + self._parse_request_tags(raw_data), + ] + # Top-level ``session_id`` on LiteLLM spend logs is LiteLLM's own + # trace session, not a benchmark session, so it is not a join source. + top_level = {key: value for key, value in raw_data.items() if key != "session_id"} + sources.append(top_level) + return sources + + @staticmethod + def _parse_request_tags(raw_data: dict[str, Any]) -> dict[str, Any]: + parsed: dict[str, Any] = {} + for tag in parse_json_list(raw_data.get("request_tags")): + if not isinstance(tag, str): + continue + for separator in (":", "="): + if separator in tag: + key, value = tag.split(separator, 1) + parsed.setdefault(key.strip(), value.strip()) + break + return parsed + + def _extract_join_fields( + self, + raw_data: dict[str, Any], + record_metadata: dict[str, Any], + diagnostics: UsageRowDiagnostics, + ) -> tuple[dict[str, Any], UUID | None]: + sources = self._metadata_sources(raw_data, record_metadata) + join_fields: dict[str, Any] = {} + for canonical, aliases in USAGE_SESSION_JOIN_FIELDS.items(): + for source in sources: + value = first_present(source, *aliases) + sanitized = self._sanitize_scalar(value) + if sanitized is not None: + join_fields[canonical] = sanitized + break + for trace_field in USAGE_TRACE_FIELDS: + for source in sources: + sanitized = self._sanitize_scalar(source.get(trace_field)) + if sanitized is not None: + join_fields[trace_field] = sanitized + break + + session_candidate: UUID | None = None + session_value = join_fields.get("benchmark_session_id") + if session_value is not None: + try: + session_candidate = UUID(str(session_value)) + except ValueError: + diagnostics.notes.append( + "benchmark_session_id is not a UUID; preserved in request_metadata only" + ) + return join_fields, session_candidate + + def _build_metadata( + self, + raw_data: dict[str, Any], + record_metadata: dict[str, Any], + join_fields: dict[str, Any], + stream: bool | None, + ) -> dict[str, Any]: + metadata: dict[str, Any] = dict(join_fields) + if stream is not None: + metadata["stream"] = stream + for target, source_key in ( + ("end_user", "end_user"), + ("team_id", "team_id"), + ("litellm_session_id", "session_id"), + ("cache_key", "cache_key"), + ): + sanitized = self._sanitize_scalar(raw_data.get(source_key)) + if sanitized is not None and sanitized != "": + metadata[target] = sanitized + for target, source_key in ( + ("team_alias", "user_api_key_team_alias"), + ("litellm_overhead_time_ms", "litellm_overhead_time_ms"), + ("attempted_retries", "attempted_retries"), + ("attempted_fallbacks", "attempted_fallbacks"), + ): + sanitized = self._sanitize_scalar(record_metadata.get(source_key)) + if sanitized is not None: + metadata[target] = sanitized + tags = [ + self._sanitize_scalar(tag) + for tag in parse_json_list(raw_data.get("request_tags"))[:MAX_REQUEST_TAGS] + if isinstance(tag, str) + ] + if tags: + metadata["request_tags"] = tags + return metadata + + def _sanitize_scalar(self, value: Any) -> Any: + """Return a redacted scalar suitable for metadata, or None.""" + if value is None: + return None + if isinstance(value, bool | int | float): + return value + if isinstance(value, str): + if value == "": + return None + return self._redaction.redact_string(value)[:MAX_METADATA_STRING_LENGTH] + if isinstance(value, UUID): + return str(value) + return None diff --git a/tests/unit/test_usage_normalizer.py b/tests/unit/test_usage_normalizer.py new file mode 100644 index 0000000..caa22c0 --- /dev/null +++ b/tests/unit/test_usage_normalizer.py @@ -0,0 +1,660 @@ +"""Fixture-driven tests for sessionless usage normalization and collection.""" + +import json +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import httpx +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +from benchmark_core.db.models import Base, Experiment, ProxyKey, TaskCard, Variant +from benchmark_core.db.models import Session as SessionORM +from benchmark_core.db.models import UsageRequest as UsageRequestORM +from benchmark_core.repositories.proxy_key_repository import SQLProxyKeyRepository +from benchmark_core.repositories.session_repository import SQLSessionRepository +from benchmark_core.repositories.usage_request_repository import SQLUsageRequestRepository +from collectors import ( + LiteLLMUsageCollector, + ProxyKeyAttributionResolver, + UsageReconciliationReport, + UsageRequestNormalizer, +) +from collectors.extraction import ( + extract_cache_fields, + extract_error_fields, + extract_latency_ms, + extract_ttft_ms, + parse_timestamp, +) +from collectors.key_attribution import key_fingerprint +from collectors.normalize_requests import SKIP_MISSING_STABLE_ID, apply_key_attribution + +FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "litellm_spend_logs" + + +def load_fixture(name: str) -> dict[str, Any]: + with open(FIXTURES_DIR / f"{name}.json") as handle: + data: dict[str, Any] = json.load(handle) + return data + + +def all_fixtures() -> list[dict[str, Any]]: + return [json.loads(path.read_text()) for path in sorted(FIXTURES_DIR.glob("*.json"))] + + +def orm_values(row: UsageRequestORM) -> dict[str, Any]: + return {column.name: getattr(row, column.name) for column in UsageRequestORM.__table__.columns} + + +@pytest.fixture +def normalizer() -> UsageRequestNormalizer: + return UsageRequestNormalizer() + + +@pytest.fixture +def db_session(): + engine = create_engine("sqlite:///:memory:", echo=False) + Base.metadata.create_all(bind=engine) + session = sessionmaker(bind=engine)() + yield session + session.close() + + +class TestFixtureMapping: + """Normalizer maps key, model, tokens, cost, latency, TTFT, errors, timestamps.""" + + def test_successful_request(self, normalizer): + result = normalizer.normalize(load_fixture("successful_request"), row_index=0) + row = result.usage_request + + assert row is not None + assert result.diagnostics.outcome == "mapped" + assert result.diagnostics.missing_fields == [] + assert row.litellm_call_id == "req-success-001" + assert row.request_id == "req-success-001" + assert row.key_alias == "bench-session-alpha" + assert row.provider == "openai" + assert row.provider_route == "openai" + assert row.requested_model == "gpt-4o" + assert row.resolved_model == "gpt-4o" + assert row.input_tokens == 50 + assert row.output_tokens == 100 + assert row.cached_input_tokens == 0 + assert row.cache_write_tokens == 0 + assert row.cost_usd == pytest.approx(0.00525) + assert row.latency_ms == pytest.approx(2500.0) + assert row.ttft_ms is None + assert row.started_at == datetime(2025, 4, 21, 10, 0, 0, tzinfo=UTC) + assert row.finished_at == datetime(2025, 4, 21, 10, 0, 2, 500000, tzinfo=UTC) + assert row.status == "success" + assert row.error_code is None + assert row.error_message is None + assert row.cache_hit is False + assert row.request_metadata["stream"] is False + + def test_failed_request(self, normalizer): + row = normalizer.normalize(load_fixture("failed_request")).usage_request + assert row is not None + assert row.status == "failure" + assert row.error_code == "429" + assert row.error_message == "Rate limit exceeded" + assert row.input_tokens == 0 + assert row.cost_usd == 0.0 + assert row.latency_ms == pytest.approx(1200.0) + + def test_streaming_request(self, normalizer): + result = normalizer.normalize(load_fixture("streaming_request")) + row = result.usage_request + assert row is not None + assert result.diagnostics.outcome == "mapped" + assert row.ttft_ms == pytest.approx(650.0) + assert row.latency_ms == pytest.approx(8750.0) + assert row.output_tokens == 400 + assert row.request_metadata["stream"] is True + + def test_cached_request(self, normalizer): + row = normalizer.normalize(load_fixture("cached_request")).usage_request + assert row is not None + assert row.cache_hit is True + assert row.cached_input_tokens == 20 + assert row.cache_write_tokens == 0 + assert row.latency_ms == pytest.approx(150.0) + + def test_non_streaming_ttft_derived_from_completion_start(self, normalizer): + row = normalizer.normalize( + load_fixture("non_streaming_with_completion_start") + ).usage_request + assert row is not None + assert row.ttft_ms == pytest.approx(1500.0) + + def test_sparse_request_is_partially_mapped(self, normalizer): + result = normalizer.normalize(load_fixture("sparse_request"), row_index=4) + row = result.usage_request + assert row is not None + assert row.litellm_call_id == "req-sparse-001" + assert row.key_alias is None + assert row.finished_at is None + assert row.cost_usd is None + assert row.latency_ms == pytest.approx(1200.0) + assert result.diagnostics.outcome == "partial" + assert set(result.diagnostics.missing_fields) == {"endTime", "spend", "api_key_alias"} + assert result.key_reference == "sk-litellm-hash-e5f6g7h8i9j0" + + def test_call_id_fallback_is_stable_id(self, normalizer): + result = normalizer.normalize(load_fixture("fallback_to_call_id")) + row = result.usage_request + assert row is not None + assert row.litellm_call_id == "call-fallback-001" + assert row.request_id is None + + @pytest.mark.parametrize("fixture", all_fixtures(), ids=lambda f: str(f.get("call_id"))) + def test_raw_api_key_never_persisted(self, normalizer, fixture): + row = normalizer.normalize(fixture).usage_request + assert row is not None + serialized = json.dumps(orm_values(row), default=str) + assert fixture["api_key"] not in serialized + assert row.litellm_key_id is None + + +class TestStableIdentifiers: + def test_missing_request_and_call_id_is_skipped_with_diagnostics(self, normalizer): + raw = load_fixture("successful_request") + del raw["request_id"] + del raw["call_id"] + + result = normalizer.normalize(raw, row_index=7) + + assert result.usage_request is None + assert result.diagnostics.outcome == "skipped" + assert result.diagnostics.reason == SKIP_MISSING_STABLE_ID + assert result.diagnostics.missing_fields == ["request_id", "call_id"] + assert result.diagnostics.row_index == 7 + assert "model" in result.diagnostics.raw_keys + + def test_empty_request_id_falls_back_to_call_id(self, normalizer): + raw = load_fixture("failed_request") + raw["request_id"] = "" + row = normalizer.normalize(raw).usage_request + assert row is not None + assert row.litellm_call_id == "call-failed-001" + + def test_non_object_record_is_skipped(self, normalizer): + result = normalizer.normalize(["not", "a", "row"], row_index=1) + assert result.usage_request is None + assert result.diagnostics.reason == "invalid_record_type" + + +class TestSessionMetadata: + def test_records_without_session_metadata_are_accepted(self, normalizer): + result = normalizer.normalize(load_fixture("sparse_request")) + row = result.usage_request + assert row is not None + assert row.benchmark_session_id is None + assert result.benchmark_session_candidate is None + for join_field in ("benchmark_session_id", "experiment_id", "variant_id"): + assert join_field not in row.request_metadata + + def test_fixture_session_metadata_preserves_join_fields(self, normalizer): + result = normalizer.normalize(load_fixture("successful_request")) + row = result.usage_request + assert row is not None + assert row.request_metadata["benchmark_session_id"] == "session-alpha-001" + assert row.request_metadata["experiment_id"] == "model-comparison-q2" + assert row.request_metadata["variant_id"] == "openai-gpt4o" + assert row.request_metadata["task_card_id"] == "repo-auth-analysis" + assert row.request_metadata["harness_profile"] == "openhands" + # Non-UUID session identifiers cannot populate the FK column. + assert row.benchmark_session_id is None + assert result.benchmark_session_candidate is None + assert any("not a UUID" in note for note in result.diagnostics.notes) + + def test_credential_service_tags_in_spend_logs_metadata(self, normalizer): + session_id = uuid4() + raw = load_fixture("sparse_request") + raw["metadata"] = { + "spend_logs_metadata": { + "benchmark_session_id": str(session_id), + "benchmark_experiment_id": "exp-1", + "benchmark_variant_id": "var-1", + "benchmark_harness_profile": "claude-code", + "benchmark_task_card_id": "task-1", + "trace_id": "trace-abc", + } + } + result = normalizer.normalize(raw) + row = result.usage_request + assert row is not None + assert result.benchmark_session_candidate == session_id + assert row.request_metadata["benchmark_session_id"] == str(session_id) + assert row.request_metadata["experiment_id"] == "exp-1" + assert row.request_metadata["variant_id"] == "var-1" + assert row.request_metadata["harness_profile"] == "claude-code" + assert row.request_metadata["task_card_id"] == "task-1" + assert row.request_metadata["trace_id"] == "trace-abc" + + def test_session_metadata_from_request_tags(self, normalizer): + session_id = uuid4() + raw = load_fixture("sparse_request") + raw["request_tags"] = json.dumps( + [f"benchmark_session_id:{session_id}", "benchmark_variant_id=v2", "team-a"] + ) + result = normalizer.normalize(raw) + row = result.usage_request + assert row is not None + assert result.benchmark_session_candidate == session_id + assert row.request_metadata["variant_id"] == "v2" + assert "team-a" in row.request_metadata["request_tags"] + + def test_top_level_litellm_session_id_is_not_benchmark_session(self, normalizer): + raw = load_fixture("sparse_request") + raw["session_id"] = str(uuid4()) + result = normalizer.normalize(raw) + row = result.usage_request + assert row is not None + assert result.benchmark_session_candidate is None + assert "benchmark_session_id" not in row.request_metadata + assert row.request_metadata["litellm_session_id"] == raw["session_id"] + + +class TestMetadataRedaction: + def test_content_and_unknown_metadata_are_dropped(self, normalizer): + raw = load_fixture("successful_request") + raw["messages"] = [{"role": "user", "content": "secret prompt"}] + raw["response"] = {"choices": [{"message": {"content": "secret answer"}}]} + raw["proxy_server_request"] = {"body": {"messages": "secret prompt"}} + raw["metadata"]["headers"] = {"authorization": "Bearer abc"} + raw["metadata"]["arbitrary_blob"] = {"nested": "secret prompt"} + + row = normalizer.normalize(raw).usage_request + assert row is not None + serialized = json.dumps(row.request_metadata) + assert "secret prompt" not in serialized + assert "secret answer" not in serialized + assert "headers" not in row.request_metadata + assert "arbitrary_blob" not in row.request_metadata + + def test_secret_values_in_metadata_are_redacted(self, normalizer): + raw = load_fixture("successful_request") + raw["metadata"]["experiment"] = "sk-abcdefghijklmnopqrstuvwx1234" + raw["metadata"]["user_api_key"] = "sk-abcdefghijklmnopqrstuvwx1234" + row = normalizer.normalize(raw).usage_request + assert row is not None + assert row.request_metadata["experiment_id"] == "[REDACTED]" + assert "sk-abcdefghijklmnopqrstuvwx1234" not in json.dumps(orm_values(row), default=str) + + def test_secret_values_in_error_message_are_redacted(self, normalizer): + raw = load_fixture("failed_request") + raw["error"] = "Invalid key sk-abcdefghijklmnopqrstuvwx1234 provided" + row = normalizer.normalize(raw).usage_request + assert row is not None + assert row.error_message is not None + assert "sk-abcdefghijklmnopqrstuvwx1234" not in row.error_message + assert "[REDACTED]" in row.error_message + + +class TestLiteLLMNativeShape: + """Records shaped like LiteLLM's SpendLogs table rows.""" + + def test_json_string_metadata_and_string_cache_hit(self, normalizer): + raw = { + "request_id": "chatcmpl-123", + "api_key": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", + "model": "anthropic/claude-sonnet", + "model_group": "claude-sonnet", + "custom_llm_provider": "anthropic", + "call_type": "acompletion", + "spend": 0.01, + "prompt_tokens": 100, + "completion_tokens": 20, + "startTime": "2025-04-21T10:00:00+00:00", + "endTime": "2025-04-21T10:00:01+00:00", + "completionStartTime": "2025-04-21T10:00:00.400000+00:00", + "cache_hit": "True", + "session_id": "litellm-trace", + "messages": [{"role": "user", "content": "hello"}], + "metadata": json.dumps( + { + "user_api_key_alias": "team-alpha-key", + "status": "success", + "additional_usage_values": { + "cache_read_input_tokens": 80, + "cache_creation_input_tokens": 5, + }, + } + ), + } + result = normalizer.normalize(raw) + row = result.usage_request + assert row is not None + assert result.diagnostics.outcome == "mapped" + assert row.key_alias == "team-alpha-key" + assert row.provider == "anthropic" + assert row.requested_model == "claude-sonnet" + assert row.route == "acompletion" + assert row.latency_ms == pytest.approx(1000.0) + assert row.ttft_ms == pytest.approx(400.0) + assert row.cache_hit is True + assert row.cached_input_tokens == 80 + assert row.cache_write_tokens == 5 + assert "messages" not in row.request_metadata + + def test_failure_from_metadata_error_information(self, normalizer): + raw = { + "request_id": "chatcmpl-err", + "model": "gpt-4o", + "metadata": { + "status": "failure", + "error_information": {"error_code": "500", "error_message": "upstream boom"}, + }, + } + row = normalizer.normalize(raw).usage_request + assert row is not None + assert row.status == "failure" + assert row.error_code == "500" + assert row.error_message == "upstream boom" + + +class TestExtractionHelpers: + def test_latency_prefers_seconds_then_duration_then_timestamps(self): + assert extract_latency_ms({"latency": 1.5}) == pytest.approx(1500.0) + assert extract_latency_ms({"request_duration_ms": 42}) == pytest.approx(42.0) + assert extract_latency_ms( + {"startTime": "2025-01-01T00:00:00Z", "endTime": "2025-01-01T00:00:00.250Z"} + ) == pytest.approx(250.0) + assert extract_latency_ms({}) is None + + def test_ttft_ignores_negative_derivation(self): + raw = {"startTime": "2025-01-01T00:00:01Z", "completion_start_time": "2025-01-01T00:00:00Z"} + assert extract_ttft_ms(raw) is None + + def test_parse_timestamp_variants(self): + assert parse_timestamp("2025-01-01T00:00:00") == datetime(2025, 1, 1, tzinfo=UTC) + assert parse_timestamp(0) == datetime(1970, 1, 1, tzinfo=UTC) + assert parse_timestamp("not-a-date") is None + + def test_cache_and_error_defaults(self): + assert extract_cache_fields({"cache_hit": "False"}).cache_hit is False + assert extract_cache_fields({}).cache_hit is None + assert extract_error_fields({}).status == "success" + assert extract_error_fields({"error": "boom"}).status == "failure" + + +class FakeKeyRepo: + def __init__(self, keys: list[ProxyKey]) -> None: + self.keys = keys + self.calls = 0 + + async def get_by_litellm_key_id(self, litellm_key_id: str) -> ProxyKey | None: + self.calls += 1 + return next((k for k in self.keys if k.litellm_key_id == litellm_key_id), None) + + async def get_by_alias(self, key_alias: str) -> ProxyKey | None: + self.calls += 1 + return next((k for k in self.keys if k.key_alias == key_alias), None) + + +def make_key(alias: str, litellm_key_id: str | None = None) -> ProxyKey: + return ProxyKey(id=uuid4(), key_alias=alias, litellm_key_id=litellm_key_id, owner="o", team="t") + + +class TestKeyAttribution: + async def test_resolve_by_litellm_key_id(self): + key = make_key("registry-alias", "sk-litellm-hash-e5f6g7h8i9j0") + resolver = ProxyKeyAttributionResolver(FakeKeyRepo([key])) + attribution = await resolver.resolve("sk-litellm-hash-e5f6g7h8i9j0", None) + assert attribution is not None + assert attribution.proxy_key_id == key.id + assert attribution.matched_by == "litellm_key_id" + + async def test_resolve_by_alias(self): + key = make_key("bench-session-alpha") + resolver = ProxyKeyAttributionResolver(FakeKeyRepo([key])) + attribution = await resolver.resolve("unknown-hash", "bench-session-alpha") + assert attribution is not None + assert attribution.matched_by == "key_alias" + + async def test_resolve_by_ephemeral_virtual_key_mapping(self): + key = make_key("mapped-alias") + resolver = ProxyKeyAttributionResolver( + FakeKeyRepo([key]), virtual_key_aliases={"hash-1": "mapped-alias"} + ) + attribution = await resolver.resolve("hash-1", None) + assert attribution is not None + assert attribution.matched_by == "virtual_key_mapping" + + async def test_unresolved_and_cached(self): + repo = FakeKeyRepo([]) + resolver = ProxyKeyAttributionResolver(repo) + assert await resolver.resolve("hash-x", "alias-x") is None + calls = repo.calls + assert await resolver.resolve("hash-x", "alias-x") is None + assert repo.calls == calls + + def test_apply_unresolved_attribution_marks_partial(self, normalizer): + result = normalizer.normalize(load_fixture("successful_request")) + apply_key_attribution(result, None) + row = result.usage_request + assert row is not None + assert row.proxy_key_id is None + assert row.key_alias is None + assert row.request_metadata["litellm_key_alias"] == "bench-session-alpha" + assert row.request_metadata["key_attribution"] == "unresolved" + assert result.diagnostics.outcome == "partial" + assert "proxy_key_id" in result.diagnostics.missing_fields + + def test_fingerprint_is_not_reversible_value(self): + fingerprint = key_fingerprint("sk-secret-value") + assert fingerprint.startswith("sha256:") + assert "sk-secret-value" not in fingerprint + + +def spend_logs_transport( + records: Any, status_code: int = 200, captured: list[httpx.Request] | None = None +) -> httpx.MockTransport: + def handler(request: httpx.Request) -> httpx.Response: + if captured is not None: + captured.append(request) + return httpx.Response(status_code, json=records) + + return httpx.MockTransport(handler) + + +WINDOW_START = datetime(2025, 4, 21, 10, 0, tzinfo=UTC) +WINDOW_END = datetime(2025, 4, 21, 11, 0, tzinfo=UTC) + + +class TestLiteLLMUsageCollector: + async def test_fetches_all_keys_without_session(self): + captured: list[httpx.Request] = [] + collector = LiteLLMUsageCollector( + "http://litellm:4000/", + "master-key", + transport=spend_logs_transport(all_fixtures(), captured=captured), + ) + result = await collector.collect(WINDOW_START, WINDOW_END, dry_run=True) + + request = captured[0] + assert request.url.path == "/spend/logs" + assert request.url.params["summarize"] == "false" + assert request.url.params["start_date"] == "2025-04-21" + assert request.url.params["end_date"] == "2025-04-22" + assert "session_id" not in request.url.params + assert request.headers["authorization"] == "Bearer master-key" + assert len(result.usage_requests) == len(all_fixtures()) + assert result.report.dry_run is True + assert result.report.written_count == 0 + + async def test_collects_and_writes_idempotently(self, db_session): + db_session.add(make_key("bench-session-alpha")) + db_session.commit() + repo = SQLUsageRequestRepository(db_session) + missing_id = load_fixture("successful_request") + del missing_id["request_id"], missing_id["call_id"] + records = [*all_fixtures(), missing_id] + + def build() -> LiteLLMUsageCollector: + return LiteLLMUsageCollector( + "http://litellm:4000", + "master-key", + repository=repo, + key_resolver=ProxyKeyAttributionResolver(SQLProxyKeyRepository(db_session)), + transport=spend_logs_transport(records), + ) + + first = await build().collect(WINDOW_START, WINDOW_END) + report = first.report + assert report.total_rows == len(records) + assert report.skipped_count == 1 + assert report.skip_reason_counts == {SKIP_MISSING_STABLE_ID: 1} + assert report.written_count == len(records) - 1 + assert report.attributed_count == 1 + assert report.unattributed_count == len(records) - 2 + assert report.missing_field_counts["proxy_key_id"] == report.unattributed_count + assert db_session.query(UsageRequestORM).count() == len(records) - 1 + + stored = await repo.get_by_litellm_call_id("req-success-001") + assert stored is not None + assert stored.proxy_key_id is not None + assert stored.request_metadata["key_attribution"] == "key_alias" + assert stored.request_metadata["owner"] == "o" + + second = await build().collect(WINDOW_START, WINDOW_END) + assert second.report.written_count == 0 + assert second.report.duplicate_count == len(records) - 1 + assert db_session.query(UsageRequestORM).count() == len(records) - 1 + + async def test_out_of_window_rows_are_counted_not_collected(self): + collector = LiteLLMUsageCollector( + "http://litellm:4000", "k", transport=spend_logs_transport(all_fixtures()) + ) + result = await collector.collect( + WINDOW_START, datetime(2025, 4, 21, 10, 10, tzinfo=UTC), dry_run=True + ) + assert {row.litellm_call_id for row in result.usage_requests} == { + "req-success-001", + "req-failed-001", + } + assert result.report.out_of_window_count == len(all_fixtures()) - 2 + + async def test_links_known_benchmark_session(self, db_session): + experiment = Experiment(name="exp") + variant = Variant( + name="var", provider="openai", model_alias="gpt-4o", harness_profile="default" + ) + task_card = TaskCard(name="task", goal="g", starting_prompt="s", stop_condition="c") + db_session.add_all([experiment, variant, task_card]) + db_session.flush() + session = SessionORM( + experiment_id=experiment.id, + variant_id=variant.id, + task_card_id=task_card.id, + harness_profile="default", + repo_path="/tmp/repo", + git_branch="main", + git_commit="abc1234", + status="active", + ) + db_session.add(session) + db_session.commit() + + known = load_fixture("sparse_request") + known["metadata"] = {"benchmark_session_id": str(session.id)} + unknown = load_fixture("fallback_to_call_id") + unknown["metadata"] = {"benchmark_session_id": str(uuid4())} + + collector = LiteLLMUsageCollector( + "http://litellm:4000", + "k", + repository=SQLUsageRequestRepository(db_session), + session_lookup=SQLSessionRepository(db_session), + transport=spend_logs_transport([known, unknown]), + ) + result = await collector.collect(WINDOW_START, WINDOW_END) + + by_id = {row.litellm_call_id: row for row in result.usage_requests} + assert by_id["req-sparse-001"].benchmark_session_id == session.id + assert by_id["call-fallback-001"].benchmark_session_id is None + assert isinstance(by_id["call-fallback-001"].request_metadata["benchmark_session_id"], str) + assert result.report.session_linked_count == 1 + + @pytest.mark.parametrize( + ("status_code", "category"), [(401, "auth_error"), (500, "http_error")] + ) + async def test_http_errors_are_categorized(self, status_code, category): + collector = LiteLLMUsageCollector( + "http://litellm:4000", + "k", + transport=spend_logs_transport({"error": "nope"}, status_code=status_code), + ) + result = await collector.collect(WINDOW_START, WINDOW_END) + assert result.usage_requests == [] + assert result.report.errors[0]["category"] == category + + async def test_connection_error_is_categorized(self): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + collector = LiteLLMUsageCollector( + "http://litellm:4000", "k", transport=httpx.MockTransport(handler) + ) + result = await collector.collect(WINDOW_START, WINDOW_END) + assert result.report.errors[0]["category"] == "connection_error" + + async def test_wrapped_and_invalid_payloads(self): + wrapped = LiteLLMUsageCollector( + "http://litellm:4000", + "k", + transport=spend_logs_transport({"data": [load_fixture("sparse_request")]}), + ) + assert len((await wrapped.collect(WINDOW_START, WINDOW_END)).usage_requests) == 1 + + invalid = LiteLLMUsageCollector( + "http://litellm:4000", "k", transport=spend_logs_transport({"spend": 1.0}) + ) + result = await invalid.collect(WINDOW_START, WINDOW_END) + assert result.report.errors[0]["category"] == "invalid_response" + + async def test_rejects_empty_window(self): + collector = LiteLLMUsageCollector("http://litellm:4000", "k") + with pytest.raises(ValueError): + await collector.collect(WINDOW_END, WINDOW_START) + + +class TestUsageReconciliationReport: + async def test_report_lists_actionable_field_names(self): + missing_id = load_fixture("successful_request") + del missing_id["request_id"], missing_id["call_id"] + collector = LiteLLMUsageCollector( + "http://litellm:4000", + "k", + key_resolver=ProxyKeyAttributionResolver(FakeKeyRepo([])), + transport=spend_logs_transport([load_fixture("sparse_request"), missing_id]), + ) + result = await collector.collect(WINDOW_START, WINDOW_END, dry_run=True) + report = result.report + + markdown = report.to_markdown() + assert "missing_stable_request_id" in markdown + assert "endTime" in markdown + assert "spend" in markdown + assert "proxy_key_id" in markdown + assert "sk-litellm-hash-e5f6g7h8i9j0" not in markdown + assert key_fingerprint("sk-litellm-hash-e5f6g7h8i9j0") in report.unattributed_keys + + payload = report.to_dict() + json.dumps(payload) + assert payload["summary"]["accepted_count"] == 1 + assert payload["summary"]["skipped_count"] == 1 + assert {row["outcome"] for row in payload["rows"]} == {"partial", "skipped"} + + def test_report_caps_row_diagnostics(self): + report = UsageReconciliationReport() + normalizer = UsageRequestNormalizer() + for index in range(150): + report.record_row(normalizer.normalize({"model": "m"}, row_index=index).diagnostics) + assert report.skipped_count == 150 + assert len(report.rows) == 100 From c5a3dcd2228a815cccec9f2d6ad5676a0528dac3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 7 Oct 2026 07:20:01 +0000 Subject: [PATCH 2/2] fix(collectors): drop LiteLLM cache_key from usage metadata Live /spend/logs dry run showed cache_key is the placeholder 'Cache OFF' on every row when caching is disabled, adding noise without join value. Co-Authored-By: bot_apk --- src/collectors/normalize_requests.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/collectors/normalize_requests.py b/src/collectors/normalize_requests.py index 70fb041..ba80c2f 100644 --- a/src/collectors/normalize_requests.py +++ b/src/collectors/normalize_requests.py @@ -1036,7 +1036,6 @@ def _build_metadata( ("end_user", "end_user"), ("team_id", "team_id"), ("litellm_session_id", "session_id"), - ("cache_key", "cache_key"), ): sanitized = self._sanitize_scalar(raw_data.get(source_key)) if sanitized is not None and sanitized != "":