From b18188dee2044de608d6b4389bbff7cf13366ff4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Nov 2026 09:55:53 -0800 Subject: [PATCH 001/178] test(e2e): retire the quiet_stack lock and batch the claude_code model registration (#49984) * test(e2e): retire the quiet_stack lock quiet_stack held an exclusive host-wide lock for six tests, draining every other xdist worker while they ran. On the full e2e lane that left 35-38% of worker time idle. Each of those tests now lives where it can run without pausing the suite - the two memory tests move to perf/ behind the perf marker, so the perf lane runs them on its own fresh stack. Their refusing deployments use a literal key because the perf stack carries no provider keys - the stored-request size check from the memory test becomes an integration test on an owned proxy that stores prompts - the discount and margin spend-row tests are dropped from e2e. Their breakdown assertions join the existing integration pricing test, which already owns the discount and margin repricing flow - the Bedrock guardrail cost test stays in e2e without the pricing reset fixture, since nothing else on the stack writes pricing config now - the MCP semantic tool filter test turns on a proxy-wide setting, so it gets a new proxy_wide_config opt-in marker and runs after the main run * test(e2e): register the claude_code deployments as one batch The claude_code session fixture registered its 24 deployments one at a time on every xdist worker, and each one waited out the 15 s propagation window on its own, about 6 minutes per worker. ProxyClient.register_models posts the whole batch, polls each model until every replica serves it, then waits out propagation once from the last write. A model that fails to post or never becomes servable comes back as ModelNotRegistered, and an unservable one is deleted, so the fixture keeps reporting per-deployment failures the way it did --- tests/e2e/AGENTS.md | 8 +- tests/e2e/claude_code/conftest.py | 35 ++-- tests/e2e/conftest.py | 21 +-- .../coverage_registry/quota_management.yaml | 2 - tests/e2e/e2e_config.py | 2 +- .../mcp/test_mcp_semantic_tool_filter_e2e.py | 5 +- tests/e2e/models.py | 21 --- tests/e2e/perf/test_database_upgrade.py | 2 +- tests/e2e/perf/test_perf_cases.py | 2 +- .../test_reliability_memory_e2e.py | 111 ++++++------- tests/e2e/perf/test_ui_journeys.py | 1 - tests/e2e/proxy_client.py | 96 +++++++---- tests/e2e/pytest.ini | 2 +- ...2e.py => test_guardrail_cost_spend_e2e.py} | 152 +----------------- tests/e2e/stack_lock.py | 45 ------ tests/e2e_harness/AGENTS.md | 2 +- tests/e2e_harness/test_proxy_client.py | 90 ++++++++++- tests/e2e_harness/test_stack_lock.py | 117 -------------- .../pricing/test_cost_margin_discount.py | 32 +++- .../test_failing_request_stored_snapshot.py | 84 ++++++++++ 20 files changed, 359 insertions(+), 471 deletions(-) rename tests/e2e/{router => perf}/test_reliability_memory_e2e.py (77%) rename tests/e2e/quota_management/spend_tracking/{test_pricing_config_spend_e2e.py => test_guardrail_cost_spend_e2e.py} (52%) delete mode 100644 tests/e2e/stack_lock.py delete mode 100644 tests/e2e_harness/test_stack_lock.py create mode 100644 tests/integration/spend/test_failing_request_stored_snapshot.py diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 781d8fd051b..b4ca8bd3d50 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -40,11 +40,11 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below) +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below). The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). The memory tests (`test_reliability_memory_e2e.py`, run by the perf lane on its own fresh stack: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches, the check for the v1.100.0 retry-breadcrumb leak). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests @@ -100,7 +100,7 @@ Each suite provides its own `client` fixture (see `llm_translation/passthrough_c Request and response bodies are typed pydantic models in `models.py`. Only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass -Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. Add `@pytest.mark.quiet_stack` to a test that measures the proxy itself (RSS, latency): the shared stack lock in `stack_lock.py` then runs it while no other test on the host is hitting the stack, marked or not, so the reading depends only on the test's own traffic. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache +Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. A test that measures the proxy itself (RSS, latency) cannot share the stack with the rest of the suite, so it goes in `perf/`, marked `perf`. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache ## Record and replay fixtures diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index b2442d492cf..14201ef6846 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -597,14 +597,21 @@ def _build_control_plane_client(proxy_config: ProxyConfig): ) -def _register_deployment(proxy, deployment: CompatDeployment) -> str: - """Register one deployment and return its proxy-assigned model_id - once it is servable on the data plane. The aliases are shared by every - cell and, under xdist, by every worker, so no call to them belongs to - one test and none is cached: the matrix exists to reach real providers.""" - return proxy.create_model( - deployment.model_name, - deployment.litellm_params, +def _register_deployments(proxy, deployments: tuple[CompatDeployment, ...]): + """The aliases are shared by every cell and, under xdist, by every worker, + so no call to them belongs to one test and none is cached: the matrix + exists to reach real providers.""" + from models import ModelInfoBody, ModelNewBody + + return proxy.register_models( + tuple( + ModelNewBody( + model_name=deployment.model_name, + litellm_params=deployment.litellm_params, + model_info=ModelInfoBody(), + ) + for deployment in deployments + ), provider_live=True, ) @@ -632,18 +639,14 @@ def _compat_models_registered() -> Any: yield return + from proxy_client import ModelNotRegistered, ModelRegistered from requests import RequestException proxy = _build_control_plane_client(proxy_config) - registered_ids: list[str] = [] - failures: list[tuple[str, str]] = [] + registrations = _register_deployments(proxy, load_all_deployments()) + registered_ids = tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) + failures = tuple((r.model_name, r.reason) for r in registrations if isinstance(r, ModelNotRegistered)) try: - for deployment in load_all_deployments(): - try: - model_id = _register_deployment(proxy, deployment) - registered_ids.append(model_id) - except (AssertionError, RequestException) as exc: - failures.append((deployment.model_name, str(exc))) if failures: summary = "\n".join( f" - {name}: {reason}" for name, reason in failures diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4c7c57ca18f..8fb96bef2f9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -35,6 +35,7 @@ from e2e_config import ( PROMPT_CACHING_OPT_IN_ENV, PROVIDER_EDGE_HOST_OPT_IN_ENV, PROXY_BASE_URL, + PROXY_WIDE_CONFIG_OPT_IN_ENV, SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, SAGEMAKER_NOVA_OPT_IN_ENV, SECRET_MANAGER_OPT_IN_ENV, @@ -54,7 +55,6 @@ from models import TeamNewBody, UserNewBody, UserNewResponse from provider_cache_routing import LIVE_PROVIDER_REQUIRED from provider_edge import replay_leftover_error from proxy_client import ProxyClient, build_proxy_client -from stack_lock import stack_lock _E2E_TEST_RAN = pytest.StashKey[bool]() _CALL_PASSED = pytest.StashKey[bool]() @@ -77,6 +77,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "secret_manager": SECRET_MANAGER_OPT_IN_ENV, "sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV, "sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, + "proxy_wide_config": PROXY_WIDE_CONFIG_OPT_IN_ENV, } ) @@ -162,11 +163,6 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set", ) - config.addinivalue_line( - "markers", - "quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; " - "every other test waits for it to finish", - ) config.addinivalue_line( "markers", "provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the " @@ -200,6 +196,11 @@ def pytest_configure(config: pytest.Config) -> None: "sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; " "deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it", ) + config.addinivalue_line( + "markers", + "proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite " + "is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: @@ -269,7 +270,7 @@ def pytest_collection_finish(session: pytest.Session) -> None: ahead of xdist's own collection-finish report, and the controller schedules no test until every worker has reported, so this is the idle footprint of a stack that just passed its readiness gate. The fixture hands the capture to the - idle-budget test in router/test_reliability_memory_e2e.py.""" + idle-budget test in perf/test_reliability_memory_e2e.py.""" if session.config.getoption("collectonly") or not any(_uses_idle_rss(item) for item in session.items): return session.config.stash[_IDLE_RSS] = read_rss_everywhere(build_proxy_client(), timeout=IDLE_RSS_READ_TIMEOUT_SECONDS) @@ -299,12 +300,6 @@ def _proxy_fail_reason() -> str | None: return None -@pytest.hookimpl(wrapper=True) -def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, object]: - with stack_lock(exclusive=item.get_closest_marker("quiet_stack") is not None): - return (yield) - - @pytest.hookimpl(tryfirst=True) def pytest_runtest_setup(item: pytest.Item) -> None: """Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe. diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index cbaf46130d6..49997aef162 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -44,8 +44,6 @@ - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/messages/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} - {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"} - {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"} -- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"} -- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"} - {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"} - {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"} - {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 724aa2362e7..6316d22d3fa 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -210,6 +210,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT" SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER" SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT" SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT" +PROXY_WIDE_CONFIG_OPT_IN_ENV: Final = "E2E_PROXY_WIDE_CONFIG_STACK" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) @@ -234,7 +235,6 @@ MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES", MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1")) MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200")) MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768")) -MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64")) def ws_base_url() -> str: diff --git a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py index 6e564f4625a..1ec2a6822e3 100644 --- a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py +++ b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py @@ -5,7 +5,8 @@ embedding deployment, and turns on `litellm_settings.mcp_semantic_tool_filter` t /config/update. A chat completion that attaches the server and asks to search logs must report, in the x-litellm-semantic-filter headers, that the catalog was cut to at most top_k tools, all of them tools the key lists. Requires DD_API_KEY + DD_APP_KEY and the proxy's -OPENAI_API_KEY for the embedding deployment. +OPENAI_API_KEY for the embedding deployment. The filter is proxy-wide, so the test is +marked `proxy_wide_config` and runs on a stack the main suite is not driving. """ from __future__ import annotations @@ -23,7 +24,7 @@ from mcp_client import McpClient from models import ChatBody, ChatMessage, LiteLLMParamsBody, McpChatTool from pydantic import BaseModel -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = [pytest.mark.e2e, pytest.mark.proxy_wide_config] TOP_K: Final = 3 diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9c12e2146c9..25683d64313 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1262,27 +1262,6 @@ class CostMap(RootModel[dict[str, CostMapEntry]]): pass -class ConfigPatchResponse(BaseModel): - status: str - values: dict[str, float | dict[str, float]] - - -class CostDiscountConfig(RootModel[dict[str, float]]): - pass - - -class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]): - pass - - -class CostDiscountConfigResponse(BaseModel): - values: dict[str, float] - - -class CostMarginConfigResponse(BaseModel): - values: dict[str, float | dict[str, float]] - - class BedrockGuardrailParams(BaseModel): guardrail: Literal["bedrock"] = "bedrock" mode: Literal["pre_call"] = "pre_call" diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py index fba6af30014..a9f5cc43fe6 100644 --- a/tests/e2e/perf/test_database_upgrade.py +++ b/tests/e2e/perf/test_database_upgrade.py @@ -28,7 +28,7 @@ from setup.run import run_plan from setup.upgrade import assert_finished, monitor_for, upgrade_case, upgrade_properties from setup.upgrade_worker import connect, read_state -pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] +pytestmark = pytest.mark.perf @pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only") diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index db657a04b00..0a137e84572 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -48,7 +48,7 @@ from setup.ramp import step_properties from setup.run import run_plan from setup.telemetry import read_gateway_telemetry, telemetry_sampler -pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] +pytestmark = pytest.mark.perf _SPEND_SAMPLE: Final = 40 _WARMUP_SECONDS: Final = 60 diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/perf/test_reliability_memory_e2e.py similarity index 77% rename from tests/e2e/router/test_reliability_memory_e2e.py rename to tests/e2e/perf/test_reliability_memory_e2e.py index fb114092303..2066137bb72 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/perf/test_reliability_memory_e2e.py @@ -1,17 +1,19 @@ """Live e2e: a few hundred requests that fail before any provider answers must not -grow the proxy's resident memory past a loose ceiling once the proxy is warm. +grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads +the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack. The regression this guards shipped in v1.100.0: every retry breadcrumb copied the whole request and the copies nested into one router-global list, so a proxy under retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that shape: a model group whose deployments refuse at the socket (an unreachable base -URL) with cooldown_time 0 so the router keeps retrying them, per-request retries, -and a fallback group that refuses the same way, each request carrying a long chat +URL, with a literal key since the perf stack carries no provider keys) with +cooldown_time 0 so the router keeps retrying them, per-request retries, and a +fallback group that refuses the same way, each request carrying a long chat transcript so every whole-request copy costs hundreds of containers instead of a -handful. Under the stack's cooldown policy a -deployment that fails a handful of times in a row is benched (a bad-credential 401 -included), the router answers "No deployments available" without retrying, and the -retry loop that leaks stops running; cooldown_time 0 keeps it running. +handful. Under the stack's cooldown policy a deployment that fails a handful of +times in a row is benched (a bad-credential 401 included), the router answers "No +deployments available" without retrying, and the retry loop that leaks stops running; +cooldown_time 0 keeps it running. Four identical batches run back to back. The first two are the warmup that grows the proxy's caches and allocator arenas, the last two are the measured window, and @@ -35,19 +37,11 @@ single load-balanced address reaches the workers behind it one answer at a time, a worker that answered only one checkpoint would otherwise drop out of the comparison, which is where a leaking worker could hide. -RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, -json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS -by only about 15 MB per 300 failing requests, while every failing request's stored -request snapshot carried a copy of the request per failed attempt, over 100 KB on -the first call and a couple of MB once the copies nested, against tens of KB with -the fix. So the first check sends one failing request before the phases, reads its -spend log back through /spend/logs, and holds the stored request body to a fixed -size budget: the deterministic catch for a breadcrumb that copies the whole -request. It runs before the phases because the leaking writer drops its own rows -under the phases' traffic (a queue budget hit, a recursion limit on the nested -copies), which would turn the size check into a missing-row check. +The stored-request half of that regression (every failing request's spend log carrying a +copy of the request per failed attempt) is a deterministic size check, so it lives in +tests/integration/spend/test_failing_request_stored_snapshot.py. -A second, cheaper check holds the idle footprint: every worker's RSS as the harness +A cheaper check holds the idle footprint: every worker's RSS as the harness read it at collection time, before this pytest process sent any traffic (see conftest.pytest_collection_finish), must sit under a fixed budget. On the release gate that is a fresh stack right after its readiness gate, one worker per @@ -74,7 +68,6 @@ from typing import Final import pytest -from complexity_router_client import ComplexityRouterClient from e2e_config import ( MEMORY_CONCURRENCY, MEMORY_IDLE_RSS_BUDGET_MB, @@ -83,7 +76,6 @@ from e2e_config import ( MEMORY_RSS_CEILING_MB, MEMORY_RSS_SAMPLE_INTERVAL_SECONDS, MEMORY_RSS_SETTLE_SAMPLES, - MEMORY_STORED_REQUEST_BUDGET_KB, MEMORY_TRANSCRIPT_TURNS, PROXY_REPLICA_URLS, unique_marker, @@ -91,11 +83,12 @@ from e2e_config import ( from e2e_metadata import Domain, Mode, Provider, Subject, meta from lifecycle import ResourceManager from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere -from models import ChatMessage, RouterSettingsOverride, SpendLogRow +from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride from proxy_client import ProxyClient -from reliability_support import REAL_MODEL, chat_override, create_never_benched_refusing_deployment -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = [pytest.mark.e2e, pytest.mark.perf] + +REFUSING_MODEL: Final = "openai/gpt-5.5" DEPLOYMENTS_PER_GROUP: Final = 2 RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES @@ -110,7 +103,6 @@ class FailedCall: status_code: int seconds: float body_head: str - call_id: str | None @dataclass(frozen=True, slots=True) @@ -123,8 +115,17 @@ class WorkerGrowth: return self.after.ram_usage_mb - self.warm.ram_usage_mb +def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str: + return proxy.create_model( + name, + LiteLLMParamsBody( + model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0 + ), + ) + + def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None: - for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): + for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): resources.defer(lambda model_id=model_id: proxy.delete_model(model_id)) @@ -141,10 +142,18 @@ TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS) def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall: started: Final = time.perf_counter() - resp: Final = chat_override( - proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT + resp: Final = proxy.transport.send( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=model, + messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")], + max_tokens=512, + router_settings_override=override, + cache={"no-cache": True}, + ), ) - return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id) + return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300]) def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]: @@ -259,24 +268,6 @@ def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fall ) -def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float: - assert call.call_id, ( - f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}" - ) - rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id) - assert rows, ( - f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack " - "writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored " - "request outgrew the writer's queue budget" - ) - snapshot: Final = rows[0].proxy_server_request - assert snapshot, ( - f"spend log {call.call_id} stored no request body, so the stack is not running with " - "general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously" - ) - return len(json.dumps(snapshot).encode()) / 1024 - - class TestReliabilityMemory: @pytest.mark.covers("reliability.perf.idle_memory.under_slo") @meta( @@ -309,13 +300,13 @@ class TestReliabilityMemory: Subject( domain=Domain.ROUTING, providers=(Provider.OPENAI,), - models=(REAL_MODEL,), + models=(REFUSING_MODEL,), mode=Mode.NONSTREAM, ) ) - def test_failing_requests_do_not_grow_rss_or_stored_request( + def test_failing_requests_do_not_grow_rss( self, - client: ComplexityRouterClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, record_property: Callable[[str, object], None], @@ -323,35 +314,25 @@ class TestReliabilityMemory: marker: Final = unique_marker() primary: Final = f"reliability-memory-{marker}" fallback: Final = f"reliability-memory-fb-{marker}" - _register_refusing_group(client.proxy, resources, primary) - _register_refusing_group(client.proxy, resources, fallback) + _register_refusing_group(proxy, resources, primary) + _register_refusing_group(proxy, resources, fallback) override: Final = RouterSettingsOverride( num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}] ) - probe: Final = _fail_once(client.proxy, scoped_key, primary, override) - _assert_every_call_failed_through_fallback((probe,), fallback) - stored_kb: Final = _stored_request_kb(client.proxy, probe) - assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, ( - f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the " - f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with " - f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole " - f"request into the stored snapshot the way the v1.100.0 ones did" - ) - warm: Final = _settled_rss_after_batches( - client.proxy, + proxy, scoped_key, primary, override, fallback, record_property, - _settled_rss_per_worker(client.proxy), + _settled_rss_per_worker(proxy), 1, WARMUP_BATCHES, ) after: Final = _settled_rss_after_batches( - client.proxy, + proxy, scoped_key, primary, override, diff --git a/tests/e2e/perf/test_ui_journeys.py b/tests/e2e/perf/test_ui_journeys.py index 63eefedec68..11b2cd1051d 100644 --- a/tests/e2e/perf/test_ui_journeys.py +++ b/tests/e2e/perf/test_ui_journeys.py @@ -33,7 +33,6 @@ from transport import HttpTransport pytestmark = [ pytest.mark.perf, - pytest.mark.quiet_stack, pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"), ] diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 8cc7bf0c2b4..9b842a723d8 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -11,7 +11,7 @@ from __future__ import annotations import os import time import warnings -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field, replace from datetime import datetime from functools import reduce @@ -155,6 +155,21 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None +@dataclass(frozen=True, slots=True) +class ModelRegistered: + model_name: str + model_id: str + + +@dataclass(frozen=True, slots=True) +class ModelNotRegistered: + model_name: str + reason: str + + +ModelRegistration = ModelRegistered | ModelNotRegistered + + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -461,6 +476,7 @@ class ProxyClient: model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT + settle: Callable[[float], None] = settle_propagation def with_caller(self, caller: Caller) -> ProxyClient: return replace(self, caller=caller) @@ -694,31 +710,59 @@ class ProxyClient: balancer address is configured (every request opens a fresh connection, so the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is what makes the model safe to use anywhere.""" - model_id = unwrap( - self.transport.post( - "/model/new", - headers=self.management_headers(), - json=body.model_copy(update={"litellm_params": route_cache_model( - body.litellm_params, provider_edge_base, - enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live, - mode=body.model_info.mode, - )}), - response_type=ModelNewResponse, - ) - ).model_id + model_id = unwrap(self._post_model(body, provider_live=provider_live)).model_id written_at = time.monotonic() try: self._await_model_servable(body.model_name, listed_for) except BaseException: self.delete_model(model_id) raise - settle_propagation(written_at) + self.settle(written_at) return model_id + @step("Add a batch of deployments, then wait once for all of them to propagate") + def register_models( + self, bodies: Sequence[ModelNewBody], *, provider_live: bool = False + ) -> tuple[ModelRegistration, ...]: + posted: Final = tuple((body.model_name, self._post_model(body, provider_live=provider_live)) for body in bodies) + written_at: Final = time.monotonic() + registrations: Final = tuple(self._servable_registration(name, result) for name, result in posted) + self.settle(written_at) + return registrations + + def _post_model(self, body: ModelNewBody, *, provider_live: bool) -> Result[ModelNewResponse]: + return self.transport.post( + "/model/new", + headers=self.management_headers(), + json=body.model_copy(update={"litellm_params": route_cache_model( + body.litellm_params, provider_edge_base, + enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live, + mode=body.model_info.mode, + )}), + response_type=ModelNewResponse, + ) + + def _servable_registration(self, model_name: str, posted: Result[ModelNewResponse]) -> ModelRegistration: + if not isinstance(posted, Success): + return ModelNotRegistered(model_name=model_name, reason=f"/model/new failed: {posted!r}") + match self._servable_outcome(model_name, None): + case Servable(): + return ModelRegistered(model_name=model_name, model_id=posted.data.model_id) + case NotServableOn() as missing: + self.delete_model(posted.data.model_id) + return ModelNotRegistered(model_name=model_name, reason=self._not_servable_message(model_name, missing)) + def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: """Block until every replica lists `model_name`, or fail at model_servable_timeout.""" + match self._servable_outcome(model_name, listed_for): + case Servable(): + return + case NotServableOn() as missing: + raise AssertionError(self._not_servable_message(model_name, missing)) + + def _servable_outcome(self, model_name: str, listed_for: str | None) -> Servable | NotServableOn: headers: Final = self.management_headers(listed_for) - outcome: Final = await_servable_everywhere( + return await_servable_everywhere( {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()}, model_name=model_name, timeout=self.model_servable_timeout, @@ -728,19 +772,15 @@ class ProxyClient: now=time.monotonic, sleep=time.sleep, ) - match outcome: - case Servable(): - return - case NotServableOn(replica=replica, last_result=last_result): - raise AssertionError( - servable_timeout_message( - model_name=model_name, - replica=replica, - timeout=self.model_servable_timeout, - db_sync_seconds=self.model_servable_db_sync_seconds, - last_result=last_result, - ) - ) + + def _not_servable_message(self, model_name: str, missing: NotServableOn) -> str: + return servable_timeout_message( + model_name=model_name, + replica=missing.replica, + timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + last_result=missing.last_result, + ) @staticmethod def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller: diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index f15cfd4e98b..891c5fc282c 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -13,7 +13,6 @@ markers = prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set - quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set @@ -21,3 +20,4 @@ markers = secret_manager: needs a proxy booted from gateway/secret_manager__ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py) sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it + proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py similarity index 52% rename from tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py rename to tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py index 8f3bed3f89b..a77ca91ad3a 100644 --- a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py @@ -17,11 +17,6 @@ from models import ( ChatBody, ChatMessage, ChatResponse, - ConfigPatchResponse, - CostDiscountConfig, - CostDiscountConfigResponse, - CostMarginConfig, - CostMarginConfigResponse, GuardrailCreateBody, GuardrailCreateResponse, KeyDeleteBody, @@ -31,14 +26,11 @@ from models import ( ) from spend_e2e_client import SpendClient -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = pytest.mark.e2e BACKEND_MODEL: Final = "openai/gpt-6.1-sol" INPUT_RATE: Final = 0.00005 OUTPUT_RATE: Final = 0.0001 -DISCOUNT: Final = 0.25 -MARGIN_PERCENT: Final = 0.1 -MARGIN_FIXED: Final = 0.0005 def _register_bedrock_guardrail( @@ -92,52 +84,6 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea ) -def _set_discount(client: SpendClient, values: dict[str, float]) -> float: - unwrap( - client.proxy.transport.patch( - "/config/cost_discount_config", - headers=client.proxy.transport.master, - json=CostDiscountConfig(values), - response_type=ConfigPatchResponse, - ) - ) - return time.monotonic() - - -def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float: - unwrap( - client.proxy.transport.patch( - "/config/cost_margin_config", - headers=client.proxy.transport.master, - json=CostMarginConfig(values), - response_type=ConfigPatchResponse, - ) - ) - return time.monotonic() - - -def _get_discount(client: SpendClient) -> dict[str, float]: - return unwrap( - client.proxy.transport.get( - "/config/cost_discount_config", - headers=client.proxy.transport.master, - params=NoBody(), - response_type=CostDiscountConfigResponse, - ) - ).values - - -def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]: - return unwrap( - client.proxy.transport.get( - "/config/cost_margin_config", - headers=client.proxy.transport.master, - params=NoBody(), - response_type=CostMarginConfigResponse, - ) - ).values - - def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str: model: Final = f"{prefix}-{unique_marker()}" model_id: Final = client.proxy.create_model( @@ -199,100 +145,7 @@ def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str: return key -@pytest.fixture -def restored_pricing_config(client: SpendClient) -> Iterator[None]: - discount: Final = _get_discount(client) - margin: Final = _get_margin(client) - try: - _set_discount(client, {}) - reset_at: Final = _set_margin(client, {}) - settle_propagation(reset_at) - yield - finally: - try: - _set_discount(client, discount) - finally: - restored_at: Final = _set_margin(client, margin) - settle_propagation(restored_at) - - -class TestPricingConfigSpend: - @pytest.mark.covers( - "quota_management.spend_tracking.discount_config.logs_cost", - exercised_on=["chat_completions"], - ) - @meta( - Subject( - domain=Domain.SPEND_BUDGETS, - route=Route.SPEND_REPORTING, - providers=(Provider.OPENAI,), - models=(BACKEND_MODEL,), - mode=Mode.NONSTREAM, - ) - ) - def test_configured_discount_reaches_persisted_spend_row( - self, - client: SpendClient, - strict_resources: ResourceManager, - scoped_key: str, - restored_pricing_config: None, - ) -> None: - discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT}) - settle_propagation(discount_set_at) - model: Final = _register_model(client, strict_resources, "discount-priced") - chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage - assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 - assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 - - row: Final = poll_cost_row(client.proxy, chat.id) - assert row is not None - base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) - breakdown: Final = row.breakdown - assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT) - assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT) - assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT)) - - @pytest.mark.covers( - "quota_management.spend_tracking.margin_config.logs_cost", - exercised_on=["chat_completions"], - ) - @meta( - Subject( - domain=Domain.SPEND_BUDGETS, - route=Route.SPEND_REPORTING, - providers=(Provider.OPENAI,), - models=(BACKEND_MODEL,), - mode=Mode.NONSTREAM, - ) - ) - def test_configured_margin_reaches_persisted_spend_row( - self, - client: SpendClient, - strict_resources: ResourceManager, - scoped_key: str, - restored_pricing_config: None, - ) -> None: - margin_set_at: Final = _set_margin( - client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}} - ) - settle_propagation(margin_set_at) - model: Final = _register_model(client, strict_resources, "margin-priced") - chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage - assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 - assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 - - row: Final = poll_cost_row(client.proxy, chat.id) - assert row is not None - base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) - breakdown: Final = row.breakdown - expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED - assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT) - assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED) - assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin) - assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin) - +class TestGuardrailCostSpend: @pytest.mark.covers( "quota_management.spend_tracking.guardrail_cost.logs_cost", exercised_on=["chat_completions"], @@ -311,7 +164,6 @@ class TestPricingConfigSpend: client: SpendClient, strict_resources: ResourceManager, scoped_key: str, - restored_pricing_config: None, ) -> None: identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"] diff --git a/tests/e2e/stack_lock.py b/tests/e2e/stack_lock.py deleted file mode 100644 index 06df7a20b6a..00000000000 --- a/tests/e2e/stack_lock.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Cross-process reader/writer lock over the proxy stack every xdist worker shares. -Every collected test holds it shared, marker or not, since the Claude Code cells and -other unmarked suites drive the same stack; a `quiet_stack` test holds it exclusive, -and the `gate` file makes a waiting exclusive holder win over readers that arrive -after it.""" - -from __future__ import annotations - -import fcntl -import hashlib -import tempfile -from collections.abc import Generator -from contextlib import ExitStack, contextmanager -from pathlib import Path -from typing import Final - -from e2e_config import PROXY_BASE_URL - -STACK_DIGEST: Final = hashlib.sha256(PROXY_BASE_URL.encode()).hexdigest()[:12] -LOCK_DIR: Final = Path(tempfile.gettempdir()) / f"litellm-e2e-stack-{STACK_DIGEST}" -GATE_FILE: Final = LOCK_DIR / "gate" -STACK_FILE: Final = LOCK_DIR / "stack" - - -@contextmanager -def _flock(path: Path, operation: int) -> Generator[None]: - with path.open("a") as handle: - fcntl.flock(handle, operation) - try: - yield - finally: - fcntl.flock(handle, fcntl.LOCK_UN) - - -@contextmanager -def stack_lock(exclusive: bool) -> Generator[None]: - LOCK_DIR.mkdir(parents=True, exist_ok=True) - if exclusive: - with _flock(GATE_FILE, fcntl.LOCK_EX), _flock(STACK_FILE, fcntl.LOCK_EX): - yield - return - with ExitStack() as held: - with _flock(GATE_FILE, fcntl.LOCK_SH): - held.enter_context(_flock(STACK_FILE, fcntl.LOCK_SH)) - yield diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index 83e573ae249..4f507c2fe7b 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -1,6 +1,6 @@ # e2e harness tests -Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network +Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, `claude_code/` covers the driver, builder and probe, `cli_harness/` covers the version resolver, the CLI env lockdown and the rate limiter, and `codex/` covers how the Codex driver reads the CLI's JSONL events and the compactions in a thread's rollout file. Put a new harness test under the folder that mirrors the suite folder whose module it covers diff --git a/tests/e2e_harness/test_proxy_client.py b/tests/e2e_harness/test_proxy_client.py index e1615ec5f65..d352e5c8d36 100644 --- a/tests/e2e_harness/test_proxy_client.py +++ b/tests/e2e_harness/test_proxy_client.py @@ -12,10 +12,11 @@ injected clock, so nothing here monkeypatches anything. from __future__ import annotations import json +import time from builtins import ExceptionGroup from collections.abc import Callable, Generator, Iterable, Mapping from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain, repeat from queue import SimpleQueue @@ -40,7 +41,10 @@ from models import ( LiteLLMParamsBody, McpServerCreateBody, McpServerUpdateBody, + ModelDeleteBody, + ModelInfoBody, ModelListEntry, + ModelNewBody, ModelsListResponse, OrgNewBody, OrgUpdateBody, @@ -59,6 +63,8 @@ from proxy_client import ( ConvergeOutcome, CredentialKind, EverywhereConverged, + ModelNotRegistered, + ModelRegistered, ModelsPoller, NeverConvergedOn, NotConverged, @@ -691,3 +697,85 @@ def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: ) assert json.loads(bodies.get_nowait()) == expected assert bodies.empty() + + +@contextmanager +def model_registry( + listed: frozenset[str], refused: frozenset[str] +) -> Generator[tuple[ProxyClient, SimpleQueue[float], SimpleQueue[str]]]: + writes: Final[SimpleQueue[float]] = SimpleQueue() + deleted: Final[SimpleQueue[str]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def reply(self, status: int, body: object) -> None: + self.send_response(status) + self.end_headers() + self.wfile.write(json.dumps(body).encode()) + + def do_GET(self) -> None: + self.reply(200, {"data": [{"id": name} for name in sorted(listed)]}) + + def do_POST(self) -> None: + raw: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if self.path == "/model/delete": + deleted.put(ModelDeleteBody.model_validate_json(raw).id) + self.reply(200, {}) + return + writes.put(time.monotonic()) + name: Final = ModelNewBody.model_validate_json(raw).model_name + if name in refused: + self.reply(500, {"error": "refused"}) + return + self.reply(200, {"model_id": f"id-{name}"}) + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + proxy: Final = build_proxy_client( + base_url=url, control_plane_base_url=url, replica_urls=(url,), control_replica_urls=(url,), master_key="bootstrap" + ) + try: + yield replace(proxy, model_servable_timeout=0.3, model_servable_interval=0.01), writes, deleted + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class TestRegisterModels: + def test_a_batch_waits_once_after_its_last_write_and_reports_each_model(self) -> None: + names: Final = ("first", "refused", "ghost", "last") + settled: Final[SimpleQueue[float]] = SimpleQueue() + with ( + model_registry(listed=frozenset({"first", "last"}), refused=frozenset({"refused"})) as ( + proxy, + writes, + deleted, + ), + without_retries(), + ): + registrations: Final = replace(proxy, settle=settled.put).register_models( + tuple( + ModelNewBody( + model_name=name, litellm_params=LiteLLMParamsBody(model="synthetic"), model_info=ModelInfoBody() + ) + for name in names + ) + ) + last_write: Final = max(writes.get_nowait() for _ in names) + assert writes.empty() + assert settled.get_nowait() >= last_write + assert settled.empty(), "a batch must wait out propagation once, not once per model" + assert tuple((type(r), r.model_name) for r in registrations) == ( + (ModelRegistered, "first"), + (ModelNotRegistered, "refused"), + (ModelNotRegistered, "ghost"), + (ModelRegistered, "last"), + ) + assert tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) == ("id-first", "id-last") + assert deleted.get_nowait() == "id-ghost" + assert deleted.empty() diff --git a/tests/e2e_harness/test_stack_lock.py b/tests/e2e_harness/test_stack_lock.py deleted file mode 100644 index b275ca8c755..00000000000 --- a/tests/e2e_harness/test_stack_lock.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Cross-process behavior of the stack lock: readers share it, an exclusive holder waits for -every reader and keeps them out, and a reader arriving behind a waiting exclusive holder -queues behind it instead of starving it.""" - -from __future__ import annotations - -import fcntl -import inspect -import os -import subprocess -import sys -import time -from contextlib import ExitStack -from pathlib import Path -from typing import Final - -import pytest -import stack_lock - -HARNESS_DIR: Final = Path(inspect.getfile(stack_lock)).resolve().parent -DEADLINE_SECONDS: Final = 30.0 -SETTLE_SECONDS: Final = 0.5 -HOLDER_SCRIPT: Final = """ -import sys, time -from pathlib import Path -from stack_lock import stack_lock -name, mode, release_path, log_path = sys.argv[1:] - - -def record(event): - with Path(log_path).open("a") as log: - log.write(f"{name} {event}\\n") - - -record("waiting") -with stack_lock(exclusive=mode == "exclusive"): - record("enter") - while not Path(release_path).exists(): - time.sleep(0.02) - record("exit") -""" - - -def _events(log_path: Path) -> tuple[str, ...]: - return tuple(log_path.read_text().splitlines()) if log_path.exists() else () - - -def _wait_for_event(log_path: Path, event: str) -> None: - deadline: Final = time.monotonic() + DEADLINE_SECONDS - while event not in _events(log_path): - if time.monotonic() > deadline: - pytest.fail(f"{event!r} never appeared; events so far: {_events(log_path)}") - time.sleep(0.02) - - -def _wait_until_gate_is_held_exclusively(gate_path: Path) -> None: - deadline: Final = time.monotonic() + DEADLINE_SECONDS - with gate_path.open("a") as handle: - while True: - try: - fcntl.flock(handle, fcntl.LOCK_SH | fcntl.LOCK_NB) - except BlockingIOError: - return - fcntl.flock(handle, fcntl.LOCK_UN) - if time.monotonic() > deadline: - pytest.fail("no exclusive holder ever took the gate") - time.sleep(0.02) - - -def _start_holder(held: ExitStack, tmp_path: Path, name: str, mode: str) -> subprocess.Popen[bytes]: - holder: Final = held.enter_context( - subprocess.Popen( - ( - sys.executable, - "-P", - "-c", - HOLDER_SCRIPT, - name, - mode, - str(tmp_path / f"release-{name}"), - str(tmp_path / "events"), - ), - cwd=HARNESS_DIR, - env={**os.environ, "TMPDIR": str(tmp_path), "PYTHONPATH": str(HARNESS_DIR)}, - ) - ) - held.callback(holder.kill) - return holder - - -def test_readers_share_exclusive_waits_and_a_waiting_exclusive_beats_later_readers(tmp_path: Path) -> None: - lock_dir: Final = tmp_path / f"litellm-e2e-stack-{stack_lock.STACK_DIGEST}" - lock_dir.mkdir() - log_path: Final = tmp_path / "events" - with ExitStack() as held: - first_reader: Final = _start_holder(held, tmp_path, "A", "shared") - _wait_for_event(log_path, "A enter") - second_reader: Final = _start_holder(held, tmp_path, "R", "shared") - _wait_for_event(log_path, "R enter") - (tmp_path / "release-R").touch() - _wait_for_event(log_path, "R exit") - writer: Final = _start_holder(held, tmp_path, "W", "exclusive") - _wait_until_gate_is_held_exclusively(lock_dir / "gate") - late_reader: Final = _start_holder(held, tmp_path, "B", "shared") - _wait_for_event(log_path, "B waiting") - time.sleep(SETTLE_SECONDS) - (tmp_path / "release-A").touch() - _wait_for_event(log_path, "W enter") - (tmp_path / "release-W").touch() - _wait_for_event(log_path, "B enter") - (tmp_path / "release-B").touch() - for holder in (first_reader, second_reader, writer, late_reader): - assert holder.wait(timeout=DEADLINE_SECONDS) == 0 - events: Final = _events(log_path) - assert events.index("R enter") < events.index("A exit") - assert events.index("W enter") > events.index("A exit") - assert events.index("B enter") > events.index("W exit") diff --git a/tests/integration/pricing/test_cost_margin_discount.py b/tests/integration/pricing/test_cost_margin_discount.py index b6cda8f125f..613fe72ed70 100644 --- a/tests/integration/pricing/test_cost_margin_discount.py +++ b/tests/integration/pricing/test_cost_margin_discount.py @@ -7,7 +7,7 @@ import httpx import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue class CostEstimate(BaseModel): @@ -110,6 +110,25 @@ def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None: assert actual == pytest.approx(dict(expectations)), rows +class LoggedBreakdown(BaseModel): + model_config = ConfigDict(frozen=True) + + original_cost: float + discount_percent: float | None = None + discount_amount: float | None = None + margin_percent: float | None = None + margin_fixed_amount: float | None = None + margin_total_amount: float | None = None + + +def _logged_breakdown(request_id: str) -> LoggedBreakdown: + rows: Final = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)) + assert len(rows) == 1, rows + metadata: Final = rows[0]["metadata"] + parsed: Final = object_value(json.loads(metadata) if isinstance(metadata, str) else metadata) + return LoggedBreakdown.model_validate(parsed["cost_breakdown"]) + + def _estimated(gateway: Gateway, model: str) -> CostEstimate: response: Final = gateway.request( "POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20} @@ -236,6 +255,17 @@ def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Ga (streamed_global_request_id, 0.12), ) ) + discounted: Final = _logged_breakdown(discounted_request_id) + assert discounted.original_cost == pytest.approx(0.06), discounted + assert discounted.discount_percent == pytest.approx(0.05), discounted + assert discounted.discount_amount == pytest.approx(0.06 * 0.05), discounted + margined: Final = _logged_breakdown(provider_margin_request_id) + assert margined.original_cost == pytest.approx(0.06), margined + assert margined.discount_percent == pytest.approx(0.05), margined + assert margined.discount_amount == pytest.approx(0.06 * 0.05), margined + assert margined.margin_percent == pytest.approx(0.1), margined + assert margined.margin_fixed_amount == pytest.approx(0.001), margined + assert margined.margin_total_amount == pytest.approx(0.057 * 0.1 + 0.001), margined rejections: Final = ( ( diff --git a/tests/integration/spend/test_failing_request_stored_snapshot.py b/tests/integration/spend/test_failing_request_stored_snapshot.py new file mode 100644 index 00000000000..4f016c33f37 --- /dev/null +++ b/tests/integration/spend/test_failing_request_stored_snapshot.py @@ -0,0 +1,84 @@ +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from integration._support.process import owned_proxy +from pydantic import JsonValue + +UNREACHABLE_BASE: Final = "http://127.0.0.1:9/v1" +TRANSCRIPT_TURNS: Final = 40 +RETRIES: Final = 2 +STORED_REQUEST_BUDGET_KB: Final = 64 + + +def _config_storing_prompts(directory: Path) -> Path: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + general_settings: Final = object_value(config["general_settings"]) + path: Final = directory / "store-prompts.yaml" + path.write_text( + yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}}) + ) + return path + + +@pytest.fixture(scope="module") +def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with gateway_from_environment() as parent: + directory: Final = tmp_path_factory.mktemp("failing_request_stored_snapshot") + with owned_proxy(parent, directory, {}, config=_config_storing_prompts(directory)) as owned: + yield owned + + +def _transcript(marker: str) -> list[JsonValue]: + turns: Final[list[JsonValue]] = [ + {"role": role, "content": f"turn {turn // 2} {role}"} + for turn, role in enumerate(("user", "assistant") * TRANSCRIPT_TURNS) + ] + return [*turns, {"role": "user", "content": marker}] + + +def _stored_request(gateway: Gateway, call_id: str) -> JsonValue: + rows: Final = eventually( + lambda: gateway.rows('SELECT proxy_server_request FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=40, + ) + stored: Final = rows[0]["proxy_server_request"] + return json.loads(stored) if isinstance(stored, str) else stored + + +def test_failing_request_with_retries_and_fallback_stores_one_copy_of_the_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + primary: Final = f"stored-snapshot-{uuid.uuid4().hex}" + fallback: Final = f"stored-snapshot-fb-{uuid.uuid4().hex}" + for name in (primary, primary, fallback, fallback): + scenario.model(name=name, api_base=UNREACHABLE_BASE, cooldown_time=0) + key: Final = scenario.key(models=[primary, fallback]) + marker: Final = f"stored snapshot {uuid.uuid4().hex}" + + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": primary, + "max_tokens": 16, + "messages": _transcript(marker), + "router_settings_override": {"num_retries": RETRIES, "fallbacks": [{primary: [fallback]}]}, + }, + key=key, + ) + assert response.status_code != 200, response.text + assert fallback in response.text, response.text + + serialized: Final = json.dumps(_stored_request(gateway, response.headers["x-litellm-call-id"])) + stored_kb: Final = len(serialized.encode()) / 1024 + assert marker in serialized, serialized[:400] + assert stored_kb <= STORED_REQUEST_BUDGET_KB, ( + f"one failing request stored a {stored_kb:.0f} KB request snapshot for a " + f"{2 * TRANSCRIPT_TURNS + 1}-message transcript with {RETRIES} retries and a fallback" + ) From 5bc7f6434dc95e565ba4fbca399df3a95e3c3675 Mon Sep 17 00:00:00 2001 From: Jon Richmond <4224495+alchemicjon@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:31:09 -0500 Subject: [PATCH 002/178] fix(logging): don't leak redacted reasoning via provider_specific_fields (#47398) * fix: prevent reasoning field from being copied into provider_specific_fields * fix: ensure reasoning field in provider_specific_fields is redacted * fix: redact reasoning_details in provider_specific_fields * fix: format redact_messages.py * ci: retrigger * ci: retrigger * ci: retrigger * fix(logging): keep provider_specific_fields.reasoning in chat responses Redaction already blanks provider_specific_fields.reasoning and reasoning_details in logs, so chat responses can keep the field as they did before. * fix(logging): redact reasoning_text, citations, and web search results in provider_specific_fields * fix(logging): redact vector store, RAG, Groq tool, Databricks agent and tool search results in provider_specific_fields * fix(logging): redact Bedrock streamed reasoningContent and file annotations in provider_specific_fields --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/redact_messages.py | 16 ++++++++++++++++ .../litellm_core_utils/test_redact_messages.py | 16 ++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 9da0e35bdda..ac5c0720eff 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -272,15 +272,31 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ return { **fields, **({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}), + **({"reasoning": REDACTED_BY_LITELLM} if fields.get("reasoning") is not None else {}), + **({"reasoning_text": REDACTED_BY_LITELLM} if fields.get("reasoning_text") is not None else {}), **{ key: None for key in ( "thinking_blocks", "reasoningContentBlocks", + "reasoningContent", + "annotations", "citations", + "citation", + "citationsContent", "compaction_blocks", "compaction_start", "compaction_delta", + "reasoning_details", + "web_search_calls", + "web_search_results", + "tool_results", + "code_interpreter_results", + "search_results", + "rerank_results", + "executed_tools", + "custom_outputs", + "tool_search_items", "token_ids", ) if key in fields diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 2fd3f16ab17..491ea5315ff 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -54,9 +54,25 @@ SIGNATURE: Final = "opaque-thought-signature" def _provider_fields() -> dict[str, object]: return { "reasoning_content": SECRET, + "reasoning": SECRET, + "reasoning_text": SECRET, + "reasoning_details": [{"type": "reasoning.text", "text": SECRET, "index": 0, "format": "unknown"}], "thinking_blocks": [{"type": "thinking", "thinking": SECRET}], "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], + "reasoningContent": {"reasoningText": {"text": SECRET}}, + "annotations": [{"type": "file", "file": {"name": "notes.pdf", "content": [{"type": "text", "text": SECRET}]}}], "citations": [{"cited_text": SECRET}], + "citationsContent": [{"citations": [{"location": {}}], "content": [{"text": SECRET}]}], + "citation": {"type": "char_location", "cited_text": SECRET}, + "web_search_calls": [{"type": "web_search_call", "action": {"type": "search", "query": SECRET}}], + "web_search_results": [{"type": "web_search_tool_result", "content": [{"title": SECRET}]}], + "tool_results": [{"type": "web_fetch_tool_result", "content": {"text": SECRET}}], + "code_interpreter_results": [{"stdout": SECRET}], + "search_results": [{"object": "vector_store.search_results.page", "data": [{"content": [{"text": SECRET}]}]}], + "rerank_results": {"results": [{"document": {"text": SECRET}}]}, + "executed_tools": [{"type": "search", "search_results": {"results": [{"content": SECRET}]}}], + "custom_outputs": {"answer": SECRET}, + "tool_search_items": [{"type": "tool_search_output", "tools": [{"description": SECRET}]}], "compaction_blocks": [{"type": "compaction", "content": SECRET, "signature": "compaction-signature"}], "compaction_start": {"type": "compaction", "content": SECRET}, "compaction_delta": {"type": "compaction_delta", "content": SECRET}, From 02b8da2f169b6e30a7ade9e1b7df4a9099294916 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 4 Nov 2026 10:37:07 -0800 Subject: [PATCH 003/178] refactor: replace accounting event scaffolding with SpendLogs projection (#49843) * refactor(accounting): replace event scaffolding with SpendLogs projection * test: preserve spend timestamps without removed capture callback * fix: guard accounting event removal across schema update paths * refactor(accounting): retain schema and prepare durable payloads * fix(accounting): preserve dates in serialized payloads * fix(accounting): retain schema push compatibility * fix(accounting): retain span registration for the legacy schema * test(accounting): remove obsolete capture acknowledgement test * fix(accounting): date every daily projection row the way the model rollup does A space-separated startTime kept the whole timestamp as the daily key, while the model usage rollup took the calendar date. Every daily row now uses the same date helper Co-Authored-By: Claude Opus 5.5 * chore(ui): regenerate dashboard API types after rebasing on main Co-Authored-By: Claude Opus 5.5 * refactor(accounting): project only from the stored accounting payload project() now reads entities only from the payload's accounting_entities, so the raw SpendLogs fallback branches are gone. build_accounting_payload() builds the metadata once instead of dumping it to JSON and parsing it back. Guardrail usage dates use a tolerant UTC parse, so a malformed startTime drops only the guardrail rows instead of failing the whole projection. Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .../common_utils/check_responses_cost.py | 1 - litellm/accounting/__init__.py | 5 +- litellm/accounting/projection.py | 304 +++ litellm/accounting/protocols.py | 74 +- litellm/accounting/snapshot.py | 97 + litellm/litellm_core_utils/litellm_logging.py | 165 +- .../litellm_core_utils/realtime_streaming.py | 2 - litellm/proxy/_types.py | 17 - litellm/proxy/common_request_processing.py | 4 - litellm/proxy/db/accounting_ledger.py | 352 ---- litellm/proxy/db/model_usage_rollup.py | 4 +- .../proxy/hooks/proxy_track_cost_callback.py | 125 +- litellm/proxy/proxy_server.py | 33 - .../spend_tracking/spend_tracking_utils.py | 12 +- litellm/proxy/spend_tracking/usage_capture.py | 1084 ---------- litellm/proxy/spend_tracking/usage_meter.py | 330 ---- litellm/proxy/utils.py | 38 - litellm/router_python_core.py | 5 - litellm/utils.py | 7 - .../database/test_accounting_ledger.py | 610 ------ tests/unit/accounting/test_projection.py | 431 ++++ tests/unit/accounting/test_types.py | 2 - ...responses_transformation_transformation.py | 49 - tests/unit/proxy/db/test_accounting_ledger.py | 149 -- .../hooks/test_proxy_track_cost_callback.py | 137 -- .../spend_tracking/test_cost_finalization.py | 24 - .../test_spend_tracking_utils.py | 6 +- .../spend_tracking/test_usage_capture.py | 1738 ----------------- .../proxy/spend_tracking/test_usage_meter.py | 473 ----- tests/unit/proxy/test_proxy_server.py | 335 ---- ..._utils_model_creation_and_error_logging.py | 90 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 44 - 32 files changed, 845 insertions(+), 5902 deletions(-) create mode 100644 litellm/accounting/projection.py create mode 100644 litellm/accounting/snapshot.py delete mode 100644 litellm/proxy/db/accounting_ledger.py delete mode 100644 litellm/proxy/spend_tracking/usage_capture.py delete mode 100644 litellm/proxy/spend_tracking/usage_meter.py delete mode 100644 tests/integration/database/test_accounting_ledger.py create mode 100644 tests/unit/accounting/test_projection.py delete mode 100644 tests/unit/proxy/db/test_accounting_ledger.py delete mode 100644 tests/unit/proxy/spend_tracking/test_usage_capture.py delete mode 100644 tests/unit/proxy/spend_tracking/test_usage_meter.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index cfe4ca98145..69e36a0c076 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -200,7 +200,6 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi unpriced: Final = _without_read_price(billed.response) cost: Final = logging_obj.response_cost_calculator(result=unpriced) end_time: Final = datetime.now() - logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time) GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( async_coroutine=logging_obj.async_success_handler( result=unpriced, start_time=logging_obj.start_time, end_time=end_time diff --git a/litellm/accounting/__init__.py b/litellm/accounting/__init__.py index d4e56f61279..5d59a35a987 100644 --- a/litellm/accounting/__init__.py +++ b/litellm/accounting/__init__.py @@ -1,7 +1,7 @@ """Public accounting contracts. Importing these values does not enable accounting.""" from litellm.accounting.measurements import measurements_for -from litellm.accounting.protocols import BudgetEnforcer, CostAggregator, CostCompiler, UsageMeter +from litellm.accounting.protocols import CostAggregator from litellm.accounting.types import ( AdmissionDecision, AggregateQuery, @@ -33,9 +33,7 @@ __all__ = [ "AdmissionDecision", "AggregateQuery", "AggregateRow", - "BudgetEnforcer", "CostAggregator", - "CostCompiler", "CostEvent", "DeliveryState", "DeliveryTarget", @@ -56,7 +54,6 @@ __all__ = [ "PricingSnapshot", "Rate", "Unpriced", - "UsageMeter", "measurements_for", "validate_aggregate_rows", ] diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py new file mode 100644 index 00000000000..fc3a89ae306 --- /dev/null +++ b/litellm/accounting/projection.py @@ -0,0 +1,304 @@ +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.constants import ( + BILLED_RETRY_ATTEMPT_METADATA_KEY, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + MODEL_INSIGHTS_DEFAULT_TASK, + MODEL_INSIGHTS_TASK_TAG_PREFIX, +) +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import GuardrailCostByUnitEntry +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, daily_spend_entity_ids +from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks +from litellm.proxy.db.model_usage_rollup import date_from_start_time +from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING +from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens +from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +_FIELDS: Final = TypeAdapter(dict[str, object]) +_STRINGS: Final = TypeAdapter(tuple[str, ...]) +_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...]) +_TASK_TYPES: Final = frozenset(load_model_insight_tasks()) + + +@dataclass(frozen=True, slots=True) +class Contribution: + table: str + key: tuple[tuple[str, str], ...] + increments: tuple[tuple[str, int | float], ...] + values: tuple[tuple[str, str | None], ...] = () + + +def _fields(value: object) -> Mapping[str, object]: + if isinstance(value, str): + return _FIELDS.validate_json(value) + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else MappingProxyType({}) + + +def _text(value: object) -> str: + return "" if value is None else str(value) + + +def _number(value: object) -> int | float: + return value if isinstance(value, (int, float)) else 0 + + +def _tags(row: Mapping[str, object]) -> tuple[str, ...]: + value: Final = row.get("request_tags") + if value is None: + return () + return _STRINGS.validate_json(value) if isinstance(value, str) else _STRINGS.validate_python(value) + + +def _date(row: Mapping[str, object]) -> str | None: + start: Final = row.get("startTime") + return date_from_start_time(start) if isinstance(start, (datetime, str)) else None + + +def _spend(table: str, key: tuple[tuple[str, str], ...], spend: float) -> Contribution: + return Contribution(table, key, (("spend", spend),)) + + +def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + spend: Final = _number(row.get("spend")) + entities: Final = _fields(metadata.get("accounting_entities")) + identities: Final = ( + ("LiteLLM_VerificationToken", "token", entities.get("api_key")), + ("LiteLLM_UserTable", "user_id", entities.get("user")), + ("LiteLLM_UserTable", "user_id", entities.get("proxy_budget_name")), + ("LiteLLM_TeamTable", "team_id", entities.get("team_id")), + ("LiteLLM_OrganizationTable", "organization_id", entities.get("organization_id")), + ("LiteLLM_EndUserTable", "user_id", entities.get("end_user")), + ("LiteLLM_AgentsTable", "agent_id", entities.get("agent_id")), + ("LiteLLM_ProjectTable", "project_id", entities.get("project_id")), + ) + groups: Final = _STRINGS.validate_python(metadata.get("accounting_model_access_groups", ())) + return ( + *( + _spend(table, ((column, identity),), spend) + for table, column, identity in identities + if isinstance(identity, str) + ), + *(_spend("LiteLLM_TagTable", (("tag_name", tag),), spend) for tag in _tags(row) if tag), + *(_spend("LiteLLM_ModelAccessGroupBudgetTable", (("access_group_name", group),), spend) for group in groups), + *_membership_contributions(entities, spend), + ) + + +def _membership_contributions(row: Mapping[str, object], spend: float) -> tuple[Contribution, ...]: + user: Final = row.get("user") + if not isinstance(user, str): + return () + owners: Final = ( + ("LiteLLM_TeamMembership", "team_id", row.get("team_id")), + ("LiteLLM_OrganizationMembership", "organization_id", row.get("organization_id")), + ) + return tuple( + _spend(table, ((column, identity), ("user_id", user)), spend) + for table, column, identity in owners + if isinstance(identity, str) + ) + + +def _daily_counters(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[tuple[str, int | float], ...]: + usage: Final = _fields(metadata.get("usage_object")) + savings: Final = _fields(metadata.get("accounting_savings")) + internal: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) + caller: Final = not internal and metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is not True + success: Final = metadata.get("status") != "failure" + duration: Final = row.get("request_duration_ms") + timed: Final = not internal and success and isinstance(duration, int) and duration >= 0 + return ( + ("spend", _number(row.get("spend"))), + ("prompt_tokens", _number(row.get("prompt_tokens"))), + ("completion_tokens", _number(row.get("completion_tokens"))), + ("api_requests", int(caller)), + ("successful_requests", int(caller and success)), + ("failed_requests", int(caller and not success)), + ("cache_read_input_tokens", extract_cache_read_tokens(usage)), + ("cache_creation_input_tokens", extract_cache_creation_tokens(usage)), + ("compression_saved_tokens", extract_compression_saved_tokens(metadata)), + ("compression_savings_spend", _number(savings.get("compression"))), + ("prompt_caching_savings_spend", _number(savings.get("prompt_caching"))), + ("gateway_injected_caching_savings_spend", _number(savings.get("gateway_injected_caching"))), + ("autorouter_savings_spend", 0.0 if internal else _number(savings.get("autorouter"))), + ("total_response_time_ms", _number(duration) if timed else 0), + ("timed_requests", int(timed)), + ) + + +def _daily_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _date(row) + if date is None or "api_key" not in row: + return () + if "mcp_namespaced_tool_name" not in row and ( + "model" not in row or "custom_llm_provider" not in row or "model_group" not in row + ): + return () + counters: Final = _daily_counters(row, metadata) + endpoint: Final = ROUTE_ENDPOINT_MAPPING.get(_text(row.get("call_type"))) + shared_key: Final = ( + ("date", date), + ("api_key", _text(row.get("api_key"))), + ("model", _text(row.get("model"))), + ("custom_llm_provider", _text(row.get("custom_llm_provider"))), + ("mcp_namespaced_tool_name", _text(row.get("mcp_namespaced_tool_name"))), + ("endpoint", _text(endpoint)), + ) + # The legacy daily organization writer uses the charged org argument, not + # SpendLogs' normalized empty-string identity for a request with no org. + daily_row: Final = {**row, "organization_id": _fields(metadata.get("accounting_entities")).get("organization_id")} + return tuple(_daily_rows(daily_row, shared_key, counters)) + + +def _daily_rows( + row: Mapping[str, object], + shared_key: tuple[tuple[str, str], ...], + counters: tuple[tuple[str, int | float], ...], +) -> Iterator[Contribution]: + for entity, table in DAILY_SPEND_TABLES.items(): + for identity in daily_spend_entity_ids(row, entity): + yield Contribution( + table.name, + ((table.entity_id_column, _text(identity)), *shared_key), + counters, + ( + ("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))), + *((("request_id", _text(row.get("request_id"))),) if table.carries_request_id else ()), + ), + ) + + +def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + windows: Final = _ROWS.validate_python(metadata.get("accounting_windows", ())) + return tuple( + Contribution( + "LiteLLM_BudgetWindowSpend", + tuple((column, _text(window[column])) for column in ("entity_type", "entity_id", "window_duration")), + (("spend", _number(window.get("spend", row.get("spend")))),), + (("window_start", _text(window["window_start"])), ("started_at", _text(row.get("startTime")))), + ) + for window in windows + ) + + +def project(row: Mapping[str, object]) -> tuple[Contribution, ...]: + metadata: Final = _fields(row.get("metadata")) + return ( + *_entity_contributions(row, metadata), + *_daily_contributions(row, metadata), + *_window_contributions(row, metadata), + *_model_contributions(row, metadata), + *_guardrail_contributions(row, metadata), + ) + + +def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _date(row) + origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) + if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): + return () + model: Final = _text(row.get("model")) or "unknown" + caller: Final = not origin + success: Final = row.get("status") == "success" + return ( + Contribution( + "LiteLLM_DailyModelUsage", + ( + ("date", date), + ("model_group", _text(row.get("model_group")) or model), + ("model", model), + ("custom_llm_provider", _text(row.get("custom_llm_provider")) or "unknown"), + ("task_type", _model_task(row)), + ), + ( + ("spend", _number(row.get("spend"))), + ("prompt_tokens", _number(row.get("prompt_tokens"))), + ("completion_tokens", _number(row.get("completion_tokens"))), + ("request_count", int(caller)), + ("successful_requests", int(caller and success)), + ("failed_requests", int(caller and not success)), + ), + ), + ) + + +def _utc_date(row: Mapping[str, object]) -> str | None: + start: Final = row.get("startTime") + try: + parsed: Final = datetime.fromisoformat(start.replace("Z", "+00:00")) if isinstance(start, str) else start + except ValueError: + return None + if not isinstance(parsed, datetime): + return None + return ( + (parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed) + .astimezone(timezone.utc) + .date() + .isoformat() + ) + + +def _guardrail_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _utc_date(row) + info: Final = metadata.get("guardrail_information") or metadata.get("standard_logging_guardrail_information") + if not row.get("request_id") or date is None or not isinstance(info, list): + return () + return tuple(chain.from_iterable(_guardrail_entry(row, date, entry) for entry in _ROWS.validate_python(info))) + + +def _guardrail_entry(row: Mapping[str, object], date: str, entry: Mapping[str, object]) -> tuple[Contribution, ...]: + identity: Final = _text(entry.get("guardrail_id") or entry.get("guardrail_name")) + usage: Final = entry.get("guardrail_usage") + if not identity or not isinstance(usage, Mapping): + return () + costs: Final = _guardrail_unit_costs(entry) + return tuple( + _guardrail_unit(row, (identity, date, unit), units, None if costs is None else costs.get(unit)) + for unit, units in _fields(usage).items() + if isinstance(units, int) and not isinstance(units, bool) and units > 0 + ) + + +def _guardrail_unit( + row: Mapping[str, object], identity: tuple[str, str, str], units: int, cost: float | None +) -> Contribution: + guardrail_id, date, unit = identity + return Contribution( + "LiteLLM_DailyGuardrailUsageUnits", + ( + ("guardrail_id", guardrail_id), + ("date", date), + ("team_id", _text(row.get("team_id"))), + ("api_key", _text(row.get("api_key"))), + ("usage_unit", unit), + ), + (("units", units), ("cost", cost or 0.0), ("untracked_units", units if cost is None else 0)), + ) + + +def _model_task(row: Mapping[str, object]) -> str: + return next( + ( + task + for tag in _tags(row) + if tag.startswith(MODEL_INSIGHTS_TASK_TAG_PREFIX) + and (task := tag.removeprefix(MODEL_INSIGHTS_TASK_TAG_PREFIX)) in _TASK_TYPES + ), + MODEL_INSIGHTS_DEFAULT_TASK, + ) + + +def _guardrail_unit_costs(entry: Mapping[str, object]) -> Mapping[str, float | None] | None: + try: + cost: Final = GuardrailCostByUnitEntry.model_validate(entry) + except ValidationError: + return None + return None if cost.guardrail_cost_in_spend is False else cost.guardrail_cost_by_unit diff --git a/litellm/accounting/protocols.py b/litellm/accounting/protocols.py index 0208e37fbdf..f6ec61983a0 100644 --- a/litellm/accounting/protocols.py +++ b/litellm/accounting/protocols.py @@ -1,58 +1,14 @@ -"""Contracts only. Implementations own persistence, pricing, projections and live budgets. - -The event row's ``payload`` is a private envelope versioned by the row's ``schema_version``. -Usage envelopes retain the public event and request-scoped pricing context: quantities with units, -model/provider, requested and billed tier/region when known, and per-request/deployment overrides. -These inputs, including the effective cost-map rates, must survive request teardown. -Cost envelopes retain the public CostEvent, including its required finalized pricing snapshot. -T5/T6 own envelope serialization; this public API does not introduce a second persisted schema. -""" +from __future__ import annotations +from collections.abc import Mapping from typing import Protocol -from litellm.accounting.types import ( - AdmissionDecision, - AggregateQuery, - AggregateRow, - CostEvent, - Dimension, - InboundUsageEvent, - Measurement, - Metric, - OutboundUsageEvent, - Unpriced, -) - - -class UsageMeter(Protocol): - async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: - """Acknowledge commit, no-op identical IDs and reject conflicting content. - - Persist the event and its request-scoped pricing context before acknowledging. - HTTP acknowledgement is a separate policy. Unknown usage remains None. - """ - ... - - -class CostCompiler(Protocol): - async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced: - """Reuse the committed charge or finalize one charge and its delivery intent. - - Resolve captured request-scoped context and rates by event identity. - Persist the resolved inputs in the required snapshot. Replay - reuses the committed result and never consumes free allowance again. - Missing pricing is Unpriced, never an inferred zero. - """ - ... +from litellm.accounting.types import AggregateQuery, AggregateRow class CostAggregator(Protocol): - async def apply(self, measurements: tuple[Measurement, ...]) -> None: - """Apply one source event's complete contribution set per target/version. - - Each target commits all contributions and its receipt atomically. A receipt - for the first metric must not suppress the rest. Usage need not wait for cost. - """ + async def apply(self, rows: tuple[Mapping[str, object], ...]) -> None: + """Apply stored request contributions atomically with their delivery receipts.""" ... async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]: @@ -63,23 +19,3 @@ class CostAggregator(Protocol): combinations. Different groups may overlap, including non-additive tags. """ ... - - -class BudgetEnforcer(Protocol): - async def reserve( - self, reservation_id: str, dimensions: tuple[Dimension, ...], estimated: tuple[Metric, ...] - ) -> AdmissionDecision: - """Apply existing admission policies and estimates; duplicate IDs add no hold.""" - ... - - async def settle(self, reservation_id: str | None, cost_event: CostEvent) -> None: - """Settle the committed charge with atomic deduplication and hold mutation. - - Preserve scope/window attribution and retry unfinished settlement after - restart. SQL projection receipts do not provide live-counter atomicity. - """ - ... - - async def release(self, reservation_id: str) -> None: - """Idempotently release unused exposure at the request boundary, preserving charges.""" - ... diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py new file mode 100644 index 00000000000..51c7bbad81b --- /dev/null +++ b/litellm/accounting/snapshot.py @@ -0,0 +1,97 @@ +from collections.abc import Callable, Mapping, Sequence +from datetime import datetime +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure +from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens +from litellm.proxy.spend_tracking.savings import compute_savings_spend, marks_gateway_injection + +if TYPE_CHECKING: + from litellm.router import Router + +_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def _mapping(value: object) -> Mapping[str, object] | None: + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else None + + +_METADATA_FIELDS: Final = frozenset( + { + "usage_object", + "status", + "guardrail_information", + "standard_logging_guardrail_information", + "compression_savings", + "internal_call_origin", + "billed_retry_attempt", + } +) + + +_ACCOUNTING_FIELDS: Final = frozenset( + { + "request_id", + "call_type", + "api_key", + "spend", + "prompt_tokens", + "completion_tokens", + "startTime", + "model", + "model_group", + "custom_llm_provider", + "user", + "team_id", + "organization_id", + "end_user", + "agent_id", + "request_tags", + "status", + "request_duration_ms", + "mcp_namespaced_tool_name", + } +) + + +def build_accounting_payload( + payload: SpendLogsPayload, + *, + entities: Mapping[str, str | None], + model_access_groups: Sequence[str], + windows: Sequence[Mapping[str, object]], + router: "Callable[[], Router | None]", +) -> Mapping[str, object]: + metadata: Final = _FIELDS.validate_json(payload["metadata"]) + savings: Final = compute_savings_spend( + model=payload["model"], + custom_llm_provider=payload["custom_llm_provider"], + compression_saved_tokens=extract_compression_saved_tokens(metadata), + gateway_injected_cache=marks_gateway_injection(metadata, payload["model_id"]), + routing_decision=_mapping(metadata.get("routing_decision")), + model_id=payload["model_id"], + llm_router=router, + usage_object=_mapping(metadata.get("usage_object")), + cost_breakdown=_mapping(metadata.get("cost_breakdown")), + recorded_autorouter_savings=metadata.get("autorouter_savings"), + recorded_autorouter_savings_estimate=_mapping(metadata.get("autorouter_savings_estimate")), + billed_at=payload["endTime"], + ) + accounting_metadata: Final = { + **{key: value for key, value in metadata.items() if key in _METADATA_FIELDS}, + "accounting_entities": entities, + "accounting_savings": savings._asdict(), + "accounting_model_access_groups": model_access_groups, + "accounting_windows": windows, + } + return { + **{ + key: value.isoformat() if isinstance(value, datetime) else value + for key, value in payload.items() + if key in _ACCOUNTING_FIELDS + }, + "metadata": _FIELDS.validate_python(safe_json_structure(accounting_metadata, non_finite_as_null=True)), + } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f8080531961..5494d0d82b7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,7 @@ from datetime import datetime as dt_object from enum import Enum from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, TypeVar, cast from httpx import Response, Timeout from openai import AsyncStream as OpenAIAsyncStream @@ -283,8 +283,6 @@ if TYPE_CHECKING: from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation - from litellm.proxy.spend_tracking.usage_capture import CapturedUsage - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -662,54 +660,7 @@ class UsageObservation(Enum): ESTIMATED = "estimated" -@dataclass(frozen=True, slots=True) -class UsageCaptureObservation: - usage_json: str | None - provenance: UsageObservation - response_cost: float | None = None - - @classmethod - def from_usage( - cls, usage: Usage | None, provenance: UsageObservation, response_cost: float | None = None - ) -> "UsageCaptureObservation | UsageCaptureError": - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError - - try: - serialized: Final = ( - None if usage is None or provenance is UsageObservation.UNKNOWN else usage.model_dump_json() - ) - except Exception as error: # noqa: BLE001 # observation failures must not interrupt legacy retry handling - return UsageCaptureError(str(error)) - else: - return cls(serialized, provenance, response_cost) - - class Logging(LiteLLMLoggingBaseClass): - usage_capture_callback: ClassVar[ - Callable[ - ["Logging", object, datetime.datetime, datetime.datetime, bool], "CapturedUsage | UsageCaptureError | None" - ] - | None - ] = None - - def capture_usage_before_logging( - self, - response: object, - start_time: datetime.datetime | None = None, - end_time: datetime.datetime | None = None, - *, - freeze_only: bool = False, - ) -> None: - capture: Final = type(self).usage_capture_callback - if capture is not None and (not self.stream or self._is_assembled_stream_success(response)): - capture( - self, - response, - start_time or self.start_time, - end_time or datetime.datetime.now(datetime.timezone.utc), - freeze_only, - ) - global \ supabaseClient, \ promptLayerLogger, \ @@ -949,7 +900,6 @@ class Logging(LiteLLMLoggingBaseClass): self._stash_owed_on_success: bool = False self._stash_billed_by_continuation: bool = False self._failed_attempts_owed_on_success: tuple[ServedRound, ...] = () - self._failed_attempt_ids_owed_on_success: tuple[str, ...] = () self.failed_attempts_billed_on_success: tuple[ServedRound, ...] = () self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None @@ -2582,11 +2532,6 @@ class Logging(LiteLLMLoggingBaseClass): ) return - self.capture_usage_before_logging( - TypeAdapter(object).validate_python(result), - TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(start_time), - TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(end_time), - ) await self.async_success_handler( result, start_time=start_time, @@ -2698,20 +2643,12 @@ class Logging(LiteLLMLoggingBaseClass): self.finalized_cost = finalized self.model_call_details["response_cost"] = finalized.amount - def record_usage_capture(self, captured: "CapturedUsage") -> None: - self.model_call_details["_accounting_capture"] = captured - - def record_usage_pricing(self, pricing: object) -> None: - self.model_call_details["_accounting_pricing_snapshot"] = pricing - def record_partial_usage_for_failure( self, usage: Usage, response_cost: float, *, observation: UsageObservation = UsageObservation.KNOWN ) -> None: """Stash what an interrupted stream already consumed so the failure log bills it instead of zero.""" self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost - if type(self).usage_capture_callback is not None: - self.model_call_details["_accounting_usage_observed"] = observation def record_failure_after_earlier_rounds( self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float @@ -2795,12 +2732,6 @@ class Logging(LiteLLMLoggingBaseClass): owed_on_success: Final = self._stash_owed_on_success self._stash_billed_by_continuation = False self._stash_owed_on_success = False - captured_attempt_ids: Final = ( - self.capture_failed_attempt_usage(usage) - if isinstance(usage, Usage) and type(self).usage_capture_callback is not None - else () - ) - self.model_call_details.pop("_accounting_capture", None) self.model_call_details.pop("_accounting_pricing_snapshot", None) self.model_call_details.pop("_accounting_usage_observed", None) self.model_call_details.pop("combined_usage_object", None) @@ -2816,60 +2747,10 @@ class Logging(LiteLLMLoggingBaseClass): if not owed_on_success or billed_by_continuation: return self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt) - self._failed_attempt_ids_owed_on_success = (*self._failed_attempt_ids_owed_on_success, *captured_attempt_ids) def mark_stash_billed_by_continuation(self) -> None: self._stash_billed_by_continuation = True - def capture_failed_attempt_usage(self, usage: Usage) -> tuple[str, ...]: - from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import CapturedUsage - - capture: Final = type(self).usage_capture_callback - if capture is None: - return () - attempt: Final = self.attempt_scoped_copy() - attempt.prepared_cost = None - attempt.finalized_cost = None - identity: Final = ( - f"{self.litellm_call_id}-attempt-{len(self.model_call_details.get('failed_attempts_usage', ()))}" - ) - params: Final = _LITELLM_PARAMS.validate_python(self.model_call_details.get("litellm_params") or {}) - metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(params) - observed: Final[object] = self.model_call_details.get("_accounting_usage_observed") - cost: Final[object] = self.model_call_details.get("response_cost") - attempt.model_call_details = { - **attempt.model_call_details, - "litellm_call_id": identity, - "parent_request_id": self.litellm_call_id, - "_accounting_capture": None, - "_accounting_prior_captures": (), - "_accounting_usage_snapshot": UsageCaptureObservation.from_usage( - usage, - observed if isinstance(observed, UsageObservation) else UsageObservation.KNOWN, - float(cost) if isinstance(cost, (int, float)) else None, - ), - "litellm_params": { - **params, - "metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, - "litellm_metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, - }, - } - captured: Final = capture( - attempt, - {"id": identity, "usage": usage.model_dump(), "status": "failed"}, - self.start_time, - datetime.datetime.now(datetime.timezone.utc), - False, - ) - if not isinstance(captured, CapturedUsage): - return () - self.model_call_details["_accounting_prior_captures"] = ( - *self.model_call_details.get("_accounting_prior_captures", ()), - captured, - ) - return tuple(write.event.event_id for write in captured.writes) - def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" self.model_call_details["continued_partial_usage"] = usage @@ -2909,8 +2790,6 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details.pop("failed_attempts_usage", None) self.failed_attempts_billed_on_success = failed_attempts self.model_call_details["failed_attempts_billed_on_success"] = failed_attempts - if self._failed_attempt_ids_owed_on_success: - self.model_call_details["_accounting_included_attempt_ids"] = self._failed_attempt_ids_owed_on_success self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects( usage_objects=[own_usage, *(usage for usage, _ in failed_attempts)] ) @@ -3073,45 +2952,6 @@ class Logging(LiteLLMLoggingBaseClass): return self._handle_a2a_response_logging(result=result) return result - def usage_capture_response(self, result: object) -> object: - if self.stream and not self._is_assembled_stream_success(result): - return None - copied: Final = result.model_copy(deep=True) if isinstance(result, BaseModel) else result - normalized: Final = self.normalize_logging_result(self.normalize_endpoint_logging_result(copied)) - if self.model_call_details.get("call_type") == CallTypes.arealtime.value and isinstance(result, list): - frames: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(result) - measured: Final = any( - isinstance(_LITELLM_PARAMS.validate_python(frame.get("response") or {}).get("usage"), Mapping) - for frame in frames - if frame.get("type") == "response.done" - ) - if not measured: - metered: Final = self.model_call_details.get(REALTIME_METERED_USAGE_KEY) - return { - "id": self.litellm_call_id, - "usage": {"metered_usage": TypeAdapter(object).dump_python(metered, mode="json")} - if metered is not None - else None, - } - assembled: Final = self._get_assembled_streaming_response( - normalized, self.start_time, self.start_time, True, [] - ) - return normalized if assembled is None else assembled - - def usage_capture_pricing(self, result: object) -> Mapping[str, JsonValue]: - from litellm.utils import count_characters - - return TypeAdapter(dict[str, JsonValue]).validate_python( - { - "custom_pricing": self.custom_pricing_for(result), - "custom_cost_per_token": self.custom_cost_per_token, - "custom_cost_per_second": self.custom_cost_per_second, - "prompt_characters": count_characters(self._prompt_for_cost_calculation()) - if self.model_call_details.get("call_type") in (CallTypes.speech.value, CallTypes.aspeech.value) - else None, - } - ) - def normalize_logging_result(self, result: object) -> object: """ Some endpoints return a different type of result than what is expected by the logging system. @@ -3826,7 +3666,6 @@ class Logging(LiteLLMLoggingBaseClass): self._as_served_chat_response(complete_streaming_response) ) self.model_call_details["complete_streaming_response"] = logged_streaming_response - self.capture_usage_before_logging(logged_streaming_response, start_time, end_time) ## STANDARDIZED LOGGING PAYLOAD if not for_async_request: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( @@ -4352,8 +4191,6 @@ class Logging(LiteLLMLoggingBaseClass): cost_for_built_in_tools_cost_usd_dollar=0.0, ) - self.capture_usage_before_logging(result, start_time, end_time) - self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( kwargs=self.model_call_details, messages=self.model_call_details.get("messages") diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 8e6e71a31d3..b6af181980e 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -836,8 +836,6 @@ class RealTimeStreaming: if self.session_tools or self.tool_calls: self.logging_obj.model_call_details["realtime_tools"] = self.session_tools self.logging_obj.model_call_details["realtime_tool_calls"] = self.tool_calls - if LiteLLMLogging.usage_capture_callback is not None: - self.logging_obj.capture_usage_before_logging(self.messages) # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index de1096701af..9d84770b288 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -38,7 +38,6 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_no_callback_env_reference, ) from litellm.proxy.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled -from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings from litellm.types.agents import ( AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export ) @@ -2481,22 +2480,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ - accounting_capture: UsageCaptureSettings | None = Field( - None, - description="Opt-in durable usage capture. Legacy spend updates remain active. Requires local Prisma writes.", - ) - accounting_sync_ack: bool = Field( - False, - description="Wait for usage commit on foreground non-stream chat, text, embedding and Responses requests. " - "Other requests remain supported with asynchronous usage capture.", - ) - accounting_sync_ack_timeout_seconds: float = Field( - 20.0, gt=0, description="Maximum response wait for usage commit, including the batch flush interval" - ) - accounting_sync_ack_failure_policy: Literal["fail_closed", "log"] = Field( - "fail_closed", description="Return 503 or log the persistence failure after the response wait expires" - ) - request_tag_mappings: Sequence[RequestTagMapping] | None = Field( None, description="Map W3C baggage values to request tags. Exact value/tag rules take precedence over tag_prefix fallbacks.", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5ae6c628324..17e6d3ee411 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -140,7 +140,6 @@ from litellm.proxy.route_llm_request import ( route_request, ) from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost -from litellm.proxy.spend_tracking.usage_capture import capture_proxy_response from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ProxyLogging, @@ -3397,9 +3396,6 @@ class ProxyBaseLLMRequestProcessing: ) ), ) - await capture_proxy_response( - logging_obj, response, background=_REQUEST_FIELDS.validate_python(self.data).get("background") is True - ) except (Exception, asyncio.CancelledError) as exc: _exception_raised = True if logging_obj.prepared_cost is not None: diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py deleted file mode 100644 index 0a9c5db1619..00000000000 --- a/litellm/proxy/db/accounting_ledger.py +++ /dev/null @@ -1,352 +0,0 @@ -from __future__ import annotations - -import hashlib -import json -from collections.abc import AsyncGenerator, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager -from dataclasses import dataclass -from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final, Literal, TypeAlias - -from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter -from typing_extensions import assert_never - -from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent -from litellm.proxy.db.db_span import db_span -from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper - -if TYPE_CHECKING: - from prisma import Prisma - - from litellm.proxy.utils import PrismaClient - -AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent -_JSON: Final = TypeAdapter(dict[str, JsonValue]) -_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime) -_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...]) -_TIMEOUT: Final = timedelta(seconds=10) - - -class AccountingConflictError(ValueError): - pass - - -@dataclass(frozen=True, slots=True) -class UsageWrite: - event: InboundUsageEvent | OutboundUsageEvent - envelope: str - targets: tuple[DeliveryTarget, ...] = () - - -_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"] - - -@dataclass(frozen=True, slots=True) -class _AcceptedEvent: - event_id: str - kind: EventKind - source_event_id: str | None - occurred_at: datetime - payload: str - payload_hash: str - targets: tuple[str, ...] - schema_version: int - version: int - - -def _json(value: object) -> str: - return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) - - -def _prepare( - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...], - envelope: str, - schema_version: int, - version: int, -) -> _AcceptedEvent: - if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1: - raise ValueError("Accounting versions must be positive integers") - validated: Final = type(event).model_validate_json(event.model_dump_json()) - payload: Final = _JSON.validate_json(envelope) - if payload.get("event") != validated.model_dump(mode="json"): - raise ValueError("Accounting envelope must contain the matching public event") - destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets))) - serialized: Final = _json(payload) - identity: Final = _json((schema_version, version, destinations, payload)) - kind: Final = ( - EventKind.COST - if isinstance(validated, CostEvent) - else EventKind.INBOUND_USAGE - if isinstance(validated, InboundUsageEvent) - else EventKind.OUTBOUND_USAGE - ) - return _AcceptedEvent( - event_id=validated.event_id, - kind=kind, - source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None, - occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at, - payload=serialized, - payload_hash=hashlib.sha256(identity.encode()).hexdigest(), - targets=destinations, - schema_version=schema_version, - version=version, - ) - - -_INSERT_EVENT: Final = """ -INSERT INTO "LiteLLM_AccountingEvent" - (event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload) -VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb) -ON CONFLICT (event_id) DO NOTHING RETURNING event_id -""" -_INSERT_DELIVERIES: Final = """ -INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) -SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb) -""" -_LOCK_SOURCE: Final = """ -SELECT event_id FROM "LiteLLM_AccountingEvent" -WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE -""" -_EXISTING_COST: Final = """ -SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1 -""" -_LOCK_EVENT: Final = """ -SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE -""" -_INSERT_USAGE_BATCH: Final = """ -INSERT INTO "LiteLLM_AccountingEvent" - (event_id, kind, schema_version, payload_hash, occurred_at, payload) -SELECT event_id, kind::"AccountingEventKind", 1, payload_hash, occurred_at::timestamp, payload -FROM jsonb_to_recordset($1::jsonb) AS records( - event_id text, kind text, payload_hash text, occurred_at text, payload jsonb) -ORDER BY event_id -ON CONFLICT (event_id) DO NOTHING RETURNING event_id -""" -_LOCK_USAGE_BATCH: Final = """ -SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent" -WHERE event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) -ORDER BY event_id FOR UPDATE -""" -_LOAD_USAGE_BATCH: Final = """ -SELECT event.payload, ARRAY( - SELECT target::text FROM "LiteLLM_AccountingDelivery" delivery - WHERE delivery.event_id = event.event_id ORDER BY target::text -) AS targets -FROM "LiteLLM_AccountingEvent" event -WHERE event.event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) - AND event.kind = 'outbound_usage' AND event.schema_version = 1 -""" -_INSERT_USAGE_DELIVERIES: Final = """ -INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) -SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1 -FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb) -CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination -""" -_COMPLETED_FAMILY: Final = """ -root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') -AND root.recorded_at < $1::timestamp -AND (root.kind = 'inbound_usage' OR EXISTS ( - SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost' -)) -AND NOT EXISTS ( - SELECT 1 FROM "LiteLLM_AccountingEvent" AS member - WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id)) - AND (member.recorded_at >= $1::timestamp - OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery - WHERE delivery.event_id = member.event_id AND delivery.state != 'applied')) -) -""" -_RETENTION_CANDIDATES: Final = ( - 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE ' - + _COMPLETED_FAMILY - + " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED" -) -_ELIGIBLE_GROUP: Final = ( - 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY -) -_DELETE_DELIVERIES: Final = """ -DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN ( - SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) -) -""" -_DELETE_EVENTS: Final = """ -DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) -""" - - -_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"] - - -async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict: - if record.source_event_id is not None: - source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id)) - if not source: - return "missing_source" - costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id))) - if any(cost["event_id"] != record.event_id for cost in costs): - return "competing_cost" - inserted: Final = _STRINGS.validate_python( - await db.query_raw( - _INSERT_EVENT, - record.event_id, - record.kind.value, - record.source_event_id, - record.schema_version, - record.payload_hash, - record.occurred_at.astimezone(timezone.utc).isoformat(), - record.payload, - ) - ) - if not inserted: - existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id))) - if not existing or existing[0]["payload_hash"] != record.payload_hash: - return "identity_conflict" - return "duplicate" - await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version) - return "inserted" - - -async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int: - if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)): - return 0 - await db.execute_raw(_DELETE_DELIVERIES, event_id) - return await db.execute_raw(_DELETE_EVENTS, event_id) - - -class AccountingLedger: - def __init__(self, transaction: Callable[[], _Transaction]) -> None: - self._transaction: Final = transaction - - @classmethod - def for_client(cls, client: PrismaClient) -> AccountingLedger: - from prisma.client import TransactionManager - - managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True)) - - @asynccontextmanager - async def transaction() -> AsyncGenerator[Prisma]: - primary: Final = writer_wrapper(client.db) - manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT)) - async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db: - await db.execute_raw("SET LOCAL statement_timeout = 5000") - await db.execute_raw("SET LOCAL lock_timeout = 1000") - yield db - - return cls(transaction) - - async def insert_usage_batch(self, writes: tuple[UsageWrite, ...]) -> None: - records: Final = tuple(_prepare(write.event, write.targets, write.envelope, 1, 1) for write in writes) - if not records: - return - unique: Final = {record.event_id: record for record in records} - if any(record.payload_hash != unique[record.event_id].payload_hash for record in records): - raise AccountingConflictError("Batch contains conflicting copies of a usage identity") - batch: Final = _json( - tuple( - { - "event_id": record.event_id, - "kind": record.kind.value, - "payload_hash": record.payload_hash, - "occurred_at": record.occurred_at.astimezone(timezone.utc).isoformat(), - "payload": _JSON.validate_json(record.payload), - } - for record in unique.values() - ) - ) - async with self._transaction() as db: - inserted: Final = _STRINGS.validate_python(await db.query_raw(_INSERT_USAGE_BATCH, batch)) - accepted: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_USAGE_BATCH, _json(tuple(unique)))) - if len(accepted) != len(unique) or any( - row["payload_hash"] != unique[row["event_id"]].payload_hash for row in accepted - ): - raise AccountingConflictError("Usage identity conflicts with accepted content or delivery set") - deliveries: Final = _json( - tuple({"event_id": row["event_id"], "targets": unique[row["event_id"]].targets} for row in inserted) - ) - await db.execute_raw(_INSERT_USAGE_DELIVERIES, deliveries) - - async def load_usage_batch(self, event_ids: tuple[str, ...]) -> tuple[UsageWrite, ...]: - if not event_ids: - return () - async with self._transaction() as db: - rows: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( - await db.query_raw(_LOAD_USAGE_BATCH, _json(event_ids)) - ) - return tuple( - UsageWrite( - OutboundUsageEvent.model_validate_json(_json(_JSON.validate_python(row["payload"])["event"])), - _json(row["payload"]), - TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]), - ) - for row in rows - ) - - async def insert_event( - self, - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...], - *, - envelope: str, - schema_version: int = 1, - version: int = 1, - ) -> Literal["inserted", "duplicate"]: - """Accept T5/T6's serialized envelope, containing its public event under `event`. - - The immutable acceptance hash includes the target set and delivery version. - Changing either for an accepted event conflicts. T7 must provide a separate - add-target or new-version workflow that preserves receipts and prevents - duplicate contributions; redelivery cannot silently create or reset work. - - The Prisma adapter opens one interactive transaction per event, including - two SET LOCAL statements and a source-row lock for costs. I/O is async, - but awaiting this method inline adds commit latency per event, potentially - several times per request. T5 owns bounded asynchronous capture and an - explicit batching strategy; looping over this method does not batch commits. - Success acknowledges commit, independently of HTTP response acknowledgement. - """ - record: Final = _prepare(event, targets, envelope, schema_version, version) - async with self._transaction() as db: - result: Final = await _accept(db, record) - match result: - case "inserted" | "duplicate": - return result - case "missing_source": - raise AccountingConflictError("Cost source must be a durably accepted outbound usage event") - case "competing_cost": - raise AccountingConflictError("Usage event already has an authoritative cost") - case "identity_conflict": - raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set") - case _: - assert_never(result) - - async def purge_completed( - self, - *, - redelivery_before: datetime | None = None, - recovery_before: datetime | None = None, - limit: int = 100, - ) -> int: - """Delete at most `limit` completed families once both retention owners permit it. - - T8 must supply a coherent recovery watermark including outstanding holds. - Omitting either watermark disables deletion. Each usage row and its cost - remain together so a retained charge never loses its source or vice versa. - Applied receipts are deleted with the family, ending its retained dedup - protection. The caller must guarantee events recorded before - redelivery_before cannot be redelivered through supported paths. - T8's recovery_before must also protect checkpoints, - outstanding holds and unfinished settlements; age alone is insufficient. - """ - if type(limit) is not int or not 1 <= limit <= 1000: - raise ValueError("Accounting retention batch size must be between 1 and 1000") - if redelivery_before is None or recovery_before is None: - return 0 - cutoff: Final = min( - _INSTANT.validate_python(redelivery_before).astimezone(timezone.utc), - _INSTANT.validate_python(recovery_before).astimezone(timezone.utc), - ).isoformat() - async with self._transaction() as db: - roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit))) - deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots]) - return sum(deleted) diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index 4566b10b260..e6fbc299504 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -88,14 +88,14 @@ def _internal_call_origin(metadata: str) -> object: return decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) -def _date_from_start_time(start_time: datetime | str) -> str | None: +def date_from_start_time(start_time: datetime | str) -> str | None: if isinstance(start_time, datetime): return start_time.date().isoformat() return start_time[:10] if len(start_time) >= 10 else None def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransaction | None: - date: Final = _date_from_start_time(payload["startTime"]) + date: Final = date_from_start_time(payload["startTime"]) origin: Final = _internal_call_origin(payload["metadata"]) if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): return None diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 031bdfa19cd..10b6ec2f661 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -24,7 +24,6 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, - UsageCaptureObservation, UsageObservation, ) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost @@ -67,7 +66,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export should_store_prompts_and_responses_in_spend_logs, ) -from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, capture_callback, captured_callback_kwargs from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, @@ -375,26 +373,13 @@ def _attempt_spend_row_kwargs( "combined_usage_object": attempt.usage, "response_cost": attempt.response_cost, "_accounting_usage_observed": UsageObservation.ESTIMATED, - "_accounting_usage_snapshot": UsageCaptureObservation.from_usage( - attempt.usage, UsageObservation.ESTIMATED, attempt.response_cost - ), - "_accounting_prior_captures": (), } row.pop("billed_failed_attempts", None) - row.pop("_accounting_capture", None) call_id: Final = request_data.get("litellm_call_id") or litellm_params.get("litellm_call_id") if isinstance(call_id, str): row["litellm_call_id"] = f"{call_id}{suffix}" row["parent_request_id"] = call_id metadata["parent_request_id"] = call_id - captured_attempts: Final = request_data.get("_accounting_billed_attempts") - if isinstance(captured_attempts, tuple): - for captured in TypeAdapter(tuple[object, ...]).validate_python(captured_attempts): - if ( - isinstance(captured, CapturedUsage) - and captured.payload["litellm_call_id"] == row["litellm_call_id"] - ): - row["_accounting_capture"] = captured standard_logging_object: Final = request_data.get("standard_logging_object") if isinstance(standard_logging_object, dict): payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict @@ -480,9 +465,6 @@ class _ProxyDBLogger(CustomLogger): "Failed to invalidate budget reservation counters after failure release failed" ) - capture_metadata: Final = await self._try_capture_failure_usage( - request_data, original_exception, user_api_key_dict - ) if completed_cost and tracking_disabled: return @@ -519,10 +501,8 @@ class _ProxyDBLogger(CustomLogger): _error_information.pop("traceback", None) _metadata["error_information"] = _error_information - _metadata = ( # rebind-ok: enriched identity is shared with the earlier immutable capture - {**_metadata, **capture_metadata} - if capture_metadata is not None - else await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( + _metadata = ( # rebind-ok: enrich identity before building failure spend rows + await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( metadata=_metadata, original_exception=original_exception ) ) @@ -629,106 +609,6 @@ class _ProxyDBLogger(CustomLogger): ), ) - async def _failure_capture_metadata( - self, request_data: Mapping[str, object], exception: Exception, auth: UserAPIKeyAuth - ) -> Mapping[str, object] | None: - from litellm.proxy.proxy_server import usage_capture - - if usage_capture.meter is None: - return None - captured_identity: Final = ( - _batch_accounting_for(request_data.get("litellm_call_id")) is not None - and request_data.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is not REJECTED_BEFORE_REQUEST_SETUP - ) - current_identity: Final = { - **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth), - "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth), - } - metadata: Final = ( - { - **dict.fromkeys(current_identity), - **{ - key: value - for key, value in _REQUEST_DATA.validate_python( - get_litellm_metadata_from_kwargs(dict(request_data)) - ).items() - if key in current_identity or key == "tags" - }, - } - if captured_identity - else current_identity - ) - return await self._enrich_failure_metadata_unless_db_stalled( - metadata=metadata, original_exception=exception, resolve_missing_key_identity=not captured_identity - ) - - async def _try_capture_failure_usage( - self, - request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer - original_exception: Exception, - auth: UserAPIKeyAuth, - ) -> Mapping[str, object] | None: - try: - metadata: Final = await self._failure_capture_metadata(request_data, original_exception, auth) - self._capture_failure_usage(request_data, original_exception, auth, metadata) - except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing or reservation release - verbose_proxy_logger.exception("Usage capture failed on the request failure path") - return None - return metadata - - def _capture_failure_usage( - self, - request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer - original_exception: Exception, - auth: UserAPIKeyAuth, - metadata: Mapping[str, object] | None = None, - ) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.db.exception_handler import exception_chain - from litellm.proxy.proxy_server import usage_capture - from litellm.proxy.spend_tracking.usage_capture import UsageAcknowledgementError - - if any(isinstance(error, UsageAcknowledgementError) for error in exception_chain(original_exception)): - return - usage_meter: Final = usage_capture.meter - - if usage_meter is None: - return - logging_obj: Final = request_data.get("litellm_logging_obj") - details: Final[Mapping[str, object]] = ( - _str_keyed_fields(logging_obj.model_call_details) if isinstance(logging_obj, Logging) else {} - ) - params: Final = _str_keyed_fields(request_data.get("litellm_params") or details.get("litellm_params")) - failure_data: Final = { - **details, - **request_data, - "litellm_params": { - **params, - "metadata": { - **_str_keyed_fields(params.get("metadata")), - **( - metadata - if metadata is not None - else { - **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth), - "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth), - } - ), - "status": "failure", - }, - }, - } - start: Final = logging_obj.start_time if isinstance(logging_obj, Logging) else request_data.get("start_time") - if not isinstance(start, datetime): - return - end: Final = datetime.now(timezone.utc) - captured: Final = tuple( - capture_callback(usage_meter, row, original_exception, start, end) - for _, row in _failure_spend_rows(failure_data, 0.0, None, self._deployment_lookup) - ) - request_data["_accounting_capture"] = captured[0] # rebind-ok: share one receipt with the legacy writer - request_data["_accounting_billed_attempts"] = captured[1:] # rebind-ok: share attempt receipts - async def _write_failure_spend_rows( self, user_api_key_dict: UserAPIKeyAuth, @@ -914,7 +794,6 @@ class _ProxyDBLogger(CustomLogger): resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES, ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) - kwargs.update(await captured_callback_kwargs(kwargs, completion_response, start_time, end_time)) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) if accounting is not None and (budget_reservation is not None or litellm.apply_search_free_tier): accounting.write_started = True diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c00783854a7..87bc78730b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -877,8 +877,6 @@ from litellm.proxy.spend_tracking.spend_event_producer import ( SpendEventProducer, build_spend_event_producer, ) -from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, UsageCaptureRuntime -from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings from litellm.proxy.telemetry.endpoints import router as telemetry_router from litellm.proxy.telemetry.middleware import TelemetryMiddleware from litellm.proxy.telemetry.runtime import TelemetryRuntime @@ -1392,7 +1390,6 @@ async def _flush_spend_around_logging_drain_on_shutdown( finally: await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock)) await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock)) - await usage_capture.stop(GracefulShutdownManager.remaining(clock=clock)) await flush_spend() @@ -3309,39 +3306,9 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): spend_event_producer: SpendEventProducer | None = None -usage_capture: Final = UsageCaptureRuntime() - - -def _configure_usage_capture() -> None: - settings: Final = _GENERAL_SETTINGS_VIEW.validate_python(general_settings) - capture_settings: Final = settings.get("accounting_capture") - if settings.get("accounting_sync_ack") and capture_settings is None: - raise ValueError("accounting_sync_ack requires accounting_capture") - if capture_settings is None: - if usage_capture.meter is not None: - raise ValueError("Disabling accounting_capture requires a worker restart") - return - if ( - prisma_client is None - or CollectorSettings().enabled - or os.getenv("SPEND_LOGS_URL") - or litellm.apply_search_free_tier - ): - raise ValueError( - "Usage capture requires local Prisma writes and disabled collector, remote writer and search free tier" - ) - from litellm.proxy.db.accounting_ledger import AccountingLedger - - ledger: Final = AccountingLedger.for_client(prisma_client) - writer: Final = UsageBatchWriter(ledger.insert_usage_batch, ledger.load_usage_batch) - usage_capture.start( - writer.write, UsageCaptureSettings.model_validate(capture_settings), load=ledger.load_usage_batch - ) - def cost_tracking(): global prisma_client, spend_event_producer - _configure_usage_capture() if prisma_client is not None: from litellm.integrations.shadow_eval_logger import ShadowEvalLogger diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 696f4e1ec48..eb8c28fcbd7 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -3,7 +3,7 @@ import json import os import re import secrets -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt from functools import reduce @@ -58,7 +58,6 @@ from litellm.litellm_core_utils.internal_call_metadata import ( ) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, - UsageObservation, coerce_model_access_groups, is_valid_sha256_hash, request_model_access_groups_from_litellm_params, @@ -592,7 +591,6 @@ def get_logging_payload( start_time: datetime, end_time: datetime, llm_router: "Router | None" = None, - capture_usage: Callable[[SpendLogsPayload, str | None], None] | None = None, ) -> SpendLogsPayload: if kwargs is None: kwargs = {} @@ -982,14 +980,6 @@ def get_logging_payload( payload.get("spend"), ) - if capture_usage is not None: - observed_usage: Final = ( - StandardLoggingPayloadSetup.get_usage_as_dict({"usage": usage}) - if usage and kwargs.get("_accounting_usage_observed") is not UsageObservation.UNKNOWN - else None - ) - capture_usage(payload, safe_dumps(observed_usage) if observed_usage is not None else None) - # Explicitly clear large intermediate objects to reduce memory pressure del response_obj_dict, usage, clean_metadata, additional_usage_values diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py deleted file mode 100644 index 19480a66678..00000000000 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ /dev/null @@ -1,1084 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import re -from collections.abc import Awaitable, Callable, Iterator, Mapping -from dataclasses import dataclass, replace -from datetime import datetime, timezone -from decimal import Decimal -from hashlib import sha256 -from itertools import chain, groupby, islice, repeat -from typing import TYPE_CHECKING, Final - -from pydantic import BaseModel, JsonValue, TypeAdapter - -from litellm._logging import verbose_proxy_logger -from litellm.accounting import ( - Dimension, - DimensionName, - InboundUsageEvent, - Metric, - MetricName, - OutboundUsageEvent, - Outcome, -) -from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY -from litellm.proxy.spend_tracking.usage_meter import ( - DurableUsageMeter, - UsageCaptureError, - UsageCaptureSettings, - usage_delivery_targets, -) -from litellm.types.utils import CallTypes, CustomPricingLiteLLMParams - -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy._types import SpendLogsPayload - - -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite - - -class UsageAcknowledgementError(UsageCaptureError): - def __init__(self, cause: Exception) -> None: - super().__init__("Usage persistence did not acknowledge commit") - self.__cause__ = cause - - -@dataclass(frozen=True, slots=True) -class CapturedUsage: - payload: SpendLogsPayload - writes: tuple[UsageWrite, ...] - receipt: asyncio.Future[None] - - -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) - - -_FIELDS: Final = TypeAdapter(dict[str, object]) - - -_INSTANT: Final = TypeAdapter(datetime) - - -_DIMENSIONS: Final = ( - ("api_key", DimensionName.KEY), - ("user", DimensionName.USER), - ("team_id", DimensionName.TEAM), - ("organization_id", DimensionName.ORGANIZATION), - ("end_user", DimensionName.END_USER), - ("model", DimensionName.MODEL), - ("model_group", DimensionName.MODEL_GROUP), - ("model_id", DimensionName.DEPLOYMENT), - ("custom_llm_provider", DimensionName.PROVIDER), - ("mcp_namespaced_tool_name", DimensionName.TOOL), - ("agent_id", DimensionName.AGENT), -) - - -_USAGE_METRICS: Final = tuple( - metric - for metric in MetricName - if metric not in (MetricName.SPEND, MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS) -) - - -def _fields(value: object) -> Mapping[str, object]: - return _FIELDS.validate_python(value) if isinstance(value, Mapping) else {} - - -def _text(value: object) -> str | None: - return value if isinstance(value, str) and value else None - - -def _utc_time(value: datetime | str) -> datetime: - parsed: Final = _INSTANT.validate_python(value) - return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc) - - -def _json(value: object) -> str: - return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) - - -def _usage_metrics(usage: Mapping[str, JsonValue] | None) -> tuple[Metric, ...] | None: - if usage is None: - return None - return tuple( - Metric(name=metric, value=Decimal(str(value))) - for metric in _USAGE_METRICS - if isinstance(value := usage.get(metric.value), (int, float)) and not isinstance(value, bool) - ) - - -@dataclass(frozen=True, slots=True) -class UsageCaptureContext: - caller_id: str | None - attempt_id: str - parent_request_id: str | None - retry_of_attempt_id: str | None - job_id: str | None - charge_eligible: bool - request_count_eligible: bool - ownership: str - pricing: str - finalized_cost: str | None = None - allocation: str | None = None - outcome: Outcome | None = None - started_at: datetime | None = None - ended_at: datetime | None = None - phase: str = "execution" - extra_dimensions: tuple[Dimension, ...] = () - provider_timestamps: str = "{}" - usage_source: str = "measured" - included_attempt_ids: tuple[str, ...] = () - related_attempt_ids: tuple[str, ...] = () - - -def _database_usage_write(event: InboundUsageEvent | OutboundUsageEvent, envelope: Mapping[str, object]) -> UsageWrite: - from sys import maxsize - - from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure - - normalized: Final = _FIELDS.validate_python(safe_json_structure(envelope, max_depth=maxsize)) - if normalized.get("event") != event.model_dump(mode="json"): - raise UsageCaptureError("Usage event text cannot be normalized without changing identity or attribution") - serialized: Final = _json(normalized) - _JSON_OBJECT.validate_json(serialized) - return UsageWrite(event, serialized, usage_delivery_targets(event.dimensions)) - - -def usage_writes( - payload: SpendLogsPayload, usage_json: str | None, context: UsageCaptureContext -) -> tuple[UsageWrite, ...]: - fields: Final = _FIELDS.validate_python(payload) - tags: Final = TypeAdapter(tuple[str, ...]).validate_json(payload["request_tags"]) - dimensions: Final = ( - *(Dimension(name=name, value=value) for field, name in _DIMENSIONS if (value := _text(fields.get(field)))), - *(Dimension(name=DimensionName.TAG, value=tag) for tag in dict.fromkeys(tags)), - *context.extra_dimensions, - ) - outcome: Final = context.outcome or (Outcome.FAILED if payload["status"] == "failure" else Outcome.SUCCEEDED) - started: Final = context.started_at or _utc_time(payload["startTime"]) - ended: Final = context.ended_at or _utc_time(payload["endTime"]) - usage: Final = None if usage_json is None else _JSON_OBJECT.validate_json(usage_json) - outbound: Final = OutboundUsageEvent( - event_id=context.attempt_id, - attempt_id=context.attempt_id, - request_id=context.caller_id, - retry_of_attempt_id=context.retry_of_attempt_id, - call_type=CallTypes(payload["call_type"]), - started_at=started, - ended_at=ended, - outcome=outcome, - dimensions=dimensions, - metrics=_usage_metrics(None if context.usage_source == "unknown" else usage), - ) - envelope: Final = { - "event": outbound.model_dump(mode="json"), - "usage": usage, - "usage_source": "unknown" if usage is None else context.usage_source, - "included_attempt_ids": context.included_attempt_ids, - "related_attempt_ids": context.related_attempt_ids, - "pricing": _JSON_OBJECT.validate_json(context.pricing), - "charge_eligible": context.charge_eligible, - "request_count_eligible": context.request_count_eligible, - "ownership": context.ownership, - "parent_request_id": context.parent_request_id, - "job_id": context.job_id, - "phase": context.phase, - "provider_timestamps": _JSON_OBJECT.validate_json(context.provider_timestamps), - "finalized_cost": None - if context.finalized_cost is None - else _JSON_OBJECT.validate_json(context.finalized_cost), - "allocation": None if context.allocation is None else _JSON_OBJECT.validate_json(context.allocation), - } - outbound_write: Final = _database_usage_write(outbound, envelope) - if not context.request_count_eligible or context.caller_id is None: - return (outbound_write,) - inbound: Final = InboundUsageEvent( - event_id=f"{context.caller_id}:inbound", - request_id=context.caller_id, - call_type=outbound.call_type, - started_at=started, - ended_at=ended, - outcome=outcome, - dimensions=dimensions, - metrics=( - Metric(name=MetricName.REQUESTS, value=Decimal(1)), - Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(int(outcome is Outcome.SUCCEEDED))), - Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(int(outcome is not Outcome.SUCCEEDED))), - ), - ) - return ( - _database_usage_write(inbound, {"event": inbound.model_dump(mode="json")}), - outbound_write, - ) - - -_INCLUDED_IN_PARENT: Final = object() - - -_PRICING_FIELDS: Final = frozenset( - { - "service_tier", - "data_residency", - "vertex_location", - "region_name", - "base_model", - "custom_pricing", - "input_cost_per_token", - "output_cost_per_token", - "input_cost_per_second", - "output_cost_per_second", - "input_cost_per_character", - "output_cost_per_character", - "size", - "quality", - "n", - "duration", - "seconds", - "search_context_size", - "additional_response_cost", - } -) - - -def _selected_pricing(fields: Mapping[str, object]) -> Mapping[str, JsonValue]: - return _JSON_OBJECT.validate_python( - { - key: value - for key, value in fields.items() - if key in _PRICING_FIELDS or key in CustomPricingLiteLLMParams.model_fields - } - ) - - -def _response_fields(response: object) -> Mapping[str, object]: - if isinstance(response, BaseModel): - return _FIELDS.validate_python(response.model_dump()) - return _fields(response) - - -def _job_time(response: Mapping[str, object], key: str) -> datetime | None: - stamp: Final = response.get(key) - if isinstance(stamp, (int, float)) and not isinstance(stamp, bool): - return datetime.fromtimestamp(stamp, tz=timezone.utc) - return None - - -def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], response: object) -> UsageCaptureContext: - from dataclasses import asdict - - from litellm.litellm_core_utils.hidden_params import get_hidden_params - from litellm.litellm_core_utils.internal_call_metadata import is_background_response, is_unbilled_non_inference_call - from litellm.litellm_core_utils.litellm_logging import ( - Logging, - UsageCaptureObservation, - UsageObservation, - mask_api_base_credentials, - ) - from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost - from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation - from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups - - metadata: Final = _JSON_OBJECT.validate_json(payload["metadata"]) - params: Final = _fields(kwargs.get("litellm_params")) - request_metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) - standard: Final = _fields(kwargs.get("standard_logging_object")) - logging_obj: Final = kwargs.get("litellm_logging_obj") - raw_response: Final = _response_fields(response) - hidden: Final = _fields(get_hidden_params(response)) - tool: Final = _fields(kwargs.get("mcp_tool_call_metadata")) - tool_pricing: Final = _fields(tool.get("mcp_server_cost_info")) - caller: Final = payload["litellm_call_id"] - cached: Final = payload["cache_hit"].lower() == "true" - included_in_parent: Final = kwargs.get("_accounting_internal_aggregate") is _INCLUDED_IN_PARENT - billed_id: Final = payload["request_id"] - terminal: Final = raw_response.get("status") in { - "completed", - "complete", - "failed", - "cancelled", - "expired", - "incomplete", - } - deferred: Final = billed_id.endswith(("_batch_cost", "_background_cost")) and terminal - management: Final = ( - (billed_id.endswith(("_batch_cost", "_background_cost")) and not terminal) - or (is_background_response(response) and not terminal) - or is_unbilled_non_inference_call(payload["call_type"], request_metadata, response) - ) - attempt: Final = ( - f"{caller}_cache_hit" - if cached and caller - else f"{caller}:status" - if management - else f"{billed_id}:included" - if included_in_parent - else billed_id - ) - if not attempt or attempt == "None" or (cached and not caller): - raise UsageCaptureError("Usage capture requires a stable execution identity") - retry: Final = metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is True - origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) - parent: Final = _text(kwargs.get("parent_request_id")) or _text(request_metadata.get("parent_request_id")) - finalized_value: Final = ( - logging_obj.finalized_cost if isinstance(logging_obj, Logging) else kwargs.get("_accounting_finalized_cost") - ) - finalized: Final = finalized_value if isinstance(finalized_value, FinalizedCost) else None - allocation_value: Final = ( - logging_obj.cost_allowance_allocation - if isinstance(logging_obj, Logging) - else kwargs.get("_accounting_allocation") - ) - allocation: Final = allocation_value if isinstance(allocation_value, SearchAllowanceAllocation) else None - pricing_snapshot: Final = kwargs.get("_accounting_pricing_snapshot") - observation: Final = kwargs.get("_accounting_usage_snapshot") - captured_pricing: Final = { - "model_map_information": ( - _JSON_OBJECT.validate_json(pricing_snapshot.model_map_information) - if isinstance(pricing_snapshot, PricingCapture) - else metadata.get("model_map_information") - ), - "cost_breakdown": metadata.get("cost_breakdown"), - **( - _JSON_OBJECT.validate_json(pricing_snapshot.inputs) - if isinstance(pricing_snapshot, PricingCapture) - else _pricing_inputs(kwargs, logging_obj, response) - ), - "tool": { - "name": tool.get("name"), - "cost_per_query": _fields(tool_pricing.get("tool_name_to_cost_per_query")).get( - _text(tool.get("name")) or "", tool_pricing.get("default_cost_per_query") - ), - } - if tool - else None, - "response_inputs": { - "audio_transcription_duration": hidden.get("audio_transcription_duration", raw_response.get("duration")), - "billed_units": _fields(raw_response.get("meta")).get("billed_units"), - "batch_models": hidden.get("batch_models"), - "batch_successful_requests": hidden.get("batch_successful_requests"), - "batch_failed_requests": hidden.get("batch_failed_requests"), - "service_tier": raw_response.get("service_tier"), - "region_name": hidden.get("region_name"), - }, - "reported_cost": observation.response_cost - if isinstance(observation, UsageCaptureObservation) - else standard.get( - "response_cost", kwargs.get("response_cost", _fields(get_hidden_params(response)).get("response_cost")) - ), - "adjustments": ( - _JSON_OBJECT.validate_json(pricing_snapshot.adjustments) - if isinstance(pricing_snapshot, PricingCapture) - else None - ), - } - return UsageCaptureContext( - caller_id=None if deferred else parent if retry or included_in_parent else caller, - attempt_id=attempt, - parent_request_id=parent, - retry_of_attempt_id=_text(request_metadata.get("retry_of_attempt_id")), - job_id=_text(raw_response.get("id")) if deferred or management else None, - charge_eligible=not cached and not included_in_parent and not management, - request_count_eligible=bool(caller) and not origin and not retry and not deferred and not included_in_parent, - ownership="included_in_parent" - if included_in_parent - else "aggregate" - if kwargs.get("additional_response_cost") is not None - else "attempt", - pricing=_json(captured_pricing), - finalized_cost=_json( - { - "response_cost": finalized.amount, - "cost_breakdown": finalized.logging_breakdown(), - "allowance_discount": finalized.allowance_discount, - } - ) - if finalized is not None - else ( - _json({"response_cost": standard.get("response_cost"), "cost_breakdown": standard.get("cost_breakdown")}) - if standard.get("cost_finalized") is True - else None - ), - allocation=_json(asdict(allocation)) if allocation is not None else None, - started_at=_job_time(raw_response, "created_at") if deferred else None, - ended_at=( - _job_time(raw_response, "completed_at") - or _job_time(raw_response, "failed_at") - or _job_time(raw_response, "cancelled_at") - or _job_time(raw_response, "expired_at") - ) - if deferred - else None, - outcome=( - Outcome.CANCELLED - if raw_response.get("status") == "cancelled" - or isinstance(response, asyncio.CancelledError) - or getattr(response, "status_code", None) == 499 - else Outcome.FAILED - if raw_response.get("status") in {"failed", "expired", "incomplete"} or isinstance(response, BaseException) - else None - ), - phase="terminal" if deferred else "management" if management else "execution", - usage_source="estimated" - if kwargs.get("_accounting_usage_observed") is UsageObservation.ESTIMATED - else "measured", - included_attempt_ids=TypeAdapter(tuple[str, ...]).validate_python( - kwargs.get("_accounting_included_attempt_ids", ()) - ), - related_attempt_ids=tuple(write.event.event_id for write in _prior_writes(kwargs)), - provider_timestamps=_json( - { - key: raw_response.get(key) - for key in ("created_at", "completed_at", "failed_at", "cancelled_at", "expired_at") - if raw_response.get(key) is not None - } - ) - if deferred - else "{}", - extra_dimensions=( - *( - (Dimension(name=DimensionName.TOOL, value=tool_name),) - if not payload.get("mcp_namespaced_tool_name") - and (tool_name := _text(tool.get("namespaced_tool_name"))) - else () - ), - *( - (Dimension(name=DimensionName.API_BASE, value=mask_api_base_credentials(api_base)),) - if (api_base := _text(payload["api_base"])) - else () - ), - *( - Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value=group) - for group in get_request_model_access_groups(kwargs) - ), - *( - (Dimension(name=DimensionName.PROJECT, value=project),) - if (project := _text(request_metadata.get("user_api_key_project_id"))) - else () - ), - *( - (Dimension(name=DimensionName.SERVICE_TIER, value=tier),) - if (tier := _text(raw_response.get("service_tier") or kwargs.get("service_tier"))) - else () - ), - ), - ) - - -class _CaptureBuilder: - def __init__(self, kwargs: Mapping[str, object], response: object) -> None: - self.kwargs: Final = kwargs - self.response: Final = response - self.writes: tuple[UsageWrite, ...] = () - - def observe(self, payload: SpendLogsPayload, usage: str | None) -> None: - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation, UsageObservation - - snapshot: Final = self.kwargs.get("_accounting_usage_snapshot") - if isinstance(snapshot, UsageCaptureError): - raise snapshot - effective: Final = snapshot.usage_json if isinstance(snapshot, UsageCaptureObservation) else usage - parsed: Final = _JSON_OBJECT.validate_json(effective) if effective is not None else {} - batch: Final = _fields(parsed.get("batch_usage_by_model")) - unknown: Final = (bool(batch) or bool(parsed.get("batch_unknown_usage_requests"))) and all( - value is None for value in batch.values() - ) - observed: Final = ( - _json( - { - "batch_usage_by_model": batch, - "batch_unknown_usage_requests": parsed.get("batch_unknown_usage_requests"), - **( - {"batch_usage_by_request": parsed["batch_usage_by_request"]} - if "batch_usage_by_request" in parsed - else {} - ), - } - ) - if unknown - else effective - ) - context: Final = capture_context(payload, self.kwargs, self.response) - provenance: Final = ( - "unknown" - if unknown - else "estimated" - if isinstance(snapshot, UsageCaptureObservation) and snapshot.provenance is UsageObservation.ESTIMATED - else context.usage_source - ) - self.writes = usage_writes(payload, observed, replace(context, usage_source=provenance)) - - -def _prior_captures(kwargs: Mapping[str, object]) -> tuple[CapturedUsage, ...]: - prior: Final = kwargs.get("_accounting_prior_captures", ()) - return ( - tuple( - item for item in TypeAdapter(tuple[object, ...]).validate_python(prior) if isinstance(item, CapturedUsage) - ) - if isinstance(prior, tuple) - else () - ) - - -def _prior_writes(kwargs: Mapping[str, object]) -> tuple[UsageWrite, ...]: - return tuple(chain.from_iterable(item.writes for item in _prior_captures(kwargs))) - - -def capture_callback( - meter: DurableUsageMeter, - kwargs: Mapping[str, object], - response: object, - start_time: datetime, - end_time: datetime, -) -> CapturedUsage: - from litellm.proxy.db.db_spend_update_writer import get_llm_router - from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload - - existing: Final = kwargs.get("_accounting_capture") - if isinstance(existing, CapturedUsage): - return existing - builder: Final = _CaptureBuilder(kwargs, response) - payload: Final = get_logging_payload( - dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe - ) - prerequisites: Final = tuple(item.receipt for item in _prior_captures(kwargs)) - return CapturedUsage(payload, builder.writes, meter.submit(builder.writes, prerequisites=prerequisites)) - - -def _batch_rows(response: object) -> tuple[Mapping[str, JsonValue], ...]: - usage: Final = _response_fields(_response_fields(response).get("usage")) - rows: Final = usage.get("batch_usage_by_request") - return TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(rows) if rows else () - - -def _batch_root(write: UsageWrite, rows: tuple[Mapping[str, JsonValue], ...]) -> UsageWrite: - event: Final = write.event.model_copy(update={"metrics": None}) - envelope: Final = _JSON_OBJECT.validate_json(write.envelope) - return _database_usage_write( - event, - { - **envelope, - "event": event.model_dump(mode="json"), - "usage": { - "batch_row_count": len(rows), - "batch_rows_sha256": sha256(_json(rows).encode()).hexdigest(), - "batch_charge_eligible": envelope["charge_eligible"], - }, - "usage_source": "unknown", - "charge_eligible": False, - "ownership": "aggregate", - }, - ) - - -def _batch_row_write(root: UsageWrite, row: Mapping[str, JsonValue], index: int) -> UsageWrite: - envelope: Final = _JSON_OBJECT.validate_json(root.envelope) - pricing: Final = _JSON_OBJECT.validate_python(envelope["pricing"]) - model: Final = _text(row["model"]) - usage: Final = _JSON_OBJECT.validate_python(row["usage"]) if row["usage"] is not None else None - identity: Final = f"{root.event.event_id}:row:{index}" - event: Final = OutboundUsageEvent.model_validate( - { - **root.event.model_dump(), - "event_id": identity, - "attempt_id": identity, - "metrics": _usage_metrics(usage), - "dimensions": ( - *(dimension for dimension in root.event.dimensions if dimension.name is not DimensionName.MODEL), - *((Dimension(name=DimensionName.MODEL, value=model),) if model else ()), - ), - } - ) - return _database_usage_write( - event, - { - **envelope, - "event": event.model_dump(mode="json"), - "usage": usage, - "usage_source": "unknown" if usage is None else "measured", - "charge_eligible": _fields(envelope["usage"])["batch_charge_eligible"], - "ownership": "attempt", - "parent_request_id": root.event.event_id, - "finalized_cost": None, - "allocation": None, - "pricing": { - **pricing, - "model_map_information": _fields(pricing.get("batch_models")).get(model or ""), - "batch_models": {}, - "reported_cost": None, - "cost_breakdown": None, - "response_inputs": { - key: value - for key, value in _fields(pricing.get("response_inputs")).items() - if key not in {"batch_models", "batch_successful_requests"} - }, - }, - }, - ) - - -@dataclass(frozen=True, slots=True) -class _BatchCapture: - payload: SpendLogsPayload - root: UsageWrite - rows: bytes - prerequisites: tuple[asyncio.Future[None], ...] - - -def _prepare_batch_capture( - kwargs: Mapping[str, object], response: object, start_time: datetime, end_time: datetime -) -> _BatchCapture: - from litellm.proxy.db.db_spend_update_writer import get_llm_router - from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload - - rows: Final = tuple(sorted(_batch_rows(response), key=_json)) - builder: Final = _CaptureBuilder(kwargs, response) - payload: Final = get_logging_payload( - dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe - ) - if len(builder.writes) != 1 or _terminal_facts(builder.writes[0]) is None: - raise UsageCaptureError("Batch rows require one terminal job observation") - return _BatchCapture( - payload, - _batch_root(builder.writes[0], rows), - b"\n".join(_json((row, sum(1 for _ in group))).encode("utf-8") for row, group in groupby(map(_json, rows))), - tuple(prior.receipt for prior in _prior_captures(kwargs)), - ) - - -def _batch_groups(rows: bytes) -> Iterator[tuple[str, int]]: - return (TypeAdapter(tuple[str, int]).validate_json(match.group()) for match in re.finditer(rb"[^\n]+", rows)) - - -def _batch_writes(root: UsageWrite, rows: bytes) -> Iterator[UsageWrite]: - expanded: Final = chain.from_iterable(repeat(row, count) for row, count in _batch_groups(rows)) - return (_batch_row_write(root, _JSON_OBJECT.validate_json(row), index) for index, row in enumerate(expanded)) - - -def _batch_write_bound(root: UsageWrite, rows: bytes) -> int: - last_index: Final = sum(count for _, count in _batch_groups(rows)) - 1 - return max( - len(root.envelope.encode("utf-8")), - max( - len(_batch_row_write(root, _JSON_OBJECT.validate_json(row), last_index).envelope.encode("utf-8")) - for row, _ in _batch_groups(rows) - ), - ) - - -def _batch_capacity( - meter: DurableUsageMeter, root: UsageWrite, rows: bytes, *, extra_bytes: int = 0 -) -> tuple[int, int, int]: - retained: Final = len(rows) + len(root.envelope.encode("utf-8")) + extra_bytes - bound: Final = _batch_write_bound(root, rows) - chunk_bytes: Final = max( - bound, - min( - bound * min(meter.settings.batch_size, sum(count for _, count in _batch_groups(rows))), - meter.settings.max_pending_bytes // 4, - (meter.settings.max_pending_bytes - retained) // 2, - ), - ) - return retained + chunk_bytes, chunk_bytes, bound - - -async def _persist_batch_capture( - meter: DurableUsageMeter, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]], - root: UsageWrite, - rows: bytes, - prerequisites: tuple[asyncio.Future[None], ...], - chunk_bytes: int, - reservation: object | None = None, -) -> None: - for receipt in prerequisites: - await asyncio.shield(receipt) - await meter.record_writes((root,)) - accepted: Final = await load((root.event.event_id,)) - if len(accepted) != 1: - raise UsageCaptureError("Committed batch snapshot could not be loaded") - frozen: Final = _reuse_terminal_snapshot(root, accepted[0]) - retained, capacity, bound = _batch_capacity(meter, frozen, rows, extra_bytes=len(root.envelope.encode("utf-8"))) - if reservation is not None: - meter.resize_producer(reservation, retained, capacity) - write_capacity: Final = capacity if reservation is not None else chunk_bytes - if bound > write_capacity: - raise UsageCaptureError("Batch observation exceeds reserved producer write capacity") - group_size: Final = min(meter.settings.batch_size, max(1, write_capacity // bound)) - remaining: Final = _batch_writes(frozen, rows) - while group := tuple(islice(remaining, group_size)): - await meter.record_writes(group) - - -async def capture_batch_callback( - meter: DurableUsageMeter, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None, - kwargs: Mapping[str, object], - response: object, - start_time: datetime, - end_time: datetime, -) -> CapturedUsage: - if load is None: - raise UsageCaptureError("Batch capture requires the existing ledger snapshot reader") - capture: Final = _prepare_batch_capture(kwargs, response, start_time, end_time) - await _persist_batch_capture( - meter, load, capture.root, capture.rows, capture.prerequisites, meter.settings.max_pending_bytes - ) - receipt: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - receipt.set_result(None) - return CapturedUsage(capture.payload, (capture.root,), receipt) - - -async def captured_callback_kwargs( - kwargs: Mapping[str, object], response: object, start_time: datetime | None, end_time: datetime | None -) -> Mapping[str, object]: - from litellm.proxy.proxy_server import usage_capture - - usage_meter: Final = usage_capture.meter - - if usage_meter is None: - return kwargs - if start_time is None or end_time is None: - verbose_proxy_logger.error("Usage capture requires callback start and end timestamps") - return kwargs - try: - captured: Final = ( - usage_capture.submit_batch(kwargs, response, start_time, end_time) - if _batch_rows(response) - else capture_callback(usage_meter, kwargs, response, start_time, end_time) - ) - except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing - verbose_proxy_logger.exception("Usage capture failed before legacy cost tracking") - return kwargs - return {**kwargs, "_accounting_capture": captured} - - -async def capture_proxy_response(logging_obj: object, response: object, *, background: bool = False) -> None: - from fastapi import HTTPException - - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.proxy_server import general_settings, usage_capture - - settings: Final = _FIELDS.validate_python(general_settings) - usage_meter: Final = usage_capture.meter - - if usage_meter is None or not isinstance(logging_obj, Logging): - return - sync_ack: Final = ( - settings.get("accounting_sync_ack") is True - and logging_obj.stream is not True - and not background - and logging_obj.call_type in {"acompletion", "atext_completion", "aembedding", "aresponses"} - ) - try: - capture_logging_usage(usage_meter, logging_obj, response, logging_obj.start_time, datetime.now(timezone.utc)) - captured: Final = logging_obj.model_call_details.get("_accounting_capture") - if sync_ack: - if not isinstance(captured, CapturedUsage): - raise UsageCaptureError("Usage capture was not admitted") - timeout: Final = TypeAdapter(float).validate_python( - settings.get("accounting_sync_ack_timeout_seconds", 20.0) - ) - await asyncio.wait_for(asyncio.shield(captured.receipt), timeout=timeout) - except Exception as error: # noqa: BLE001 # the configured response policy owns every capture failure - verbose_proxy_logger.exception("Usage capture failed before the response boundary") - if sync_ack and settings.get("accounting_sync_ack_failure_policy", "fail_closed") == "fail_closed": - acknowledgement: Final = UsageAcknowledgementError(error) - raise HTTPException(503, str(acknowledgement)) from acknowledgement - - -def capture_logging_usage( - meter: DurableUsageMeter, - logging_obj: object, - response: object, - start_time: datetime, - end_time: datetime, - freeze_only: bool = False, -) -> CapturedUsage | UsageCaptureError | None: - from litellm._internal_context import internal_parent_request_id, is_internal_call - from litellm.batches.batch_utils import batch_cost_is_final - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import LiteLLMBatch - - if not isinstance(logging_obj, Logging): - return - internal: Final = is_internal_call.get() - if isinstance(response, LiteLLMBatch) and batch_cost_is_final(response) and response.usage is None: - return - try: - if not internal and isinstance(logging_obj.model_call_details.get("_accounting_capture"), CapturedUsage): - return capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time) - normalized_result: Final = logging_obj.usage_capture_response(response) - prepared: Final = logging_obj.prepared_cost - normalized: Final = ( - {**_response_fields(normalized_result), "usage": prepared.usage.model_dump()} - if prepared is not None - else normalized_result - ) - if normalized is None: - return - details: Final = _FIELDS.validate_python(logging_obj.model_call_details) - params: Final = _fields(details.get("litellm_params")) - metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) - if metadata.get("batch_ignore_default_logging") is True: - return - existing_pricing: Final = details.get("_accounting_pricing_snapshot") - pricing: Final = ( - existing_pricing - if isinstance(existing_pricing, PricingCapture) and not internal - else _capture_pricing(logging_obj, normalized, details) - ) - if freeze_only or _batch_rows(normalized): - if not internal: - logging_obj.record_usage_pricing(pricing) - return - captured: Final = capture_callback( - meter, - { - **details, - "litellm_logging_obj": logging_obj, - "_accounting_capture": None if internal else details.get("_accounting_capture"), - "_accounting_pricing_snapshot": pricing, - "_accounting_internal_aggregate": _INCLUDED_IN_PARENT if internal else None, - "parent_request_id": internal_parent_request_id.get() if internal else details.get("parent_request_id"), - }, - normalized, - start_time, - end_time, - ) - if not internal: - logging_obj.record_usage_capture(captured) - except Exception as error: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging - verbose_proxy_logger.exception("Usage capture failed at the execution boundary") - return UsageCaptureError(str(error)) - else: - return captured - - -@dataclass(frozen=True, slots=True) -class PricingCapture: - model_map_information: str - adjustments: str - inputs: str - - -def _capture_pricing(logging_obj: Logging, normalized: object, details: Mapping[str, object]) -> PricingCapture: - import litellm - from litellm.litellm_core_utils.cost_discount import resolve_cost_discount - from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup - - params: Final = _fields(details.get("litellm_params")) - prepared: Final = logging_obj.prepared_cost - provider: Final = _text(details.get("custom_llm_provider")) - builtins: Final = logging_obj.standard_built_in_tools_params - pricing: Final = ( - prepared.payload["model_map_information"] - if prepared is not None - else StandardLoggingPayloadSetup.get_model_cost_information( - base_model=logging_obj.get_deployment_model_for_cost(), - custom_pricing=bool(params.get("custom_pricing")), - custom_llm_provider=_text(details.get("custom_llm_provider")), - init_response_obj=normalized, - api_base=_text(params.get("api_base")), - router_model_id=logging_obj.get_router_model_id(), - ) - ) - adjustments: Final = _json( - { - "discount": resolve_cost_discount( - litellm.cost_discount_config, - provider, - logging_obj.get_deployment_model_for_cost(), - _text(params.get("region_name") or params.get("vertex_location")), - ), - "margin": litellm.cost_margin_config.get(provider or "", litellm.cost_margin_config.get("global")), - "web_search_options": ( - _selected_pricing(_fields(builtins.get("web_search_options"))) - if builtins.get("web_search_options") is not None - else None - ), - "file_search": builtins.get("file_search") is not None, - } - ) - batch_usage: Final = _fields(_fields(_response_fields(normalized).get("usage")).get("batch_usage_by_model")) - return PricingCapture( - _json(pricing), - adjustments, - _json( - { - **_pricing_inputs(details, logging_obj, normalized), - "batch_models": { - model: StandardLoggingPayloadSetup.get_model_cost_information( - base_model=model, - custom_pricing=bool(params.get("custom_pricing")), - custom_llm_provider=provider, - init_response_obj={"model": model}, - api_base=_text(params.get("api_base")), - router_model_id=logging_obj.get_router_model_id(), - ) - for model in batch_usage - }, - } - ), - ) - - -def _pricing_inputs(kwargs: Mapping[str, object], logging_obj: object, response: object) -> Mapping[str, object]: - from litellm.litellm_core_utils.litellm_logging import Logging - - return { - "request": _selected_pricing(kwargs), - "deployment": _selected_pricing(_fields(kwargs.get("litellm_params"))), - "optional_params": _selected_pricing(_fields(kwargs.get("optional_params"))), - "logging": logging_obj.usage_capture_pricing(response) if isinstance(logging_obj, Logging) else None, - "model_info": _selected_pricing( - _fields( - _fields( - _fields(kwargs.get("litellm_params")).get("litellm_metadata") - or _fields(kwargs.get("litellm_params")).get("metadata") - ).get("model_info") - ) - ), - } - - -def _terminal_facts(write: UsageWrite) -> Mapping[str, JsonValue] | None: - payload: Final = _JSON_OBJECT.validate_json(write.envelope) - if payload.get("phase") != "terminal": - return None - event: Final = _JSON_OBJECT.validate_python(payload["event"]) - return { - **{ - key: value - for key, value in payload.items() - if key not in {"event", "pricing", "finalized_cost", "allocation"} - }, - "event": {key: value for key, value in event.items() if key not in {"started_at", "ended_at"}}, - } - - -def _reuse_terminal_snapshot(candidate: UsageWrite, accepted: UsageWrite) -> UsageWrite: - facts: Final = _terminal_facts(candidate) - if facts is None or facts != _terminal_facts(accepted) or set(candidate.targets) != set(accepted.targets): - raise AccountingConflictError("Repeated terminal job disagrees with its accepted usage or attribution") - return accepted - - -class UsageBatchWriter: - def __init__( - self, - insert: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]], - ) -> None: - self._insert: Final = insert - self._load: Final = load - - async def _restore(self, writes: tuple[UsageWrite, ...]) -> tuple[UsageWrite, ...]: - terminal: Final = tuple(write for write in writes if _terminal_facts(write) is not None) - if not terminal: - return writes - identities: Final = tuple(dict.fromkeys(write.event.event_id for write in terminal)) - stored: Final = await self._load(identities) - # Re-read the first accepted observation and pricing before the ledger's unchanged id/hash check. - snapshots: Final = {write.event.event_id: write for write in (*reversed(terminal), *stored)} - return tuple( - _reuse_terminal_snapshot(write, snapshots[write.event.event_id]) - if write.event.event_id in snapshots - else write - for write in writes - ) - - async def write(self, writes: tuple[UsageWrite, ...]) -> None: - restored: Final = await self._restore(writes) - try: - await self._insert(restored) - except AccountingConflictError: - if not any(_terminal_facts(write) is not None for write in writes): - raise - await self._insert(await self._restore(writes)) - - -class UsageCaptureRuntime: - def __init__(self) -> None: - self._batches: Final[set[asyncio.Task[None]]] = set() # mutable-ok: owned producer lifecycle - self._stopping = False - self.meter: DurableUsageMeter | None = None - self.load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None - - def start( - self, - write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - settings: UsageCaptureSettings, - *, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None, - ) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - - if self.meter is not None: - if self.meter.settings != settings: - raise ValueError("Changing accounting_capture requires a worker restart") - return - meter: Final = DurableUsageMeter(write_batch, settings) - self.meter = meter - self.load = load - self._stopping = False - Logging.usage_capture_callback = lambda logger, response, start, end, freeze: capture_logging_usage( - meter, logger, response, start, end, freeze - ) - - def submit_batch( - self, kwargs: Mapping[str, object], response: object, start_time: datetime, end_time: datetime - ) -> CapturedUsage: - meter: Final = self.meter - load: Final = self.load - if meter is None or load is None or self._stopping: - raise UsageCaptureError("Batch capture requires an active ledger writer and snapshot reader") - if len(self._batches) >= meter.settings.max_pending_records: - raise UsageCaptureError("Batch capture producer capacity exhausted") - capture: Final = _prepare_batch_capture(kwargs, response, start_time, end_time) - retained, chunk_bytes, _ = _batch_capacity( - meter, capture.root, capture.rows, extra_bytes=len(capture.root.envelope.encode("utf-8")) - ) - reservation: Final = meter.reserve_producer(retained, chunk_bytes) - task: Final = asyncio.create_task( - _persist_batch_capture( - meter, load, capture.root, capture.rows, capture.prerequisites, chunk_bytes, reservation - ), - name="accounting batch capture", - ) - task.add_done_callback(lambda _: meter.release_producer(reservation)) - self._batches.add(task) - task.add_done_callback(self._batch_finished) - return CapturedUsage(capture.payload, (capture.root,), task) - - def _batch_finished(self, task: asyncio.Task[None]) -> None: - self._batches.discard(task) - if task.cancelled(): - verbose_proxy_logger.error("Batch usage capture stopped before complete persistence") - elif (error := task.exception()) is not None: - verbose_proxy_logger.error("Batch usage capture failed", exc_info=error) - - async def stop(self, timeout: float) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - - if self.meter is None: - return - self._stopping = True - deadline: Final = asyncio.get_running_loop().time() + timeout - try: - if self._batches: - await asyncio.wait(tuple(self._batches), timeout=max(0, timeout)) - finally: - remaining: Final = tuple(self._batches) - for task in remaining: - task.cancel() - await asyncio.gather(*remaining, return_exceptions=True) - try: - await self.meter.stop(max(0, deadline - asyncio.get_running_loop().time())) - finally: - Logging.usage_capture_callback = None - self.meter = None - self.load = None diff --git a/litellm/proxy/spend_tracking/usage_meter.py b/litellm/proxy/spend_tracking/usage_meter.py deleted file mode 100644 index 431e4c2c3c3..00000000000 --- a/litellm/proxy/spend_tracking/usage_meter.py +++ /dev/null @@ -1,330 +0,0 @@ -from __future__ import annotations - -import asyncio -from collections import deque -from collections.abc import Awaitable, Callable, Mapping -from concurrent.futures import Future -from dataclasses import dataclass -from itertools import accumulate, chain, count, takewhile -from threading import Lock -from typing import Final - -from pydantic import BaseModel, ConfigDict, Field - -from litellm._logging import verbose_proxy_logger -from litellm.accounting import ( - DeliveryTarget, - Dimension, - DimensionName, - InboundUsageEvent, - OutboundUsageEvent, -) -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite - -_DAILY_TARGETS: Final = ( - (DimensionName.KEY, DeliveryTarget.DAILY_USER_SPEND), - (DimensionName.USER, DeliveryTarget.DAILY_USER_SPEND), - (DimensionName.TEAM, DeliveryTarget.DAILY_TEAM_SPEND), - (DimensionName.ORGANIZATION, DeliveryTarget.DAILY_ORGANIZATION_SPEND), - (DimensionName.END_USER, DeliveryTarget.DAILY_END_USER_SPEND), - (DimensionName.AGENT, DeliveryTarget.DAILY_AGENT_SPEND), - (DimensionName.TAG, DeliveryTarget.DAILY_TAG_SPEND), - (DimensionName.MODEL, DeliveryTarget.DAILY_MODEL_USAGE), - (DimensionName.TOOL, DeliveryTarget.DAILY_TOOL_SPEND), -) - - -def usage_delivery_targets(dimensions: tuple[Dimension, ...]) -> tuple[DeliveryTarget, ...]: - names: Final = frozenset(dimension.name for dimension in dimensions) - return tuple(dict.fromkeys(target for dimension, target in _DAILY_TARGETS if dimension in names)) - - -class UsageCaptureSettings(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - batch_size: int = Field(default=1000, ge=2, le=1000) - flush_interval_seconds: float = Field(default=5.0, gt=0, le=10) - max_pending_bytes: int = Field(default=16 * 1024 * 1024, ge=1024) - max_pending_records: int = Field(default=10000, ge=2) - - -class UsageCaptureError(RuntimeError): - pass - - -@dataclass(frozen=True, slots=True) -class _Pending: - writes: tuple[UsageWrite, ...] - size: int - committed: Future[None] - finished: asyncio.Event - prerequisites: tuple[asyncio.Future[None], ...] - - -async def _wait_for_batch(wakeup: asyncio.Event, timeout: float) -> None: - try: - await asyncio.wait_for(wakeup.wait(), timeout) - except asyncio.TimeoutError: - return - - -def _report_failure(receipt: Future[None]) -> None: - if not receipt.cancelled() and receipt.exception() is not None: - verbose_proxy_logger.error("Usage capture did not acknowledge commit", exc_info=receipt.exception()) - - -def _retryable_persistence_failure(error: Exception) -> bool: - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain - - outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(error) - return (outage is not None and not PrismaDBExceptionHandler.is_permanent_database_fault(outage)) or any( - isinstance(cause, Exception) - and ( - PrismaDBExceptionHandler.is_deadlock_error(cause) - or PrismaDBExceptionHandler.is_lock_timeout_error(cause) - or PrismaDBExceptionHandler.postgres_sqlstate(cause) == "40001" - ) - for cause in exception_chain(error) - ) - - -def _isolatable_persistence_failure(error: Exception) -> bool: - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain - - return isinstance(error, AccountingConflictError) or any( - isinstance(cause, Exception) and (PrismaDBExceptionHandler.postgres_sqlstate(cause) or "").startswith("22") - for cause in exception_chain(error) - ) - - -class DurableUsageMeter: - def __init__( - self, - write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - settings: UsageCaptureSettings, - *, - wait_for_batch: Callable[[asyncio.Event, float], Awaitable[None]] = _wait_for_batch, - retry_delay: Callable[[float], Awaitable[None]] = asyncio.sleep, - ) -> None: - self._write_batch: Final = write_batch - self.settings: Final = settings - self._wait_for_batch: Final = wait_for_batch - self._retry_delay: Final = retry_delay - self._loop: Final = asyncio.get_running_loop() - self._admission: Final = Lock() - self._queue: Final[deque[_Pending]] = deque() # mutable-ok: one bounded producer/consumer queue - self._wakeup: Final = asyncio.Event() - self._capacity_changed: Final = asyncio.Event() - self._reservations: Mapping[object, tuple[int, int]] = {} - self._reserved_bytes = 0 - self._task: asyncio.Task[None] | None = None - self._closed = False - self.pending_bytes = 0 - self.pending_records = 0 - - def submit( - self, writes: tuple[UsageWrite, ...], *, prerequisites: tuple[asyncio.Future[None], ...] = () - ) -> asyncio.Future[None]: - size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes) - with self._admission: - if ( - self._closed - or not writes - or len(writes) > self.settings.batch_size - or self.pending_bytes + size > self.settings.max_pending_bytes - or self.pending_records + len(writes) > self.settings.max_pending_records - ): - raise UsageCaptureError("Usage capture admission rejected: closed, empty or capacity exhausted") - receipt: Final = Future[None]() - receipt.set_running_or_notify_cancel() - receipt.add_done_callback(_report_failure) - self._queue.append(_Pending(writes, size, receipt, asyncio.Event(), prerequisites)) - self.pending_bytes += size - self.pending_records += len(writes) - self._loop.call_soon_threadsafe(self._ensure_worker) - wrapped: Final = asyncio.wrap_future(receipt, loop=self._loop) - wrapped.add_done_callback(lambda done: None if done.cancelled() else done.exception()) - return wrapped - - def reserve_producer(self, size: int, headroom: int) -> object: - token: Final = object() - self.resize_producer(token, size, headroom) - return token - - def resize_producer(self, token: object, size: int, headroom: int) -> None: - with self._admission: - previous: Final = self._reservations.get(token, (0, 0))[0] - delta: Final = size - previous - others: Final = (value[1] for key, value in self._reservations.items() if key is not token) - required: Final = max((headroom, *others)) - if ( - self._closed - or self.pending_bytes + delta > self.settings.max_pending_bytes - or self._reserved_bytes + delta + required > self.settings.max_pending_bytes - ): - raise UsageCaptureError("Batch capture producer byte capacity exhausted") - self._reservations = {**self._reservations, token: (size, headroom)} - self._reserved_bytes += delta - self.pending_bytes += delta - self._capacity_changed.set() - - def release_producer(self, token: object) -> None: - with self._admission: - size, _ = self._reservations[token] - self._reservations = {key: value for key, value in self._reservations.items() if key is not token} - self._reserved_bytes -= size - self.pending_bytes -= size - self._capacity_changed.set() - - def _ensure_worker(self) -> None: - if self.pending_records >= self.settings.batch_size or self._queue and self._queue[0].prerequisites: - self._wakeup.set() - if self._task is None or self._task.done(): - self._task = asyncio.create_task(self._run(), name="accounting usage capture") - - async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: - envelope: Final = '{"event":' + event.model_dump_json() + "}" - await asyncio.shield(self.submit((UsageWrite(event, envelope, usage_delivery_targets(event.dimensions)),))) - - async def record_writes(self, writes: tuple[UsageWrite, ...]) -> None: - size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes) - capacity: Final = self.settings.max_pending_bytes - self._reserved_bytes - if len(writes) > min(self.settings.batch_size, self.settings.max_pending_records) or ( - size > capacity and len(writes) > 1 - ): - middle: Final = len(writes) // 2 - await self.record_writes(writes[:middle]) - await self.record_writes(writes[middle:]) - return - if size > capacity: - raise UsageCaptureError("One usage observation exceeds capture capacity") - while (receipt := self._submit_available(writes)) is None: - self._wakeup.set() - await self._capacity_changed.wait() - self._wakeup.set() - await asyncio.shield(receipt) - - def _submit_available(self, writes: tuple[UsageWrite, ...]) -> asyncio.Future[None] | None: - self._capacity_changed.clear() - try: - receipt: Final = self.submit(writes) - except UsageCaptureError: - if self._closed or not self.pending_records or not writes: - raise - return None - return receipt - - async def flush(self) -> None: - self._ensure_worker() - self._wakeup.set() - if self._task is not None: - await asyncio.shield(self._task) - - async def close(self) -> None: - with self._admission: - self._closed = True - await self.flush() - - async def stop(self, timeout: float) -> None: - with self._admission: - self._closed = True - try: - await asyncio.wait_for(self.close(), timeout) - except asyncio.TimeoutError: - if self._task is not None: - self._task.cancel() - await asyncio.gather(self._task, return_exceptions=True) - while self._queue: - self._finish( - self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") - ) - verbose_proxy_logger.error("Usage capture shutdown deadline exhausted; commit acknowledgement incomplete") - - def _take_batch(self) -> tuple[_Pending, ...]: - with self._admission: - if self._queue[0].prerequisites: - return (self._queue.popleft(),) - independent: Final = takewhile(lambda item: not item.prerequisites, self._queue) - sizes: Final = accumulate(len(item.writes) for item in independent) - count: Final = sum(1 for _ in takewhile(lambda size: size <= self.settings.batch_size, sizes)) - return tuple(self._queue.popleft() for _ in range(count)) - - async def _persist(self, writes: tuple[UsageWrite, ...]) -> None: - for attempt in count(): - try: - await self._write_batch(writes) - except Exception as error: # noqa: BLE001 # retry only classified database failures, preserving frozen writes - if not _retryable_persistence_failure(error): - raise - verbose_proxy_logger.warning("Usage persistence unavailable; retaining admitted batch for retry") - await self._retry_delay(min(10.0, 0.2 * 2 ** min(attempt, 6))) - else: - return - - async def _commit(self, batch: tuple[_Pending, ...]) -> None: - writes: Final = tuple(chain.from_iterable(item.writes for item in batch)) - try: - if batch[0].prerequisites: - outcomes: Final = await asyncio.gather( - *(asyncio.shield(receipt) for receipt in batch[0].prerequisites), return_exceptions=True - ) - for outcome in outcomes: - if isinstance(outcome, BaseException): - raise UsageCaptureError("Usage dependency did not acknowledge commit") from outcome - await self._persist(writes) - except asyncio.CancelledError: - self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) - raise - except Exception as error: # noqa: BLE001 # isolate data faults, surface systemic failures to all producers - if len(batch) > 1 and _isolatable_persistence_failure(error): - try: - for item in batch: - await self._commit((item,)) - except asyncio.CancelledError: - self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) - raise - return - self._finish( - batch, - UsageCaptureError("Usage identity conflicts with accepted content") - if isinstance(error, AccountingConflictError) - else error, - ) - else: - self._finish(batch, None) - - def _finish(self, batch: tuple[_Pending, ...], error: Exception | None) -> None: - for item in batch: - if item.finished.is_set(): - continue - item.finished.set() - with self._admission: - self.pending_bytes -= item.size - self.pending_records -= len(item.writes) - self._capacity_changed.set() - if item.committed.done(): - continue - if error is None: - item.committed.set_result(None) - else: - item.committed.set_exception(error) - - async def _run(self) -> None: - try: - while self._queue: - await self._wait_for_batch(self._wakeup, self.settings.flush_interval_seconds) - self._wakeup.clear() - await self._commit(self._take_batch()) - if ( - self._closed - or self.pending_records >= self.settings.batch_size - or self._queue - and self._queue[0].prerequisites - ): - self._wakeup.set() - except (Exception, asyncio.CancelledError): # noqa: BLE001 # reject queued receipts if the writer exits abnormally - while self._queue: - self._finish( - self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") - ) - raise diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e68d348c16c..145c8e38229 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -168,7 +168,6 @@ from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int from litellm.litellm_core_utils.litellm_logging import ( BilledRoundsDispatch, Logging, - UsageObservation, deployment_pricing_model_info, ) from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token @@ -1805,39 +1804,6 @@ def _failure_fields_to_lift( if _prompt_estimate is not None else () ) - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation - - _recorded_observation: Final[object] = _model_call_details.get("_accounting_usage_observed") - _current_usage: Final = ( - _current_failure_usage_to_lift( - _model_call_details, - request_data, - _model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None, - _billed, - _deployment_id, - _base_model, - _estimate_call_type, - ) - if Logging.usage_capture_callback is not None - and (_billed or _model_call_details.get("combined_usage_object") is not None) - else None - ) - _current_observation: Final = ( - _recorded_observation - if isinstance(_recorded_observation, UsageObservation) - else UsageObservation.UNKNOWN - if _current_usage is None - else UsageObservation.KNOWN - if _model_call_details.get("combined_usage_object") is not None - else UsageObservation.ESTIMATED - ) - _capture_observation: Final = UsageCaptureObservation.from_usage( - _current_usage[0] if _current_usage is not None and isinstance(_current_usage[0], Usage) else None, - _current_observation, - float(_current_usage[1]) - if _current_usage is not None and isinstance(_current_usage[1], (int, float)) - else None, - ) _entries: Final = ( ("start_time", _model_call_details.get("start_time")), ("first_api_call_start_time", _first_handoff), @@ -1846,16 +1812,12 @@ def _failure_fields_to_lift( ("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)), ("standard_logging_object", _model_call_details.get("standard_logging_object")), ("billed_failed_attempts", _earlier_attempts or None), - ("_accounting_capture", _model_call_details.get("_accounting_capture")), ("_accounting_pricing_snapshot", _model_call_details.get("_accounting_pricing_snapshot")), ("_accounting_finalized_cost", _logging_obj.finalized_cost if isinstance(_logging_obj, Logging) else None), ( "_accounting_allocation", _logging_obj.cost_allowance_allocation if isinstance(_logging_obj, Logging) else None, ), - ("_accounting_usage_observed", _current_observation if Logging.usage_capture_callback is not None else None), - ("_accounting_usage_snapshot", _capture_observation if Logging.usage_capture_callback is not None else None), - ("_accounting_prior_captures", _model_call_details.get("_accounting_prior_captures")), ) return MappingProxyType({key: value for key, value in _entries if value is not None}) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index e28ebf439b5..3ea900642bd 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -3842,11 +3842,6 @@ class PythonCore(DeprecatedRouterMembers): priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" } - if LiteLLMLogging.usage_capture_callback is not None: - from litellm._internal_context import internal_sub_call - - with internal_sub_call(cut_logging.litellm_call_id): - cut_logging.capture_usage_before_logging(priced_partial) return ContinuedPartial( usage=BaseTokenUsageProcessor.combine_usage_objects( usage_objects=[partial_usage, inherited_usage] diff --git a/litellm/utils.py b/litellm/utils.py index b457da058ad..e210b2fe026 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1636,7 +1636,6 @@ def dispatch_success_logging( is_litellm_internal_call: bool, ) -> None: if is_litellm_internal_call: - logging_obj.capture_usage_before_logging(result, start_time, end_time) return _schedule_async_success_logging( @@ -1675,14 +1674,9 @@ def _schedule_async_success_logging( if not logging_obj.claim_async_success_log(): return - if logging_obj.defer_async_logging and not is_completion_with_fallbacks: - logging_obj.capture_usage_before_logging(result, start_time, end_time, freeze_only=True) - def _enqueue_async_logging() -> None: from litellm.responses.streaming_iterator import spawn_logging_task - if not is_completion_with_fallbacks: - logging_obj.capture_usage_before_logging(result, start_time, end_time) if logging_obj.defer_async_logging: logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, end_time=end_time @@ -2280,7 +2274,6 @@ def client(original_function): # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") - logging_obj.capture_usage_before_logging(TypeAdapter(object).validate_python(result), start_time, end_time) if not is_internal_call.get(): # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py deleted file mode 100644 index 15cd5413149..00000000000 --- a/tests/integration/database/test_accounting_ledger.py +++ /dev/null @@ -1,610 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import os -import subprocess -import sys -from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence -from contextlib import asynccontextmanager -from datetime import datetime, timezone -from decimal import Decimal -from itertools import chain -from pathlib import Path -from typing import TYPE_CHECKING, Final, TypeAlias - -import psycopg -import pytest -import pytest_asyncio -from psycopg import sql -from pydantic import JsonValue - -from litellm.accounting import ( - CostEvent, - DeliveryState, - DeliveryTarget, - EventKind, - InboundUsageEvent, - Metric, - MetricName, - OutboundUsageEvent, - Outcome, - PricingSnapshot, -) -from litellm.proxy.db import accounting_ledger -from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite -from litellm.types.utils import CallTypes -from tests.integration._support.database import scratch_database -from tests.integration._support.prisma_clients import connected_proxy_prisma - -if TYPE_CHECKING: - from prisma import Prisma - - from litellm.proxy.utils import PrismaClient - -_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger] - -_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma" -_EVENT_INSERT: Final = ( - 'INSERT INTO "LiteLLM_AccountingEvent" ' - "(event_id, kind, schema_version, payload_hash, occurred_at, payload) " - "VALUES ('event', %s, %s, 'hash', '2026-01-01', '{}') RETURNING kind::text, schema_version" -) -_DELIVERY_INSERT: Final = ( - 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state, attempts) ' - "VALUES ('event', %s, %s, %s, %s) RETURNING target::text, version, state::text, attempts" -) - - -@pytest.fixture(scope="module") -def ledger_database() -> Iterator[str]: - with scratch_database("accounting") as database_url: - result: Final = subprocess.run( - [sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(_SCHEMA)], - env={**os.environ, "DATABASE_URL": database_url}, - capture_output=True, - text=True, - timeout=300, - ) - assert result.returncode == 0, result.stdout + result.stderr - yield database_url - - -@pytest.fixture -def ledger(ledger_database: str) -> Iterator[psycopg.Connection[tuple[object, ...]]]: - with psycopg.connect(ledger_database) as connection: - try: - yield connection - finally: - connection.rollback() - - -@pytest.mark.parametrize("kind", tuple(EventKind)) -def test_ledger_accepts_every_public_event_kind( - ledger: psycopg.Connection[tuple[object, ...]], kind: EventKind -) -> None: - row: Final = ledger.execute(_EVENT_INSERT, (kind.value, 1)).fetchone() - assert row == (kind.value, 1) - - -@pytest.mark.parametrize("target", tuple(DeliveryTarget)) -@pytest.mark.parametrize("state", tuple(DeliveryState)) -def test_delivery_accepts_every_public_target_and_state( - ledger: psycopg.Connection[tuple[object, ...]], target: DeliveryTarget, state: DeliveryState -) -> None: - row: Final = ledger.execute(_DELIVERY_INSERT, (target.value, 1, state.value, 0)).fetchone() - assert row == (target.value, 1, state.value, 0) - - -@pytest.mark.parametrize( - ("enum_name", "values"), - [ - ("AccountingEventKind", tuple(value.value for value in EventKind)), - ("AccountingDeliveryState", tuple(value.value for value in DeliveryState)), - ("AccountingDeliveryTarget", tuple(value.value for value in DeliveryTarget)), - ], -) -def test_migrated_enums_match_the_public_contract( - ledger: psycopg.Connection[tuple[object, ...]], enum_name: str, values: tuple[str, ...] -) -> None: - rows: Final = ledger.execute( - "SELECT e.enumlabel FROM pg_enum e JOIN pg_type t ON t.oid = e.enumtypid WHERE t.typname = %s", - (enum_name,), - ).fetchall() - assert frozenset(row[0] for row in rows) == frozenset(values) - - -@pytest.mark.parametrize("version", [0, -1]) -def test_event_rejects_nonpositive_schema_versions( - ledger: psycopg.Connection[tuple[object, ...]], version: int -) -> None: - with pytest.raises(psycopg.errors.CheckViolation, match="AccountingEvent_schema_version_check"): - ledger.execute(_EVENT_INSERT, (EventKind.OUTBOUND_USAGE.value, version)) - - -@pytest.mark.parametrize( - ("version", "attempts", "constraint"), - [ - (0, 0, "AccountingDelivery_version_check"), - (-1, 0, "AccountingDelivery_version_check"), - (1, -1, "AccountingDelivery_attempts_check"), - ], -) -def test_delivery_rejects_invalid_versions_and_attempts( - ledger: psycopg.Connection[tuple[object, ...]], version: int, attempts: int, constraint: str -) -> None: - with pytest.raises(psycopg.errors.CheckViolation, match=constraint): - ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, version, "pending", attempts)) - - -@pytest.mark.parametrize("column", ["state", "target"]) -def test_delivery_rejects_unknown_enum_values(ledger: psycopg.Connection[tuple[object, ...]], column: str) -> None: - ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, 1, "pending", 0)) - with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): - ledger.execute( - sql.SQL('UPDATE "LiteLLM_AccountingDelivery" SET {} = %s').format(sql.Identifier(column)), ("invalid",) - ) - - -def test_event_rejects_unknown_kind(ledger: psycopg.Connection[tuple[object, ...]]) -> None: - with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): - ledger.execute(_EVENT_INSERT, ("invalid", 1)) - - -def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Connection[tuple[object, ...]]) -> None: - row: Final = ledger.execute( - 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) ' - "VALUES ('event', %s, 1) RETURNING state::text, attempts", - (DeliveryTarget.KEY_SPEND.value,), - ).fetchone() - assert row == ("pending", 0) - - -_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) -_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc) -_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND) - - -def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent: - fields: Final = dict( - event_id=event_id, - request_id="request", - call_type=CallTypes.completion, - started_at=_AT, - ended_at=_AT, - outcome=Outcome.SUCCEEDED, - dimensions=(), - ) - if inbound: - return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),)) - return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None) - - -def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent: - return CostEvent( - event_id=event_id, - usage_event_id=source, - occurred_at=_AT, - dimensions=(), - charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),), - snapshot=PricingSnapshot( - calculator_version="fixture", - model="fixture", - provider=None, - pricing_source="cost_map", - priced_model_key="fixture", - total_cost=Decimal(amount), - ), - ) - - -def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str: - return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}}) - - -async def _insert( - store: AccountingLedger, - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...] = _TARGETS, - *, - context: Mapping[str, JsonValue] | None = None, -) -> str: - return await store.insert_event(event, targets, envelope=_envelope(event, context=context)) - - -def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]: - with psycopg.connect(database_url) as connection: - cursor: Final = connection.execute(statement, parameters) - return cursor.fetchall() if cursor.description is not None else [] - - -def _complete_and_age(database_url: str) -> None: - with psycopg.connect(database_url) as connection: - connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,)) - connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,)) - - -@pytest_asyncio.fixture(loop_scope="function") -async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]: - monkeypatch.setenv("DATABASE_URL", ledger_database) - monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) - _rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"') - async with connected_proxy_prisma(ledger_database) as client: - yield client, AccountingLedger.for_client(client) - - -@pytest.mark.asyncio -async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None: - _, store = writer - event: Final = _usage() - assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted" - rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"') - assert rows == [ - ("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"}))) - ] - deliveries: Final = _rows( - ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"' - ) - assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS} - - -@pytest.mark.asyncio -async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None: - _, store = writer - event: Final = _usage(inbound=True) - results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8))) - assert results.count("inserted") == 1 and results.count("duplicate") == 7 - _complete_and_age(ledger_database) - before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') - reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2) - assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate" - assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"]) -async def test_conflicting_acceptance_cannot_overwrite_or_add_work( - writer: _Writer, ledger_database: str, change: str -) -> None: - _, store = writer - event: Final = _usage() - await _insert(store, event) - altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event - targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS - context: Final = {"override": "0.5"} if change == "context" else {} - options: Final = {change: 2} if change in ("schema_version", "version") else {} - with pytest.raises(AccountingConflictError, match="conflicts"): - await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options) - assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)] - assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == { - (t.value,) for t in _TARGETS - } - - -@pytest.mark.asyncio -async def test_competing_cost_identities_accept_only_one_authoritative_charge( - writer: _Writer, ledger_database: str -) -> None: - _, store = writer - await _insert(store, _usage()) - results: Final = await asyncio.gather( - _insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True - ) - assert sum(result == "inserted" for result in results) == 1 - assert sum(isinstance(result, AccountingConflictError) for result in results) == 1 - accepted: Final = _rows( - ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" - ) - assert len(accepted) == 1 - event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"])) - assert await _insert(store, event) == "duplicate" - assert event.snapshot.total_cost == Decimal("0.125") - with pytest.raises(AccountingConflictError): - await _insert(store, _cost(event_id=event.event_id, amount="0.5")) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("source", ["missing", "inbound", "cost"]) -async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None: - _, store = writer - if source == "inbound": - await _insert(store, _usage("source", inbound=True)) - if source == "cost": - await _insert(store, _usage()) - await _insert(store, _cost(event_id="source")) - with pytest.raises(AccountingConflictError, match="outbound usage"): - await _insert(store, _cost(source="source", event_id="invalid")) - assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [ - (0,) - ] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"]) -async def test_interruption_rolls_back_or_retry_recognizes_commit( - writer: _Writer, ledger_database: str, failure: str -) -> None: - client, store = writer - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - if failure == "before_commit": - raise RuntimeError("injected pre-commit failure") - if failure == "cancellation": - raise asyncio.CancelledError - raise RuntimeError("injected lost acknowledgement") - - error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError - with pytest.raises(error): - await _insert(AccountingLedger(interrupted), _usage()) - expected: Final = 1 if failure == "after_commit" else 0 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)] - assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted") - - -@pytest.mark.asyncio -async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None: - from prisma.errors import RawQueryError - - _, store = writer - with psycopg.connect(ledger_database) as connection: - connection.execute( - 'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)' - ) - try: - with pytest.raises(RawQueryError, match="reject_delivery"): - await _insert(store, _usage()) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] - finally: - with psycopg.connect(ledger_database) as connection: - connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery') - - -@pytest.mark.asyncio -@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"]) -async def test_retention_preserves_incomplete_or_recovery_needed_families( - writer: _Writer, ledger_database: str, blocked: str -) -> None: - _, store = writer - await _insert(store, _usage()) - if blocked != "unpriced": - await _insert(store, _cost()) - _complete_and_age(ledger_database) - with psycopg.connect(ledger_database) as connection: - if blocked in ("pending", "parked"): - connection.execute( - "UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,) - ) - if blocked == "recent_cost": - connection.execute( - "UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,) - ) - before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') - assert ( - await store.purge_completed( - redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF - ) - == 0 - ) - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before - - -@pytest.mark.asyncio -async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage()) - await _insert(store, _cost()) - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] - - -@pytest.mark.asyncio -async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None: - client, store = writer - for name in ("a", "b", "c"): - await _insert(store, _usage(name, inbound=True)) - _complete_and_age(ledger_database) - async with client.writer_db.tx() as tx: - await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE") - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("a",), - ("c",), - ] - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - - -@pytest.mark.asyncio -async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage("a", inbound=True)) - await _insert(store, _usage("b")) - await _insert(store, _usage("c", inbound=True)) - _complete_and_age(ledger_database) - _rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'") - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("a",), - ("b",), - ] - - -@pytest.mark.asyncio -async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage(), targets=()) - event: Final = _cost(amount="0") - assert await _insert(store, event, targets=()) == "inserted" - assert await _insert(store, event, targets=()) == "duplicate" - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - - -@pytest.mark.asyncio -async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted( - writer: _Writer, ledger_database: str -) -> None: - client, store = writer - await _insert(store, _usage()) - _complete_and_age(ledger_database) - inserted: Final = asyncio.Event() - release: Final = asyncio.Event() - - @asynccontextmanager - async def paused_commit() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - inserted.set() - await release.wait() - - task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost())) - try: - await asyncio.wait_for(inserted.wait(), timeout=5) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0 - finally: - release.set() - await task - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("cost",), - ("usage",), - ] - - -@pytest.mark.asyncio -async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None: - client, store = writer - await _insert(store, _usage(inbound=True)) - _complete_and_age(ledger_database) - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - raise RuntimeError("injected cleanup failure") - - with pytest.raises(RuntimeError, match="cleanup failure"): - await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1 - - -@pytest.mark.asyncio -async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured( - writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch -) -> None: - with scratch_database("replica") as replica_url: - monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url) - async with connected_proxy_prisma(ledger_database) as client: - store: Final = AccountingLedger.for_client(client) - await _insert(store, _usage()) - assert await _insert(store, _usage()) == "duplicate" - assert await _insert(store, _cost()) == "inserted" - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)] - - -def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]: - if isinstance(plan, dict): - own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else () - children: Final = (plan.get("Plan"), *plan.get("Plans", ())) - return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None)) - return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else () - - -@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES]) -def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None: - explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s") - with psycopg.connect(ledger_database) as connection: - connection.execute("SET enable_seqscan = off") - plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone() - assert plan is not None - conditions: Final = _index_conditions(plan[0]) - assert any( - name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition) - for name, condition in conditions - ), conditions - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "conflict"]) -async def test_usage_batch_is_atomic_and_replay_preserves_receipts( - writer: _Writer, ledger_database: str, failure: str -) -> None: - client, store = writer - events: Final = ( - _usage("caller", inbound=True), - _usage("attempt").model_copy( - update={"metrics": (Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),)} - ), - ) - writes: Final = tuple(UsageWrite(event, _envelope(event), _TARGETS) for event in events) - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - if failure == "before_commit": - raise RuntimeError("injected pre-commit failure") - if failure == "after_commit": - raise RuntimeError("injected lost acknowledgement") - - if failure == "conflict": - await store.insert_usage_batch((writes[1],)) - _complete_and_age(ledger_database) - rejected: Final = (writes[0], UsageWrite(events[1], _envelope(events[1], context={"changed": True}), _TARGETS)) - with pytest.raises(AccountingConflictError): - await store.insert_usage_batch(rejected) - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("attempt",)] - else: - with pytest.raises(RuntimeError, match="injected"): - await AccountingLedger(interrupted).insert_usage_batch(writes) - count: Final = 2 if failure == "after_commit" else 0 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(count,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(count * len(_TARGETS),)] - await store.insert_usage_batch(writes) - recovered: Final = await store.load_usage_batch(tuple(event.event_id for event in events)) - assert len(recovered) == 1 - assert recovered[0].event == events[1] - assert json.loads(recovered[0].envelope) == json.loads(writes[1].envelope) - assert set(recovered[0].targets) == set(_TARGETS) - _complete_and_age(ledger_database) - before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') - await store.insert_usage_batch((*writes, *writes)) - assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(2,)] - - -@pytest.mark.asyncio -async def test_invalid_jsonb_submission_cannot_reject_another_producers_commit( - writer: _Writer, ledger_database: str -) -> None: - from prisma.errors import DataError - - from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings - - _, store = writer - invalid: Final = _usage("invalid") - valid: Final = _usage("valid") - meter: Final = DurableUsageMeter(store.insert_usage_batch, UsageCaptureSettings(batch_size=2)) - rejected: Final = meter.submit( - (UsageWrite(invalid, _envelope(invalid, context={"parent": "bad\x00text"}), _TARGETS),) - ) - accepted: Final = meter.submit((UsageWrite(valid, _envelope(valid), _TARGETS),)) - try: - with pytest.raises(DataError): - await rejected - await accepted - finally: - await meter.close() - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("valid",)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] - assert (meter.pending_bytes, meter.pending_records) == (0, 0) diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py new file mode 100644 index 00000000000..ec0ab44e505 --- /dev/null +++ b/tests/unit/accounting/test_projection.py @@ -0,0 +1,431 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.accounting.projection import project +from litellm.accounting.snapshot import build_accounting_payload +from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.proxy_server import proxy_logging_obj +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.utils import PrismaClient +from litellm.types.utils import ModelResponse, Usage + +_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) +_FIELDS: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.parametrize("status", ["success", "failure"]) +@pytest.mark.parametrize("origin", [None, "auto_router", "billed_retry"]) +@pytest.mark.parametrize("spend", [0.0, 0.25]) +@pytest.mark.asyncio +async def test_daily_projection_matches_legacy_writer(status: str, origin: str | None, spend: float) -> None: + payload: Final = get_logging_payload( + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "projection-request", + "litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user", "status": status}}, + }, + response_obj=ModelResponse( + id="projection-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ) + metadata: Final = { + **_FIELDS.validate_json(payload["metadata"]), + **({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if origin == "billed_retry" else {}), + **({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin} if origin not in (None, "billed_retry") else {}), + } + row: Final = {**payload, "spend": spend, "metadata": json.dumps(metadata), "request_duration_ms": 25} + legacy: Final = await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction( + row, PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj), "user" + ) + assert legacy is not None + accounting_input: Final[SpendLogsPayload] = {**row} + accounting: Final = build_accounting_payload( + accounting_input, + entities={"api_key": payload["api_key"], "user": payload["user"]}, + model_access_groups=(), + windows=(), + router=lambda: None, + ) + restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) + stored_daily: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend") + assert dict(stored_daily.increments) == {name: legacy[name] for name, _ in stored_daily.increments} + assert dict(stored_daily.key)["date"] == legacy["date"] + assert dict(stored_daily.key)["api_key"] == legacy["api_key"] + assert not {"messages", "response", "proxy_server_request", "requester_ip_address"}.intersection(accounting) + + +def test_projection_uses_frozen_savings_and_keeps_duplicate_entity_tags() -> None: + row: Final[Mapping[str, object]] = { + "request_id": "saved", + "startTime": _NOW, + "api_key": "key", + "user": "user", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "request_tags": '["tag", "tag"]', + "spend": 0.5, + "metadata": json.dumps( + { + "accounting_savings": { + "compression": 1.0, + "prompt_caching": -2.0, + "gateway_injected_caching": 3.0, + "autorouter": 4.0, + }, + "autorouter_savings": 9000, + } + ), + } + contributions: Final = project(row) + tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_TagTable") + daily_tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_DailyTagSpend") + assert sum(dict(item.increments)["spend"] for item in tags) == 1.0 + assert len(daily_tags) == 1 + assert {key: value for key, value in daily_tags[0].increments if key.endswith("savings_spend")} == { + "compression_savings_spend": 1.0, + "prompt_caching_savings_spend": -2.0, + "gateway_injected_caching_savings_spend": 3.0, + "autorouter_savings_spend": 4.0, + } + + +def test_projection_carries_windows_and_does_not_reprice_guardrails() -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "api_key": "key", + "spend": 0.75, + "metadata": json.dumps( + { + "accounting_windows": [ + { + "entity_type": "key", + "entity_id": "key", + "window_duration": "1d", + "window_start": _NOW.isoformat(), + } + ], + "guardrail_information": [ + { + "guardrail_id": "guard", + "guardrail_usage": {"text": 3, "images": 2}, + "guardrail_cost_by_unit": {"text": 0.3, "images": None}, + } + ], + } + ), + } + contributions: Final = project(row) + window: Final = next(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend") + units: Final = { + dict(item.key)["usage_unit"]: dict(item.increments) + for item in contributions + if item.table == "LiteLLM_DailyGuardrailUsageUnits" + } + assert dict(window.increments) == {"spend": 0.75} + assert dict(window.values)["window_start"] == _NOW.isoformat() + assert units == { + "text": {"units": 3, "cost": 0.3, "untracked_units": 0}, + "images": {"units": 2, "cost": 0.0, "untracked_units": 2}, + } + + +@pytest.mark.asyncio +async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.responses.streaming_iterator import drain_logging_tasks + + client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(client, "spend_log_transactions", []) + monkeypatch.setattr(client, "model_usage_transactions", []) + writer: Final = DBSpendUpdateWriter() + charged: Final = await writer.update_database( + token="sk-projection", + user_id="user", + end_user_id="end-user", + team_id="team", + org_id="org", + project_id="project", + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "entity-request", + "agent_id": "agent", + "user": "end-user", + "litellm_params": { + "metadata": { + "user_api_key": "sk-projection", + "user_api_key_user_id": "user", + "user_api_key_team_id": "team", + "user_api_key_org_id": "org", + "user_api_key_project_id": "project", + "user_api_key_end_user_id": "end-user", + "tags": ["tag", "tag"], + }, + "user": "end-user", + }, + }, + completion_response=ModelResponse( + id="entity-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + response_cost=0.25, + ) + await drain_logging_tasks() + assert charged + assert len(client.spend_log_transactions) == 1 + queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + logged: Final[SpendLogsPayload] = client.spend_log_transactions[0] + projected: Final = project( + _FIELDS.validate_json( + json.dumps( + build_accounting_payload( + logged, + entities={ + "api_key": logged["api_key"], + "user": "user", + "end_user": "end-user", + "team_id": "team", + "organization_id": "org", + "project_id": "project", + "agent_id": "agent", + }, + model_access_groups=(), + windows=(), + router=lambda: None, + ), + default=str, + ) + ) + ) + legacy: Final = {(item["entity_type"].value, item["entity_id"]): item["response_cost"] for item in queued} + assert legacy[("key", proxy_server.hash_token("sk-projection"))] == next( + dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_VerificationToken" + ) + for table, entity, identity in ( + ("LiteLLM_UserTable", "user", "user"), + ("LiteLLM_TeamTable", "team", "team"), + ("LiteLLM_OrganizationTable", "organization", "org"), + ("LiteLLM_EndUserTable", "end_user", "end-user"), + ("LiteLLM_AgentsTable", "agent", "agent"), + ("LiteLLM_ProjectTable", "project", "project"), + ): + assert legacy[(entity, identity)] == next( + dict(item.increments)["spend"] for item in projected if item.table == table + ) + assert sum(dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_TagTable") == sum( + item["response_cost"] or 0 for item in queued if item["entity_type"].value == "tag" + ) + + daily_queues: Final = ( + ("LiteLLM_DailyUserSpend", writer.daily_spend_update_queue), + ("LiteLLM_DailyTeamSpend", writer.daily_team_spend_update_queue), + ("LiteLLM_DailyOrganizationSpend", writer.daily_org_spend_update_queue), + ("LiteLLM_DailyEndUserSpend", writer.daily_end_user_spend_update_queue), + ("LiteLLM_DailyAgentSpend", writer.daily_agent_spend_update_queue), + ("LiteLLM_DailyTagSpend", writer.daily_tag_spend_update_queue), + ) + for table, queue in daily_queues: + transactions = await queue.flush_and_get_aggregated_daily_spend_update_transactions() + matching = tuple(item for item in projected if item.table == table) + assert matching, table + contribution = matching[0] + assert len(transactions) == 1 + transaction = next(iter(transactions.values())) + assert dict(contribution.increments) == {name: transaction[name] for name, _ in contribution.increments} + assert len(client.model_usage_transactions) == 1 + model: Final = client.model_usage_transactions[0] + model_projection: Final = next(item for item in projected if item.table == "LiteLLM_DailyModelUsage") + assert dict(model_projection.increments) == { + "spend": model.spend, + "prompt_tokens": model.prompt_tokens, + "completion_tokens": model.completion_tokens, + "request_count": int(model.counts_as_request), + "successful_requests": int(model.successful), + "failed_requests": int(not model.successful), + } + + +@pytest.mark.parametrize("served_id", ["requested", "fallback", "missing"]) +@pytest.mark.asyncio +async def test_projection_preserves_fallback_narrowed_model_access_groups(served_id: str) -> None: + from litellm import Router + from litellm.proxy.db.db_spend_update_writer import debitable_model_access_groups + + router: Final = Router( + model_list=[ + { + "model_name": "first", + "litellm_params": {"model": "openai/projection-test", "api_key": "unused"}, + "model_info": {"id": "requested", "access_groups": ["first-group"]}, + }, + { + "model_name": "second", + "litellm_params": {"model": "openai/projection-test", "api_key": "unused"}, + "model_info": {"id": "fallback", "access_groups": ["second-group", "unattributed"]}, + }, + ] + ) + attributed: Final = ("first-group", "second-group", "first-group") + writer: Final = DBSpendUpdateWriter() + client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) + await writer._update_model_access_group_db(0.5, attributed, served_id, client, router) + queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + frozen_groups: Final = debitable_model_access_groups(attributed, served_id, router) + projected: Final = project( + {"spend": 0.5, "metadata": json.dumps({"accounting_model_access_groups": frozen_groups})} + ) + assert {dict(item.key)["access_group_name"]: dict(item.increments)["spend"] for item in projected} == { + item["entity_id"]: item["response_cost"] for item in queued + } + assert "unattributed" not in frozen_groups + + +def test_guardrail_projection_skips_unmeasured_counters() -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "metadata": json.dumps( + { + "guardrail_information": [ + {"guardrail_usage": {"text": 3}}, + { + "guardrail_id": "guard", + "guardrail_usage": {"zero": 0, "negative": -1, "bool": True, "unknown": None, "text": 2}, + }, + ] + } + ), + } + units: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyGuardrailUsageUnits") + assert len(units) == 1 + assert dict(units[0].key)["usage_unit"] == "text" + assert dict(units[0].increments) == {"units": 2, "cost": 0.0, "untracked_units": 2} + + +def test_lean_row_without_model_does_not_fabricate_daily_entity_usage() -> None: + projected: Final = project( + { + "request_id": "lean", + "api_key": "key", + "startTime": _NOW, + "spend": 0.5, + "metadata": {"accounting_entities": {"api_key": "key"}}, + } + ) + assert tuple(item.table for item in projected) == ("LiteLLM_VerificationToken", "LiteLLM_DailyModelUsage") + + +@pytest.mark.parametrize( + "cost_fields", + [ + {"guardrail_cost_in_spend": False, "guardrail_cost_by_unit": {"text": 1.0}}, + {"guardrail_cost_by_unit": {"text": -1.0}}, + ], +) +def test_report_only_or_invalid_guardrail_prices_remain_untracked(cost_fields: Mapping[str, object]) -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "request_tags": '["task:debugging"]', + "metadata": json.dumps( + { + "guardrail_information": [ + {"guardrail_id": "guard", "guardrail_usage": {"text": 2}, **cost_fields}, + ] + } + ), + } + contributions: Final = project(row) + guardrail: Final = next(item for item in contributions if item.table == "LiteLLM_DailyGuardrailUsageUnits") + model: Final = next(item for item in contributions if item.table == "LiteLLM_DailyModelUsage") + assert dict(guardrail.increments) == {"units": 2, "cost": 0.0, "untracked_units": 2} + assert dict(model.key)["task_type"] == "debugging" + + +@pytest.mark.parametrize("organization_id", [None, "", "charged-org"]) +def test_daily_organization_uses_frozen_charged_identity(organization_id: str | None) -> None: + row: Final[Mapping[str, object]] = { + "startTime": _NOW, + "api_key": "key", + "model": "projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "organization_id": "stored-org", + "spend": 0.25, + "metadata": {"accounting_entities": {"organization_id": organization_id}}, + } + daily: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyOrganizationSpend") + if organization_id is None: + assert daily == () + else: + assert len(daily) == 1 + assert dict(daily[0].key)["organization_id"] == organization_id + assert dict(daily[0].increments)["spend"] == 0.25 + + +@pytest.mark.parametrize("start", [_NOW, "2026-11-02T12:00:00+00:00", "2026-11-02 12:00:00+00:00"]) +def test_every_daily_row_uses_the_calendar_date_of_the_request(start: object) -> None: + row: Final[Mapping[str, object]] = { + "request_id": "dated", + "startTime": start, + "api_key": "key", + "user": "user", + "team_id": "team", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "spend": 0.5, + "metadata": "{}", + } + dates: Final = { + item.table: dict(item.key)["date"] for item in project(row) if item.table.startswith("LiteLLM_Daily") + } + assert {"LiteLLM_DailyUserSpend", "LiteLLM_DailyModelUsage"} <= dates.keys() + assert set(dates.values()) == {"2026-11-02"} + + +def test_unparseable_start_time_skips_only_guardrail_usage() -> None: + projected: Final = project( + { + "request_id": "bad-time", + "startTime": "yesterday at noon", + "spend": 0.5, + "metadata": { + "accounting_entities": {"api_key": "key"}, + "guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}], + }, + } + ) + tables: Final = {item.table for item in projected} + assert "LiteLLM_VerificationToken" in tables + assert "LiteLLM_DailyGuardrailUsageUnits" not in tables + + +def test_guardrail_usage_is_dated_in_utc() -> None: + projected: Final = project( + { + "request_id": "offset", + "startTime": "2026-11-02T22:30:00-05:00", + "metadata": {"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 1}}]}, + } + ) + assert [dict(item.key)["date"] for item in projected if item.table == "LiteLLM_DailyGuardrailUsageUnits"] == [ + "2026-11-03" + ] diff --git a/tests/unit/accounting/test_types.py b/tests/unit/accounting/test_types.py index f645ee6c2cc..c1be39ab111 100644 --- a/tests/unit/accounting/test_types.py +++ b/tests/unit/accounting/test_types.py @@ -19,7 +19,6 @@ from litellm.accounting import ( Dimension, DimensionFilter, DimensionName, - EventKind, Granularity, InboundUsageEvent, Metric, @@ -442,7 +441,6 @@ def test_subset_token_counts_round_trip_without_changing_totals() -> None: @pytest.mark.parametrize( ("enum_name", "python_enum"), [ - ("AccountingEventKind", EventKind), ("AccountingDeliveryState", DeliveryState), ("AccountingDeliveryTarget", DeliveryTarget), ], diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 2d02468fd8c..04ba823d08f 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -484,55 +484,6 @@ async def test_bridged_text_completion_is_logged_as_the_text_completion_it_serve assert logged.payload["response"]["choices"][0]["text"] == "pong" -@pytest.mark.asyncio -async def test_a_bridged_chat_response_waits_on_the_usage_ack_like_any_chat_response( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - logging_obj: Final = Logging( - model="openai/gpt-5.3-codex", - messages=[], - stream=False, - call_type="acompletion", - start_time=datetime.datetime.now(), - litellm_call_id="bridged-chat", - function_id="bridged-chat", - ) - with respx.mock(assert_all_called=True) as router: - router.post(url__startswith="https://api.openai.com/v1/responses").mock( - return_value=httpx.Response(200, json=_searched_responses_body("openai/gpt-5.3-codex")) - ) - served: Final = await litellm.acompletion( - model="openai/gpt-5.3-codex", - messages=[{"role": "user", "content": "ping"}], - api_key="test-key", - litellm_logging_obj=logging_obj, - ) - response: Final = asyncio.create_task(capture_proxy_response(logging_obj, served)) - try: - await database.entered.wait() - assert not response.done() - assert database.committed == () - database.release.set() - await response - assert "bridged-chat:inbound" in database.committed - finally: - database.release.set() - await runtime.stop(1) - await response - - def _served_message_start_id(events: Sequence[bytes]) -> object: lines: Final = b"".join(events).decode().splitlines() starts: Final = tuple( diff --git a/tests/unit/proxy/db/test_accounting_ledger.py b/tests/unit/proxy/db/test_accounting_ledger.py deleted file mode 100644 index 6b301146301..00000000000 --- a/tests/unit/proxy/db/test_accounting_ledger.py +++ /dev/null @@ -1,149 +0,0 @@ -import json -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final, Literal, NoReturn - -import pytest - -from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome -from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingLedger, UsageWrite -from litellm.types.utils import CallTypes - - -def _no_transaction() -> NoReturn: - raise AssertionError("Invalid accounting input must not open a database transaction") - - -@pytest.fixture -def event() -> InboundUsageEvent: - return InboundUsageEvent( - event_id="caller", - request_id="caller", - call_type=CallTypes.completion, - started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"]) -async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None: - with pytest.raises(ValueError, match=r"Accounting envelope|validation error"): - await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("version", [0, -1, True, 1.5]) -@pytest.mark.parametrize("field", ["version", "schema_version"]) -async def test_invalid_versions_fail_before_database_access( - event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"] -) -> None: - with pytest.raises(ValueError, match="positive integers"): - await AccountingLedger(_no_transaction).insert_event( - event, - (), - envelope="{}", - **{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions - ) - - -@pytest.mark.asyncio -async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None: - invalid: Final = event.model_copy(update={"event_id": ""}) - with pytest.raises(ValueError, match="event_id"): - await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5]) -async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None: - with pytest.raises(ValueError, match="batch size"): - await AccountingLedger(_no_transaction).purge_completed( - limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"]) -async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None: - cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - deleted: Final = await AccountingLedger(_no_transaction).purge_completed( - redelivery_before=cutoff if missing == "recovery" else None, - recovery_before=cutoff if missing == "redelivery" else None, - ) - assert deleted == 0 - - -@pytest.mark.asyncio -async def test_retention_rejects_ambiguous_naive_watermarks() -> None: - with pytest.raises(ValueError, match="timezone"): - await AccountingLedger(_no_transaction).purge_completed( - redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc) - ) - - -@pytest.mark.asyncio -async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}}) - with pytest.raises(ValueError, match="Out of range float"): - await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) - - -@pytest.mark.asyncio -async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - envelope: Final = json.dumps({"event": event.model_dump(mode="json")}) - with pytest.raises(ValueError, match="not a valid DeliveryTarget"): - await AccountingLedger(_no_transaction).insert_event( - event, - ("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target - envelope=envelope, - ) - - -@pytest.mark.asyncio -async def test_empty_usage_batches_do_not_open_a_transaction() -> None: - ledger: Final = AccountingLedger(_no_transaction) - await ledger.insert_usage_batch(()) - assert await ledger.load_usage_batch(()) == () - - -@pytest.mark.asyncio -async def test_conflicting_usage_within_one_batch_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - original: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) - changed: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json"), "pricing": {"rate": 2}})) - with pytest.raises(AccountingConflictError, match="conflicting copies"): - await AccountingLedger(_no_transaction).insert_usage_batch((original, changed)) - - -@pytest.mark.asyncio -async def test_loaded_usage_preserves_decimal_metrics_from_jsonb(event: InboundUsageEvent) -> None: - from collections.abc import AsyncIterator - from contextlib import asynccontextmanager - from typing import cast - from unittest.mock import AsyncMock - - from prisma import Prisma - - from litellm.accounting import OutboundUsageEvent - - outbound: Final = OutboundUsageEvent( - **event.model_dump(exclude={"metrics"}), - attempt_id="attempt", - retry_of_attempt_id=None, - metrics=(Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),), - ) - payload: Final = {"event": outbound.model_dump(mode="json"), "usage": {"prompt_tokens": 120000}} - database: Final = AsyncMock(spec=Prisma) - database.query_raw.return_value = [{"payload": payload, "targets": []}] - - @asynccontextmanager - async def transaction() -> AsyncIterator[Prisma]: - yield cast(Prisma, database) - - recovered: Final = await AccountingLedger(transaction).load_usage_batch((outbound.event_id,)) - assert len(recovered) == 1 - assert recovered[0].event == outbound - assert json.loads(recovered[0].envelope) == payload diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 21b5bd300c1..c8434007a12 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -3981,140 +3981,3 @@ async def test_async_post_call_failure_hook_disable_pre_routing_error_logs( user_api_key_dict=user_api_key_dict, ) assert len(writer.rows) == rows_written - - -@pytest.mark.asyncio -async def test_failure_capture_overload_does_not_skip_budget_release( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _write - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - meter: Final = runtime.meter - assert meter is not None - receipt: Final = meter.submit((_write("occupied-one"), _write("occupied-two"))) - auth: Final = UserAPIKeyAuth(budget_reservation={"reserved_cost": 0.5, "entries": [], "finalized": False}) - request: Final = { - "call_type": "acompletion", - "litellm_call_id": "failed", - "model": "capture-model", - "start_time": _CALL_START, - "litellm_params": {}, - } - try: - await database.entered.wait() - await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), auth) - assert auth.budget_reservation is not None and auth.budget_reservation["finalized"] is True - assert "Usage capture failed on the request failure path" in caplog.text - assert "_accounting_capture" not in request - finally: - database.release.set() - await runtime.stop(1) - await receipt - assert database.committed == ("occupied-one", "occupied-two") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", [False, True]) -@pytest.mark.parametrize("complete", [False, True]) -@pytest.mark.parametrize("deferred", [False, True]) -async def test_capture_and_legacy_share_resolved_key_attribution( - failure: bool, complete: bool, deferred: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - from pydantic import TypeAdapter - - from litellm.accounting import DeliveryTarget, DimensionName - from litellm.proxy import proxy_server - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage - - key: Final = UserAPIKeyAuth( - api_key="a" * 64, - user_id="owner", - team_id="team", - org_id="org", - project_id="project", - key_alias="key", - team_alias="team-alias", - request_route="/v1/chat/completions", - ) - cache: Final = UserApiKeyCache() - await cache.async_set_cache(key="a" * 64, value=key) - monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) - monkeypatch.setattr(proxy_server, "prisma_client", object()) - writer: Final = _RecordingSpendWriter() - monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(db_spend_update_writer=writer)) - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - rejected: Final = failure and deferred and not complete - owner: Final = "original-owner" if deferred and not rejected else "owner" - metadata: Final = { - "user_api_key": "b" * 64 if rejected else "a" * 64, - **( - { - "user_api_key_user_id": "forged-owner" if rejected else owner, - "user_api_key_team_id": "forged-team" if rejected else "team", - "user_api_key_org_id": "forged-org" if rejected else "org", - "user_api_key_project_id": None if deferred and not complete else "project", - "user_api_key_alias": "key", - "user_api_key_team_alias": "team-alias", - } - if complete or deferred - else {} - ), - } - kwargs: Final = { - "call_type": "aretrieve_batch" if deferred else "call_mcp_tool", - "litellm_call_id": "attributed", - "model": "MCP: echo", - "response_cost": 0.25, - "start_time": _CALL_START, - "litellm_params": {"metadata": metadata}, - } - if rejected: - from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY - - kwargs[REJECTED_BEFORE_REQUEST_SETUP_KEY] = REJECTED_BEFORE_REQUEST_SETUP - logger: Final = ProxyDBLogger(spend_writer=lambda: cast(DBSpendUpdateWriter, writer)) - try: - with track_batch_accounting("attributed" if deferred and not rejected else "other"): - if failure: - await logger.async_post_call_failure_hook( - kwargs, - HTTPException(401, "model not allowed"), - key if complete else UserAPIKeyAuth(api_key="a" * 64, request_route="/v1/chat/completions"), - ) - else: - await logger.async_log_success_event(kwargs, {"id": "tool"}, _CALL_START, _CALL_START) - finally: - await runtime.stop(1) - assert len(writer.rows) == 1 - assert len(database.rows) == 2 - dimensions: Final = {item.name: item.value for item in database.rows[-1].event.dimensions} - if rejected: - assert dimensions.get(DimensionName.KEY) == key.api_key - assert { - name: dimensions.get(name) - for name in (DimensionName.USER, DimensionName.TEAM, DimensionName.ORGANIZATION, DimensionName.PROJECT) - } == { - DimensionName.USER: owner, - DimensionName.TEAM: "team", - DimensionName.ORGANIZATION: "org", - DimensionName.PROJECT: None if deferred and not complete and not rejected else "project", - } - assert set(database.rows[-1].targets) >= {DeliveryTarget.DAILY_TEAM_SPEND, DeliveryTarget.DAILY_ORGANIZATION_SPEND} - legacy: Final = TypeAdapter(dict[str, object]).validate_python(writer.rows[0]["kwargs"]) - payload: Final = get_logging_payload(legacy, {"id": "tool"}, _CALL_START, _CALL_START) - assert (payload["user"], payload["team_id"], payload["organization_id"]) == (owner, "team", "org") diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index a426d8f4e14..ef70a193280 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -92,12 +92,6 @@ def response(cost: float = 0.5) -> ModelResponse: async def test_header_and_logging_share_one_allowance_allocation( logging_obj: Logging, used: int, expected: float ) -> None: - import json - - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings - counter: Final = Counter(used) result: Final = response() costs: Final = await asyncio.gather(*(finalize_response_cost(logging_obj, result, counter, NOW) for _ in range(3))) @@ -116,24 +110,6 @@ async def test_header_and_logging_share_one_allowance_allocation( assert payload["cost_finalized"] is True assert payload["cost_breakdown"] == costs[0].logging_breakdown() - committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - committed.put_nowait(write) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings()) - assert logging_obj.prepared_cost is not None - prepared_usage: Final = logging_obj.prepared_cost.usage.model_dump() - result.usage = Usage(prompt_tokens=999) - capture_logging_usage(meter, logging_obj, result, NOW, NOW) - await meter.close() - committed.get_nowait() - event: Final = json.loads(committed.get_nowait().envelope) - assert event["finalized_cost"]["response_cost"] == pytest.approx(costs[0].amount) - assert event["allocation"]["discount"] == pytest.approx(costs[0].allowance_discount) - assert event["usage"] == prepared_usage - assert counter.calls == 1 @pytest.mark.asyncio diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 4e113ce03a4..32f6eea91ac 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7695,13 +7695,11 @@ def test_optional_mcp_metadata_does_not_change_legacy_spend_payload(tool_metadat assert actual == expected -@pytest.mark.parametrize("capture_enabled", [False, True]) @pytest.mark.parametrize("completion_start", [None, "absent", "reported"]) -def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, completion_start: str | None) -> None: +def test_spend_payload_preserves_unknown_completion_start(completion_start: str | None) -> None: start: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc) end: Final = start + datetime.timedelta(seconds=2) reported: Final = start + datetime.timedelta(seconds=1) - observations: Final[list[object]] = [] kwargs: Final = { "call_type": "acompletion", "litellm_call_id": "timestamp-call", @@ -7713,9 +7711,7 @@ def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, {"id": "timestamp-response"}, start, end, - capture_usage=(lambda row, usage: observations.append(row)) if capture_enabled else None, ) assert payload["completionStartTime"] == ( end if completion_start == "absent" else reported if completion_start else None ) - assert observations == ([payload] if capture_enabled else []) diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py deleted file mode 100644 index 0c59473114a..00000000000 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ /dev/null @@ -1,1738 +0,0 @@ -import asyncio -import json -from collections.abc import Awaitable, Mapping -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import Metric, MetricName, Outcome -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite -from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings -from litellm.types.utils import CallTypes -from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _manual_flush, _write - - -@pytest.mark.asyncio -@pytest.mark.parametrize("usage", [None, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}]) -async def test_capture_keeps_unknown_usage_distinct_from_measured_zero(usage: Mapping[str, int] | None) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - captured: Final = capture_callback( - meter, - {"call_type": "acompletion", "litellm_call_id": "caller", "model": "unknown-model", "litellm_params": {}}, - {"id": "provider", "usage": usage}, - timestamp, - timestamp, - ) - await meter.close() - outbound: Final = captured.writes[1].event - assert outbound.metrics == ( - None - if usage is None - else tuple(Metric(name=MetricName(key), value=Decimal(value)) for key, value in usage.items()) - ) - envelope: Final = json.loads(captured.writes[1].envelope) - assert envelope["usage"] == usage - assert envelope["pricing"]["reported_cost"] is None - assert database.committed == ("caller:inbound", "provider") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "call_type", - [ - call.value - for call in ( - CallTypes.acompletion, - CallTypes.arealtime, - CallTypes.call_mcp_tool, - CallTypes.asend_message, - CallTypes.pass_through, - ) - ], -) -async def test_shared_callback_capture_retains_usage_and_attribution_without_bodies(call_type: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - captured: Final = capture_callback( - meter, - { - "call_type": call_type, - "litellm_call_id": "caller", - "model": "test-model", - "messages": "private prompt", - "litellm_params": {"api_key": "private-secret", "metadata": {"user_api_key_team_id": "team"}}, - }, - { - "id": "provider", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, - "output": "private answer", - }, - timestamp, - timestamp, - ) - await meter.close() - envelope: Final = captured.writes[1].envelope - assert all(secret not in envelope for secret in ("private prompt", "private-secret", "private answer")) - event: Final = captured.writes[1].event - assert event.call_type.value == call_type - assert event.metrics == ( - Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(3)), - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), - Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(5)), - ) - assert any(dimension.value == "team" for dimension in event.dimensions) - - -@pytest.mark.asyncio -async def test_cache_duplicate_uses_caller_identity_despite_timestamped_diagnostic_ids() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "test-model", - "cache_hit": True, - "litellm_params": {}, - } - response: Final = {"id": "cached-provider-id", "usage": {"total_tokens": 5}} - first: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) - second: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) - await meter.close() - assert first.writes == second.writes - assert first.writes[1].event.event_id == "caller_cache_hit" - assert json.loads(first.writes[1].envelope)["charge_eligible"] is False - - -NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - - -class StoredUsage: - def __init__(self) -> None: - self.rows: tuple[UsageWrite, ...] = () - - async def load(self, identities: tuple[str, ...]) -> tuple[UsageWrite, ...]: - return tuple(row for row in self.rows if row.event.event_id in identities) - - async def insert(self, writes: tuple[UsageWrite, ...]) -> None: - existing: Final = {row.event.event_id: row for row in self.rows} - if any(existing.get(row.event.event_id, row) != row for row in writes): - raise AccountingConflictError("Same identity has different content") - unique: Final = {row.event.event_id: row for row in writes} - if any(unique[row.event.event_id] != row for row in writes): - raise AccountingConflictError("Submission contains conflicting content") - self.rows += tuple(row for identity, row in unique.items() if identity not in existing) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("text", ["parent\x00value", "pärent-✓-日本", r"\u0000"]) -async def test_capture_normalizes_database_text_without_poisoning_neighbors(text: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import PricingCapture, capture_callback - - database: Final = StoredUsage() - batches: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - assert all("\x00" not in (json.loads(write.envelope).get("parent_request_id") or "") for write in writes) - batches.put_nowait(writes) - await database.insert(writes) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(), wait_for_batch=_manual_flush) - nested: Final = json.loads("[" * 100 + "3" + "]" * 100) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "search_context_size": text, - "_accounting_pricing_snapshot": PricingCapture( - "{}", "{}", json.dumps({"request": {"search_context_size": text}, "nested": nested}) - ), - "litellm_params": {"metadata": {"parent_request_id": text}}, - } - response: Final = {"id": "provider", "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}} - captured: Final = capture_callback(meter, kwargs, response, NOW, NOW) - neighbor: Final = capture_callback( - meter, - {**kwargs, "litellm_call_id": "neighbor", "litellm_params": {}}, - {**response, "id": "neighbor-provider"}, - NOW, - NOW, - ) - try: - await meter.flush() - await asyncio.gather(captured.receipt, neighbor.receipt) - assert batches.qsize() == 1 - assert len(batches.get_nowait()) == 4 - envelope: Final = json.loads(captured.writes[-1].envelope) - assert envelope["parent_request_id"] == text.replace("\x00", "") - assert envelope["pricing"]["request"]["search_context_size"] == text.replace("\x00", "") - assert envelope["pricing"]["nested"] == nested - assert envelope["event"] == captured.writes[-1].event.model_dump(mode="json") - assert envelope["usage"] == response["usage"] - assert kwargs["litellm_params"] == {"metadata": {"parent_request_id": text}} - repeated: Final = capture_callback(meter, kwargs, response, NOW, NOW) - await meter.flush() - await repeated.receipt - assert repeated.writes == captured.writes - assert len(database.rows) == 4 - finally: - await meter.stop(0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("invalid_field", ["identity", "context"]) -async def test_invalid_capture_text_is_rejected_before_admission(invalid_field: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "caller", "litellm_params": {}} - try: - accepted: Final = capture_callback(meter, kwargs, {"id": "provider"}, NOW, NOW) - invalid: Final = ( - {**kwargs, "litellm_call_id": "caller\x00"} - if invalid_field == "identity" - else {**kwargs, "litellm_params": {"metadata": {"parent_request_id": "parent\ud800"}}} - ) - with pytest.raises((UsageCaptureError, ValueError)): - capture_callback(meter, invalid, {"id": "other"}, NOW, NOW) - assert meter.pending_records == 2 - await meter.flush() - await accepted.receipt - assert tuple(write.event.event_id for write in database.rows) == ("caller:inbound", "provider") - finally: - await meter.stop(0) - - -def _logger(call_type: str = "acompletion", *, stream: bool = False) -> Logging: - from litellm.litellm_core_utils.litellm_logging import Logging - - logger: Final = Logging( - model="capture-model", - messages=[], - stream=stream, - call_type=call_type, - start_time=NOW, - litellm_call_id="caller", - function_id="function", - ) - logger.update_environment_variables( - model="capture-model", custom_llm_provider="openai", litellm_params={}, optional_params={} - ) - return logger - - -@pytest.mark.asyncio -async def test_terminal_retrieval_reuses_accepted_snapshot_but_rejects_changed_usage() -> None: - from datetime import timedelta - - from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback - from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = { - "call_type": "aget_responses", - "litellm_call_id": "poll-one", - "model": "capture-model", - "response_cost": 2.0, - "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, - } - response: Final = { - "id": "job", - "status": "completed", - "background": True, - "created_at": NOW.timestamp(), - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - } - first: Final = capture_callback(meter, kwargs, response, NOW, NOW) - await meter.flush() - later: Final = capture_callback( - meter, {**kwargs, "litellm_call_id": "poll-two", "response_cost": 9.0}, response, NOW, NOW + timedelta(hours=2) - ) - await meter.flush() - await later.receipt - assert database.rows == first.writes - assert len(database.rows) == 1 - assert database.rows[0].event.event_id == "job_background_cost" - changed: Final = capture_callback(meter, kwargs, {**response, "usage": {"total_tokens": 8}}, NOW, NOW) - await meter.close() - with pytest.raises(UsageCaptureError, match="identity"): - await changed.receipt - assert database.rows == first.writes - - -@pytest.mark.asyncio -@pytest.mark.parametrize("terminal", [False, True]) -async def test_competing_writers_reuse_only_terminal_pricing_snapshots(terminal: bool) -> None: - from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback - from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - - database: Final = StoredUsage() - attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - competing: Final = asyncio.Event() - - async def insert(writes: tuple[UsageWrite, ...]) -> None: - attempts.put_nowait(writes) - if attempts.qsize() >= 2: - competing.set() - await competing.wait() - await database.insert(writes) - - meters: Final = tuple( - DurableUsageMeter(UsageBatchWriter(insert, database.load).write, UsageCaptureSettings()) for _ in range(2) - ) - kwargs: Final = { - "call_type": "aget_responses" if terminal else "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "litellm_params": { - "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} - if terminal - else {} - }, - } - response: Final = { - "id": "job", - "status": "completed", - "background": terminal, - "created_at": NOW.timestamp(), - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - } - captured: Final = tuple( - capture_callback(meter, {**kwargs, "response_cost": price}, response, NOW, NOW) - for meter, price in zip(meters, (2.0, 9.0), strict=True) - ) - try: - await asyncio.gather(*(meter.flush() for meter in meters)) - results: Final = await asyncio.gather(*(capture.receipt for capture in captured), return_exceptions=True) - assert sum(result is None for result in results) == (2 if terminal else 1) - assert sum(isinstance(result, UsageCaptureError) for result in results) == (0 if terminal else 1) - assert database.rows in tuple(capture.writes for capture in captured) - assert len(database.rows) == (1 if terminal else 2) - assert attempts.qsize() == (3 if terminal else 2) - finally: - await asyncio.gather(*(meter.stop(0) for meter in meters)) - - -@pytest.mark.asyncio -async def test_batch_poll_has_distinct_identity_and_only_terminal_usage_is_chargeable() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = { - "call_type": "aretrieve_batch", - "litellm_call_id": "poll", - "model": "capture-model", - "litellm_params": {}, - } - pending: Final = capture_callback(meter, kwargs, {"id": "batch", "status": "in_progress"}, NOW, NOW) - completed: Final = capture_callback( - meter, kwargs, {"id": "batch", "status": "completed", "usage": {"total_tokens": 9}}, NOW, NOW - ) - await meter.close() - assert pending.writes[-1].event.event_id == "poll:status" - assert json.loads(pending.writes[-1].envelope)["charge_eligible"] is False - assert completed.writes[-1].event.event_id == "batch_batch_cost" - assert json.loads(completed.writes[-1].envelope)["charge_eligible"] is True - assert len(completed.writes) == 1 - - -@pytest.mark.asyncio -async def test_internal_capture_does_not_claim_reused_parent_logger(monkeypatch: pytest.MonkeyPatch) -> None: - import litellm - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.utils import ModelResponse, Usage - - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - { - "input_cost_per_token": 0.1, - "output_cost_per_token": 0.2, - "litellm_provider": "openai", - "mode": "chat", - }, - ) - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger() - response: Final = ModelResponse( - id="response", model="capture-model", usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3) - ) - with internal_sub_call("caller"): - capture_logging_usage(meter, logger, response, NOW, NOW) - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - { - "input_cost_per_token": 0.7, - "output_cost_per_token": 0.9, - "litellm_provider": "openai", - "mode": "chat", - }, - ) - litellm.register_model({"capture-model": litellm.model_cost["capture-model"]}) - capture_logging_usage(meter, logger, response, NOW, NOW) - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - assert tuple(row.event.event_id for row in database.rows) == ("response:included", "caller:inbound", "response") - child: Final = json.loads(database.rows[0].envelope) - parent: Final = json.loads(database.rows[-1].envelope) - assert (child["ownership"], child["charge_eligible"], child["parent_request_id"]) == ( - "included_in_parent", - False, - "caller", - ) - assert parent["charge_eligible"] is True - assert child["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.1 - assert parent["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.7 - - -@pytest.mark.asyncio -async def test_stream_capture_ignores_chunks_and_preserves_terminal_usage_extras() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.utils import ModelResponse, ModelResponseStream, Usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(stream=True) - capture_logging_usage(meter, logger, ModelResponseStream(id="chunk"), NOW, NOW) - assert meter.pending_records == 0 - capture_logging_usage( - meter, logger, ModelResponse(id="stream", usage=Usage(prompt_tokens=3, provider_units=4)), NOW, NOW - ) - await meter.close() - assert json.loads(database.rows[-1].envelope)["usage"]["provider_units"] == 4 - assert database.rows[-1].event.event_id == "stream" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["acompletion", "atext_completion", "aembedding", "aresponses"]) -async def test_supported_sync_ack_waits_for_commit(call_type: str, monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy._types import ConfigGeneralSettings - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - deadlines: Final = asyncio.Queue[float | None]() - original_wait: Final = asyncio.wait_for - - async def wait_with_deadline(awaitable: Awaitable[object], timeout: float | None) -> object: - deadlines.put_nowait(timeout) - return await original_wait(awaitable, timeout) - - monkeypatch.setattr(asyncio, "wait_for", wait_with_deadline) - response: Final = asyncio.create_task( - capture_proxy_response( - _logger(call_type), - ModelResponse(id="provider", usage=Usage(prompt_tokens=3)), - ) - ) - try: - await database.entered.wait() - deadline: Final = deadlines.get_nowait() - assert deadline == ConfigGeneralSettings().accounting_sync_ack_timeout_seconds - assert deadline is not None and deadline > UsageCaptureSettings().flush_interval_seconds + 10 - assert not response.done() - assert database.committed == () - database.release.set() - await response - assert database.committed == ("caller:inbound", "provider") - finally: - database.release.set() - await runtime.stop(1) - await response - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure_policy", ["fail_closed", "log"]) -async def test_sync_ack_commit_failure_follows_response_policy( - failure_policy: str, monkeypatch: pytest.MonkeyPatch -) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier(failure=ValueError("invalid envelope")) - database.release.set() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - meter: Final = runtime.meter - assert meter is not None - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, - "general_settings", - {"accounting_sync_ack": True, "accounting_sync_ack_failure_policy": failure_policy}, - ) - logger: Final = _logger() - response: Final = ModelResponse(id="provider", usage=Usage(prompt_tokens=4)) - if failure_policy == "fail_closed": - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, response) - assert error.value.status_code == 503 - else: - await capture_proxy_response(logger, response) - await runtime.stop(1.0) - assert database.committed == () - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_failed_attempt_keeps_its_usage_without_claiming_the_callers_success() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - for _ in range(2): - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4) - logger.carry_failed_attempt_usage() - logger.capture_usage_before_logging( - ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, total_tokens=3)), NOW, NOW - ) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ( - "caller-attempt-0", - "caller-attempt-1", - "caller:inbound", - "fallback", - ) - failed: Final = json.loads(database.rows[0].envelope) - assert failed["usage"]["total_tokens"] == 7 - assert failed["parent_request_id"] == "caller" - assert failed["request_count_eligible"] is False - assert failed["charge_eligible"] is True - assert failed["event"]["retry_of_attempt_id"] is None - assert failed["event"]["outcome"] == "failed" - assert failed["pricing"]["reported_cost"] == 0.4 - assert all(json.loads(row.envelope)["included_attempt_ids"] == [] for row in database.rows[:2]) - successful: Final = json.loads(database.rows[-1].envelope) - assert successful["related_attempt_ids"] == ["caller-attempt-0", "caller-attempt-1"] - assert successful["included_attempt_ids"] == [] - assert successful["usage"]["total_tokens"] == 3 - assert successful["charge_eligible"] is True - - -@pytest.mark.asyncio -async def test_success_that_bills_a_failed_follow_up_names_that_attempt_as_included() -> None: - from litellm.litellm_core_utils.hidden_params import set_hidden_param - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - fallback: Final = ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4)) - set_hidden_param(fallback, "response_cost", 0.1) - try: - logger.record_served_rounds_for_failure(((Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4),)) - logger.carry_failed_attempt_usage() - logger.prepare_response_cost(fallback, NOW) - logger.capture_usage_before_logging(fallback, NOW, NOW) - finally: - await runtime.stop(1) - envelopes: Final = {row.event.event_id: json.loads(row.envelope) for row in database.rows} - assert envelopes["caller-attempt-0"]["usage"]["total_tokens"] == 7 - assert envelopes["caller-attempt-0"]["charge_eligible"] is True - assert envelopes["fallback"]["usage"]["total_tokens"] == 4 + 7 - assert envelopes["fallback"]["included_attempt_ids"] == ["caller-attempt-0"] - assert database.rows[2].event.outcome is Outcome.SUCCEEDED - - -@pytest.mark.asyncio -@pytest.mark.parametrize("estimated", [False, True]) -async def test_cancelled_usage_keeps_measured_or_estimated_provenance(estimated: bool) -> None: - from fastapi import HTTPException - - from litellm.litellm_core_utils.litellm_logging import UsageObservation - from litellm.proxy.spend_tracking.usage_capture import capture_callback - from litellm.types.utils import Usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - capture_callback( - meter, - { - "litellm_call_id": "cancelled", - "call_type": "acompletion", - "model": "capture-model", - "litellm_params": {}, - "combined_usage_object": Usage(prompt_tokens=4, total_tokens=4), - "_accounting_usage_observed": UsageObservation.ESTIMATED if estimated else UsageObservation.KNOWN, - }, - HTTPException(499, "Client disconnected"), - NOW, - NOW, - ) - await meter.close() - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["event"]["outcome"] == "cancelled" - assert payload["usage"]["total_tokens"] == 4 - assert payload["usage_source"] == ("estimated" if estimated else "measured") - assert payload["pricing"]["reported_cost"] is None - - -@pytest.mark.asyncio -async def test_realtime_capture_keeps_metered_units_and_provider_usage_without_frames() -> None: - from litellm.constants import REALTIME_METERED_USAGE_KEY - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.llms.xai import RealtimeMeteredUsage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(CallTypes.arealtime.value) - logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 1.5, 2.5, 3) - capture_logging_usage( - meter, - logger, - [ - { - "type": "response.done", - "response": { - "id": "turn", - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - "output": "private realtime transcript", - }, - } - ], - NOW, - NOW, - ) - await meter.close() - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["usage"]["total_tokens"] == 7 - assert payload["usage"]["metered_usage"] == { - "push_to_talk": True, - "audio_seconds_sent": 1.5, - "audio_seconds_received": 2.5, - "text_messages": 3, - } - assert "private realtime transcript" not in database.rows[-1].envelope - - -@pytest.mark.asyncio -async def test_default_http_capture_returns_before_database_commit(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {}) - try: - await capture_proxy_response( - _logger(), - ModelResponse(id="asynchronous", usage=Usage(prompt_tokens=1)), - ) - await database.entered.wait() - assert database.committed == () - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("caller:inbound", "asynchronous") - - -@pytest.mark.asyncio -async def test_deferred_capture_freezes_pricing_before_configuration_changes(monkeypatch: pytest.MonkeyPatch) -> None: - import litellm - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_logging_usage - from litellm.types.utils import ModelResponse, Usage - from litellm.utils import dispatch_success_logging - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - logger.defer_async_logging = True - logger.custom_cost_per_second = 0.25 - monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.1}) - response: Final = ModelResponse(id="deferred", usage=Usage(prompt_tokens=3)) - try: - dispatch_success_logging(logger, response, NOW, NOW, False, False) - assert database.rows == () - logger.custom_cost_per_second = 9 - monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.9}) - assert runtime.meter is not None - capture_logging_usage(runtime.meter, logger, response, NOW, NOW) - await runtime.meter.flush() - finally: - await runtime.stop(1) - pricing: Final = json.loads(database.rows[-1].envelope)["pricing"] - assert pricing["logging"]["custom_cost_per_second"] == 0.25 - assert pricing["adjustments"]["discount"] == 0.1 - - -@pytest.mark.asyncio -async def test_synchronous_internal_capture_reserves_capacity_before_scheduling() -> None: - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) - logger: Final = _logger() - - def capture_child() -> None: - with internal_sub_call("caller"): - logger.capture_usage_before_logging(ModelResponse(id="child", usage=Usage(prompt_tokens=2)), NOW, NOW) - - try: - await asyncio.to_thread(capture_child) - assert runtime.meter is not None - assert runtime.meter.pending_records == 1 - receipt: Final = await asyncio.to_thread(runtime.meter.submit, (_write("second"),)) - await database.entered.wait() - with pytest.raises(UsageCaptureError, match="capacity"): - await asyncio.to_thread(runtime.meter.submit, (_write("overflow"),)) - database.release.set() - await receipt - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("child:included", "second") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("has_metered_usage", [False, True]) -async def test_realtime_without_provider_token_usage_does_not_invent_zero(has_metered_usage: bool) -> None: - from litellm.constants import REALTIME_METERED_USAGE_KEY - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.llms.xai import RealtimeMeteredUsage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(CallTypes.arealtime.value) - if has_metered_usage: - logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 0, 0, 0) - capture_logging_usage(meter, logger, [], NOW, NOW) - await meter.close() - usage: Final = json.loads(database.rows[-1].envelope)["usage"] - assert usage == ( - { - "metered_usage": { - "push_to_talk": True, - "audio_seconds_sent": 0, - "audio_seconds_received": 0, - "text_messages": 0, - } - } - if has_metered_usage - else None - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("committed_before_failure", [False, True]) -async def test_database_recovery_retries_frozen_usage_without_callback_or_duplicate_rows( - committed_before_failure: bool, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - failures: Final = asyncio.Queue[None]() - failures.put_nowait(None) - retrying: Final = asyncio.Event() - recovered: Final = asyncio.Event() - delays: Final = asyncio.Queue[float]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - if not failures.empty(): - failures.get_nowait() - if committed_before_failure: - await database.insert(writes) - raise ConnectionError("database unavailable or acknowledgement lost") - await database.insert(writes) - - async def retry_delay(delay: float) -> None: - delays.put_nowait(delay) - retrying.set() - await recovered.wait() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(max_pending_records=2), wait_for_batch=_manual_flush, retry_delay=retry_delay - ) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "litellm_params": {}, - } - response: Final = {"id": "provider", "usage": {"total_tokens": 3}} - first: Final = capture_callback(meter, kwargs, response, NOW, NOW) - reserved: Final = meter.pending_bytes - flushing: Final = asyncio.create_task(meter.flush()) - await retrying.wait() - assert not first.receipt.done() - assert (meter.pending_records, meter.pending_bytes) == (2, reserved) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("overflow"),)) - recovered.set() - await flushing - await first.receipt - duplicate: Final = capture_callback(meter, {**kwargs, "_accounting_capture": first}, response, NOW, NOW) - await meter.close() - assert duplicate is first - assert database.rows == first.writes - assert delays.get_nowait() == 0.2 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -async def test_cut_stream_keeps_original_model_usage_as_part_of_the_parent_aggregate() -> None: - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.router import Router - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger(stream=True) - stream: Final = CustomStreamWrapper(completion_stream=iter(()), model="capture-model", logging_obj=logger) - partial: Final = ModelResponse( - id="cut", model="capture-model", usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) - ) - try: - Router._price_cut_partial(stream, partial) - finally: - await runtime.stop(1) - assert len(database.rows) == 1 - payload: Final = json.loads(database.rows[0].envelope) - assert payload["event"]["event_id"] == "cut:included" - assert payload["parent_request_id"] == "caller" - assert payload["charge_eligible"] is False - assert payload["usage"]["total_tokens"] == 5 - assert payload["pricing"]["model_map_information"]["model_map_key"] == "capture-model" - - -@pytest.mark.asyncio -async def test_direct_mcp_logging_captures_tool_pricing_and_identity_without_arguments() -> None: - from mcp.types import CallToolResult, TextContent - - from litellm.accounting import DimensionName - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger(CallTypes.call_mcp_tool.value) - logger.model_call_details["mcp_tool_call_metadata"] = { - "name": "lookup", - "namespaced_tool_name": "server.lookup", - "arguments": {"query": "private query"}, - "mcp_server_cost_info": {"default_cost_per_query": 0.2, "tool_name_to_cost_per_query": {"lookup": 0.3}}, - } - try: - await logger.async_success_handler( - CallToolResult(content=[TextContent(type="text", text="private result")]), NOW, NOW - ) - finally: - await runtime.stop(1) - outbound: Final = database.rows[-1] - assert any(d.name is DimensionName.TOOL and d.value == "server.lookup" for d in outbound.event.dimensions) - payload: Final = json.loads(outbound.envelope) - assert payload["pricing"]["tool"] == {"name": "lookup", "cost_per_query": 0.3} - assert payload["usage"] is None - assert all(secret not in outbound.envelope for secret in ("private query", "private result")) - - -@pytest.mark.asyncio -async def test_callback_overload_is_visible_without_suppressing_legacy_tracking( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, captured_callback_kwargs - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "one", "litellm_params": {}} - try: - for start, end in ((None, NOW), (NOW, None)): - missing_timestamp: Final = await captured_callback_kwargs(kwargs, {"id": "one"}, start, end) - assert missing_timestamp == kwargs - assert "Usage capture requires callback start and end timestamps" in caplog.text - admitted: Final = await captured_callback_kwargs(kwargs, {"id": "one"}, NOW, NOW) - await database.entered.wait() - rejected: Final = await captured_callback_kwargs({**kwargs, "litellm_call_id": "two"}, {"id": "two"}, NOW, NOW) - assert "_accounting_capture" in admitted - assert "_accounting_capture" not in rejected - assert "Usage capture failed before legacy cost tracking" in caplog.text - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("one:inbound", "one") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("call_type", "streaming", "background"), - [("acompletion", True, False), ("aresponses", False, True), ("aocr", False, False)], -) -@pytest.mark.parametrize("overloaded", [False, True]) -async def test_async_capture_preserves_response_without_commit_ack( - call_type: str, streaming: bool, background: bool, overloaded: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, - "general_settings", - { - "accounting_sync_ack": True, - "accounting_sync_ack_timeout_seconds": 0, - }, - ) - meter: Final = runtime.meter - assert meter is not None - if overloaded: - meter.submit((_write("occupied-one"), _write("occupied-two"))) - try: - await capture_proxy_response( - _logger(call_type, stream=streaming), - ModelResponse(id="provider", usage=Usage(prompt_tokens=3)), - background=background, - ) - assert database.committed == () - assert meter.pending_records == 2 - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == (("occupied-one", "occupied-two") if overloaded else ("caller:inbound", "provider")) - - -@pytest.mark.asyncio -async def test_new_connection_turn_does_not_reuse_the_previous_turns_capture() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.capture_usage_before_logging(ModelResponse(id="first", usage=Usage(prompt_tokens=2)), NOW, NOW) - next_turn: Final = logger.copy_for_new_call("second-caller", NOW) - next_turn.capture_usage_before_logging(ModelResponse(id="second", usage=Usage(prompt_tokens=4)), NOW, NOW) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ( - "caller:inbound", - "first", - "second-caller:inbound", - "second", - ) - assert json.loads(database.rows[-1].envelope)["usage"]["prompt_tokens"] == 4 - - -@pytest.mark.asyncio -async def test_billed_failure_is_captured_before_legacy_persistence_gates(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - request: Final = { - "call_type": "acompletion", - "litellm_call_id": "failed", - "model": "capture-model", - "start_time": NOW, - "combined_usage_object": Usage(prompt_tokens=3, total_tokens=3), - "response_cost": 0.2, - "litellm_params": {}, - } - try: - await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), UserAPIKeyAuth()) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("failed:inbound", "failed") - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["event"]["outcome"] == "failed" - assert payload["usage"]["total_tokens"] == 3 - assert payload["charge_eligible"] is True - assert payload["pricing"]["reported_cost"] == 0.2 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("pages", [None, 0, 3]) -async def test_prepared_ocr_capture_preserves_reported_pages(pages: int | None) -> None: - from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger("aocr") - response: Final = OCRResponse(pages=[], model="capture-model", usage_info=OCRUsageInfo(pages_processed=pages)) - logger.prepare_response_cost(response, NOW) - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["pages_processed"]) == pages - assert envelope["usage_source"] == ("unknown" if pages is None else "measured") - assert response.usage_info is not None and response.usage_info.pages_processed == pages - - -@pytest.mark.asyncio -@pytest.mark.parametrize("measured", [None, {"total_tokens": 0}, {"total_tokens": 3}]) -async def test_batch_capture_provenance_tracks_observations_not_envelope_presence( - measured: Mapping[str, int] | None, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - usage: Final = { - "total_tokens": 0 if measured is None else measured["total_tokens"], - "batch_usage_by_model": {"missing-model": None, "other-model": measured}, - "batch_unknown_usage_requests": 2 if measured is None else 1, - } - capture_callback( - meter, - {"call_type": "aretrieve_batch", "litellm_call_id": "caller", "litellm_params": {}}, - {"id": "batch", "usage": usage}, - NOW, - NOW, - ) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert envelope["usage_source"] == ("unknown" if measured is None else "measured") - assert envelope["usage"]["batch_usage_by_model"] == usage["batch_usage_by_model"] - assert database.rows[-1].event.metrics == ( - None if measured is None else (Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(measured["total_tokens"])),) - ) - - -@pytest.mark.asyncio -async def test_ack_admission_failure_does_not_turn_provider_success_into_failed_usage( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - monkeypatch.setattr(proxy_server, "prisma_client", None) - meter: Final = runtime.meter - assert meter is not None - logger: Final = _logger() - try: - meter.submit((_write("occupied-one"), _write("occupied-two"))) - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, ModelResponse(id="caller", usage=Usage(prompt_tokens=3))) - assert error.value.status_code == 503 - await meter.flush() - request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} - await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("occupied-one", "occupied-two") - - -@pytest.mark.asyncio -async def test_response_ack_timeout_preserves_success_and_later_commits(monkeypatch: pytest.MonkeyPatch) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - logger: Final = _logger() - - async def expired_deadline(future: object, timeout: float) -> None: - raise TimeoutError("controlled response deadline") - - try: - with monkeypatch.context() as clock: - clock.setattr(asyncio, "wait_for", expired_deadline) - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, ModelResponse(id="provider", usage=Usage(prompt_tokens=3))) - assert error.value.status_code == 503 - await database.entered.wait() - captured: Final = logger.model_call_details["_accounting_capture"] - request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} - await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) - database.release.set() - await captured.receipt - assert tuple(row.event.outcome for row in captured.writes) == (Outcome.SUCCEEDED, Outcome.SUCCEEDED) - assert database.committed == ("caller:inbound", "provider") - assert database.transactions == 1 - finally: - await runtime.stop(1) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["generate_content", "agenerate_content"]) -@pytest.mark.parametrize("tokens", [None, 0, 7]) -async def test_prepared_native_google_capture_preserves_reported_usage( - call_type: str, tokens: int | None, monkeypatch: pytest.MonkeyPatch -) -> None: - import httpx - - import litellm - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(call_type) - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - {"input_cost_per_token": 0.1, "output_cost_per_token": 0.2, "litellm_provider": "openai", "mode": "chat"}, - ) - reported: Final = ( - {"trafficType": "ON_DEMAND"} - if tokens is None - else { - "promptTokenCount": tokens, - "candidatesTokenCount": tokens, - "totalTokenCount": 2 * tokens, - "cachedContentTokenCount": tokens, - "promptTokensDetails": [{"modality": "DOCUMENT", "tokenCount": tokens}], - "cacheTokensDetails": [{"modality": "document", "token_count": tokens}], - } - ) - response: Final = { - "responseId": "native-google", - "candidates": [{"content": {"parts": [{"text": "answer"}], "role": "model"}, "finishReason": "STOP"}], - "usageMetadata": reported, - } - native_response: Final = httpx.Response(200, json=response) - logger.model_call_details["httpx_response"] = native_response - prepared: Final = logger.prepare_response_cost(response, NOW) - assert prepared is not None - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == ( - None if tokens is None else 2 * tokens - ) - assert (None if envelope["usage"] is None else envelope["usage"]["prompt_tokens_details"]["text_tokens"]) == tokens - assert ( - None - if envelope["usage"] is None - else envelope["usage"]["prompt_tokens_details"]["cached_tokens_details"]["text_tokens"] - ) == tokens - assert envelope["usage_source"] == ("unknown" if tokens is None else "measured") - assert response == native_response.json() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("rejected", [False, True]) -async def test_retry_links_only_admitted_usage_and_preserves_later_measurement(rejected: bool) -> None: - from litellm.litellm_core_utils.litellm_logging import UsageObservation - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings(max_pending_records=4)) - logger: Final = _logger() - assert runtime.meter is not None - try: - if rejected: - runtime.meter.submit(tuple(_write(f"occupied-{index}") for index in range(4))) - logger.model_call_details["_accounting_usage_observed"] = UsageObservation.UNKNOWN - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - await runtime.meter.flush() - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - finally: - await runtime.stop(1) - successful: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "success")) - assert successful.get("related_attempt_ids", successful["included_attempt_ids"]) == ( - [] if rejected else ["caller-attempt-0"] - ) - assert successful["usage"]["prompt_tokens"] == 3 - if not rejected: - failed: Final = json.loads( - next(row.envelope for row in database.rows if row.event.event_id == "caller-attempt-0") - ) - assert failed["usage"]["prompt_tokens"] == 5 - assert failed["usage_source"] == "measured" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("current_tokens", [None, 0, 3]) -async def test_final_failure_captures_current_attempt_without_legacy_retry_aggregate( - current_tokens: int | None, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.utils import _failure_fields_to_lift - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - if current_tokens is not None: - logger.record_partial_usage_for_failure( - Usage(prompt_tokens=current_tokens, total_tokens=current_tokens), 0.2 - ) - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert lifted["combined_usage_object"].total_tokens == 5 + (current_tokens or 0) - assert runtime.meter is not None - capture_callback( - runtime.meter, - {**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"}, - None, - NOW, - NOW, - ) - finally: - await runtime.stop(1) - final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) - assert (None if final["usage"] is None else final["usage"]["total_tokens"]) == current_tokens - assert final["usage_source"] == ("unknown" if current_tokens is None else "measured") - assert final["pricing"]["reported_cost"] == (None if current_tokens is None else 0.2) - assert final["included_attempt_ids"] == [] - assert final["related_attempt_ids"] == ["caller-attempt-0"] - - -@pytest.mark.asyncio -async def test_estimated_final_attempt_keeps_its_provenance_after_measured_retry() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.utils import _failure_fields_to_lift - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - logger.messages = [{"role": "user", "content": "hello"}] - logger.model_call_details["messages"] = logger.messages - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - logger.model_call_details["current_attempt_api_call_start_time"] = NOW - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert runtime.meter is not None - capture_callback(runtime.meter, {**logger.model_call_details, **lifted}, None, NOW, NOW) - finally: - await runtime.stop(1) - final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) - assert final["usage_source"] == "estimated" - assert final["usage"]["total_tokens"] > 0 - assert lifted["combined_usage_object"].total_tokens == final["usage"]["total_tokens"] + 5 - assert final["pricing"]["reported_cost"] == 0 - - -@pytest.mark.asyncio -async def test_parent_receipt_cannot_acknowledge_when_a_linked_attempt_conflicts() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - assert runtime.meter is not None - await runtime.meter.flush() - other: Final = _logger() - other.record_partial_usage_for_failure(Usage(prompt_tokens=8), 0.8) - other.carry_failed_attempt_usage() - other.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - captured: Final = other.model_call_details["_accounting_capture"] - await runtime.meter.flush() - with pytest.raises(UsageCaptureError, match="dependency"): - await captured.receipt - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("caller-attempt-0",) - - -@pytest.mark.asyncio -async def test_unserializable_observation_does_not_interrupt_retry_or_legacy_usage() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - usage: Final = Usage(prompt_tokens=5, total_tokens=5, opaque_usage=object()) - try: - logger.record_partial_usage_for_failure(usage, 0.4) - assert logger.model_call_details["combined_usage_object"] is usage - logger.carry_failed_attempt_usage() - assert logger.model_call_details["failed_attempts_usage"] == ((usage, 0.4),) - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("caller:inbound", "success") - assert json.loads(database.rows[-1].envelope)["related_attempt_ids"] == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prior_committed", [False, True]) -@pytest.mark.parametrize("batch_size", [2, 3]) -async def test_retry_dependencies_do_not_resubmit_usage_or_expand_final_submission( - prior_committed: bool, batch_size: int -) -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - database.release.set() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=batch_size, max_pending_records=3)) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - assert runtime.meter is not None - if prior_committed: - await runtime.meter.flush() - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - captured: Final = logger.model_call_details.get("_accounting_capture") - assert captured is not None, "Final usage must be admitted independently of prior usage batch size" - await runtime.meter.close() - await captured.receipt - assert runtime.meter.pending_records == 0 - finally: - await runtime.stop(1) - assert database.committed == ("caller-attempt-0", "caller:inbound", "success") - final: Final = json.loads(captured.writes[-1].envelope) - assert final["related_attempt_ids"] == ["caller-attempt-0"] - assert final["usage"]["prompt_tokens"] == 3 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prepared", [False, True]) -@pytest.mark.parametrize( - ("call_type", "response", "field", "quantity"), - [ - ("arerank", {"id": "rerank", "results": [], "meta": {"tokens": {"input_tokens": 7}}}, "prompt_tokens", 7), - ("aocr", {"id": "ocr", "usage_info": {"pages_processed": 3}}, "pages_processed", 3), - ( - "acompletion", - {"id": "google", "usageMetadata": {"promptTokenCount": 7, "totalTokenCount": 7}}, - "prompt_tokens", - 7, - ), - ], -) -async def test_capture_preserves_native_usage_before_and_after_preparation( - prepared: bool, call_type: str, response: dict[str, object], field: str, quantity: int -) -> None: - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(call_type) - logger.model_call_details["response_cost"] = 0.25 - if prepared: - assert logger.prepare_response_cost(response, NOW) is not None - with internal_sub_call("parent"): - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - assert len(database.rows) == 1 - envelope: Final = json.loads(database.rows[0].envelope) - assert envelope["usage"][field] == quantity - assert envelope["usage_source"] == "measured" - assert envelope["ownership"] == "included_in_parent" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("tokens", [(120000, 120000), (210000, 30000)]) -@pytest.mark.parametrize("priced", [False, True]) -async def test_batch_capture_retains_request_basis_for_later_tiered_pricing( - tokens: tuple[int, int], priced: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - import litellm - from litellm.batches.batch_utils import batch_cost_and_usage - from litellm.cost_calculator import batch_cost_calculator - from litellm.proxy.spend_tracking.usage_capture import ( - UsageBatchWriter, - capture_batch_callback, - capture_logging_usage, - ) - from litellm.types.utils import Usage - - rates: Final = { - "input_cost_per_token_batches": 0.001, - "input_cost_per_token_above_200k_tokens_batches": 0.002, - "output_cost_per_token_batches": 0.0, - "litellm_provider": "openai", - "mode": "chat", - } - monkeypatch.setitem(litellm.model_cost, "capture-model", rates if priced else {"litellm_provider": "openai"}) - rows: Final = [ - { - "response": { - "status_code": 200, - "body": { - "model": "capture-model", - "usage": {"prompt_tokens": n, "completion_tokens": 0, "total_tokens": n}, - }, - } - } - for n in tokens - ] - result: Final = batch_cost_and_usage(rows, "openai", model_info=rates if priced else {}) - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger("aretrieve_batch") - logger.model_call_details["response_cost"] = result.cost - capture_logging_usage(meter, logger, {"id": "batch", "status": "completed", "usage": result.usage}, NOW, NOW) - await capture_batch_callback( - meter, - database.load, - logger.model_call_details, - {"id": "batch", "status": "completed", "usage": result.usage}, - NOW, - NOW, - ) - await meter.close() - retained: Final = tuple( - {"model": "capture-model", "usage": json.loads(write.envelope)["usage"]} for write in database.rows[1:] - ) - assert sorted(row["usage"]["prompt_tokens"] for row in retained) == sorted(tokens) - reconstructed: Final = sum( - sum(batch_cost_calculator(Usage(**row["usage"]), row["model"], "openai", model_info=rates)) for row in retained - ) - expected: Final = sum(n * (0.002 if n > 200000 else 0.001) for n in tokens) - assert reconstructed == pytest.approx(expected) - assert result.cost == pytest.approx(expected if priced else 0) - assert database.rows[0].event.metrics is None - assert json.loads(database.rows[0].envelope)["charge_eligible"] is False - assert all(json.loads(write.envelope)["finalized_cost"] is None for write in database.rows[1:]) - - -@pytest.mark.asyncio -async def test_large_batch_recovers_partial_commit_with_frozen_pricing_and_rejects_changed_usage() -> None: - from litellm.proxy.spend_tracking.usage_capture import ( - PricingCapture, - UsageBatchWriter, - capture_batch_callback, - capture_callback, - ) - - database: Final = StoredUsage() - transactions: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - failure: Final = asyncio.Queue[bool]() - failure.put_nowait(True) - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - assert sum(len(write.envelope.encode()) for write in writes) <= 65536 - assert len(writes) <= 32 - if len(database.rows) > 1 and not failure.empty(): - failure.get_nowait() - raise RuntimeError("Interrupted batch persistence") - await database.insert(writes) - transactions.put_nowait(writes) - - writer: Final = UsageBatchWriter(persist, database.load) - meter: Final = DurableUsageMeter( - writer.write, - UsageCaptureSettings(batch_size=32, max_pending_records=32, max_pending_bytes=65536), - wait_for_batch=_manual_flush, - ) - rows: Final = tuple( - { - "model": "capture-model" if index % 7 else None, - "usage": {"prompt_tokens": index % 5 + 1} if index % 7 else None, - } - for index in range(1000) - ) - response: Final = {"id": "large-batch", "status": "completed", "usage": {"batch_usage_by_request": rows}} - pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"capture-model": {"rate": 1}}})) - kwargs: Final = { - "call_type": "aretrieve_batch", - "litellm_call_id": "poll", - "litellm_params": {}, - "_accounting_pricing_snapshot": pricing, - } - failed_database: Final = CommitBarrier(failure=RuntimeError("Earlier capture failed")) - failed_database.release.set() - prior_meter: Final = DurableUsageMeter(failed_database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - prior: Final = capture_callback( - prior_meter, - {"call_type": "acompletion", "litellm_call_id": "earlier", "litellm_params": {}}, - {"id": "earlier"}, - NOW, - NOW, - ) - await prior_meter.close() - with pytest.raises(RuntimeError, match="Earlier capture failed"): - await capture_batch_callback( - meter, - database.load, - {**kwargs, "_accounting_prior_captures": (prior,)}, - response, - NOW, - NOW, - ) - assert database.rows == () - with pytest.raises(RuntimeError, match="Interrupted batch persistence"): - await capture_batch_callback(meter, database.load, kwargs, response, NOW, NOW) - partial: Final = database.rows - assert 1 < len(partial) < len(rows) + 1 - changed_pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"capture-model": {"rate": 2}}})) - retried: Final = await capture_batch_callback( - meter, - database.load, - {**kwargs, "_accounting_pricing_snapshot": changed_pricing}, - {**response, "usage": {"batch_usage_by_request": tuple(reversed(rows))}}, - NOW, - NOW, - ) - assert retried.receipt.done() and retried.receipt.exception() is None - assert len(database.rows) == len(rows) + 1 - assert database.rows[: len(partial)] == partial - envelopes: Final = tuple(json.loads(write.envelope) for write in database.rows[1:]) - assert all( - envelope["pricing"]["model_map_information"] == ({"rate": 1} if envelope["usage"] is not None else None) - for envelope in envelopes - ) - assert sum(envelope["usage"] is None for envelope in envelopes) == sum(row["usage"] is None for row in rows) - assert all(envelope["request_count_eligible"] is False for envelope in envelopes) - assert sum(envelope["usage"]["prompt_tokens"] for envelope in envelopes if envelope["usage"] is not None) == sum( - row["usage"]["prompt_tokens"] for row in rows if row["usage"] is not None - ) - before_conflict: Final = database.rows - with pytest.raises(UsageCaptureError, match="identity conflicts"): - await capture_batch_callback( - meter, - database.load, - kwargs, - {**response, "usage": {"batch_usage_by_request": rows[:-1]}}, - NOW, - NOW, - ) - await meter.close() - assert database.rows == before_conflict - assert transactions.qsize() > 2 - assert meter.pending_bytes == meter.pending_records == 0 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("outcome", ["drain", "cancel", "failure"]) -async def test_batch_callback_handoff_is_bounded_and_runtime_owns_persistence( - outcome: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, UsageCaptureRuntime, captured_callback_kwargs - - database: Final = StoredUsage() - entered: Final = asyncio.Event() - release: Final = asyncio.Event() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - entered.set() - await release.wait() - if outcome == "failure": - raise RuntimeError("Persistence unavailable") - await database.insert(writes) - - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), load=database.load) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - response: Final = { - "id": "batch", - "status": "completed", - "usage": {"batch_usage_by_request": ({"model": "model", "usage": {"prompt_tokens": 2}},) * 3}, - } - callback: Final = asyncio.create_task(captured_callback_kwargs(kwargs, response, NOW, NOW)) - try: - await entered.wait() - assert callback.done(), "Batch persistence must not consume the legacy callback deadline" - captured: Final = (await callback)["_accounting_capture"] - assert isinstance(captured, CapturedUsage) and not captured.receipt.done() - second: Final = await captured_callback_kwargs(kwargs, response, NOW, NOW) - assert isinstance(second.get("_accounting_capture"), CapturedUsage) - overflow: Final = await captured_callback_kwargs(kwargs, response, NOW, NOW) - assert "_accounting_capture" not in overflow - assert "Batch capture producer capacity exhausted" in caplog.text - if outcome != "cancel": - release.set() - await runtime.stop(0 if outcome == "cancel" else 1) - assert captured.receipt.done() - assert captured.receipt.cancelled() is (outcome == "cancel") - if outcome == "failure": - assert isinstance(captured.receipt.exception(), RuntimeError) - assert "Batch usage capture failed" in caplog.text - assert len(database.rows) == (4 if outcome == "drain" else 0) - with pytest.raises(UsageCaptureError, match="active ledger writer"): - runtime.submit_batch(kwargs, response, NOW, NOW) - finally: - release.set() - await asyncio.gather(callback, return_exceptions=True) - await runtime.stop(1) - - -@pytest.mark.asyncio -async def test_batch_producers_share_byte_capacity_and_release_it_after_drain() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - database: Final = StoredUsage() - release: Final = asyncio.Event() - runtime: Final = UsageCaptureRuntime() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - await release.wait() - assert runtime.meter is not None - assert runtime.meter.pending_bytes <= 65536 - await database.insert(writes) - - runtime.start(persist, UsageCaptureSettings(batch_size=32, max_pending_bytes=65536), load=database.load) - from litellm.batches.batch_utils import batch_cost_and_usage - - parsed: Final = batch_cost_and_usage( - [ - { - "response": { - "status_code": 200, - "body": { - "model": "capture-model", - "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, - }, - } - } - for _ in range(1000) - ], - "openai", - model_info={}, - ) - rows: Final = ({"model": "capture-model", "usage": {"prompt_tokens": 2}},) - response: Final = {"id": "batch", "status": "completed", "usage": parsed.usage} - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - try: - large: Final = asyncio.Queue[asyncio.Future[None]]() - for index in range(100): - try: - captured: Final = runtime.submit_batch(kwargs, {**response, "id": f"batch-{index}"}, NOW, NOW) - except UsageCaptureError: - break - large.put_nowait(captured.receipt) - large_count: Final = large.qsize() - assert 0 < large_count < 100 - assert runtime.meter is not None and 0 < runtime.meter.pending_bytes <= 65536 - release.set() - await asyncio.gather(*(large.get_nowait() for _ in range(large_count))) - assert len(database.rows) == 1001 * large_count - observations: Final = tuple(json.loads(write.envelope) for write in database.rows) - assert sum(row["usage"].get("prompt_tokens", 0) for row in observations) == 2000 * large_count - assert sum(row["usage"].get("completion_tokens", 0) for row in observations) == 1000 * large_count - assert runtime.meter.pending_bytes == 0 - admitted: Final = asyncio.Queue[asyncio.Future[None]]() - for index in range(100): - try: - small: Final = runtime.submit_batch( - kwargs, - {**response, "id": f"small-{index}", "usage": {"batch_usage_by_request": rows}}, - NOW, - NOW, - ) - except UsageCaptureError: - break - admitted.put_nowait(small.receipt) - count: Final = admitted.qsize() - assert 1 < count < 100 - await asyncio.gather(*(admitted.get_nowait() for _ in range(count))) - assert len(database.rows) == 1001 * large_count + 2 * count - assert runtime.meter.pending_bytes == 0 - finally: - release.set() - await runtime.stop(1) - - -@pytest.mark.asyncio -async def test_batch_replay_reserves_the_accepted_pricing_snapshot_size() -> None: - from litellm.proxy.spend_tracking.usage_capture import PricingCapture, UsageBatchWriter, UsageCaptureRuntime - - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - runtime: Final = UsageCaptureRuntime() - runtime.start(writer.write, UsageCaptureSettings(max_pending_bytes=16384), load=database.load) - response: Final = { - "id": "batch", - "status": "completed", - "usage": {"batch_usage_by_request": ({"model": "model", "usage": {"prompt_tokens": 2}},)}, - } - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"model": {"source": "x" * 1200}}})) - try: - initial: Final = runtime.submit_batch({**kwargs, "_accounting_pricing_snapshot": pricing}, response, NOW, NOW) - await asyncio.shield(initial.receipt) - accepted: Final = database.rows - replay: Final = runtime.submit_batch(kwargs, response, NOW, NOW) - await asyncio.shield(replay.receipt) - assert database.rows == accepted and len(accepted) == 2 - assert runtime.meter is not None and runtime.meter.pending_bytes == 0 - finally: - await runtime.stop(1) diff --git a/tests/unit/proxy/spend_tracking/test_usage_meter.py b/tests/unit/proxy/spend_tracking/test_usage_meter.py deleted file mode 100644 index bd3a0b9ed70..00000000000 --- a/tests/unit/proxy/spend_tracking/test_usage_meter.py +++ /dev/null @@ -1,473 +0,0 @@ -import asyncio -import json -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite -from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings -from litellm.types.utils import CallTypes - - -def _write(identity: str) -> UsageWrite: - event: Final = InboundUsageEvent( - event_id=identity, - request_id=identity, - call_type=CallTypes.acompletion, - started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - return UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) - - -class CommitBarrier: - def __init__(self, *, conflict: str | None = None, failure: Exception | None = None) -> None: - self.entered: Final = asyncio.Event() - self.release: Final = asyncio.Event() - self.committed: tuple[str, ...] = () - self.transactions = 0 - self.conflict: Final = conflict - self.failure: Final = failure - - async def write(self, writes: tuple[UsageWrite, ...]) -> None: - self.transactions += 1 - self.entered.set() - await self.release.wait() - if self.failure is not None: - raise self.failure - identities: Final = tuple(write.event.event_id for write in writes) - if self.conflict in identities: - raise AccountingConflictError("Conflicting test event") - self.committed += identities - - -async def _manual_flush(wakeup: asyncio.Event, timeout: float) -> None: - await wakeup.wait() - - -@pytest.mark.asyncio -async def test_receipts_wait_for_one_batched_commit_and_keep_inflight_capacity_reserved() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await database.entered.wait() - assert not first.done() and not second.done() - assert meter.pending_records == 2 - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("three"),)) - database.release.set() - await asyncio.gather(first, second) - await meter.close() - assert database.committed == ("one", "two") - assert database.transactions == 1 - assert (meter.pending_bytes, meter.pending_records) == (0, 0) - - -@pytest.mark.asyncio -async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> None: - database: Final = CommitBarrier() - write: Final = _write("large" * 300) - size: Final = len(write.envelope.encode()) - meter: Final = DurableUsageMeter( - database.write, UsageCaptureSettings(max_pending_bytes=size), wait_for_batch=_manual_flush - ) - reservation: Final = meter.reserve_producer(size // 2, size // 2) - with pytest.raises(UsageCaptureError, match="producer byte capacity"): - meter.reserve_producer(1, size) - with pytest.raises(UsageCaptureError, match="producer byte capacity"): - meter.resize_producer(reservation, size, 1) - assert meter.pending_bytes == size // 2 - meter.release_producer(reservation) - receipt: Final = meter.submit((write,)) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("extra"),)) - assert meter.pending_bytes == size - database.release.set() - await meter.close() - await receipt - assert database.committed == (write.event.event_id,) - with pytest.raises(UsageCaptureError, match="closed"): - meter.submit((_write("after-close"),)) - with pytest.raises(UsageCaptureError, match="producer byte capacity"): - meter.reserve_producer(1, 1) - - -@pytest.mark.asyncio -async def test_conflicting_group_rolls_back_and_valid_neighbor_still_commits() -> None: - database: Final = CommitBarrier(conflict="conflict") - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3)) - rejected: Final = meter.submit((_write("caller"), _write("conflict"))) - accepted: Final = meter.submit((_write("valid"),)) - await database.entered.wait() - database.release.set() - with pytest.raises(UsageCaptureError, match="identity"): - await rejected - await accepted - await meter.close() - assert database.committed == ("valid",) - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_permanent_persistence_failure_reaches_each_receipt_and_releases_capacity() -> None: - failure: Final = ValueError("invalid envelope") - database: Final = CommitBarrier(failure=failure) - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await database.entered.wait() - database.release.set() - results: Final = await asyncio.gather(first, second, return_exceptions=True) - await meter.close() - assert results == [failure, failure] - assert meter.pending_bytes == 0 - - -@pytest.mark.asyncio -async def test_cancelling_record_waiter_does_not_cancel_persistence() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - waiting: Final = asyncio.create_task(meter.record(_write("caller").event)) - await database.entered.wait() - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await waiting - database.release.set() - await meter.close() - assert database.committed == ("caller",) - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_shutdown_deadline_rejects_unacknowledged_work_and_releases_capacity() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("in-flight"), _write("in-flight-two"))) - await database.entered.wait() - queued: Final = meter.submit((_write("queued"),)) - await meter.stop(0) - results: Final = await asyncio.gather(first, queued, return_exceptions=True) - assert all(isinstance(result, UsageCaptureError) for result in results) - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - assert database.committed == () - - -@pytest.mark.asyncio -async def test_shutdown_during_database_backoff_releases_reserved_capacity() -> None: - retrying: Final = asyncio.Event() - blocked: Final = asyncio.Event() - - async def unavailable(writes: tuple[UsageWrite, ...]) -> None: - raise ConnectionError("database unavailable") - - async def backoff(delay: float) -> None: - retrying.set() - await blocked.wait() - - meter: Final = DurableUsageMeter(unavailable, UsageCaptureSettings(batch_size=2), retry_delay=backoff) - receipt: Final = meter.submit((_write("one"), _write("two"))) - await retrying.wait() - await meter.stop(0) - with pytest.raises(UsageCaptureError, match="interrupted"): - await receipt - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("wrapped", [False, True]) -async def test_closed_pooled_connection_retains_batch_until_automatic_recovery(wrapped: bool) -> None: - import httpx - from prisma.engine.errors import EngineRequestError - from prisma.http import Response - - database: Final = CommitBarrier() - attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - retrying: Final = asyncio.Event() - recovered: Final = asyncio.Event() - closure: Final = EngineRequestError( - Response(httpx.Response(500)), '{"error_code":"P1017","message":"Server has closed the connection."}' - ) - failure: Final = RuntimeError("ledger transaction failed") if wrapped else closure - if wrapped: - failure.__cause__ = closure - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - attempts.put_nowait(writes) - if attempts.qsize() == 1: - raise failure - await database.write(writes) - - async def backoff(delay: float) -> None: - retrying.set() - await recovered.wait() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), retry_delay=backoff - ) - writes: Final = (_write("caller"), _write("attempt")) - receipt: Final = meter.submit(writes) - retry_signal: Final = asyncio.create_task(retrying.wait()) - reserved: Final = meter.pending_bytes - try: - await asyncio.wait((receipt, retry_signal), return_when=asyncio.FIRST_COMPLETED) - assert not receipt.done(), "A dropped pooled connection must retain admitted usage for retry" - assert (meter.pending_records, meter.pending_bytes) == (2, reserved) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("overflow"),)) - database.release.set() - recovered.set() - await receipt - assert attempts.get_nowait() == attempts.get_nowait() == writes - assert database.committed == ("caller", "attempt") - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - retry_signal.cancel() - await asyncio.gather(retry_signal, return_exceptions=True) - await meter.stop(0) - await asyncio.gather(receipt, return_exceptions=True) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("first_commits", [False, True]) -async def test_shutdown_during_conflict_isolation_completes_every_receipt(first_commits: bool) -> None: - isolating: Final = asyncio.Event() - blocked: Final = asyncio.Event() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - if len(writes) > 1: - raise AccountingConflictError("one producer conflicts") - if first_commits and writes[0].event.event_id == "one": - return - isolating.set() - await blocked.wait() - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await isolating.wait() - await meter.stop(0) - results: Final = await asyncio.gather(first, second, return_exceptions=True) - assert results[0] is None if first_commits else isinstance(results[0], UsageCaptureError) - assert isinstance(results[1], UsageCaptureError) - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -async def test_elapsed_flush_deadline_commits_a_partial_batch() -> None: - from litellm.proxy.spend_tracking.usage_meter import _wait_for_batch - - database: Final = CommitBarrier() - database.release.set() - - async def expired_flush(wakeup: asyncio.Event, timeout: float) -> None: - assert timeout == 5.0 - await _wait_for_batch(wakeup, 0) - - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=10), wait_for_batch=expired_flush) - receipt: Final = meter.submit((_write("partial-batch"),)) - try: - await receipt - assert database.committed == ("partial-batch",) - assert database.transactions == 1 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - await meter.stop(0) - - -@pytest.mark.asyncio -async def test_burst_drains_full_batches_without_waiting_for_partial_batch_deadline() -> None: - database: Final = CommitBarrier() - database.release.set() - waits: Final = asyncio.Queue[bool]() - - async def observe_wait(wakeup: asyncio.Event, timeout: float) -> None: - waits.put_nowait(wakeup.is_set()) - await wakeup.wait() - - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2), wait_for_batch=observe_wait) - receipts: Final = tuple(meter.submit((_write(str(index)),)) for index in range(5)) - try: - assert await waits.get() is True - await asyncio.gather(*receipts[:2]) - assert await waits.get() is True, "A full queued batch must not wait for the flush deadline" - await asyncio.gather(*receipts[2:4]) - assert await waits.get() is False - assert database.committed == ("0", "1", "2", "3") - assert meter.pending_records == 1 - assert not receipts[-1].done() - await meter.close() - await receipts[-1] - assert database.committed == ("0", "1", "2", "3", "4") - assert database.transactions == 3 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - await meter.stop(0) - await asyncio.gather(*receipts, return_exceptions=True) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sqlstate", ["22P05", "22003", "42P01"]) -async def test_invalid_submission_isolated_without_retrying_systemic_faults(sqlstate: str) -> None: - from prisma.errors import DataError - - database: Final = CommitBarrier() - database.release.set() - attempts: Final = asyncio.Queue[tuple[str, ...]]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - identities: Final = tuple(write.event.event_id for write in writes) - attempts.put_nowait(identities) - if "invalid" in identities or sqlstate == "42P01": - raise DataError( - data={"user_facing_error": {"message": "database rejected input", "meta": {"code": sqlstate}}} - ) - await database.write(writes) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=3)) - rejected: Final = meter.submit((_write("caller"), _write("invalid"))) - neighbor: Final = meter.submit((_write("neighbor"),)) - try: - with pytest.raises(DataError): - await rejected - if sqlstate == "42P01": - with pytest.raises(DataError): - await neighbor - assert attempts.qsize() == 1 - else: - await neighbor - assert database.committed == ("neighbor",) - assert attempts.qsize() == 3 - finally: - await meter.close() - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("dependency_outcome", ["commit", "failure", "cancel"]) -async def test_commit_dependencies_preserve_capacity_and_isolate_failed_dependents(dependency_outcome: str) -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3, max_pending_records=4)) - prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) - neighbor: Final = meter.submit((_write("neighbor"),)) - assert meter.pending_records == 3 - assert not dependent.done() - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("excess-one"), _write("excess-two"))) - if dependency_outcome == "commit": - prerequisite.set_result(None) - elif dependency_outcome == "failure": - prerequisite.set_exception(ValueError("dependency failed")) - else: - prerequisite.cancel() - database.release.set() - await meter.close() - await neighbor - if dependency_outcome == "commit": - await dependent - assert database.committed == ("parent", "final", "neighbor") - else: - with pytest.raises(UsageCaptureError, match="dependency"): - await dependent - assert database.committed == ("neighbor",) - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - - -@pytest.mark.asyncio -async def test_shutdown_does_not_cancel_dependency_receipt_or_leak_reserved_capacity() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) - await meter.stop(0) - with pytest.raises(UsageCaptureError, match="acknowledgement"): - await dependent - assert not prerequisite.done() - assert database.committed == () - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - prerequisite.set_result(None) - - -@pytest.mark.asyncio -async def test_large_write_sequence_drains_within_byte_and_record_bounds() -> None: - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter( - database.write, - UsageCaptureSettings(batch_size=8, max_pending_records=2, max_pending_bytes=1024), - wait_for_batch=_manual_flush, - ) - writes: Final = tuple(_write(str(index)) for index in range(20)) - pending: Final = meter.submit((_write("pending-a"), _write("pending-b"))) - await meter.record_writes(writes) - await pending - with pytest.raises(UsageCaptureError, match="One usage observation"): - await meter.record_writes((_write("oversized" * 300),)) - await meter.close() - assert database.committed == ("pending-a", "pending-b", *(write.event.event_id for write in writes)) - assert database.transactions > 1 - assert (meter.pending_bytes, meter.pending_records) == (0, 0) - with pytest.raises(UsageCaptureError): - await meter.record_writes((_write("closed"),)) - - -@pytest.mark.asyncio -async def test_cancelling_bounded_capture_preserves_admitted_writes_without_admitting_remaining_rows() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter( - database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2), wait_for_batch=_manual_flush - ) - recording: Final = asyncio.create_task(meter.record_writes(tuple(_write(str(index)) for index in range(8)))) - await database.entered.wait() - recording.cancel() - with pytest.raises(asyncio.CancelledError): - await recording - database.release.set() - await meter.close() - assert database.committed == ("0", "1") - assert meter.pending_bytes == meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_capacity_wait_retries_admission_before_other_producers_finish() -> None: - entered: Final = asyncio.Queue[tuple[str, ...]]() - release: Final = asyncio.Queue[None]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - entered.put_nowait(tuple(write.event.event_id for write in writes)) - await release.get() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(batch_size=2, max_pending_records=3), wait_for_batch=_manual_flush - ) - initial: Final = meter.submit((_write("first"), _write("second"))) - remaining: Final = meter.submit((_write("remaining"),)) - assert await entered.get() == ("first", "second") - recording: Final = asyncio.create_task(meter.record_writes((_write("target"),))) - try: - await asyncio.sleep(0) - release.put_nowait(None) - assert await entered.get() == ("remaining",) - await asyncio.sleep(0) - assert meter.pending_records == 2 - release.put_nowait(None) - assert await entered.get() == ("target",) - release.put_nowait(None) - await recording - await asyncio.gather(initial, remaining) - finally: - recording.cancel() - await asyncio.gather(recording, return_exceptions=True) - release.put_nowait(None) - release.put_nowait(None) - await meter.close() diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index a50db09128f..dae9ffcf7fd 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -268,76 +268,6 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm token = MASTER_KEY -@pytest.mark.asyncio -async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from datetime import datetime - - from litellm.litellm_core_utils.logging_worker import LoggingWorker - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from litellm.responses.streaming_iterator import spawn_logging_task - - worker: Final = LoggingWorker() - started: Final = asyncio.Event() - release: Final = asyncio.Event() - first_flush: Final = asyncio.Event() - pending: Final[asyncio.Queue[str]] = asyncio.Queue() - persisted: Final[asyncio.Queue[str]] = asyncio.Queue() - - captured: Final[asyncio.Queue[str]] = asyncio.Queue() - - async def persist_usage(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - captured.put_nowait(write.event.event_id) - - capture: Final = UsageCaptureRuntime() - capture.start(persist_usage, UsageCaptureSettings(flush_interval_seconds=10)) - monkeypatch.setattr(litellm.proxy.proxy_server, "usage_capture", capture) - - async def enqueue_spend() -> None: - assert capture.meter is not None - capture_callback( - capture.meter, - {"call_type": "acompletion", "litellm_call_id": "shutdown-usage", "litellm_params": {}}, - {"id": "shutdown-attempt", "usage": {"total_tokens": 3}}, - datetime(2026, 1, 1), - datetime(2026, 1, 1), - ) - pending.put_nowait("accounted request") - - async def callback() -> None: - started.set() - await release.wait() - spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting") - - async def flush_spend() -> None: - while not pending.empty(): - persisted.put_nowait(pending.get_nowait()) - first_flush.set() - - worker.ensure_initialized_and_enqueue(callback()) - await started.wait() - shutdown: Final = asyncio.create_task( - litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker) - ) - try: - await first_flush.wait() - waited_for_callback: Final = not shutdown.done() - release.set() - await shutdown - assert waited_for_callback, "Shutdown returned while accounting callbacks were still running" - assert persisted.get_nowait() == "accounted request" - assert pending.empty() - assert (captured.get_nowait(), captured.get_nowait()) == ("shutdown-usage:inbound", "shutdown-attempt") - assert capture.meter is None - finally: - release.set() - await worker.flush() - await worker.stop() - await capture.stop(0) @pytest.mark.asyncio @@ -4377,276 +4307,11 @@ async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> await worker.stop(discard_pending=True) -@pytest.mark.asyncio -@pytest.mark.parametrize("unsupported", ["database", "collector", "remote", "free_tier"]) -async def test_capture_enablement_rejects_unrecoverable_writer_modes( - unsupported: str, monkeypatch: pytest.MonkeyPatch -) -> None: - import litellm - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - monkeypatch.setattr(proxy_server, "usage_capture", UsageCaptureRuntime()) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {}}) - monkeypatch.setattr(proxy_server, "prisma_client", None if unsupported == "database" else object()) - monkeypatch.setattr(litellm, "apply_search_free_tier", unsupported == "free_tier") - monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "true" if unsupported == "collector" else "false") - monkeypatch.delenv("SPEND_LOGS_URL", raising=False) - if unsupported == "remote": - monkeypatch.setenv("SPEND_LOGS_URL", "https://writer.invalid") - with pytest.raises(ValueError, match="requires local Prisma"): - proxy_server._configure_usage_capture() - assert proxy_server.usage_capture.meter is None -@pytest.mark.asyncio -async def test_capture_configuration_requires_restart_without_replacing_the_active_writer( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.utils import PrismaClient - - runtime: Final = UsageCaptureRuntime() - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - with pytest.raises(ValueError, match="requires accounting_capture"): - proxy_server._configure_usage_capture() - client: Final = PrismaClient("postgresql://localhost/unused", ProxyLogging(UserApiKeyCache())) - monkeypatch.setattr(proxy_server, "prisma_client", client) - monkeypatch.setattr(litellm, "apply_search_free_tier", False) - monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "false") - monkeypatch.delenv("SPEND_LOGS_URL", raising=False) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 4}}) - try: - proxy_server._configure_usage_capture() - active: Final = runtime.meter - assert active is not None - assert active.settings.batch_size == 4 - proxy_server._configure_usage_capture() - assert runtime.meter is active - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 6}}) - with pytest.raises(ValueError, match="Changing accounting_capture requires a worker restart"): - proxy_server._configure_usage_capture() - monkeypatch.setattr(proxy_server, "general_settings", {}) - with pytest.raises(ValueError, match="Disabling accounting_capture requires a worker restart"): - proxy_server._configure_usage_capture() - assert runtime.meter is active - assert active.settings.batch_size == 4 - finally: - await runtime.stop(1) - assert runtime.meter is None -@pytest.mark.asyncio -@pytest.mark.parametrize("sync_ack", [False, True]) -@pytest.mark.parametrize( - ("path", "payload"), - [ - ( - "/v1/chat/completions", - {"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True}, - ), - ("/v1/responses", {"model": "openai/gpt-5-mini", "input": "hi", "background": True}), - ( - "/v1/messages", - {"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1}, - ), - ], -) -async def test_sync_ack_preserves_async_request_support( - sync_ack: bool, - path: str, - payload: dict[str, object], - monkeypatch: pytest.MonkeyPatch, - respx_mock: respx.MockRouter, -) -> None: - from fastapi import Request - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - - persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - persisted.put_nowait(write) - - async def authorize(request: Request) -> UserAPIKeyAuth: - return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path) - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key") - upstream: Final = { - "id": "provider-result", - "object": "chat.completion", - "created": 1, - "model": "gpt-5-mini", - "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, - } - if payload.get("stream"): - chunk: Final = { - **upstream, - "object": "chat.completion.chunk", - "choices": [{"index": 0, "finish_reason": "stop", "delta": {"role": "assistant", "content": "hi"}}], - } - respx_mock.post().respond( - 200, text=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n", headers={"content-type": "text/event-stream"} - ) - elif path == "/v1/responses": - respx_mock.post().respond( - 200, - json={ - "id": "resp_background", - "object": "response", - "created_at": 1, - "model": "gpt-5-mini", - "status": "in_progress", - "background": True, - "output": [], - "usage": None, - }, - ) - else: - respx_mock.post(url__regex=r".*/responses$").respond( - 200, - json={ - "id": "resp_message", - "object": "response", - "created_at": 1, - "model": "gpt-5-mini", - "status": "completed", - "output": [ - { - "id": "msg_reply", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "hi", "annotations": []}], - } - ], - "usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, - }, - ) - respx_mock.post(url__regex=r".*/chat/completions$").respond(200, json=upstream) - runtime: Final = UsageCaptureRuntime() - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, "general_settings", {"accounting_sync_ack": sync_ack, "accounting_sync_ack_timeout_seconds": 0} - ) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr( - proxy_server, - "llm_router", - litellm.Router( - model_list=[ - { - "model_name": payload["model"], - "litellm_params": {"model": payload["model"], "api_key": "test-provider-key"}, - } - ], - num_retries=0, - ), - ) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache())) - monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()]) - monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize) - runtime.start(persist, UsageCaptureSettings()) - try: - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test" - ) as client: - response: Final = await client.post(path, json=payload) - finally: - await runtime.stop(1) - assert response.status_code == 200, response.text - assert len(respx_mock.calls) == 1 - if payload.get("background"): - assert response.json()["status"] == "in_progress" - else: - assert "hi" in response.text - assert not persisted.empty(), "Terminal usage must still be captured" - -@pytest.mark.asyncio -@pytest.mark.parametrize("sync_ack", [False, True]) -async def test_text_completion_provider_rejection_captures_unknown_usage( - sync_ack: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter -) -> None: - from fastapi import Request - - from litellm.accounting import Outcome - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - - persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - persisted.put_nowait(write) - - async def authorize(request: Request) -> UserAPIKeyAuth: - return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path) - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key") - respx_mock.post().respond( - 404, - json={ - "error": {"message": "The model is unavailable", "type": "invalid_request_error", "code": "model_not_found"} - }, - ) - runtime: Final = UsageCaptureRuntime() - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": sync_ack}) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr( - proxy_server, - "llm_router", - litellm.Router( - model_list=[ - { - "model_name": "instruct", - "litellm_params": {"model": "openai/gpt-3.5-turbo-instruct", "api_key": "test-provider-key"}, - } - ], - num_retries=0, - ), - ) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache())) - monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()]) - monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize) - runtime.start(persist, UsageCaptureSettings()) - try: - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test" - ) as client: - response: Final = await client.post( - "/v1/completions", json={"model": "instruct", "prompt": "hi", "max_tokens": 1} - ) - finally: - await runtime.stop(1) - assert response.status_code == 404, response.text - assert len(respx_mock.calls) == 1 - assert persisted.qsize() == 2 - inbound: Final = persisted.get_nowait() - outbound: Final = persisted.get_nowait() - assert outbound.event.outcome is Outcome.FAILED - assert outbound.event.call_type.value == "atext_completion" - assert outbound.event.metrics is None - assert inbound.event.request_id == outbound.event.request_id - envelope: Final = json.loads(outbound.envelope) - assert envelope["usage"] is None - assert envelope["usage_source"] == "unknown" - assert envelope["pricing"]["reported_cost"] is None @pytest.mark.parametrize( diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index d28619c0938..3b000ba741f 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -3735,42 +3735,6 @@ def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) -@pytest.mark.asyncio -async def test_a_billed_responses_route_failure_captures_usage_from_the_newest_compaction_item() -> None: - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from litellm.proxy.utils import _failure_fields_to_lift - - async def persist(_: tuple[UsageWrite, ...]) -> None: - return None - - def captured_prompt_tokens(request_input: list[dict[str, str]]) -> int: - logging_obj: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params()) - logging_obj.model_call_details["call_type"] = "aresponses" - logging_obj.model_call_details["messages"] = request_input - lifted: Final = _failure_fields_to_lift( - {"litellm_logging_obj": logging_obj}, provider_billed=True, route_call_type="aresponses" - ) - snapshot: Final = lifted["_accounting_usage_snapshot"] - assert isinstance(snapshot, UsageCaptureObservation) and snapshot.usage_json is not None - return json.loads(snapshot.usage_json)["prompt_tokens"] - - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings()) - try: - sent: Final = captured_prompt_tokens( - [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] - ) - since_compaction: Final = captured_prompt_tokens( - [COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION] - ) - finally: - await runtime.stop(1) - - assert sent == since_compaction @pytest.mark.parametrize( @@ -4405,57 +4369,3 @@ def test_is_client_error_exception_existing_behaviour_unchanged() -> None: ProxyException(message="x", type="x", param=None, code="429") ) is True assert _is_client_error_exception(ValueError("x")) is False - - -@pytest.mark.asyncio -@pytest.mark.parametrize("tokens", [None, 0, 3]) -async def test_failure_capture_preserves_prepared_usage_observation(tokens: int | None) -> None: - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - - timestamp: Final = real_datetime.datetime(2026, 1, 1) - committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - committed.put_nowait(write) - - logger: Final = Logging( - model=FAILURE_USAGE_MODEL, - messages=[], - stream=False, - call_type="acompletion", - start_time=timestamp, - litellm_call_id="post-response-failure", - function_id="post-response-failure", - ) - logger.update_environment_variables( - model=FAILURE_USAGE_MODEL, custom_llm_provider="openai", litellm_params={}, optional_params={} - ) - logger.model_call_details["response_cost"] = 0.2 - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings()) - try: - prepared: Final = logger.prepare_response_cost( - {"id": "provider", "usage": None if tokens is None else {"prompt_tokens": tokens, "total_tokens": tokens}}, - timestamp, - ) - assert prepared is not None - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert lifted["combined_usage_object"] is logger.model_call_details["combined_usage_object"] - assert lifted["response_cost"] == 0.2 - assert runtime.meter is not None - capture_callback( - runtime.meter, - {**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"}, - None, - timestamp, - timestamp, - ) - finally: - await runtime.stop(1) - committed.get_nowait() - envelope: Final = json.loads(committed.get_nowait().envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == tokens - assert envelope["usage_source"] == ("unknown" if tokens is None else "measured") diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1963a0302c5..2191454a2f7 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30261,27 +30261,6 @@ export interface components { * @description Documents all the fields supported by `general_settings` in config.yaml */ ConfigGeneralSettings: { - /** @description Opt-in durable usage capture. Legacy spend updates remain active. Requires local Prisma writes. */ - accounting_capture?: components["schemas"]["UsageCaptureSettings"] | null; - /** - * Accounting Sync Ack - * @description Wait for usage commit on foreground non-stream chat, text, embedding and Responses requests. Other requests remain supported with asynchronous usage capture. - * @default false - */ - accounting_sync_ack: boolean; - /** - * Accounting Sync Ack Failure Policy - * @description Return 503 or log the persistence failure after the response wait expires - * @default fail_closed - * @enum {string} - */ - accounting_sync_ack_failure_policy: "fail_closed" | "log"; - /** - * Accounting Sync Ack Timeout Seconds - * @description Maximum response wait for usage commit, including the batch flush interval - * @default 20 - */ - accounting_sync_ack_timeout_seconds: number; /** * Admission Queue Timeout Seconds * @description maximum time a request waits for a worker slot @@ -49983,29 +49962,6 @@ export interface components { */ version: string | null; }; - /** UsageCaptureSettings */ - UsageCaptureSettings: { - /** - * Batch Size - * @default 1000 - */ - batch_size: number; - /** - * Flush Interval Seconds - * @default 5 - */ - flush_interval_seconds: number; - /** - * Max Pending Bytes - * @default 16777216 - */ - max_pending_bytes: number; - /** - * Max Pending Records - * @default 10000 - */ - max_pending_records: number; - }; /** UsageChartPoint */ UsageChartPoint: { /** Blocked */ From 0d70b9b9a1cecd8094339c0a2667b68c7e1bdce0 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 10:39:46 -0800 Subject: [PATCH 004/178] test(perf): remove the database_upgrade_5rps e2e case and its harness tests (#50076) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/perf/setup/upgrade.py | 223 -------------- tests/e2e/perf/test_database_upgrade.py | 102 ------- tests/e2e_harness/perf/test_run.py | 24 +- tests/e2e_harness/perf/test_upgrade.py | 375 ------------------------ 4 files changed, 23 insertions(+), 701 deletions(-) delete mode 100644 tests/e2e/perf/setup/upgrade.py delete mode 100644 tests/e2e/perf/test_database_upgrade.py delete mode 100644 tests/e2e_harness/perf/test_upgrade.py diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py deleted file mode 100644 index c9114ccbff2..00000000000 --- a/tests/e2e/perf/setup/upgrade.py +++ /dev/null @@ -1,223 +0,0 @@ -from __future__ import annotations - -import math -import time -from collections.abc import Callable -from dataclasses import replace -from typing import Final - -import psycopg -from cases.small_bodies_5rps import CASE -from e2e_config import MASTER_KEY -from e2e_http import send_text_external -from prometheus_client.parser import text_string_to_metric_families -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter - -from setup.performance_test_case import PerformanceTestCase -from setup.upgrade_worker import UpgradeState, fingerprints, missing_indexes, read_state, release_indexes -from setup.wire import WireRecord - -MAX_TRAFFIC_SECONDS: Final = 6 * 3600 -_ROWS: Final = TypeAdapter(tuple[tuple[str, str | None], ...]) -_ACTIVITY: Final = TypeAdapter(tuple[list[int], int]) -_COUNT: Final = TypeAdapter(tuple[int]) -_SERVICES: Final = ("postgres", "batch_write_to_db") - - -def upgrade_case() -> PerformanceTestCase: - return replace( - CASE, - name="database_upgrade_5rps", - traffic=replace(CASE.traffic, total_requests=int(CASE.traffic.rps * MAX_TRAFFIC_SECONDS)), - limits=tuple( - limit - for limit in CASE.limits - if limit.metric - in { - "unexpected_status_count", - "gateway_added_seconds", - "time_to_first_token_seconds", - "memory_growth_megabytes", - } - ), - ) - - -class MetricsHeaders(BaseModel): - model_config = ConfigDict(frozen=True) - authorization: str = Field(default=f"Bearer {MASTER_KEY}", alias="Authorization") - - -def postgres_failures(text: str) -> float: - families: Final = tuple(text_string_to_metric_families(text)) - names: Final = frozenset(family.name for family in families) - required: Final = frozenset(f"litellm_{service}_failed_requests" for service in _SERVICES) - samples: Final = tuple(sample for family in families for sample in family.samples) - has_database_activity: Final = any( - sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0 - for sample in samples - ) - assert required <= names or has_database_activity, f"Missing database service counters: {required - names}" - failures: Final = tuple(sample.value for sample in samples if sample.name in {f"{name}_total" for name in required}) - assert all(math.isfinite(value) and value >= 0 for value in failures), "Invalid database failure counter" - return sum(failures) - - -def assert_database_metrics(urls: tuple[str, ...]) -> None: - assert urls, "E2E_DATABASE_METRICS_URLS must name every gateway/backend /metrics endpoint" - for url in urls: - assert_database_metrics_url(url) - - -def assert_database_metrics_url(url: str) -> None: - result: Final = send_text_external("GET", url, headers=MetricsHeaders(), timeout=10) - assert result.status_code == 200, f"Database metrics unavailable at {url}: {result.status_code}" - assert postgres_failures(result.body) == 0, f"Failed PostgreSQL operations reported at {url}" - - -def persisted(connection: psycopg.Connection[tuple[object, ...]], records: tuple[WireRecord, ...]) -> frozenset[str]: - ids: Final = tuple(record.call_id for record in records if record.call_id) - if not ids: - return frozenset() - rows: Final = _ROWS.validate_python( - connection.execute( - 'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=ANY(%s)', (list(ids),) - ).fetchall() - ) - return matched_response_ids(records, rows) - - -def matched_response_ids(records: tuple[WireRecord, ...], rows: tuple[tuple[str, str | None], ...]) -> frozenset[str]: - call_ids: Final = frozenset(call_id for _, call_id in rows if call_id) - return frozenset(record.response_id for record in records if record.response_id and record.call_id in call_ids) - - -def activity(connection: psycopg.Connection[tuple[object, ...]]) -> tuple[list[int], int]: - return _ACTIVITY.validate_python( - connection.execute("""SELECT - ARRAY(SELECT pid FROM pg_stat_progress_create_index WHERE datname=current_database()), - (SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock') - """).fetchone() - ) - - -def assert_finished(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None: - assert state.phase == "complete", f"Upgrade did not finish: {state.phase}" - assert not missing_indexes(connection, state.required_indexes), "Required target indexes are missing or invalid" - (missing,) = _COUNT.validate_python( - connection.execute( - """ - SELECT count(*) FROM unnest(%s::text[]) name WHERE NOT EXISTS ( - SELECT 1 FROM "_prisma_migrations" WHERE migration_name=name - AND finished_at IS NOT NULL AND rolled_back_at IS NULL)""", - (list(state.required_migrations),), - ).fetchone() - ) - assert missing == 0, f"{missing} target migrations were not applied" - (unfinished,) = _COUNT.validate_python( - connection.execute(""" - SELECT count(*) FROM "_prisma_migrations" WHERE finished_at IS NULL AND rolled_back_at IS NULL - """).fetchone() - ) - assert unfinished == 0, f"{unfinished} unfinished migrations" - assert fingerprints(connection, tuple(state.baseline_rows)) == state.baseline_rows, ( - "Baseline rows changed or disappeared" - ) - - -class UpgradeMonitor: - def __init__( - self, - *, - state: Callable[[], UpgradeState], - check_metrics: Callable[[], None], - read_activity: Callable[[], tuple[list[int], int]], - read_persisted: Callable[[tuple[WireRecord, ...]], frozenset[str]], - start_indexes: Callable[[], None], - clock: Callable[[], float] = time.monotonic, - wall_clock: Callable[[], float] = time.time, - ) -> None: - self.state: Final = state - self.check_metrics: Final = check_metrics - self.read_activity: Final = read_activity - self.read_persisted: Final = read_persisted - self.start_indexes: Final = start_indexes - self.clock: Final = clock - self.wall_clock: Final = wall_clock - self.started: Final = clock() - self.status_limit: Final = next( - limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count" and limit.max is not None - ) - self.pending: tuple[WireRecord, ...] = () - self.processed = 0 - self.successes = 0 - self.persisted_count = 0 - self.overlap_writes = 0 - self.max_lock_waiters = 0 - self.previous_builds: frozenset[int] = frozenset() - self.build_seen_at: float | None = None - self.released = False - self.completed_at: float | None = None - - def poll(self, records: tuple[WireRecord, ...]) -> bool: - self.check_metrics() - current: Final = self.state() - assert current.phase != "failed", "Target migration/index worker failed" - assert self.clock() - self.started < MAX_TRAFFIC_SECONDS, "Upgrade exceeded its six-hour operational timeout" - successful: Final = tuple(record for record in records if record.status == 200) - self.processed += len(records) - self.successes += len(successful) - assert self.processed - self.successes <= self.status_limit, "Too many unexpected non-200 responses" - assert all(record.ok and record.text for record in successful), "Failed or empty inference response" - assert all(record.response_id and record.call_id for record in successful), ( - "Response missing spend-log identifiers" - ) - builds, lock_waiters = self.read_activity() - if not self.previous_builds.intersection(builds): - self.build_seen_at = self.wall_clock() if builds else None - self.max_lock_waiters = max(self.max_lock_waiters, lock_waiters) - self.pending += successful - found: Final = self.read_persisted(self.pending) - self.persisted_count += len(found) - still_building, _ = self.read_activity() - since: Final = self.build_seen_at - if since is not None and frozenset(builds).intersection(still_building): - self.overlap_writes += sum( - 1 for record in self.pending if record.response_id in found and record.started_at >= since - ) - self.previous_builds = frozenset(builds) - self.pending = tuple(record for record in self.pending if record.response_id not in found) - if self.successes and not self.released: - self.start_indexes() - self.released = True - if current.phase == "complete" and self.completed_at is None: - self.completed_at = self.clock() - return self.completed_at is not None and self.clock() - self.completed_at >= 60 and self.successes >= 300 - - def assert_complete(self, state: UpgradeState) -> None: - assert self.successes >= 300, f"Only {self.successes} requests completed" - assert not self.pending, f"{len(self.pending)} successful requests have no matching spend log" - assert self.persisted_count == self.successes, "Spend-log count does not match successful requests" - assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work" - - -def upgrade_properties(state: UpgradeState, monitor: UpgradeMonitor) -> tuple[tuple[str, float], ...]: - """The upgrade's own numbers under the `perf_` prefix the dashboard ingests, next to the perf rows.""" - return ( - ("perf_upgrade_schema_seconds", state.schema_seconds), - ("perf_upgrade_index_seconds", state.index_seconds), - ("perf_upgrade_overlap_writes", float(monitor.overlap_writes)), - ("perf_upgrade_max_lock_waiters", float(monitor.max_lock_waiters)), - ("perf_upgrade_pending_indexes", float(len(state.pending_indexes))), - ("perf_upgrade_unmatched_requests", float(len(monitor.pending))), - ) - - -def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor: - return UpgradeMonitor( - state=lambda: read_state(connection), - check_metrics=lambda: assert_database_metrics(urls), - read_activity=lambda: activity(connection), - read_persisted=lambda records: persisted(connection, records), - start_indexes=lambda: release_indexes(connection), - ) diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py deleted file mode 100644 index a9f5cc43fe6..00000000000 --- a/tests/e2e/perf/test_database_upgrade.py +++ /dev/null @@ -1,102 +0,0 @@ -from __future__ import annotations - -import os -import sys -import time -from pathlib import Path -from typing import Final - -import pytest -from idp import Keycloak -from lifecycle import ResourceManager -from memory_readings import read_rss_everywhere -from proxy_client import ProxyClient -from setup.datastore import read_db_counters, read_redis_counters -from setup.observe import ( - Observation, - evaluate, - junit_properties, - read_mock_requests, - render, - reset_mock, - set_mock_latency, -) -from setup.plan import build_plan -from setup.podcpu import gateway_cpu_sampler -from setup.provision import MOCK_PROVIDERS, build_setup, mock_url -from setup.run import run_plan -from setup.upgrade import assert_finished, monitor_for, upgrade_case, upgrade_properties -from setup.upgrade_worker import connect, read_state - -pytestmark = pytest.mark.perf - - -@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only") -def test_database_upgrade( - proxy: ProxyClient, idp: Keycloak, resources: ResourceManager, request: pytest.FixtureRequest -) -> None: - case: Final = upgrade_case() - urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS) - for url in urls: - _ = reset_mock(url) - _ = set_mock_latency(url, case.mock_latency) - setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources) - metrics: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip()) - with connect() as connection: - assert read_state(connection).phase == "schema_ready", "Expected the target schema ready, index work waiting" - rss_before: Final = read_rss_everywhere(proxy) - db_before: Final = read_db_counters() - redis_before: Final = read_redis_counters() - monitor: Final = monitor_for(connection, metrics) - with gateway_cpu_sampler() as cpu: - outcome: Final = run_plan(case, build_plan(case), setup, stop_when=monitor.poll, request_timeout=30) - assert outcome.locust_exit == 0, outcome.locust_output - remaining: Final = outcome.records[monitor.processed :] - _ = monitor.poll(remaining) - deadline: Final = time.monotonic() + 120 - while monitor.pending and time.monotonic() < deadline: - time.sleep(1) - _ = monitor.poll(()) - state: Final = read_state(connection) - _ = sys.stdout.write( - f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} " - f"unmatched_requests={len(monitor.pending)}\n" - ) - rss_after: Final = read_rss_everywhere(proxy) - observed: Final = Observation( - records=outcome.records, - mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)), - guardrail_calls=(), - rss_before=rss_before, - rss_after=rss_after, - wall_seconds=outcome.wall_seconds, - spend_logged_share=1, - db_before=db_before, - db_after=read_db_counters(), - redis_before=redis_before, - redis_after=read_redis_counters(), - gateway_cpu_seconds=cpu.seconds, - ) - rows, exceeded = evaluate(case, observed) - item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped - assert isinstance(item, pytest.Item) - item.user_properties[:] = [ - entry for entry in item.user_properties if not entry[0].startswith("perf_") - ] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set - item.user_properties.extend(junit_properties(rows)) - item.user_properties.extend(upgrade_properties(state, monitor)) - _ = sys.stdout.write(render(case, observed, rows) + "\n") - _ = sys.stdout.write( - f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} " - f"overlap_writes={monitor.overlap_writes} max_lock_waiters={monitor.max_lock_waiters} " - f"pending_indexes={state.pending_indexes}\n" - ) - artifact: Final = os.environ.get("E2E_ARTIFACT_DIR") - if artifact: - _ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8") - monitor.assert_complete(state) - assert_finished(connection, state) - assert {(r.replica, r.hostname) for r in rss_before.readings} == { - (r.replica, r.hostname) for r in rss_after.readings - }, "Gateway pods replaced during the upgrade" - assert not exceeded, "\n".join(exceeded) diff --git a/tests/e2e_harness/perf/test_run.py b/tests/e2e_harness/perf/test_run.py index 23f62c9c560..480af96a12a 100644 --- a/tests/e2e_harness/perf/test_run.py +++ b/tests/e2e_harness/perf/test_run.py @@ -8,7 +8,29 @@ from typing import Final import pytest from setup.run import run_until from setup.wire import WireRecord -from test_upgrade import record + + +def record(seq: int = 0, *, started_at: float = 100) -> WireRecord: + return WireRecord( + seq=seq, + marker=f"request-{seq}", + caller="key", + endpoint="chat_completions", + stream=False, + blocked=False, + started_at=started_at, + status=200, + header_s=1, + first_byte_s=1, + first_token_s=None, + total_s=1, + text="answer", + ok=True, + mock_delay_s=1, + call_id=f"call-{seq}", + error=None, + response_id=f"response-{seq}", + ) class StopAfterRecord: diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py deleted file mode 100644 index 3414476ba4c..00000000000 --- a/tests/e2e_harness/perf/test_upgrade.py +++ /dev/null @@ -1,375 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Final - -import pytest -from cases.small_bodies_5rps import CASE -from memory_readings import RssCapture, RssReading -from setup.observe import Observation, evaluate -from setup.upgrade import ( - MAX_TRAFFIC_SECONDS, - UpgradeMonitor, - matched_response_ids, - postgres_failures, - upgrade_case, - upgrade_properties, -) -from setup.upgrade_worker import UpgradeState -from setup.wire import PlanProgress, WireRecord, response_id - - -def record(seq: int = 0, *, started_at: float = 100) -> WireRecord: - return WireRecord( - seq=seq, - marker=f"request-{seq}", - caller="key", - endpoint="chat_completions", - stream=False, - blocked=False, - started_at=started_at, - status=200, - header_s=1, - first_byte_s=1, - first_token_s=None, - total_s=1, - text="answer", - ok=True, - mock_delay_s=1, - call_id=f"call-{seq}", - error=None, - response_id=f"response-{seq}", - ) - - -def test_spend_logs_match_call_ids_when_response_ids_are_transformed() -> None: - records: Final = (record(0), record(1), record(2)) - rows: Final = (("provider-response-id", "call-0"), ("response-1", "unrelated-call")) - assert matched_response_ids(records, rows) == frozenset({"response-0"}) - - -def state(phase: str = "schema_ready", *, pending: bool = True) -> UpgradeState: - return UpgradeState.model_validate( - { - "phase": phase, - "started_at": 0, - "required_indexes": ["index"], - "required_migrations": ["migration"], - "pending_indexes": ["index"] if pending else [], - "baseline_rows": {"seed-row": "hash"}, - } - ) - - -class Probe: - def __init__(self) -> None: - self.now = 0.0 - self.current = state() - self.builds: list[int] = [] - self.found: frozenset[str] = frozenset() - self.releases = 0 - self.metric_checks = 0 - self.errors = False - - def check_metrics(self) -> None: - self.metric_checks += 1 - assert not self.errors, "Failed PostgreSQL operation" - - def persisted(self, records: tuple[WireRecord, ...]) -> frozenset[str]: - return frozenset(row.response_id for row in records if row.response_id in self.found and row.response_id) - - def release(self) -> None: - self.releases += 1 - - def monitor(self) -> UpgradeMonitor: - return UpgradeMonitor( - state=lambda: self.current, - check_metrics=self.check_metrics, - read_activity=lambda: (self.builds, 2), - read_persisted=self.persisted, - start_indexes=self.release, - clock=lambda: self.now, - wall_clock=lambda: self.now, - ) - - -def test_stop_drains_issued_requests_without_issuing_the_rest() -> None: - progress: Final = PlanProgress(100) - assert progress.take(stop=False) - assert progress.take(stop=False) - progress.completed = 1 - assert not progress.take(stop=True) - assert not progress.drained - assert not progress.take(stop=False) - progress.completed = 2 - assert progress.drained - assert progress.issued == 2 - - -def test_finite_plan_still_waits_for_its_last_response() -> None: - progress: Final = PlanProgress(1) - assert progress.take(stop=False) - assert not progress.take(stop=False) - assert not progress.drained - progress.completed = 1 - assert progress.drained - - -@pytest.mark.parametrize( - "payload,expected", - [ - ({"id": "chat-id"}, "chat-id"), - ({"type": "response.created", "response": {"id": "response-id"}}, "response-id"), - ({"type": "message_start", "message": {"id": "message-id"}}, "message-id"), - ({"type": "response.output_item.added", "item": {"id": "item-id"}}, None), - ({"type": "content_block_delta", "delta": {"text": "hello"}}, None), - ], -) -def test_response_id_uses_response_not_output_item(payload: Mapping[str, object], expected: str | None) -> None: - assert response_id(payload) == expected - - -def metrics(*, failure: str = "0", batch: str = "0") -> str: - return f"""# TYPE litellm_postgres_failed_requests_total counter -litellm_postgres_failed_requests_total{{function_name="query",error_class="Timeout"}} {failure} -# TYPE litellm_batch_write_to_db_failed_requests_total counter -litellm_batch_write_to_db_failed_requests_total{{function_name="flush",error_class="DatabaseError"}} {batch} -""" - - -def test_failed_sql_and_async_batch_writes_are_both_counted() -> None: - assert postgres_failures(metrics(failure="2", batch="3")) == 5 - assert postgres_failures(metrics()) == 0 - - -def test_no_error_label_series_is_a_valid_zero() -> None: - assert ( - postgres_failures("""# TYPE litellm_postgres_failed_requests_total counter -# TYPE litellm_batch_write_to_db_failed_requests_total counter -""") - == 0 - ) - - -@pytest.mark.parametrize("service", [None, "postgres", "batch_write_to_db"]) -def test_multiprocess_metrics_omit_failure_families_until_a_failure(service: str | None) -> None: - total: Final = """# TYPE litellm_postgres_total_requests_total counter -litellm_postgres_total_requests_total{postgres="postgres"} 1 -""" - failure: Final = ( - f"# TYPE litellm_{service}_failed_requests_total counter\n" - f'litellm_{service}_failed_requests_total{{function_name="query",error_class="Timeout"}} 1\n' - if service is not None - else "" - ) - assert postgres_failures(total + failure) == (1 if service is not None else 0) - - -@pytest.mark.parametrize("total", ["0", "-1", "NaN", "+Inf"]) -def test_missing_failure_families_require_observed_database_activity(total: str) -> None: - with pytest.raises(AssertionError): - _ = postgres_failures( - "# TYPE litellm_postgres_total_requests_total counter\n" - f'litellm_postgres_total_requests_total{{postgres="postgres"}} {total}\n' - ) - - -@pytest.mark.parametrize("body", ["", "# TYPE unrelated counter\n", metrics(failure="NaN"), metrics(batch="-1")]) -def test_missing_or_invalid_counters_do_not_pass_as_zero(body: str) -> None: - with pytest.raises(AssertionError): - _ = postgres_failures(body) - - -def test_index_completion_stops_after_the_post_upgrade_window() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - assert not monitor.poll((record(),)) - assert probe.releases == 1 - probe.builds = [17] - probe.now = 101 - assert not monitor.poll(()) - probe.found = frozenset({"response-0", "response-1"}) - probe.now = 103 - assert not monitor.poll((record(1, started_at=102),)) - assert monitor.overlap_writes == 1 - assert monitor.persisted_count == 2 - assert not monitor.pending - assert monitor.max_lock_waiters == 2 - probe.current = state("complete") - assert not monitor.poll(()) - probe.now = 162 - assert not monitor.poll(()) - probe.now = 163 - assert not monitor.poll(()) - probe.found = frozenset(f"response-{i}" for i in range(300)) - assert monitor.poll(tuple(record(i) for i in range(2, 300))) - assert probe.releases == 1 - assert probe.metric_checks == 7 - - -def test_http_200_without_spend_row_cannot_complete() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - records: Final = tuple(record(i) for i in range(300)) - _ = monitor.poll(records) - with pytest.raises(AssertionError, match="no matching spend log"): - monitor.assert_complete(state("complete", pending=False)) - - -def test_no_pending_index_work_still_requires_300_persisted_requests() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - records: Final = tuple(record(i) for i in range(300)) - probe.found = frozenset(f"response-{i}" for i in range(300)) - _ = monitor.poll(records) - monitor.assert_complete(state("complete", pending=False)) - assert monitor.successes == monitor.persisted_count == 300 - with pytest.raises(AssertionError, match="during index work"): - monitor.assert_complete(state("complete", pending=True)) - - -def test_postgres_failure_fails_even_when_inference_succeeded() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - probe.errors = True - with pytest.raises(AssertionError, match="PostgreSQL"): - _ = monitor.poll((record(),)) - - -def test_failed_migration_and_operational_timeout_are_not_success() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - probe.current = state("failed") - with pytest.raises(AssertionError, match="worker failed"): - _ = monitor.poll(()) - probe.current = state() - probe.now = MAX_TRAFFIC_SECONDS - with pytest.raises(AssertionError, match="operational timeout"): - _ = monitor.poll(()) - - -def test_upgrade_case_uses_5rps_with_the_same_request_and_memory_limits() -> None: - case: Final = upgrade_case() - assert case.traffic.rps == 5 - assert not {"cpu_seconds_per_request", "postgres_calls_per_request", "redis_calls_per_request"}.intersection( - limit.metric for limit in case.limits - ) - assert case.limits == tuple( - limit - for limit in CASE.limits - if limit.metric - in { - "unexpected_status_count", - "gateway_added_seconds", - "time_to_first_token_seconds", - "memory_growth_megabytes", - } - ) - - -def failed_record(seq: int) -> WireRecord: - return record(seq).model_copy(update={"status": 503, "ok": False, "text": "", "response_id": None, "call_id": None}) - - -def test_three_non_200s_do_not_start_indexes_or_require_spend_rows() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - assert not monitor.poll(tuple(failed_record(i) for i in range(3))) - assert probe.releases == 0 - assert monitor.processed == 3 - assert monitor.successes == 0 - assert not monitor.pending - probe.found = frozenset(f"response-{i}" for i in range(3, 303)) - assert not monitor.poll(tuple(record(i) for i in range(3, 303))) - monitor.assert_complete(state("complete", pending=False)) - assert probe.releases == 1 - assert monitor.processed == 303 - assert monitor.successes == monitor.persisted_count == 300 - - -def test_fourth_non_200_fails_across_poll_batches() -> None: - monitor: Final = Probe().monitor() - assert not monitor.poll(tuple(failed_record(i) for i in range(3))) - with pytest.raises(AssertionError, match="Too many unexpected non-200"): - _ = monitor.poll((failed_record(3),)) - - -@pytest.mark.parametrize( - "failed", [record().model_copy(update={"text": ""}), record().model_copy(update={"ok": False})] -) -def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None: - with pytest.raises(AssertionError, match="Failed or empty inference"): - _ = Probe().monitor().poll((failed,)) - - -def _ceiling(metric: str, percentile: int | None = None) -> float: - return next( - limit.max - for limit in upgrade_case().limits - if limit.metric == metric and limit.percentile == percentile and limit.max is not None - ) - - -@pytest.mark.parametrize( - "over,expected", - [ - ("", ()), - ("unexpected_status_count", ("unexpected_status_count",)), - ("gateway_added_seconds", ("gateway_added_seconds percentile 50",)), - ("time_to_first_token_seconds", ("time_to_first_token_seconds percentile 99",)), - ("memory_growth_megabytes", ("memory_growth_megabytes",)), - ], -) -def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples( - over: str, expected: tuple[str, ...] -) -> None: - """Each sample sits exactly on the case's ceiling, so only the one nudged past it goes over.""" - failures: Final = int(_ceiling("unexpected_status_count")) + (over == "unexpected_status_count") - overhead: Final = _ceiling("gateway_added_seconds", 50) * (1.01 if over == "gateway_added_seconds" else 1) - ttft: Final = _ceiling("time_to_first_token_seconds", 99) * (1.01 if over == "time_to_first_token_seconds" else 1) - growth: Final = _ceiling("memory_growth_megabytes") + (over == "memory_growth_megabytes") - nonstream: Final = tuple(record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100)) - stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft}) - observed: Final = Observation( - records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))), - mock_requests=(), - guardrail_calls=(), - rss_before=RssCapture((RssReading("gateway", "host", 1, 1000),), ()), - rss_after=RssCapture((RssReading("gateway", "host", 1, 1000 + growth),), ()), - wall_seconds=60, - ) - rows, exceeded = evaluate(upgrade_case(), observed) - assert tuple(row.label for row in rows if row.over or row.missing) == expected - assert len(exceeded) == len(expected) - - -def test_upgrade_properties_carry_the_upgrade_timings_and_what_the_monitor_saw() -> None: - state: Final = UpgradeState( - phase="complete", - started_at=0, - schema_seconds=12.5, - index_seconds=240.25, - required_indexes=("idx_a", "idx_b"), - required_migrations=("20260101_a",), - pending_indexes=("idx_b",), - baseline_rows={}, - ) - monitor: Final = UpgradeMonitor( - state=lambda: state, - check_metrics=lambda: None, - read_activity=lambda: ([], 0), - read_persisted=lambda records: frozenset(), - start_indexes=lambda: None, - ) - monitor.overlap_writes = 37 - monitor.max_lock_waiters = 3 - monitor.pending = (record(0), record(1)) - - assert dict(upgrade_properties(state, monitor)) == { - "perf_upgrade_schema_seconds": 12.5, - "perf_upgrade_index_seconds": 240.25, - "perf_upgrade_overlap_writes": 37, - "perf_upgrade_max_lock_waiters": 3, - "perf_upgrade_pending_indexes": 1, - "perf_upgrade_unmatched_requests": 2, - } From b662f3f7a441cfd85ba89be5e0a8a78b9ea3fb09 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 10:41:16 -0800 Subject: [PATCH 005/178] feat(decisions): add EmpirioLabs AI as a decisions provider (#49977) * feat(decisions): add EmpirioLabs AI as a decisions provider empiriolabs/aplomb-1 now works with litellm.decisions and litellm.systemone, the proxy /v1/decisions and /v1/systemone routes, the decision model guardrail and the Admin UI. Requests go to POST /v1/decisions on the EmpirioLabs AI API with the System One body, keyed by EMPIRIOLABS_API_KEY with an optional EMPIRIOLABS_API_BASE Chat models keep routing through the existing OpenAI-compatible JSON provider, and a test pins that * ci: run the EmpirioLabs AI decisions tests in the llms unit shard * fix: align EmpirioLabs registration with CI gates --------- Co-authored-by: Adam Dalloul <47503782+Adam-Dalloul@users.noreply.github.com> Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + README.md | 1 + litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 2 + .../empiriolabs/decisions/transformation.py | 10 ++ .../provider_endpoints_support_backup.json | 2 +- .../provider_create_fields.json | 28 ++++ litellm/types/utils.py | 1 + litellm/utils.py | 2 + provider_endpoints_support.json | 2 +- .../providers/test_decisions_wire.py | 1 + tests/unit/decisions/test_main.py | 143 ++++++++++++++++++ tests/unit/llms/empiriolabs/__init__.py | 0 .../llms/empiriolabs/decisions/__init__.py | 0 ...st_empiriolabs_decisions_transformation.py | 40 +++++ .../openai_like/test_empiriolabs_provider.py | 53 +++++++ .../public_endpoints/test_public_endpoints.py | 1 + .../components/provider_info_helpers.test.tsx | 19 +++ .../src/components/provider_info_helpers.tsx | 3 + 19 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 litellm/llms/empiriolabs/decisions/transformation.py create mode 100644 tests/unit/llms/empiriolabs/__init__.py create mode 100644 tests/unit/llms/empiriolabs/decisions/__init__.py create mode 100644 tests/unit/llms/empiriolabs/decisions/test_empiriolabs_decisions_transformation.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 0a490bb8f45..9370de85109 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -280,6 +280,7 @@ jobs: tests/unit/llms/e2b tests/unit/llms/edenai tests/unit/llms/elevenlabs + tests/unit/llms/empiriolabs tests/unit/llms/exa_ai tests/unit/llms/fal_ai tests/unit/llms/fastcrw diff --git a/README.md b/README.md index c4aeee80fb4..31a93bf8c67 100644 --- a/README.md +++ b/README.md @@ -398,6 +398,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | | | [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | | +| [EmpirioLabs AI (`empiriolabs`)](https://docs.litellm.ai/docs/providers/empiriolabs) | ✅ | | ✅ | | | | | | | | | [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | | | [eu/jev Decisions API (`eu_jev`)](https://docs.litellm.ai/docs/pass_through/eu_jev) | | | | | | | | | | | | [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | | diff --git a/litellm/__init__.py b/litellm/__init__.py index c1f4bb6a3df..8e5dd68a186 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1796,6 +1796,9 @@ if TYPE_CHECKING: from .llms.bespoke.decisions.transformation import ( BespokeDecisionsConfig as BespokeDecisionsConfig, ) + from .llms.empiriolabs.decisions.transformation import ( + EmpirioLabsDecisionsConfig as EmpirioLabsDecisionsConfig, + ) from .llms.eu_jev.decisions.transformation import ( EuJevDecisionsConfig as EuJevDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 65074093dbb..bf903e23bc0 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = ( "OllamaDecisionsConfig", "LayaDecisionsConfig", "BespokeDecisionsConfig", + "EmpirioLabsDecisionsConfig", "EuJevDecisionsConfig", "OpenAIDecisionsConfig", "AzureDecisionsConfig", @@ -759,6 +760,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "OllamaDecisionsConfig": (".llms.ollama.decisions.transformation", "OllamaDecisionsConfig"), "LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"), "BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"), + "EmpirioLabsDecisionsConfig": (".llms.empiriolabs.decisions.transformation", "EmpirioLabsDecisionsConfig"), "EuJevDecisionsConfig": (".llms.eu_jev.decisions.transformation", "EuJevDecisionsConfig"), "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"), "AzureDecisionsConfig": (".llms.azure.decisions.transformation", "AzureDecisionsConfig"), diff --git a/litellm/llms/empiriolabs/decisions/transformation.py b/litellm/llms/empiriolabs/decisions/transformation.py new file mode 100644 index 00000000000..ce3464310c4 --- /dev/null +++ b/litellm/llms/empiriolabs/decisions/transformation.py @@ -0,0 +1,10 @@ +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class EmpirioLabsDecisionsConfig(BaseDecisionsConfig): + path = "/v1/decisions" + api_key_env = ("EMPIRIOLABS_API_KEY",) + api_base_env = ("EMPIRIOLABS_API_BASE",) + + def get_default_api_base(self) -> str | None: + return "https://api.empiriolabs.ai" diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 672a2739964..85bd18c7df5 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -5277,7 +5277,7 @@ } }, "empiriolabs": { - "display_name": "EmpirioLabs (`empiriolabs`)", + "display_name": "EmpirioLabs AI (`empiriolabs`)", "url": "https://docs.litellm.ai/docs/providers/empiriolabs", "endpoints": { "chat_completions": false, diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 9d4ede7b90d..5ded2fa5f5e 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1553,6 +1553,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "EMPIRIOLABS", + "provider_display_name": "EmpirioLabs AI", + "litellm_provider": "empiriolabs", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.empiriolabs.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "empiriolabs/aplomb-1" + }, { "provider": "EMPOWER", "provider_display_name": "Empower", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index dc71ba149da..4c47489e08f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4650,6 +4650,7 @@ class LlmProviders(str, Enum): STRANDS_DECIDER = "strands_decider" LAYA = "laya" BESPOKE = "bespoke" + EMPIRIOLABS = "empiriolabs" EU_JEV = "eu_jev" MISTRAL = "mistral" MILVUS = "milvus" diff --git a/litellm/utils.py b/litellm/utils.py index e210b2fe026..4add9d53cb3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10296,6 +10296,7 @@ class ProviderConfigManager: if provider == LlmProviders.BESPOKE: return litellm.BespokeDecisionsConfig() if provider in ( + LlmProviders.EMPIRIOLABS, LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI, @@ -10303,6 +10304,7 @@ class ProviderConfigManager: LlmProviders.PARALLEL_AI, ): configs: Final = { + LlmProviders.EMPIRIOLABS: litellm.EmpirioLabsDecisionsConfig, LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig, LlmProviders.AZURE: litellm.AzureDecisionsConfig, LlmProviders.AZURE_AI: litellm.AzureAIDecisionsConfig, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 672a2739964..85bd18c7df5 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -5277,7 +5277,7 @@ } }, "empiriolabs": { - "display_name": "EmpirioLabs (`empiriolabs`)", + "display_name": "EmpirioLabs AI (`empiriolabs`)", "url": "https://docs.litellm.ai/docs/providers/empiriolabs", "endpoints": { "chat_completions": false, diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 4e091460667..bb4973d6f2f 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -139,6 +139,7 @@ _PROVIDERS: Final = ( False, "databricks/databricks-openjev-qwen35-4b", ), + _Provider("empiriolabs", "empiriolabs/aplomb-test", "/v1/decisions", "aplomb-test", _API_KEY, False, None), ) _LOCAL_PROVIDERS: Final = ( pytest.param( diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 789d59bd613..1e42e74806d 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -132,6 +132,7 @@ _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = ( "https://openrouter.ai/api/alpha/decisions", "typesafe/jev-1.13", ), + ("empiriolabs", "empiriolabs/aplomb-1", "https://api.empiriolabs.ai/v1/decisions", "aplomb-1"), ) _OPENAI_RESPONSE: Final[Mapping[str, object]] = { @@ -1549,6 +1550,7 @@ _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset( LlmProviders.LM_STUDIO.value, LlmProviders.OLLAMA.value, LlmProviders.OLLAMA_CHAT.value, + LlmProviders.EMPIRIOLABS.value, } ) _DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple( @@ -3449,6 +3451,147 @@ async def test_fireworks_ai_bills_cached_input_tokens_at_the_cache_read_rate( assert recording_logger.standard_logging_object["response_cost"] < 300 * float(prices["input_cost_per_token"]) +_EMPIRIOLABS_MODEL: Final = "empiriolabs/aplomb-1" +_EMPIRIOLABS_DEFAULT_URL: Final = "https://api.empiriolabs.ai/v1/decisions" +_EMPIRIOLABS_STATE: Final = "I was charged twice for order 4411 and want a refund today." +_EMPIRIOLABS_QUESTIONS: Final[Mapping[str, object]] = MappingProxyType( + {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}} +) +_EMPIRIOLABS_RESPONSE: Final[Mapping[str, object]] = { + "model": "aplomb-1", + "answers": {"is_refund": {"type": "noul", "noul": 0.97, "confidence": 0.94}}, + "usage": {"input_tokens": 24, "output_tokens": 0}, +} + + +@pytest.mark.parametrize( + ("api_base", "environment_api_base", "url"), + ( + (None, None, _EMPIRIOLABS_DEFAULT_URL), + (None, "https://environment-gateway.example/v1", "https://environment-gateway.example/v1/decisions"), + ( + "https://gateway.example/v1/", + "https://environment-gateway.example/v1", + "https://gateway.example/v1/decisions", + ), + ("https://gateway.example", None, "https://gateway.example/v1/decisions"), + ), + ids=("default_base", "environment_base", "explicit_base_over_environment", "explicit_base_without_v1"), +) +def test_empiriolabs_posts_the_systemone_body_to_v1_decisions_on_the_resolved_base( + api_base: str | None, + environment_api_base: str | None, + url: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + if environment_api_base is not None: + monkeypatch.setenv("EMPIRIOLABS_API_BASE", environment_api_base) + route: Final = respx_mock.post(url).respond(json=_EMPIRIOLABS_RESPONSE) + + response: Final = litellm.systemone( + model=_EMPIRIOLABS_MODEL, + state=_EMPIRIOLABS_STATE, + questions=_EMPIRIOLABS_QUESTIONS, + api_base=api_base, + api_key="caller-key", + ) + + assert route.call_count == 1 + assert json.loads(route.calls[0].request.content) == { + "model": "aplomb-1", + "state": _EMPIRIOLABS_STATE, + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + } + assert response.answers["is_refund"] == NoulAnswer(type="noul", noul=0.97, confidence=0.94) + assert response._hidden_params["custom_llm_provider"] == "empiriolabs" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api_key", "authorization"), + ((None, "Bearer environment-key"), ("caller-key", "Bearer caller-key")), + ids=("environment_key", "caller_key_over_environment"), +) +async def test_empiriolabs_sends_the_caller_key_or_else_the_environment_key( + api_key: str | None, + authorization: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + monkeypatch.setenv("EMPIRIOLABS_API_KEY", "environment-key") + route: Final = respx_mock.post(_EMPIRIOLABS_DEFAULT_URL).respond(json=_EMPIRIOLABS_RESPONSE) + + response: Final = await litellm.adecisions( + model=_EMPIRIOLABS_MODEL, + input=_EMPIRIOLABS_STATE, + questions=[{"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"}], + api_key=api_key, + ) + + assert route.call_count == 1 + assert route.calls[0].request.headers["authorization"] == authorization + assert isinstance(response, OpenAIDecisionResponse) + assert [answer.model_dump(mode="json") for answer in response.answers] == [ + {"type": "predicate", "name": "is_refund", "probability": 0.97} + ] + + +def test_empiriolabs_without_a_key_is_refused_before_http( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_KEY", raising=False) + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + + with pytest.raises(litellm.AuthenticationError, match="Missing API key for Decisions provider 'empiriolabs'"): + litellm.systemone(model=_EMPIRIOLABS_MODEL, state=_EMPIRIOLABS_STATE, questions=_EMPIRIOLABS_QUESTIONS) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize( + ("status", "error_type", "code", "exception"), + ( + (401, "authentication_error", "invalid_api_key", litellm.AuthenticationError), + (402, "insufficient_credits", "insufficient_credits", litellm.PaymentRequiredError), + (404, "not_found_error", "model_not_found", litellm.NotFoundError), + (422, "invalid_request_error", "invalid_request", litellm.BadRequestError), + (429, "rate_limit_error", "rate_limit_exceeded", litellm.RateLimitError), + (503, "service_unavailable", "model_unavailable", litellm.ServiceUnavailableError), + ), +) +def test_empiriolabs_upstream_errors_keep_their_status_and_message( + status: int, + error_type: str, + code: str, + exception: type[Exception], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + message: Final = f"upstream refused the decision with {code}" + route: Final = respx_mock.post(_EMPIRIOLABS_DEFAULT_URL).respond( + status, + json={"error": {"message": message, "type": error_type, "code": code}}, + headers={"retry-after": "1"}, + ) + + with pytest.raises(exception, match=message) as caught: + litellm.systemone( + model=_EMPIRIOLABS_MODEL, + state=_EMPIRIOLABS_STATE, + questions=_EMPIRIOLABS_QUESTIONS, + api_key="caller-key", + num_retries=0, + ) + + assert getattr(caught.value, "status_code") == status + assert route.call_count == 1 + + _BOX_QUESTION_REQUESTS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( ( "decisions", diff --git a/tests/unit/llms/empiriolabs/__init__.py b/tests/unit/llms/empiriolabs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/empiriolabs/decisions/__init__.py b/tests/unit/llms/empiriolabs/decisions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/empiriolabs/decisions/test_empiriolabs_decisions_transformation.py b/tests/unit/llms/empiriolabs/decisions/test_empiriolabs_decisions_transformation.py new file mode 100644 index 00000000000..02e9eab00ae --- /dev/null +++ b/tests/unit/llms/empiriolabs/decisions/test_empiriolabs_decisions_transformation.py @@ -0,0 +1,40 @@ +from collections.abc import Mapping + +import pytest + +from litellm.llms.empiriolabs.decisions.transformation import EmpirioLabsDecisionsConfig + + +@pytest.mark.parametrize( + ("environment", "connection"), + ( + ({}, ("https://api.empiriolabs.ai", None)), + ( + {"EMPIRIOLABS_API_BASE": "https://gateway.example/v1", "EMPIRIOLABS_API_KEY": "environment-key"}, + ("https://gateway.example/v1", "environment-key"), + ), + ), + ids=("public_api_without_a_key", "environment"), +) +def test_connection_comes_from_the_environment_before_the_public_api( + environment: Mapping[str, str], + connection: tuple[str, str | None], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + monkeypatch.delenv("EMPIRIOLABS_API_KEY", raising=False) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + + assert EmpirioLabsDecisionsConfig().resolve_connection(None, None) == connection + + +@pytest.mark.parametrize( + "api_base", + ("https://api.empiriolabs.ai", "https://api.empiriolabs.ai/v1", "https://api.empiriolabs.ai/v1/"), +) +def test_requests_go_to_v1_decisions_whether_or_not_the_base_ends_in_v1(api_base: str) -> None: + assert ( + EmpirioLabsDecisionsConfig().get_complete_url(api_base=api_base, model="aplomb-1") + == "https://api.empiriolabs.ai/v1/decisions" + ) diff --git a/tests/unit/llms/openai_like/test_empiriolabs_provider.py b/tests/unit/llms/openai_like/test_empiriolabs_provider.py index a8acce5cc9d..6ebca61073f 100644 --- a/tests/unit/llms/openai_like/test_empiriolabs_provider.py +++ b/tests/unit/llms/openai_like/test_empiriolabs_provider.py @@ -2,13 +2,19 @@ Unit tests for the EmpirioLabs OpenAI-like provider. """ +import json import os import sys +from typing import Final + +import pytest +import respx sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) ) +import litellm from litellm.llms.openai_like.dynamic_config import create_config_class from litellm.llms.openai_like.json_loader import JSONProviderRegistry @@ -61,3 +67,50 @@ def test_empiriolabs_complete_url_appends_endpoint(): stream=False, ) assert url == f"{EMPIRIOLABS_BASE_URL}/chat/completions" + + +def test_empiriolabs_chat_models_still_route_through_the_json_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("EMPIRIOLABS_API_KEY", "environment-key") + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + + with respx.mock() as upstream: + route: Final = upstream.post(f"{EMPIRIOLABS_BASE_URL}/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_empiriolabs", + "object": "chat.completion", + "created": 1_789_550_000, + "model": "qwen3-7-plus", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello from EmpirioLabs AI"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 5, "total_tokens": 9}, + }, + ) + response: Final = litellm.completion( + model="empiriolabs/qwen3-7-plus", + messages=[{"role": "user", "content": "Say hello"}], + max_completion_tokens=16, + ) + + request: Final = route.calls.last.request + assert litellm.get_llm_provider("empiriolabs/qwen3-7-plus") == ( + "qwen3-7-plus", + "empiriolabs", + "environment-key", + EMPIRIOLABS_BASE_URL, + ) + assert route.call_count == 1 + assert request.headers["authorization"] == "Bearer environment-key" + assert json.loads(request.content) == { + "model": "qwen3-7-plus", + "messages": [{"role": "user", "content": "Say hello"}], + "max_tokens": 16, + } + assert response.choices[0].message.content == "Hello from EmpirioLabs AI" diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 29413d2b189..c17a9ba9371 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -529,6 +529,7 @@ def _decisions_provider_entry(provider: str) -> ProviderCreateInfo: ("Nace", "Nace", LlmProviders.NACE), ("PARALLEL_AI", "Parallel AI", LlmProviders.PARALLEL_AI), ("EU_JEV", "eu/jev", LlmProviders.EU_JEV), + ("EMPIRIOLABS", "EmpirioLabs AI", LlmProviders.EMPIRIOLABS), ), ) def test_decisions_provider_fields(provider_key: str, display_name: str, provider: LlmProviders) -> None: diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index c88e44f21d4..604d5923caa 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -252,6 +252,7 @@ describe("provider_info_helpers", () => { Providers.DATAROBOT, Providers.DOCKER_MODEL_RUNNER, Providers.DOTPROMPT, + Providers.EMPIRIOLABS, Providers.EMPOWER, Providers.GALADRIEL, Providers.GradientAI, @@ -616,6 +617,24 @@ describe("provider_info_helpers", () => { ]); }); + it("should list EmpirioLabs AI chat and decision models for the 'EMPIRIOLABS' provider key", () => { + const modelMap = { + "empiriolabs/aplomb-1": { + litellm_provider: "empiriolabs", + mode: "chat", + supported_endpoints: ["/v1/chat/completions", "/v1/decisions"], + }, + "empiriolabs/qwen3-7-plus": { litellm_provider: "empiriolabs", mode: "chat" }, + "typesafe/jev-latest": { litellm_provider: "typesafe", mode: "evaluation" }, + }; + expect(getProviderModels("EMPIRIOLABS" as Providers, modelMap)).toEqual([ + "empiriolabs/aplomb-1", + "empiriolabs/qwen3-7-plus", + ]); + expect(getProviderLogoAndName("empiriolabs").displayName).toBe("EmpirioLabs AI"); + expect(getPlaceholder("EMPIRIOLABS")).toBe("empiriolabs/aplomb-1"); + }); + it("should list only strands_decider models for the 'StrandsDecider' provider key", () => { const modelMap = { "strands_decider/strands-decider-2B-hobson-v19": { litellm_provider: "strands_decider", mode: "decisions" }, diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 1680258e10b..167343fc05e 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -118,6 +118,7 @@ export enum Providers { DOTPROMPT = "Dotprompt", EDENAI = "Eden AI", ElevenLabs = "ElevenLabs", + EMPIRIOLABS = "EmpirioLabs AI", EMPOWER = "Empower", FalAI = "Fal AI", FEATHERLESS_AI = "Featherless Ai", @@ -252,6 +253,7 @@ export const provider_map: Record = { DOTPROMPT: "dotprompt", EDENAI: "edenai", ElevenLabs: "elevenlabs", + EMPIRIOLABS: "empiriolabs", EMPOWER: "empower", FalAI: "fal_ai", FEATHERLESS_AI: "featherless_ai", @@ -499,6 +501,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Cursor]: "cursor/claude-4-sonnet", [Providers.DeepInfra]: "deepinfra/", [Providers.EDENAI]: "edenai/openai/gpt-mini-latest", + [Providers.EMPIRIOLABS]: "empiriolabs/aplomb-1", [Providers.FalAI]: "fal_ai/fal-ai/flux-pro/v1.1-ultra", [Providers.Google_AI_Studio]: "gemini-pro", [Providers.HOSTINGER_ROUTER]: "hostinger_router/claude-sonnet-5", From 8cbd40893eb94f12068f0cf9b22d3f0e34f57578 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Wed, 4 Nov 2026 11:08:21 -0800 Subject: [PATCH 006/178] fix(router): stamp the failed deployment on every endpoint's attempt (#49941) * fix(router): stamp the failed deployment on every endpoint's attempt Weighted failover (#27980) re-picks within a model group by excluding the deployment named in the error's failed_deployment_id. #27980 stamped that id in only three attempt functions: _completion, _acompletion and _ageneric_api_call_with_fallbacks_helper. Embedding, image generation, transcription, speech, text completion, rerank, adapter completion, the sync generic call and vector store create never stamped it, so with enable_weighted_failover on, a failure on one of their deployments skipped the in-group re-pick and went straight to the cross-group fallbacks, or raised. Each of those attempts now stamps the deployment it picked, read from the model_info that _update_kwargs_with_deployment put on the request (so a client-side-credential call names its own dynamic deployment, as on completions). A failed pick stamps nothing. File and batch creation and batch cancellation are left alone: they are bound to a resource on the deployment they picked, so re-picking another one cannot help. * test(router): narrow the failed-pick raise to the no-healthy-deployments error Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013FDwtoMjnHa4xxazGJJNSm * fix(router): exclude the configured deployment on failover for client-side credentials Weighted failover excluded the per-request id that client-supplied credentials generate, so the failed configured deployment stayed eligible for the re-pick. Stamp the configured id alongside it and exclude that. Also drop the vector store create stamp, which no fallback path reads, and redundant test comments. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013FDwtoMjnHa4xxazGJJNSm --------- Co-authored-by: Claude --- litellm/router_python_core.py | 74 ++++++++-- tests/unit/test_router_weighted_failover.py | 141 ++++++++++++++++++++ 2 files changed, 204 insertions(+), 11 deletions(-) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 3ea900642bd..e251248d675 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -5211,8 +5211,19 @@ class PythonCore(DeprecatedRouterMembers): # _update_kwargs_with_deployment and stamped into kwargs["model_info"]; stamping # the static shared deployment's id instead would let one tenant's bad credentials # cool down the deployment every other tenant sharing this config relies on. + if getattr(exception, "failed_deployment_id", None): + return effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({}) self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info})) + configured_deployment_id: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info).get( + "original_model_id" + ) + if isinstance(configured_deployment_id, str) and configured_deployment_id: + exception.failed_configured_deployment_id = configured_deployment_id # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: the exception carries the failure to _maybe_run_weighted_failover + + def _stamp_failed_attempt(self, exception: Exception, kwargs: Mapping[str, object], *, picked: bool) -> None: + if picked: + self._stamp_failed_deployment_id_with_effective_model_info(exception, MappingProxyType({}), kwargs) @staticmethod def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None: @@ -6089,9 +6100,10 @@ class PythonCore(DeprecatedRouterMembers): def _image_generation(self, prompt: str, model: str, **kwargs): model_name: Final = "" + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - deployment: Final = self.get_available_deployment( + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6127,6 +6139,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aimage_generation(self, prompt: str, model: str, **kwargs): @@ -6152,11 +6167,12 @@ class PythonCore(DeprecatedRouterMembers): async def _aimage_generation(self, prompt: str, model: str, **kwargs): model_name = model start_time: Final = datetime.now() # noqa: DTZ005 # SDK logging uses naive local timestamps + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) try: - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6217,6 +6233,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def atranscription(self, file: FileTypes, model: str, **kwargs): @@ -6263,10 +6282,11 @@ class PythonCore(DeprecatedRouterMembers): async def _atranscription(self, file: FileTypes, model: str, **kwargs): model_name: Final = model + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _atranscription()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6300,6 +6320,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): @@ -6356,10 +6379,11 @@ class PythonCore(DeprecatedRouterMembers): async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): model_name: Final = model + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _aspeech()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6393,6 +6417,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def arerank(self, model: str, **kwargs): @@ -6418,9 +6445,10 @@ class PythonCore(DeprecatedRouterMembers): async def _arerank(self, model: str, **kwargs): model_name = None + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _rerank()- model: %s", model) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, @@ -6452,6 +6480,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e def text_completion( @@ -6533,10 +6564,11 @@ class PythonCore(DeprecatedRouterMembers): raise e async def _atext_completion(self, model: str, prompt: str, **kwargs): + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _atext_completion()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": prompt}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6572,6 +6604,9 @@ class PythonCore(DeprecatedRouterMembers): if model is not None: self.fail_calls[model] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aadapter_completion( @@ -6603,10 +6638,11 @@ class PythonCore(DeprecatedRouterMembers): raise e async def _aadapter_completion(self, adapter_id: str, model: str, **kwargs): + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _aadapter_completion()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "default text"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6642,6 +6678,9 @@ class PythonCore(DeprecatedRouterMembers): if model is not None: self.fail_calls[model] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs) -> object: @@ -7630,6 +7669,7 @@ class PythonCore(DeprecatedRouterMembers): """ handler_name: Final = original_function.__name__ metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) self._update_kwargs_before_fallbacks( @@ -7638,7 +7678,7 @@ class PythonCore(DeprecatedRouterMembers): metadata_variable_name=metadata_variable_name, ) with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): - deployment: Final = self.get_available_deployment( + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=kwargs.get("messages", None), input=kwargs.get("input", None), @@ -7684,6 +7724,9 @@ class PythonCore(DeprecatedRouterMembers): verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) if model is not None: self.fail_calls[model] += 1 + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e def embedding( @@ -7705,10 +7748,11 @@ class PythonCore(DeprecatedRouterMembers): def _embedding(self, input: str | list[object], model: str, **kwargs): model_name = None + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside embedding()- model: %s", model) with embedding_deployment_selection(): - deployment: Final = self.get_available_deployment( + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init model=model, input=input, specific_deployment=kwargs.pop("specific_deployment", None), @@ -7765,6 +7809,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aembedding( @@ -7794,11 +7841,12 @@ class PythonCore(DeprecatedRouterMembers): async def _aembedding(self, input: str | list[object], model: str, **kwargs): model_name = None + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _aembedding()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) with embedding_deployment_selection(): - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, input=input, specific_deployment=kwargs.pop("specific_deployment", None), @@ -7845,6 +7893,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e #### FILES API #### @@ -9083,8 +9134,9 @@ class PythonCore(DeprecatedRouterMembers): kwargs[metadata_variable_name] = meta if not isinstance(meta, dict): return None + configured_failed_id: Final[str | None] = getattr(exception, "failed_configured_deployment_id", None) prev_excluded: Final = set(meta.get("_failover_excluded_ids") or []) - excluded: Final = prev_excluded | {failed_id} + excluded: Final = prev_excluded | {configured_failed_id or failed_id} all_ids: Final = { (d.get("model_info") or {}).get("id") diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index c00670473af..70231b620be 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -1019,3 +1019,144 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy assert while_a_is_down._hidden_params["model_id"] == "B" assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 assert after_a_recovers._hidden_params["model_id"] == "A" + + +def _two_deployment_router(model: str) -> Router: + return Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": model, "api_key": "bad", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": model, "api_key": "good", "weight": 0}, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + disable_cooldowns=True, + enable_weighted_failover=True, + ) + + +_ENDPOINTS: Final = [ + ("aembedding", "openai/text-embedding-3-small", {"input": ["hi"]}), + ("aimage_generation", "openai/dall-e-3", {"prompt": "a cat"}), + ("atranscription", "openai/whisper-1", {"file": b"audio"}), + ("aspeech", "openai/tts-1", {"input": "hi", "voice": "alloy"}), + ("atext_completion", "openai/gpt-3.5-turbo-instruct", {"prompt": "hi"}), + ("arerank", "cohere/rerank-english-v3.0", {"query": "q", "documents": ["d"]}), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("endpoint", "model", "arguments"), _ENDPOINTS, ids=[e[0] for e in _ENDPOINTS]) +async def test_weighted_failover_moves_every_endpoint_to_another_deployment( + monkeypatch: pytest.MonkeyPatch, endpoint: str, model: str, arguments: dict[str, object] +) -> None: + keys_called: list[object] = [] + answer: Final = object() + + async def provider(**kwargs: object) -> object: + keys_called.append(kwargs.get("api_key")) + if kwargs.get("api_key") == "bad": + raise litellm.InternalServerError(message="A is down", llm_provider="openai", model=model) + return answer + + monkeypatch.setattr(litellm, endpoint, provider) + router: Final = _two_deployment_router(model) + + response: Final = await getattr(router, endpoint)(model="test-model", **arguments) + + assert response is answer + assert keys_called == ["bad", "good"] + + +@pytest.mark.asyncio +async def test_weighted_failover_skips_the_failed_deployment_for_clientside_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + api_bases_called: list[object] = [] + answer: Final = object() + + async def provider(**kwargs: object) -> object: + api_bases_called.append(kwargs.get("api_base")) + if kwargs.get("api_base") == "http://a": + raise litellm.InternalServerError(message="A is down", llm_provider="openai", model="m") + return answer + + monkeypatch.setattr(litellm, "aembedding", provider) + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "http://a", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "http://b", "weight": 0}, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + disable_cooldowns=True, + enable_weighted_failover=True, + ) + + response: Final = await router.aembedding(model="test-model", input=["hi"], api_key="tenant-key") + + assert response is answer + assert api_bases_called == ["http://a", "http://b"] + + +@pytest.mark.parametrize( + ("endpoint", "attempt", "arguments"), + [ + ("embedding", "_embedding", {"input": ["hi"]}), + ("image_generation", "_image_generation", {"prompt": "a cat"}), + ], +) +def test_a_sync_attempt_stamps_the_failed_deployment( + monkeypatch: pytest.MonkeyPatch, endpoint: str, attempt: str, arguments: dict[str, object] +) -> None: + def provider(**kwargs: object) -> object: + raise litellm.InternalServerError(message="down", llm_provider="openai", model="m") + + monkeypatch.setattr(litellm, endpoint, provider) + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/m", "api_key": "k"}, + "model_info": {"id": "A"}, + } + ], + ) + + with pytest.raises(litellm.InternalServerError) as raised: + getattr(router, attempt)(model="test-model", **arguments) + + assert getattr(raised.value, "failed_deployment_id", None) == "A" + + +@pytest.mark.asyncio +async def test_a_failed_pick_stamps_no_deployment() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "k"}, + "model_info": {"id": "A"}, + } + ], + ) + + with pytest.raises(litellm.BadRequestError, match="no healthy deployments") as raised: + await router._aembedding(input=["hi"], model="no-such-model") + + assert getattr(raised.value, "failed_deployment_id", None) is None From 53740a4426990598945edf83496d2c63149839b8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Nov 2026 11:30:50 -0800 Subject: [PATCH 007/178] refactor(ui): move keys, projects and cost text onto i18n (#49994) Moves the page-authored English on the virtual keys, projects, cost tracking, cost optimization and budgets pages, plus the create key form and the shared key and team pickers, into t() and Trans with 1,318 messages under 32 new top-level keys in en.json. The English on screen is unchanged key_edit_view.tsx grows past 800 lines, so it gets a max-lines entry in eslint-suppressions.json instead of a split --- ui/litellm-dashboard/eslint-suppressions.json | 3 + .../budgets/_components/BudgetFormFields.tsx | 31 +- .../budgets/_components/BudgetModal.tsx | 16 +- .../budgets/_components/BudgetPanel.tsx | 40 +- .../budgets/_components/BudgetTable.tsx | 62 +- .../_components/BudgetTableColumns.tsx | 58 +- .../budgets/_components/EditBudgetModal.tsx | 14 +- .../_components/AutoRouterBenchmarksTab.tsx | 240 +- .../_components/AutoRouterSummaryTable.tsx | 52 +- .../_components/CacheLeakageCard.tsx | 55 +- .../_components/CostOptimizationView.tsx | 48 +- .../PromptCachingRequestsTable.tsx | 82 +- .../_components/PromptCachingTab.tsx | 8 +- .../_components/PromptCompressionTab.tsx | 93 +- .../_components/ShadowEvalSection.tsx | 327 ++- .../_components/ShadowEvalStartForm.tsx | 144 +- .../_components/TierTurnsChart.tsx | 11 +- .../_components/UsageTab.tsx | 38 +- .../_components/autoRouterBenchmarks.ts | 24 +- .../_components/costOptimizationUtils.ts | 14 +- .../_components/add_margin_form.tsx | 32 +- .../_components/add_provider_form.tsx | 26 +- .../_components/cost_tracking_settings.tsx | 94 +- .../_components/how_it_works.tsx | 54 +- .../_components/pricing_calculator/index.tsx | 29 +- .../pricing_calculator/multi_cost_results.tsx | 93 +- .../multi_export_dropdown.tsx | 8 +- .../pricing_calculator/multi_export_utils.ts | 104 +- .../use_multi_cost_estimate.ts | 9 +- .../_components/provider_discount_table.tsx | 22 +- .../_components/provider_margin_table.tsx | 29 +- .../_components/use_block_unpriced_config.ts | 8 +- .../_components/use_discount_config.ts | 28 +- .../_components/use_margin_config.ts | 31 +- .../_components/content_filter/TagsInput.tsx | 7 +- .../_components/ProjectDetailsPage.tsx | 73 +- .../_components/ProjectKeysSection.tsx | 12 +- .../projects/_components/ProjectKeysTable.tsx | 11 +- .../_components/ProjectKeysTableColumns.tsx | 23 +- .../ProjectModals/CreateProjectModal.tsx | 13 +- .../ProjectModals/EditProjectModal.tsx | 13 +- .../ProjectModals/ProjectBaseForm.tsx | 116 +- .../ProjectModals/projectFormSchema.ts | 14 +- .../projects/_components/ProjectsPage.tsx | 16 +- .../projects/_components/ProjectsTable.tsx | 13 +- .../_components/ProjectsTableColumns.tsx | 37 +- .../src/components/DocsMenu.tsx | 3 +- .../VirtualKeysPage/VirtualKeysTable.tsx | 46 +- .../VirtualKeysPage/keyTableColumns.tsx | 160 +- .../agent_management/AgentSelector.tsx | 11 +- .../common_components/AccessGroupSelector.tsx | 9 +- .../KeyLifecycleSettings.tsx | 70 +- .../common_components/ModelAliasManager.tsx | 50 +- .../OrganizationDropdown.tsx | 7 +- .../PassThroughRoutesSelector.tsx | 7 +- .../PremiumLoggingSettings.tsx | 20 +- .../common_components/ProjectDropdown.tsx | 6 +- .../RateLimitTypeFormItem.tsx | 50 +- .../RouterSettingsAccordion.tsx | 6 +- .../common_components/UserDropdown.tsx | 8 +- .../budget_duration_dropdown.tsx | 34 +- .../common_components/simple_table.tsx | 5 +- .../common_components/team_dropdown.tsx | 8 +- .../common_components/team_multi_select.tsx | 11 +- .../BudgetFallbacksEditor.tsx | 36 +- .../key_team_helpers/BudgetWindowsEditor.tsx | 6 +- .../key_team_helpers/EndUserBudgetSelect.tsx | 21 +- .../key_team_helpers/ModelMaxBudgetEditor.tsx | 35 +- .../key_team_helpers/TagRateLimitEditor.tsx | 19 +- .../components/object_permissions_view.tsx | 29 +- .../src/components/onboarding_link.tsx | 18 +- .../organisms/RegenerateKeyModal.tsx | 82 +- .../organisms/createKeyOwnerAccess.ts | 7 +- .../organisms/create_key_button.tsx | 424 ++-- .../src/components/organisms/utils.ts | 11 +- .../permissions/AgentPermissions.tsx | 20 +- .../permissions/MCPServerPermissions.tsx | 36 +- .../permissions/VectorStorePermissions.tsx | 6 +- .../components/permissions/inheritedGrants.ts | 11 +- .../LatencyBasedConfiguration.tsx | 11 +- .../ReliabilityRetriesSection.tsx | 6 +- .../router_settings/RouterSettingsForm.tsx | 6 +- .../RoutingStrategySelector.tsx | 4 +- .../router_settings/TagFilteringToggle.tsx | 6 +- .../src/components/router_settings/index.tsx | 10 +- .../templates/KeyEditViewControls.tsx | 208 +- .../components/templates/KeyInfoHeader.tsx | 83 +- .../components/templates/KeyProjectField.tsx | 12 +- .../components/templates/KeySavingsTab.tsx | 26 +- .../templates/estimatedOutputTokens.ts | 21 +- .../templates/keyEditFieldNormalizers.ts | 13 +- .../components/templates/key_edit_view.tsx | 173 +- .../components/templates/key_info_view.tsx | 332 +-- .../VectorStoreSelector.tsx | 3 +- ui/litellm-dashboard/src/i18n/locales/en.json | 1936 +++++++++++++++++ .../src/lib/autorouter_presets.ts | 7 +- 96 files changed, 4612 insertions(+), 1792 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 3d30964129d..a4a30f3021a 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2000,6 +2000,9 @@ "src/components/templates/key_edit_view.tsx": { "local/filename-pascal-case": { "count": 1 + }, + "max-lines": { + "count": 1 } }, "src/components/templates/key_info_view.tsx": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx index 9ffd23884b4..56e2c6dbecc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx @@ -1,4 +1,5 @@ import React, { useId } from "react"; +import { useTranslation } from "react-i18next"; import { Controller } from "react-hook-form"; import { z } from "zod"; import { FormField } from "@/components/shared/form/FormField"; @@ -16,6 +17,7 @@ import { import { Input } from "@/components/ui/input"; import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/components/ui/input-group"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import i18n from "@/i18n/i18n"; import { useZodForm } from "@/lib/forms/useZodForm"; const budgetShape = { @@ -54,19 +56,23 @@ const numberProps = ) => onChange(event.target.value === "" ? null : event.target.valueAsNumber), }); -const OptionalBadge = () => Optional; +const OptionalBadge = () => { + const { t } = useTranslation(); + return {t("budgets.form.optional")}; +}; function SpendLimitField({ form }: { form: BudgetForm }) { + const { t } = useTranslation(); const amountId = useId(); return ( - Spend limit + {t("budgets.form.spendLimit")} @@ -91,7 +97,7 @@ function SpendLimitField({ form }: { form: BudgetForm }) { value={field.value ?? NO_RESET} onValueChange={(next) => field.onChange(next === NO_RESET ? null : next)} > - + @@ -124,13 +130,14 @@ function RateLimitField({ form, name, label }: { form: BudgetForm; name: NumberF } export function BudgetFormFields({ form, isEditing }: { form: BudgetForm; isEditing: boolean }) { + const { t } = useTranslation(); return ( {({ ref, ...field }) => } @@ -139,15 +146,13 @@ export function BudgetFormFields({ form, isEditing }: { form: BudgetForm; isEdit
- Rate limits + {t("budgets.form.rateLimits")} - - - - - Batches count against the daily limit instead of the per-minute ones. Provider rate limits still apply. - + + + + {t("budgets.form.batchNote")}
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx index bd10837640d..c6bf7e8121e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx @@ -1,4 +1,5 @@ import React, { useId } from "react"; +import { useTranslation } from "react-i18next"; import { useCreateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { applyBudgetPrecision } from "./budgetPrecision"; import { BudgetFormFields, type BudgetFormValues, useBudgetForm } from "./BudgetFormFields"; @@ -11,6 +12,7 @@ interface BudgetModalProps { setIsModalVisible: React.Dispatch>; } const BudgetModal: React.FC = ({ isModalVisible, setIsModalVisible }) => { + const { t } = useTranslation(); const formId = useId(); const form = useBudgetForm({ budget_id: "" }); const createBudget = useCreateBudget(); @@ -22,14 +24,14 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis const handleCreate = async (formValues: BudgetFormValues) => { try { - toast.info("Making API Call"); + toast.info(t("budgets.createModal.toast.pending")); await createBudget.mutateAsync(applyBudgetPrecision(formValues)); - toast.success("Budget Created"); + toast.success(t("budgets.createModal.toast.created")); form.reset(); setIsModalVisible(false); } catch (error) { console.error("Error creating the budget:", error); - toast.fromError(`Error creating the budget: ${error}`); + toast.fromError(t("budgets.createModal.toast.error", { error: String(error) })); } }; @@ -37,15 +39,17 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis !open && handleCancel()}> - Create budget + + {t("budgets.createModal.title")} +
- }>Cancel + }>{t("budgets.createModal.cancel")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx index 785b09252af..bbc4fed94f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx @@ -6,6 +6,7 @@ import { Page, PageTabs, PageTabsList, PageTabsTrigger } from "@/components/shared/Page"; import { Plus, Wallet } from "lucide-react"; import React, { useCallback, useState } from "react"; +import { useTranslation } from "react-i18next"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { prism } from "react-syntax-highlighter/dist/esm/styles/prism"; @@ -29,6 +30,7 @@ interface BudgetSettingsPageProps { } const BudgetPanel: React.FC = ({ accessToken }) => { + const { t } = useTranslation(); const syntaxTheme = useSyntaxTheme(prism); const [isCreateModelVisible, setIsCreateModelVisible] = useState(false); const [isEditModalVisible, setIsEditModalVisible] = useState(false); @@ -65,10 +67,10 @@ const BudgetPanel: React.FC = ({ accessToken }) => { } try { await deleteBudget.mutateAsync(selectedBudget.budget_id); - toast.success("Budget deleted."); + toast.success(t("budgets.panel.toast.deleted")); } catch (error) { console.error("Error deleting budget:", error); - toast.fromError("Failed to delete budget"); + toast.fromError(t("budgets.panel.toast.deleteFailed")); } finally { setIsDeleteModalVisible(false); setSelectedBudget(null); @@ -85,22 +87,22 @@ const BudgetPanel: React.FC = ({ accessToken }) => { - Budgets + {t("budgets.panel.title")} - Spend, TPM and RPM limits you can assign to customers. + {t("budgets.panel.description")} {canModify && ( <> )} - Budgets - Examples + {t("budgets.panel.tabs.budgets")} + {t("budgets.panel.tabs.examples")} @@ -122,15 +124,15 @@ const BudgetPanel: React.FC = ({ accessToken }) => { /> = ({ accessToken }) => {
-

How to use budget id

+

{t("budgets.panel.examples.intro")}

- Assign Budget to Customer + {t("budgets.panel.examples.tabs.assignBudget")} - Test it (Curl) + {t("budgets.panel.examples.tabs.curl")} - Test it (OpenAI SDK) + {t("budgets.panel.examples.tabs.openaiSdk")} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx index 80872c58d28..9d9e32ed86e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx @@ -2,6 +2,7 @@ import { Inbox, ShieldAlert } from "lucide-react"; import React, { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { BUDGET_DURATION_FILTER_OPTIONS, @@ -21,6 +22,7 @@ import { import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import i18n from "@/i18n/i18n"; import { ApiError } from "@/lib/http/client"; import { BUDGET_TABLE_HIDDEN_COLUMNS, getBudgetTableColumns } from "./BudgetTableColumns"; @@ -49,11 +51,19 @@ const formatFilterValue = (columnId: string, value: unknown): string => { } if (columnId === "max_budget") { const { min, max, unlimitedOnly } = (value ?? {}) as MaxBudgetFilterValue; - return unlimitedOnly === true ? "Unlimited only" : `${min ? `$${min}` : "any"} to ${max ? `$${max}` : "any"}`; + return unlimitedOnly === true + ? i18n.t("budgets.table.filterChips.unlimitedOnly") + : i18n.t("budgets.table.filterChips.maxBudgetRange", { + min: min ? `$${min}` : i18n.t("budgets.table.filterChips.any"), + max: max ? `$${max}` : i18n.t("budgets.table.filterChips.any"), + }); } if (columnId === "created_at") { const { from, to } = (value ?? {}) as CreatedAtFilterValue; - return `${from || "any"} to ${to || "any"}`; + return i18n.t("budgets.table.filterChips.createdRange", { + from: from || i18n.t("budgets.table.filterChips.any"), + to: to || i18n.t("budgets.table.filterChips.any"), + }); } return String(value); }; @@ -81,22 +91,24 @@ const normalizeCreatedAt = (draft: CreatedAtFilterValue): CreatedAtFilterValue | }; function EmptyState({ hasQuery }: { hasQuery: boolean }) { + const { t } = useTranslation(); return (
-
{hasQuery ? "No matching budgets" : "No budgets yet"}
+
+ {hasQuery ? t("budgets.table.emptyState.noMatchTitle") : t("budgets.table.emptyState.title")} +
- {hasQuery - ? "No budget matches your search or filters." - : "Create a budget to set spend, TPM and RPM limits for customers."} + {hasQuery ? t("budgets.table.emptyState.noMatchDescription") : t("budgets.table.emptyState.description")}
); } function ErrorState({ error }: { error: Error }) { + const { t } = useTranslation(); const forbidden = error instanceof ApiError && error.status === 403; return (
@@ -104,10 +116,10 @@ function ErrorState({ error }: { error: Error }) {
- {forbidden ? "You do not have access to budgets" : "Could not load budgets"} + {forbidden ? t("budgets.table.errorState.forbiddenTitle") : t("budgets.table.errorState.title")}
- {forbidden ? "Ask a proxy admin to grant you the admin viewer role." : error.message} + {forbidden ? t("budgets.table.errorState.forbiddenDescription") : error.message}
); @@ -141,19 +153,20 @@ function DurationFilter({ selected, onChange }: { selected: string[]; onChange: } function BudgetFilterFields({ get, set }: FilterDraft) { + const { t } = useTranslation(); const maxBudget = (get("max_budget") as MaxBudgetFilterValue | undefined) ?? {}; const created = (get("created_at") as CreatedAtFilterValue | undefined) ?? {}; const unlimitedOnly = maxBudget.unlimitedOnly === true; return ( <> - + set("budget_duration", selected)} /> - +
set("max_budget", normalizeMaxBudget({ ...maxBudget, min: event.target.value }))} - placeholder="Min" - aria-label="Minimum max budget" + placeholder={t("budgets.table.filters.min")} + aria-label={t("budgets.table.filters.minAriaLabel")} data-testid="budget-filter-max-budget-min" /> set("max_budget", normalizeMaxBudget({ ...maxBudget, max: event.target.value }))} - placeholder="Max" - aria-label="Maximum max budget" + placeholder={t("budgets.table.filters.max")} + aria-label={t("budgets.table.filters.maxAriaLabel")} data-testid="budget-filter-max-budget-max" />
@@ -184,23 +197,23 @@ function BudgetFilterFields({ get, set }: FilterDraft) { onCheckedChange={(checked) => set("max_budget", normalizeMaxBudget({ unlimitedOnly: checked === true }))} data-testid="budget-filter-max-budget-unlimited" /> - Unlimited only + {t("budgets.table.filters.unlimitedOnly")} - +
set("created_at", normalizeCreatedAt({ ...created, from: event.target.value }))} - aria-label="Created from" + aria-label={t("budgets.table.filters.createdFrom")} data-testid="budget-filter-created-from" /> set("created_at", normalizeCreatedAt({ ...created, to: event.target.value }))} - aria-label="Created to" + aria-label={t("budgets.table.filters.createdTo")} data-testid="budget-filter-created-to" />
@@ -210,11 +223,12 @@ function BudgetFilterFields({ get, set }: FilterDraft) { } const BudgetTable: React.FC = ({ list, canModify, onEditClick, onDeleteClick }) => { + const { t } = useTranslation(); const [filtersOpen, setFiltersOpen] = useState(false); const columns = useMemo( - () => getBudgetTableColumns({ canModify, onEditClick, onDeleteClick }), - [canModify, onEditClick, onDeleteClick], + () => getBudgetTableColumns({ canModify, onEditClick, onDeleteClick }, t), + [canModify, onEditClick, onDeleteClick, t], ); const hasQuery = list.searchValue.trim() !== "" || list.columnFilters.length > 0; @@ -239,7 +253,7 @@ const BudgetTable: React.FC = ({ list, canModify, onEditClick, columnFilters={list.columnFilters} onColumnFiltersChange={list.onColumnFiltersChange} isLoading={list.isLoading} - loadingMessage="Loading budgets…" + loadingMessage={t("budgets.table.loading")} noDataMessage={emptyMessage} size="compact" toolbar={(table) => ( @@ -248,7 +262,7 @@ const BudgetTable: React.FC = ({ list, canModify, onEditClick, table={table} searchValue={list.searchValue} onSearchChange={list.onSearchChange} - searchPlaceholder="Search by budget ID…" + searchPlaceholder={t("budgets.table.searchPlaceholder")} onOpenFilters={() => setFiltersOpen(true)} onRefresh={list.refetch} isRefreshing={list.isFetching} @@ -259,8 +273,8 @@ const BudgetTable: React.FC = ({ list, canModify, onEditClick, table={table} open={filtersOpen} onOpenChange={setFiltersOpen} - title="Filters" - description="Narrow down your budgets" + title={t("budgets.table.filters.title")} + description={t("budgets.table.filters.description")} > {(draft) => } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx index 3ec291c1158..d0538e7ed1c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -1,7 +1,9 @@ "use client"; import { ColumnDef, FilterFn } from "@tanstack/react-table"; +import type { TFunction } from "i18next"; import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; +import { useTranslation } from "react-i18next"; import { DataTableSortHeader } from "@/components/shared/DataTable"; import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; @@ -27,15 +29,17 @@ const serverFilter: FilterFn = () => true; serverFilter.autoRemove = () => false; function RateLimitCell({ value }: { value: number | null | undefined }) { + const { t } = useTranslation(); if (value == null) { - return n/a; + return {t("budgets.table.cells.notApplicable")}; } return {value}; } function BudgetDurationCell({ value }: { value: string | null | undefined }) { + const { t } = useTranslation(); if (!value) { - return Not set; + return {t("budgets.table.cells.notSet")}; } return {getBudgetDurationLabel(value)}; } @@ -47,10 +51,11 @@ interface BudgetRowActionsProps { } function BudgetRowActions({ budget, onEditClick, onDeleteClick }: BudgetRowActionsProps) { + const { t } = useTranslation(); return ( @@ -59,7 +64,7 @@ function BudgetRowActions({ budget, onEditClick, onDeleteClick }: BudgetRowActio onEditClick(budget)}> - Edit budget + {t("budgets.table.rowActions.edit")} onDeleteClick(budget)} > - Delete budget + {t("budgets.table.rowActions.delete")} @@ -87,16 +92,15 @@ interface BudgetTableColumnsDeps { onDeleteClick: (budget: budgetItem) => void; } -export const getBudgetTableColumns = ({ - canModify, - onEditClick, - onDeleteClick, -}: BudgetTableColumnsDeps): ColumnDef[] => [ +export const getBudgetTableColumns = ( + { canModify, onEditClick, onDeleteClick }: BudgetTableColumnsDeps, + t: TFunction, +): ColumnDef[] => [ { id: "budget_id", accessorKey: "budget_id", - meta: { title: "Budget ID" }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.budgetId") }, + header: ({ column }) => , cell: ({ row }) => ( ), @@ -105,32 +109,34 @@ export const getBudgetTableColumns = ({ id: "max_budget", accessorKey: "max_budget", filterFn: serverFilter, - meta: { title: "Max Budget", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.maxBudget"), numeric: true }, + header: ({ column }) => , size: 120, - cell: ({ row }) => , + cell: ({ row }) => ( + + ), }, { id: "tpm_limit", accessorKey: "tpm_limit", - meta: { title: "TPM", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.tpm"), numeric: true }, + header: ({ column }) => , size: 100, cell: ({ row }) => , }, { id: "rpm_limit", accessorKey: "rpm_limit", - meta: { title: "RPM", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.rpm"), numeric: true }, + header: ({ column }) => , size: 100, cell: ({ row }) => , }, { id: "tpd_limit", accessorKey: "tpd_limit", - meta: { title: "TPD (batch)", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.tpd"), numeric: true }, + header: ({ column }) => , size: 110, cell: ({ row }) => , }, @@ -138,9 +144,9 @@ export const getBudgetTableColumns = ({ id: "budget_duration", accessorKey: "budget_duration", filterFn: serverFilter, - meta: { title: "Reset" }, + meta: { title: t("budgets.table.columns.reset") }, sortDescFirst: false, - header: ({ column }) => , + header: ({ column }) => , size: 110, cell: ({ row }) => , }, @@ -148,8 +154,8 @@ export const getBudgetTableColumns = ({ id: "created_at", accessorKey: "created_at", filterFn: serverFilter, - meta: { title: "Created" }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.created") }, + header: ({ column }) => , size: 160, cell: ({ row }) => , }, @@ -158,7 +164,7 @@ export const getBudgetTableColumns = ({ { id: "actions", meta: { className: "text-right", headerClassName: "text-right" }, - header: () => Actions, + header: () => {t("budgets.table.columns.actions")}, size: 64, enableSorting: false, enableHiding: false, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx index 4aa69c5c15e..81c1f652eab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx @@ -1,4 +1,5 @@ import React, { useEffect, useId } from "react"; +import { useTranslation } from "react-i18next"; import { useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { applyBudgetPrecision } from "./budgetPrecision"; @@ -23,6 +24,7 @@ interface EditBudgetModalProps { existingBudget: budgetItem; } const EditBudgetModal: React.FC = ({ isModalVisible, setIsModalVisible, existingBudget }) => { + const { t } = useTranslation(); const formId = useId(); const form = useBudgetForm(toFormValues(existingBudget)); const { dirtyFields } = form.formState; @@ -39,15 +41,15 @@ const EditBudgetModal: React.FC = ({ isModalVisible, setIs const handleUpdate = async (formValues: BudgetFormValues) => { try { - toast.info("Making API Call"); + toast.info(t("budgets.editModal.toast.pending")); const changedValues = pickDirty(formValues, dirtyFields); await updateBudget.mutateAsync(applyBudgetPrecision({ budget_id: formValues.budget_id, ...changedValues })); - toast.success("Budget Updated"); + toast.success(t("budgets.editModal.toast.updated")); form.reset(); setIsModalVisible(false); } catch (error) { console.error("Error updating the budget:", error); - toast.fromError(`Error updating the budget: ${error}`); + toast.fromError(t("budgets.editModal.toast.error", { error: String(error) })); } }; @@ -55,15 +57,15 @@ const EditBudgetModal: React.FC = ({ isModalVisible, setIs !open && handleCancel()}> - Edit budget + {t("budgets.editModal.title")}
- }>Cancel + }>{t("budgets.editModal.cancel")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index b3acb8f9168..a6823d50b21 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -1,6 +1,7 @@ "use client"; import React, { useState } from "react"; +import { Trans, useTranslation } from "react-i18next"; import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels"; @@ -76,20 +77,21 @@ const SpendRow: React.FC<{ label: string; value: string; hint?: string; subdued? ); const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { + const { t } = useTranslation(); const stats = view.stats; const cheaper = stats.saved_pct != null && stats.saved_pct >= 0; const classifierCost = stats.baseline_spend == null ? null : stats.savings_estimated_classifier_cost ?? null; const comparedAll = stats.savings_estimated_turns === stats.turns; return ( - +

- Total estimated savings + {t("costOptimization.autoRouter.hero.totalSavings")}

- {stats.saved_spend == null ? "Unavailable" : usd(stats.saved_spend)} + {stats.saved_spend == null ? t("costOptimization.autoRouter.unavailable") : usd(stats.saved_spend)}

{stats.saved_pct != null && ( = ({ view }) => {
{stats.baseline_spend != null && !comparedAll && (

- Compared on {stats.savings_estimated_turns.toLocaleString()} of {stats.turns.toLocaleString()} requests; - adaptive and quality routers are excluded + {t("costOptimization.autoRouter.hero.comparedOn", { + compared: stats.savings_estimated_turns.toLocaleString(), + total: stats.turns.toLocaleString(), + })}

)} {stats.unattributed_saved_spend != null && (

- Per-router records differ from recorded savings by {usd(Math.abs(stats.unattributed_saved_spend))}, for - example history from before per-router tracking, so the baseline comparison is unavailable + {t("costOptimization.autoRouter.hero.unattributed", { + amount: usd(Math.abs(stats.unattributed_saved_spend)), + })}

)}
{stats.baseline_spend != null && classifierCost == null && (

- Breakdown unavailable because some usage predates classification-cost tracking. + {t("costOptimization.autoRouter.hero.breakdownUnavailable")}

)}
@@ -157,20 +170,24 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { }; const StackedTurnBar: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => { + const { t } = useTranslation(); const segments = buckets.filter((b) => b.turns > 0); return (
{segments.map((b) => (
))}
@@ -185,47 +202,57 @@ const StackedTurnBar: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => { ); }; -const BucketTable: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => ( - - - - Bucket - Turns - - Hit rate - - - - {buckets.map((b) => ( - - - - - - {b.label} - {b.sublabel} - - - - - {b.turns.toLocaleString()} - - -
-
-
- - - {pctLabel(b.hitRatePct)} - +const BucketTable: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => { + const { t } = useTranslation(); + return ( +
+ + + + {t("costOptimization.autoRouter.caching.columns.bucket")} + + + {t("costOptimization.autoRouter.caching.columns.turns")} + + + + {t("costOptimization.autoRouter.caching.columns.hitRate")} + - ))} - -
-); + + + {buckets.map((b) => ( + + + + + + {b.label} + {b.sublabel} + + + + + {b.turns.toLocaleString()} + + +
+
+
+ + + {pctLabel(b.hitRatePct)} + + + ))} + + + ); +}; const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => { - const buckets = bucketRows(cache); + const { t } = useTranslation(); + const buckets = bucketRows(cache, t); const total = bucketTurnsTotal(cache); const expiredMissPct = expiredMissShare(cache); return ( @@ -233,7 +260,7 @@ const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => {
-

Cache hit rate

+

{t("costOptimization.autoRouter.caching.hitRate")}

{pctLabel(cache.hit_rate_pct)}

{expiredMissPct === null ? null : ( @@ -248,13 +275,12 @@ const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => { } > - Expired-miss + {t("costOptimization.autoRouter.caching.expiredMiss.label")} {pctLabel(expiredMissPct)} - share of all measured turns that missed cache because a return to an earlier tier came after its TTL - lapsed + {t("costOptimization.autoRouter.caching.expiredMiss.tooltip")} @@ -263,17 +289,27 @@ const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => {
-

Share of turns

+

+ {t("costOptimization.autoRouter.caching.shareOfTurns")} +

- {total.toLocaleString()} turns - measured + {total.toLocaleString()} + ), + }} + />

{cache.unordered_turns > 0 && (

- {cache.unordered_turns.toLocaleString()} turns arrived out of order across pods and are not bucketed + {t("costOptimization.autoRouter.caching.unorderedTurns", { + turns: cache.unordered_turns.toLocaleString(), + })}

)}
@@ -291,13 +327,14 @@ interface BenchmarksBodyProps { } const BenchmarksBody: React.FC = ({ isPending, error, data, selectedKey, autoRouters }) => { - if (isPending) return Loading auto-router usage...; + const { t } = useTranslation(); + if (isPending) return {t("costOptimization.autoRouter.body.loading")}; if (error instanceof ApiError && error.status === 403) { - return Auto-router usage is visible to proxy admin roles only; + return {t("costOptimization.autoRouter.body.forbidden")}; } - if (error || !data) return Auto-router usage is unavailable right now; + if (error || !data) return {t("costOptimization.autoRouter.body.unavailable")}; - const view = viewFor(data, selectedKey); + const view = viewFor(data, selectedKey, t); const stats = view.stats; return ( <> @@ -305,30 +342,32 @@ const BenchmarksBody: React.FC = ({ isPending, error, data, -

- Savings and spend count requests on the selected UTC days. Actual spend covers every request on complexity - routers, including LLM classification cost. Baseline is actual spend plus recorded savings, so savings can be - zero or negative. -

+

{t("costOptimization.autoRouter.body.savingsNote")}

-

- Session metrics cover every session that overlaps the range, including its turns outside the range. -

+

{t("costOptimization.autoRouter.body.sessionNote")}

@@ -338,10 +377,8 @@ const BenchmarksBody: React.FC = ({ isPending, error, data,
-

Auto-router prompt caching

-

- every turn falls in exactly one bucket, by what the router did -

+

{t("costOptimization.autoRouter.caching.title")}

+

{t("costOptimization.autoRouter.caching.subtitle")}

@@ -362,21 +399,26 @@ export const AutoRouterUsageView: React.FC = ({ apiKey, userId, }) => { + const { t } = useTranslation(); const { dateValue, onDateChange } = activity; const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey, userId); const [selectedKey, setSelectedKey] = useState(ALL_ROUTERS); const { data: autoRouters } = useAutoRouters(); const groups = data?.groups ?? []; - const selectedLabel = data ? viewFor(data, selectedKey).label : "All auto-routers"; + const selectedLabel = data ? viewFor(data, selectedKey, t).label : t("costOptimization.autoRouter.allRouters"); const rangeLabel = formatRangeLabel(dateValue.from, dateValue.to); return (
-

Auto-router usage

- {rangeLabel &&

{rangeLabel} (UTC)

} +

{t("costOptimization.autoRouter.usageTitle")}

+ {rangeLabel && ( +

+ {t("costOptimization.autoRouter.rangeUtc", { range: rangeLabel })} +

+ )}
@@ -386,7 +428,7 @@ export const AutoRouterUsageView: React.FC = ({ {selectedLabel} - All auto-routers + {t("costOptimization.autoRouter.allRouters")} {groups.map((g) => ( {groupLabel(g, groups)} @@ -398,12 +440,7 @@ export const AutoRouterUsageView: React.FC = ({
- {userId && ( -

- Usage for this user across API keys and JWT-authenticated requests. Older sessions recorded without a user ID - are not included. -

- )} + {userId &&

{t("costOptimization.autoRouter.userScopeNote")}

} = ({ }; const AutoRouterBenchmarksTab: React.FC = ({ accessToken, activity }) => { + const { t } = useTranslation(); const [visitedTabs, setVisitedTabs] = useState(["usage"]); const handleTabChange = (value: unknown) => { @@ -430,10 +468,10 @@ const AutoRouterBenchmarksTab: React.FC = ({ acces - Usage + {t("costOptimization.autoRouter.tabs.usage")} - Shadow Evals + {t("costOptimization.autoRouter.tabs.shadowEvals")} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx index 0273fae9ec1..2a5cf572008 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx @@ -1,3 +1,5 @@ +import { useTranslation } from "react-i18next"; + import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; @@ -10,33 +12,40 @@ interface AutoRouterSummaryTableProps { } export default function AutoRouterSummaryTable({ groups, selectedGroup }: AutoRouterSummaryTableProps) { + const { t } = useTranslation(); const visibleGroups = selectedGroup ? [selectedGroup] : groups; return ( - Router usage and savings - - Selected UTC days. Cost includes LLM calls and classification. Tokens count routed LLM input and output. - + {t("costOptimization.routerSummary.title")} + {t("costOptimization.routerSummary.description")} - +
- Auto-router - Tokens through router - Cost via router - Cost per 1M tokens - Saved vs. premium model - % saved + {t("costOptimization.routerSummary.columns.router")} + + {t("costOptimization.routerSummary.columns.tokens")} + + + {t("costOptimization.routerSummary.columns.cost")} + + + {t("costOptimization.routerSummary.columns.costPerMillion")} + + + {t("costOptimization.routerSummary.columns.saved")} + + {t("costOptimization.routerSummary.columns.savedPct")} {visibleGroups.length === 0 ? ( - No auto-routers in this range + {t("costOptimization.routerSummary.empty")} ) : ( @@ -44,29 +53,32 @@ export default function AutoRouterSummaryTable({ groups, selectedGroup }: AutoRo {groupLabel(group, groups)} - {group.total_tokens == null ? "Unavailable" : group.total_tokens.toLocaleString()} + {group.total_tokens == null + ? t("costOptimization.routerSummary.unavailable") + : group.total_tokens.toLocaleString()} {usd(group.spend)} {group.total_tokens != null && group.total_tokens > 0 ? usd((group.spend * 1_000_000) / group.total_tokens) - : "Unavailable"} + : t("costOptimization.routerSummary.unavailable")} - {group.saved_spend == null ? "Unavailable" : usd(group.saved_spend)} + {group.saved_spend == null + ? t("costOptimization.routerSummary.unavailable") + : usd(group.saved_spend)} - {group.saved_pct == null ? "Unavailable" : pctLabel(group.saved_pct)} + {group.saved_pct == null + ? t("costOptimization.routerSummary.unavailable") + : pctLabel(group.saved_pct)} )) )}
-

- Savings are estimated against each router's premium baseline. Token totals and unit costs are unavailable - for usage recorded before token tracking; unit cost also requires nonzero tokens. -

+

{t("costOptimization.routerSummary.footnote")}

 ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx index 7a63ccd1ee4..ab34f09f0ff 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx @@ -1,6 +1,7 @@ "use client"; import React, { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { ArrowDown, ArrowUp, ArrowUpDown, Info } from "lucide-react"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; @@ -67,6 +68,7 @@ const SortableHead = ({ sort: SortState; onSort: (column: SortColumn) => void; }) => { + const { t } = useTranslation(); const active = sort.column === column; const ActiveArrow = sort.dir === "asc" ? ArrowUp : ArrowDown; const Arrow = active ? ActiveArrow : ArrowUpDown; @@ -76,7 +78,7 @@ const SortableHead = ({