mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
test: update the spend tests that still expected the pre-ledger accounting queues (#50814)
* test: update the spend tests that still expected the pre-ledger accounting queues #49872 made the SpendLogs ledger the default accounting path on purpose. A batch with no billable request now writes no cost row, a call the guardrail fails closed during a vendor outage keeps one spend row of either status while its key is still charged once, and charges left at a graceful shutdown are applied by the next start. The daily spend batch shutdown tests exercise the old queue, so they pin use_spend_logs_ledger to false. * test(integration): price the guardrail relay model so the outage test can assert the charge The relay model had no price, so every stream cost $0 and the outage test could not tell a charged call from an uncharged one. Both rigs now price it, and the test expects each served stream to cost 9 input and 4 output tokens and the key to be charged for all sixteen. Also notes the skipped batch cost row case in the e2e rules file, next to the other explicit skips. * test: repair four tests that fail for reasons outside the product The OpenAPI docs gate test rebuilds the Rust extension with uv, and the integration lane blocks outbound traffic, so uv could not refresh its PyPI index for maturin. The test now runs uv offline against the cache the dependency install already filled. The response cache test expected a chat reply blocked with empty content to be cached. Since #50057 a reply with no output is never cached, so that case now expects the repeat to reach the upstream and the real answer to be cached. The Gemini cache hit e2e test gave the model 16 tokens, which its thinking can use up, and a reply with no text is not cached. It now allows 512 and checks the first reply has text. The Together cache read e2e test picks the cheapest model with a cache read price. Since #50491 that is a -latest alias, whose target reports no cached tokens. Aliases are now left out of the pick.
This commit is contained in:
parent
9a7f5eeec4
commit
07be67603c
12 changed files with 227 additions and 112 deletions
|
|
@ -37,7 +37,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
|
|||
- `embeddings/` - the `/embeddings` endpoint across providers
|
||||
- `batches/` - the `/batches` endpoint. GovCloud batch coverage is explicitly skipped until its infrastructure is provisioned. Other missing credentials and provider errors remain failures
|
||||
- `realtime/` - realtime websocket sessions, including the pipecat audio path
|
||||
- `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/*`)
|
||||
- `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/*`). The batch cost row attribution case in `spend_tracking/` is explicitly skipped, because the SpendLogs ledger writes no cost row for a $0 batch and a batch with a cost does not finish within one run
|
||||
- `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). 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
|
||||
|
|
|
|||
|
|
@ -1617,17 +1617,6 @@ class TestBatchFailurePaths:
|
|||
assert first_error.message, "batch error item has no message"
|
||||
assert first_error.code, "batch error item has no code"
|
||||
|
||||
rows = client.proxy.poll_logs_for_request_id(f"{fetched.id}_batch_cost")
|
||||
assert rows, (
|
||||
f"failed batch {fetched.id} wrote no spend row; retrieve must book $0 (LIT-4852)"
|
||||
)
|
||||
assert all((row.spend or 0) == 0 for row in rows), (
|
||||
f"failed batch must cost $0, got {[(r.request_id, r.spend) for r in rows]}"
|
||||
)
|
||||
assert rows[0].call_type == "aretrieve_batch", (
|
||||
f"batch cost row call_type={rows[0].call_type!r}"
|
||||
)
|
||||
|
||||
conflict = client.cancel_batch(batch.id, key=key)
|
||||
match conflict:
|
||||
case UnknownApiError(status_code=409, body=body):
|
||||
|
|
|
|||
|
|
@ -166,6 +166,7 @@ def _cheapest_together_chat_model(registry: Mapping[str, CostMapEntry], needs: _
|
|||
def qualifies(name: str, entry: CostMapEntry) -> bool:
|
||||
return (
|
||||
name.startswith("together_ai/")
|
||||
and not name.endswith("-latest")
|
||||
and entry.litellm_provider == "together_ai"
|
||||
and entry.mode == "chat"
|
||||
and (entry.deprecation_date is None or entry.deprecation_date > today)
|
||||
|
|
|
|||
|
|
@ -410,6 +410,9 @@ class TestKeyAttribution:
|
|||
rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT]
|
||||
assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}"
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="the SpendLogs ledger writes no cost row for a $0 batch, and a batch with a cost is out of one run's reach"
|
||||
)
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key",
|
||||
exercised_on=["batches"],
|
||||
|
|
|
|||
|
|
@ -287,8 +287,12 @@ def test_cache_hit_is_zero_cost_and_suffixed(
|
|||
# populated. The marker keeps each run isolated - a fixed prompt would persist
|
||||
# in the shared response cache across runs and make both calls hit (flaky).
|
||||
prompt = f"What is the capital of France? Answer in one word. {unique_marker()}"
|
||||
_ = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=16, cache=None))
|
||||
_ = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=16, cache=None))
|
||||
first = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=512, cache=None))
|
||||
answered = first.choices and first.choices[0].message and first.choices[0].message.content
|
||||
assert answered, (
|
||||
f"{GEMINI_MODEL} returned no text, and a reply with no output is never cached: {first}"
|
||||
)
|
||||
_ = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=512, cache=None))
|
||||
|
||||
rows = client.poll_logs_for_key(
|
||||
scoped_key,
|
||||
|
|
|
|||
|
|
@ -603,34 +603,45 @@ def test_embeddings_are_still_served_from_the_cache(gateway: Gateway) -> None:
|
|||
assert second is not None and second["data"] == first["data"], second
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"choice",
|
||||
(
|
||||
pytest.param(
|
||||
{"index": 0, "message": {"role": "assistant", "content": ""}, "finish_reason": "content_filter"},
|
||||
id="blocked-empty-content",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
|
||||
],
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
},
|
||||
id="tool-calls-only",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_chat_answers_with_a_choice_but_no_text_are_still_cached(
|
||||
gateway: Gateway, choice: dict[str, JsonValue]
|
||||
def test_chat_answer_blocked_with_empty_content_reaches_upstream_again_and_the_real_answer_is_cached(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
upstream: Final = Upstream(gateway.upstream_url)
|
||||
store: Final = redis_store()
|
||||
blocked: Final[dict[str, JsonValue]] = {
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": ""},
|
||||
"finish_reason": "content_filter",
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = scripted(scenario, json_response({**CHAT, "choices": [blocked]}))
|
||||
model: Final = scenario.model(api_base=handle.api_base())
|
||||
request: Final = chat_body(model, prompt())
|
||||
first: Final = _chat(gateway, request)
|
||||
assert object_value(choices(first)[0])["finish_reason"] == "content_filter", first.text
|
||||
assert upstream.calls(handle.scenario_id) == 1
|
||||
rescript(handle, json_response(CHAT))
|
||||
real: Final = _chat(gateway, request)
|
||||
assert response_id(real) != response_id(first) and content_of_first_choice(real) == "scripted", real.text
|
||||
assert upstream.calls(handle.scenario_id) == 2
|
||||
await_entry(store, response_id(real))
|
||||
hit: Final = _chat(gateway, request)
|
||||
assert response_id(hit) == response_id(real) and content_of_first_choice(hit) == "scripted", hit.text
|
||||
assert upstream.calls(handle.scenario_id) == 2
|
||||
|
||||
|
||||
def test_chat_answers_with_tool_calls_but_no_text_are_still_cached(gateway: Gateway) -> None:
|
||||
upstream: Final = Upstream(gateway.upstream_url)
|
||||
store: Final = redis_store()
|
||||
choice: Final[dict[str, JsonValue]] = {
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = scripted(scenario, json_response({**CHAT, "choices": [choice]}))
|
||||
model: Final = scenario.model(api_base=handle.api_base())
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
from typing import Final
|
||||
|
||||
import openapi_docs_gate as gate
|
||||
import pytest
|
||||
|
||||
|
||||
def test_base_counts_imports_native_extension_from_detached_worktree() -> None:
|
||||
def test_base_counts_imports_native_extension_from_detached_worktree(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("UV_OFFLINE", "1")
|
||||
expected_counts: Final = dict(gate.count_by_rule(gate.head_violations()))
|
||||
assert expected_counts
|
||||
assert dict(gate.base_counts("HEAD")) == expected_counts
|
||||
|
|
|
|||
|
|
@ -7,8 +7,10 @@ import time
|
|||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import websockets
|
||||
from _pytest.mark.structures import ParameterSet
|
||||
from integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration.cost_calculation.assertions import assert_exact
|
||||
|
|
@ -150,46 +152,54 @@ def _batch_input_lines(case: BatchCostCase, model_name: str) -> bytes:
|
|||
).encode()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"case",
|
||||
tuple(pytest.param(case, marks=pytest.mark.covers(case.covers), id=case.name) for case in BATCH_CASES),
|
||||
)
|
||||
def _batch_retrieved_twice(
|
||||
gateway: Gateway, scenario: Scenario, case: BatchCostCase
|
||||
) -> tuple[str, str, httpx.Response]:
|
||||
key: Final = scenario.key()
|
||||
model_name, identity = _register_deployment(
|
||||
scenario,
|
||||
case.litellm_model,
|
||||
_batch_response(case),
|
||||
case.name,
|
||||
realtime=False,
|
||||
)
|
||||
file_response: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": model_name},
|
||||
{"file": ("in.jsonl", _batch_input_lines(case, model_name), "application/jsonl")},
|
||||
key=key,
|
||||
)
|
||||
assert file_response.is_success, file_response.text
|
||||
file_body: Final = JSON_OBJECT.validate_json(file_response.content)
|
||||
batch_response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": string_value(file_body["id"]),
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": model_name,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert batch_response.is_success, batch_response.text
|
||||
batch_body: Final = JSON_OBJECT.validate_json(batch_response.content)
|
||||
batch_id: Final = string_value(batch_body["id"])
|
||||
first_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
|
||||
second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
|
||||
assert first_retrieval.is_success, first_retrieval.text
|
||||
assert second_retrieval.is_success, second_retrieval.text
|
||||
return key, identity, second_retrieval
|
||||
|
||||
|
||||
def _batch_params(cases: tuple[BatchCostCase, ...]) -> tuple[ParameterSet, ...]:
|
||||
return tuple(pytest.param(case, marks=pytest.mark.covers(case.covers), id=case.name) for case in cases)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", _batch_params(tuple(case for case in BATCH_CASES if case.expected.spend > 0)))
|
||||
def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
model_name, identity = _register_deployment(
|
||||
scenario,
|
||||
case.litellm_model,
|
||||
_batch_response(case),
|
||||
case.name,
|
||||
realtime=False,
|
||||
)
|
||||
file_response: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": model_name},
|
||||
{"file": ("in.jsonl", _batch_input_lines(case, model_name), "application/jsonl")},
|
||||
key=key,
|
||||
)
|
||||
assert file_response.is_success, file_response.text
|
||||
file_body: Final = JSON_OBJECT.validate_json(file_response.content)
|
||||
batch_response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": string_value(file_body["id"]),
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": model_name,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert batch_response.is_success, batch_response.text
|
||||
batch_body: Final = JSON_OBJECT.validate_json(batch_response.content)
|
||||
batch_id: Final = string_value(batch_body["id"])
|
||||
first_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
|
||||
second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
|
||||
assert first_retrieval.is_success, first_retrieval.text
|
||||
assert second_retrieval.is_success, second_retrieval.text
|
||||
key, identity, second_retrieval = _batch_retrieved_twice(gateway, scenario, case)
|
||||
retrieval_rows: Final = poll_rows_where(key, 1, lambda row: row.call_type == "aretrieve_batch")
|
||||
assert len(retrieval_rows) == 1
|
||||
rows: Final = read_rows_now(key)
|
||||
|
|
@ -203,6 +213,22 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None:
|
|||
assert len(tuple(row for row in read_rows_now(key) if row.call_type == "aretrieve_batch")) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", _batch_params(tuple(case for case in BATCH_CASES if case.expected.spend == 0)))
|
||||
def test_batch_with_no_billable_request_is_retrieved_free_and_writes_no_cost_row(
|
||||
gateway: Gateway, case: BatchCostCase
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
key, _, _ = _batch_retrieved_twice(gateway, scenario, case)
|
||||
later_model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
later: Final = gateway.chat(later_model, key=key, text="sent after the batch retrievals")
|
||||
assert later["usage"] == {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}, later
|
||||
rows: Final = poll_rows_where(key, 1, lambda row: row.call_type == "acompletion")
|
||||
assert [row.spend for row in rows] == [pytest.approx(20 * 0.001 + 20 * 0.002)], rows
|
||||
every_row: Final = read_rows_now(key)
|
||||
assert [row.call_type for row in every_row if row.call_type == "aretrieve_batch"] == [], every_row
|
||||
assert sum(row.spend or 0.0 for row in every_row) == pytest.approx(20 * 0.001 + 20 * 0.002), every_row
|
||||
|
||||
|
||||
def _realtime_response(case: RealtimeCostCase) -> RealtimeResponse:
|
||||
return RealtimeResponse(
|
||||
content_type="application/x-realtime",
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ request as before."""
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
|
|
@ -39,6 +40,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server
|
|||
pytestmark: Final = pytest.mark.shares_worker_state
|
||||
|
||||
MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
|
||||
INPUT_COST_PER_TOKEN: Final = 0.001
|
||||
OUTPUT_COST_PER_TOKEN: Final = 0.002
|
||||
ROUTER_MODEL: Final = "converse-relay"
|
||||
EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
|
||||
SAMPLED: Final = "relay-sampled"
|
||||
|
|
@ -390,6 +393,10 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
|||
"aws_secret_access_key": "synthetic-relay-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
"model_info": {
|
||||
"input_cost_per_token": INPUT_COST_PER_TOKEN,
|
||||
"output_cost_per_token": OUTPUT_COST_PER_TOKEN,
|
||||
},
|
||||
}
|
||||
]
|
||||
config_path: Final = directory / "converse-stream-relay.yaml"
|
||||
|
|
@ -651,6 +658,24 @@ def _spend_statuses(call_id: str, at_least: int) -> list[JsonValue]:
|
|||
return [row["status"] for row in rows]
|
||||
|
||||
|
||||
def _logged_spend(call_id: str) -> float:
|
||||
rows: Final = read_rows(
|
||||
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s OR request_id=%s '
|
||||
"OR metadata->>'litellm_call_id'=%s",
|
||||
(call_id, call_id, call_id),
|
||||
)
|
||||
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
|
||||
return rows[0]["spend"]
|
||||
|
||||
|
||||
def _key_spend(key: str) -> float:
|
||||
rows: Final = read_rows(
|
||||
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashlib.sha256(key.encode()).hexdigest(),)
|
||||
)
|
||||
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
|
||||
return rows[0]["spend"]
|
||||
|
||||
|
||||
def _exception(frames: Sequence[Frame]) -> tuple[str, str] | None:
|
||||
last: Final = frames[-1] if frames else None
|
||||
if last is None or last.message_type != "exception":
|
||||
|
|
@ -1001,6 +1026,7 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
|
|||
with ThreadPoolExecutor(max_workers=count) as pool:
|
||||
return list(zip(markers, pool.map(partial(_stream_with_call_id, rig, OUTAGE), markers)))
|
||||
|
||||
spend_before: Final = _key_spend(rig.keys[OUTAGE])
|
||||
rig.outage.set()
|
||||
try:
|
||||
during: Final = burst(8)
|
||||
|
|
@ -1013,5 +1039,17 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
|
|||
for marker, (status, _, body, _) in after:
|
||||
assert status == 200 and _exception(decode(body)) is None, (marker, body)
|
||||
assert body == b"".join(rig.scripts[marker].chunks or ()), (marker, decode(body))
|
||||
for marker, (_, _, _, call_id) in (*during, *after):
|
||||
assert _spend_statuses(call_id, 1).count("success") == 1, marker
|
||||
for marker, (_, _, _, call_id) in after:
|
||||
assert _spend_statuses(call_id, 1) == ["success"], marker
|
||||
for marker, (_, _, _, call_id) in during:
|
||||
assert len(_spend_statuses(call_id, 1)) == 1, marker
|
||||
stream_cost: Final = 9 * INPUT_COST_PER_TOKEN + 4 * OUTPUT_COST_PER_TOKEN
|
||||
costs: Final = [_logged_spend(call_id) for _, (_, _, _, call_id) in after]
|
||||
assert costs == [pytest.approx(stream_cost)] * len(after), costs
|
||||
charged: Final = eventually(
|
||||
lambda: _key_spend(rig.keys[OUTAGE]) - spend_before,
|
||||
lambda spend: spend == pytest.approx(len((*during, *after)) * stream_cost),
|
||||
seconds=70,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert charged == pytest.approx(len((*during, *after)) * stream_cost), charged
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ Streams of other model families keep going out as received"""
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
|
|
@ -43,6 +44,8 @@ pytestmark: Final = pytest.mark.shares_worker_state
|
|||
CLAUDE_ID: Final = "anthropic.claude-sonnet-5-v1:0"
|
||||
LLAMA_ID: Final = "meta.llama4-maverick-17b-instruct-v1:0"
|
||||
PROFILE_ID: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/relayprofile"
|
||||
INPUT_COST_PER_TOKEN: Final = 0.001
|
||||
OUTPUT_COST_PER_TOKEN: Final = 0.002
|
||||
CLAUDE: Final = "invoke-relay"
|
||||
LLAMA: Final = "invoke-relay-llama"
|
||||
PROFILE: Final = "invoke-relay-profile"
|
||||
|
|
@ -334,6 +337,7 @@ def _deployment(name: str, model_id: str, upstream: Wire) -> dict[str, JsonValue
|
|||
"aws_secret_access_key": "synthetic-relay-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
"model_info": {"input_cost_per_token": INPUT_COST_PER_TOKEN, "output_cost_per_token": OUTPUT_COST_PER_TOKEN},
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -468,6 +472,24 @@ def _stream(rig: Rig, guardrail: str, marker: str, model: str = CLAUDE) -> tuple
|
|||
return status, content_type, body
|
||||
|
||||
|
||||
def _logged_spend(call_id: str) -> float:
|
||||
rows: Final = read_rows(
|
||||
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s OR request_id=%s '
|
||||
"OR metadata->>'litellm_call_id'=%s",
|
||||
(call_id, call_id, call_id),
|
||||
)
|
||||
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
|
||||
return rows[0]["spend"]
|
||||
|
||||
|
||||
def _key_spend(key: str) -> float:
|
||||
rows: Final = read_rows(
|
||||
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashlib.sha256(key.encode()).hexdigest(),)
|
||||
)
|
||||
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
|
||||
return rows[0]["spend"]
|
||||
|
||||
|
||||
def _exception(frames: Sequence[Frame]) -> tuple[str, str] | None:
|
||||
last: Final = frames[-1] if frames else None
|
||||
if last is None or last.message_type != "exception":
|
||||
|
|
@ -807,6 +829,7 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
|
|||
with ThreadPoolExecutor(max_workers=count) as pool:
|
||||
return list(zip(markers, pool.map(partial(_stream_with_call_id, rig, OUTAGE), markers)))
|
||||
|
||||
spend_before: Final = _key_spend(rig.keys[OUTAGE])
|
||||
rig.outage.set()
|
||||
try:
|
||||
during: Final = burst(8)
|
||||
|
|
@ -819,8 +842,20 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
|
|||
for marker, (status, _, body, _) in after:
|
||||
assert status == 200 and _exception(decode(body)) is None, (marker, body)
|
||||
assert body == b"".join(rig.scripts[marker].chunks or ()), (marker, decode(body))
|
||||
for marker, (_, _, _, call_id) in (*during, *after):
|
||||
assert _spend_statuses(call_id, 1).count("success") == 1, marker
|
||||
for marker, (_, _, _, call_id) in after:
|
||||
assert _spend_statuses(call_id, 1) == ["success"], marker
|
||||
for marker, (_, _, _, call_id) in during:
|
||||
assert len(_spend_statuses(call_id, 1)) == 1, marker
|
||||
stream_cost: Final = 9 * INPUT_COST_PER_TOKEN + 4 * OUTPUT_COST_PER_TOKEN
|
||||
costs: Final = [_logged_spend(call_id) for _, (_, _, _, call_id) in after]
|
||||
assert costs == [pytest.approx(stream_cost)] * len(after), costs
|
||||
charged: Final = eventually(
|
||||
lambda: _key_spend(rig.keys[OUTAGE]) - spend_before,
|
||||
lambda spend: spend == pytest.approx(len((*during, *after)) * stream_cost),
|
||||
seconds=70,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert charged == pytest.approx(len((*during, *after)) * stream_cost), charged
|
||||
|
||||
|
||||
def test_client_disconnect_stops_the_proxy_reading_the_upstream(rig: Rig) -> None:
|
||||
|
|
|
|||
|
|
@ -94,14 +94,19 @@ def _tagged_chat(gateway: Gateway, key: str, tags: tuple[str, str]) -> None:
|
|||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
def _assert_requests_and_spend(table: str, count: str, key: str, requests: int) -> None:
|
||||
hashed: Final = hashlib.sha256(key.encode()).hexdigest()
|
||||
rows: Final = read_rows(
|
||||
f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,)
|
||||
)
|
||||
assert rows[0]["requests"] == requests, (table, rows)
|
||||
assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows)
|
||||
|
||||
|
||||
def _assert_shutdown_accounting(key: str, tags: tuple[str, str], requests: int) -> None:
|
||||
hashed: Final = hashlib.sha256(key.encode()).hexdigest()
|
||||
for table, count in (("LiteLLM_SpendLogs", "count(*)"), ("LiteLLM_DailyUserSpend", "sum(api_requests)")):
|
||||
rows: Final = read_rows(
|
||||
f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,)
|
||||
)
|
||||
assert rows[0]["requests"] == requests, (table, rows)
|
||||
assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows)
|
||||
_assert_requests_and_spend("LiteLLM_SpendLogs", "count(*)", key, requests)
|
||||
_assert_requests_and_spend("LiteLLM_DailyUserSpend", "sum(api_requests)", key, requests)
|
||||
key_rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,))
|
||||
assert key_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9)
|
||||
for tag in tags:
|
||||
|
|
@ -137,7 +142,7 @@ def test_tagged_request_survives_shutdown_before_the_tag_timer(
|
|||
|
||||
@pytest.mark.parametrize("redis_buffer", [False, True])
|
||||
@pytest.mark.parametrize("overflow", [False, True])
|
||||
def test_callback_backlog_is_accounted_before_shutdown_disconnects(
|
||||
def test_callback_backlog_is_logged_before_shutdown_disconnects_and_accounted_once_by_the_next_start(
|
||||
gateway: Gateway, tmp_path: Path, redis_buffer: bool, overflow: bool
|
||||
) -> None:
|
||||
release: Final = threading.Event()
|
||||
|
|
@ -148,19 +153,15 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects(
|
|||
wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger,
|
||||
):
|
||||
config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True)
|
||||
with owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{
|
||||
"DATABASE_URL": os.environ["DATABASE_URL"],
|
||||
"LOGGING_WORKER_CONCURRENCY": "2",
|
||||
"LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000",
|
||||
"SHUTDOWN_CALLBACK_URL": logger.url,
|
||||
"GRACEFUL_SHUTDOWN_TIMEOUT": "15",
|
||||
"LITELLM_LOG": "DEBUG",
|
||||
},
|
||||
config=config,
|
||||
) as owned:
|
||||
overrides: Final = {
|
||||
"DATABASE_URL": os.environ["DATABASE_URL"],
|
||||
"LOGGING_WORKER_CONCURRENCY": "2",
|
||||
"LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000",
|
||||
"SHUTDOWN_CALLBACK_URL": logger.url,
|
||||
"GRACEFUL_SHUTDOWN_TIMEOUT": "15",
|
||||
"LITELLM_LOG": "DEBUG",
|
||||
}
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned:
|
||||
key: Final = string_value(
|
||||
owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"]
|
||||
)
|
||||
|
|
@ -179,9 +180,13 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects(
|
|||
)
|
||||
release.set()
|
||||
owned.process.wait(timeout=90)
|
||||
_assert_shutdown_accounting(key, tags, requests)
|
||||
finally:
|
||||
release.set()
|
||||
_assert_requests_and_spend("LiteLLM_SpendLogs", "count(*)", key, requests)
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as restarted:
|
||||
restarted.begin_shutdown()
|
||||
restarted.process.wait(timeout=90)
|
||||
_assert_shutdown_accounting(key, tags, requests)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gunicorn", [False, True])
|
||||
|
|
@ -396,17 +401,18 @@ class _Shutdown:
|
|||
self.owned.process.wait(timeout=120)
|
||||
|
||||
|
||||
def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path:
|
||||
def _legacy_spend_queue_config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["general_settings"]["database_connection_pool_limit"] = pool_limit
|
||||
config["general_settings"]["database_connection_pool_timeout"] = 60
|
||||
config["general_settings"]["use_spend_logs_ledger"] = False
|
||||
path: Final = tmp_path / f"pool-{pool_limit}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _proxy_with_one_seeded_row(
|
||||
def _legacy_spend_queue_proxy_with_one_seeded_row(
|
||||
gateway: Gateway,
|
||||
tmp_path: Path,
|
||||
pool_limit: int,
|
||||
|
|
@ -429,7 +435,7 @@ def _proxy_with_one_seeded_row(
|
|||
"SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS": "1",
|
||||
"SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": str(cancel_timeout_seconds),
|
||||
},
|
||||
config=_config_with_pool_limit(tmp_path, pool_limit),
|
||||
config=_legacy_spend_queue_config_with_pool_limit(tmp_path, pool_limit),
|
||||
remove_environment=("DATABASE_URL_READ_REPLICA",),
|
||||
workers=workers,
|
||||
) as owned:
|
||||
|
|
@ -460,7 +466,7 @@ def test_daily_spend_batch_cancelled_while_waiting_for_a_pool_connection_is_writ
|
|||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with (
|
||||
_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=2) as shutdown,
|
||||
_legacy_spend_queue_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=2) as shutdown,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as models,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
|
||||
):
|
||||
|
|
@ -482,7 +488,7 @@ def test_daily_spend_batch_cancelled_while_waiting_for_a_row_lock_is_written_exa
|
|||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with (
|
||||
_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=10) as shutdown,
|
||||
_legacy_spend_queue_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=10) as shutdown,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as holder,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
|
||||
):
|
||||
|
|
@ -508,7 +514,7 @@ def test_daily_spend_batch_cancelled_while_postgres_is_committing_it_is_written_
|
|||
gateway: Gateway, tmp_path: Path, cancel_timeout_seconds: int, commit_fails_once: bool
|
||||
) -> None:
|
||||
with (
|
||||
_proxy_with_one_seeded_row(
|
||||
_legacy_spend_queue_proxy_with_one_seeded_row(
|
||||
gateway, tmp_path, pool_limit=10, cancel_timeout_seconds=cancel_timeout_seconds, settle_seconds=90
|
||||
) as shutdown,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
|
||||
|
|
@ -528,7 +534,7 @@ def test_daily_spend_batch_cancelled_while_postgres_is_committing_it_is_written_
|
|||
def test_daily_spend_burst_across_two_workers_survives_shutdown_during_commit_exactly_once(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with _proxy_with_one_seeded_row(
|
||||
with _legacy_spend_queue_proxy_with_one_seeded_row(
|
||||
gateway, tmp_path, pool_limit=10, settle_seconds=120, requests=BURST_REQUESTS, workers=2
|
||||
) as shutdown:
|
||||
suffix: Final = _install_slow_commit(shutdown.owner, fails_once=False)
|
||||
|
|
|
|||
|
|
@ -278,7 +278,7 @@ def test_member_added_without_any_budget_is_charged_on_its_membership_row(gatewa
|
|||
|
||||
|
||||
@pytest.mark.covers("spend.team_member.stale_low_redis_counter_still_blocks_member_over_budget")
|
||||
def test_member_over_budget_is_blocked_when_redis_counter_reads_stale_low(gateway: Gateway) -> None:
|
||||
def test_member_over_budget_is_blocked_when_redis_counter_reads_stale_low(gateway: Gateway, peer: Gateway) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
|
|
@ -308,7 +308,7 @@ def test_member_over_budget_is_blocked_when_redis_counter_reads_stale_low(gatewa
|
|||
assert float(counted) == pytest.approx(0.06), counted
|
||||
cache.set(counter_key, "0.01")
|
||||
upstream.delete("/__observations").raise_for_status()
|
||||
denied: Final = gateway.request(
|
||||
denied: Final = peer.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"stale counter {uuid.uuid4().hex}"}]},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue