mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fc27dfdc7d
* test(integration): give each owned proxy a private database and Redis by default owned_proxy and owned_proxy_process now boot on a scratch database and a per-proxy redis-server unless the caller opts into the worker's state with state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips DATABASE_URL. Under the replica profile the private database gets the writer and reader roles so the reader path stays covered. Gateway carries the database_url its proxy writes to, and the Scenario cleanup checks and delete_key_if_present read from it, so they keep checking something on a private database. The 189 test files with owned proxy sites that depend on the worker's database or Redis opt in with a shares_worker_state module marker, which a module-scoped conftest fixture turns into the shared state. The 7 helper modules with such sites opt in per call, since conftest fixtures and other test files reach them. * test(integration): reach session fixtures with the shared marker and keep the worker retry default The module-scoped fixture that applied shares_worker_state ran after session-scoped fixtures, so proxies those booted went private in marked files. The marker is now applied around each marked item's setup, call and teardown, which covers fixtures of every scope. Owned proxies on the worker database inherited the worker gateway's router_settings.num_retries of 0 from it. A private proxy now gets the same default in its config copy unless the config sets num_retries itself. Five providers files read the owned proxy's spend rows or response cache through the worker's database or Redis and opt into sharing.
603 lines
26 KiB
Python
603 lines
26 KiB
Python
import json
|
|
from collections.abc import Callable, Generator, Mapping
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from hashlib import sha256
|
|
from pathlib import Path
|
|
from typing import Final
|
|
from uuid import uuid4
|
|
|
|
import httpx
|
|
import pytest
|
|
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
|
|
from integration._support.database import read_rows, scratch_database
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.upstream import delete_scenario, register_scenario
|
|
from integration._support.wire import Reply, Request, Wire, wire_server
|
|
from integration.spend.test_batch_completion_accounting import batch_input_file, batch_routes
|
|
from integration.streaming.test_stream_contracts import text_stream
|
|
from pydantic import JsonValue
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
CACHED_PROMPT_TOKENS: Final = 4
|
|
|
|
|
|
def _ollama_fixed_id_reply(request: Request) -> Reply:
|
|
body: Final = JSON_OBJECT.validate_json(request.body)
|
|
assert request.target == "/v1/chat/completions", request.target
|
|
assert body["model"] == "gpt-oss:20b", body
|
|
common: Final = {
|
|
"id": "chatcmpl-0",
|
|
"model": "gpt-oss:20b",
|
|
"created": 1,
|
|
"system_fingerprint": "fp_ollama",
|
|
}
|
|
usage: Final = {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}
|
|
if body.get("stream") is True:
|
|
chunks: Final = (
|
|
{"choices": [{"index": 0, "delta": {"role": "assistant", "content": "h"}, "finish_reason": None}]},
|
|
{"choices": [{"index": 0, "delta": {"content": "i"}, "finish_reason": None}]},
|
|
{"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
|
{"choices": [], "usage": usage},
|
|
)
|
|
return Reply(
|
|
content_type="text/event-stream",
|
|
chunks=tuple(
|
|
("data: " + json.dumps({**common, "object": "chat.completion.chunk", **chunk}) + "\n\n").encode()
|
|
for chunk in chunks
|
|
)
|
|
+ (b"data: [DONE]\n\n",),
|
|
)
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
**common,
|
|
"object": "chat.completion",
|
|
"usage": usage,
|
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
|
}
|
|
).encode()
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("stream", [False, True])
|
|
def test_reused_ollama_response_ids_keep_every_spend_row(gateway: Gateway, stream: bool) -> None:
|
|
with wire_server(_ollama_fixed_id_reply) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(
|
|
model="openai/gpt-oss:20b",
|
|
api_base=wire.url + "/v1",
|
|
api_key="synthetic-key",
|
|
input_cost_per_token=0.001,
|
|
output_cost_per_token=0.002,
|
|
)
|
|
key: Final = scenario.key(models=[model])
|
|
responses: Final = tuple(
|
|
gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": f"hi {index}"}],
|
|
"stream": stream,
|
|
**({"stream_options": {"include_usage": True}} if stream else {}),
|
|
},
|
|
key=key,
|
|
)
|
|
for index in range(3)
|
|
)
|
|
assert tuple(response.status_code for response in responses) == (200, 200, 200)
|
|
returned: Final = tuple(
|
|
tuple(
|
|
JSON_OBJECT.validate_json(line[6:])
|
|
for line in response.text.splitlines()
|
|
if line.startswith("data: ") and line != "data: [DONE]"
|
|
)
|
|
if stream
|
|
else (JSON_OBJECT.validate_json(response.content),)
|
|
for response in responses
|
|
)
|
|
ids: Final = tuple(string_value(chunks[0]["id"]) for chunks in returned)
|
|
for chunks, response_id in zip(returned, ids, strict=True):
|
|
assert {chunk["id"] for chunk in chunks} == {response_id}
|
|
assert object_value(chunks[0]["provider_specific_fields"])["original_response_id"] == "chatcmpl-0"
|
|
assert object_value(chunks[-1]["usage"])["total_tokens"] == 6
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
),
|
|
lambda values: len(values) == 3,
|
|
seconds=70,
|
|
)
|
|
assert len(set(ids)) == 3
|
|
assert {row["request_id"] for row in rows} == set(ids)
|
|
assert tuple(float(str(row["spend"])) for row in rows) == pytest.approx((0.007, 0.007, 0.007))
|
|
|
|
|
|
def _config(
|
|
tmp_path: Path,
|
|
spend_logs_metadata_fields: Mapping[str, JsonValue] | None,
|
|
*,
|
|
model_list: tuple[Mapping[str, JsonValue], ...] = (),
|
|
litellm_settings: Mapping[str, JsonValue] | None = None,
|
|
guardrails: tuple[Mapping[str, JsonValue], ...] = (),
|
|
) -> Path:
|
|
retention: Final = (
|
|
{} if spend_logs_metadata_fields is None else {"spend_logs_metadata_fields": dict(spend_logs_metadata_fields)}
|
|
)
|
|
config: Final = tmp_path / f"spend-logs-metadata-fields-{uuid4()}.json"
|
|
config.write_text(
|
|
json.dumps(
|
|
{
|
|
"model_list": [dict(model) for model in model_list],
|
|
"general_settings": {
|
|
"master_key": "os.environ/LITELLM_MASTER_KEY",
|
|
"database_url": "os.environ/DATABASE_URL",
|
|
"store_model_in_db": True,
|
|
"proxy_batch_write_at": 1,
|
|
"proxy_batch_polling_interval": 1,
|
|
**retention,
|
|
},
|
|
"litellm_settings": dict(litellm_settings or {}),
|
|
"guardrails": [dict(guardrail) for guardrail in guardrails],
|
|
}
|
|
)
|
|
)
|
|
return config
|
|
|
|
|
|
def _respond(request: Request) -> Reply:
|
|
if request.method == "GET":
|
|
return Reply(body=b'{"object":"list","data":[]}')
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"id": f"chatcmpl-{uuid4()}",
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 2,
|
|
"total_tokens": 12,
|
|
"prompt_tokens_details": {"cached_tokens": CACHED_PROMPT_TOKENS},
|
|
},
|
|
}
|
|
).encode()
|
|
)
|
|
|
|
|
|
def _stored_row_after_one_chat(
|
|
gateway: Gateway, tmp_path: Path, spend_logs_metadata_fields: Mapping[str, JsonValue]
|
|
) -> tuple[dict[str, JsonValue], str]:
|
|
with (
|
|
wire_server(_respond) as wire,
|
|
owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, spend_logs_metadata_fields)) as isolated,
|
|
isolated.scenario() as scenario,
|
|
):
|
|
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key")
|
|
api_key: Final = scenario.key(key_alias=f"metadata-fields-{uuid4()}", models=[model])
|
|
response_id: Final = string_value(isolated.chat(model, key=api_key)["id"])
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT metadata, proxy_server_request, response FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
|
(response_id,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
return rows[0], sha256(api_key.encode()).hexdigest()
|
|
|
|
|
|
def test_excluded_metadata_fields_are_not_stored_but_still_reach_daily_spend(gateway: Gateway, tmp_path: Path) -> None:
|
|
row, hashed_key = _stored_row_after_one_chat(
|
|
gateway, tmp_path, {"exclude": ["model_map_information", "usage_object"]}
|
|
)
|
|
|
|
metadata: Final = object_value(row["metadata"])
|
|
assert "model_map_information" not in metadata, metadata
|
|
assert "usage_object" not in metadata, metadata
|
|
assert {"status", "cold_storage_object_key"} <= set(metadata), metadata
|
|
assert string_value(metadata["user_api_key_alias"]).startswith("metadata-fields-")
|
|
assert row["proxy_server_request"] == {}
|
|
assert row["response"] == {}
|
|
daily: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT prompt_tokens, cache_read_input_tokens FROM "LiteLLM_DailyUserSpend" WHERE api_key=%s',
|
|
(hashed_key,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert daily[0]["cache_read_input_tokens"] == CACHED_PROMPT_TOKENS, daily
|
|
|
|
|
|
def test_included_metadata_fields_are_the_only_ones_stored_besides_always_kept(
|
|
gateway: Gateway, tmp_path: Path
|
|
) -> None:
|
|
row, _ = _stored_row_after_one_chat(gateway, tmp_path, {"include": ["user_api_key_alias"]})
|
|
|
|
assert set(object_value(row["metadata"])) == {"status", "cold_storage_object_key", "user_api_key_alias"}
|
|
|
|
|
|
def test_guardrail_usage_is_tracked_when_guardrail_information_is_not_stored(gateway: Gateway, tmp_path: Path) -> None:
|
|
guardrail_name: Final = f"metadata-fields-guardrail-{uuid4()}"
|
|
with (
|
|
wire_server(lambda _: Reply(body=b'{"action":"NONE"}')) as policy,
|
|
wire_server(_respond) as wire,
|
|
owned_proxy(
|
|
gateway,
|
|
tmp_path,
|
|
{},
|
|
config=_config(
|
|
tmp_path,
|
|
{"exclude": ["guardrail_information"]},
|
|
guardrails=(
|
|
{
|
|
"guardrail_name": guardrail_name,
|
|
"litellm_params": {
|
|
"guardrail": "generic_guardrail_api",
|
|
"mode": "pre_call",
|
|
"default_on": True,
|
|
"api_base": policy.url,
|
|
"api_key": "synthetic-guardrail-key",
|
|
},
|
|
},
|
|
),
|
|
),
|
|
) as isolated,
|
|
isolated.scenario() as scenario,
|
|
):
|
|
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key")
|
|
response_id: Final = string_value(isolated.chat(model, key=scenario.key(models=[model]))["id"])
|
|
assert len(policy.drain()) == 1
|
|
rows: Final = eventually(
|
|
lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert "guardrail_information" not in object_value(rows[0]["metadata"]), rows[0]
|
|
indexed: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT guardrail_id FROM "LiteLLM_SpendLogGuardrailIndex" WHERE request_id=%s', (response_id,)
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
metrics: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT requests_evaluated, passed_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s',
|
|
(string_value(indexed[0]["guardrail_id"]),),
|
|
),
|
|
lambda values: values == [{"requests_evaluated": 1, "passed_count": 1}],
|
|
seconds=70,
|
|
)
|
|
assert metrics == [{"requests_evaluated": 1, "passed_count": 1}], metrics
|
|
|
|
|
|
EXCLUDED: Final = ("model_map_information", "user_api_key_alias")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Isolated:
|
|
proxy: Gateway
|
|
scenario: Scenario
|
|
key: str
|
|
database_url: str | None = None
|
|
|
|
@property
|
|
def hashed_key(self) -> str:
|
|
return sha256(self.key.encode()).hexdigest()
|
|
|
|
def rows(self, count: int, where: str = "TRUE", seconds: int = 70) -> tuple[dict[str, JsonValue], ...]:
|
|
return tuple(
|
|
eventually(
|
|
lambda: read_rows(
|
|
"SELECT request_id, call_type, status, spend, cache_hit, metadata "
|
|
f'FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND {where} ORDER BY "startTime"',
|
|
(self.hashed_key,),
|
|
database_url=self.database_url,
|
|
),
|
|
lambda values: len(values) == count,
|
|
seconds=seconds,
|
|
)
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def _isolated(gateway: Gateway, config: Path, tmp_path: Path, database_url: str | None = None) -> Generator[_Isolated]:
|
|
database: Final = {} if database_url is None else {"DATABASE_URL": database_url}
|
|
replica: Final = () if database_url is None else ("DATABASE_URL_READ_REPLICA",)
|
|
with (
|
|
owned_proxy(gateway, tmp_path, database, config=config, remove_environment=replica) as proxy,
|
|
proxy.scenario() as scenario,
|
|
):
|
|
key: Final = scenario.key(key_alias=f"metadata-fields-{uuid4()}")
|
|
yield _Isolated(proxy, scenario, key, database_url)
|
|
|
|
|
|
def _assert_filtered(row: Mapping[str, JsonValue], *kept: str) -> dict[str, JsonValue]:
|
|
metadata: Final = object_value(row["metadata"])
|
|
assert not set(EXCLUDED) & set(metadata), metadata
|
|
assert {"status", "cold_storage_object_key", *kept} <= set(metadata), metadata
|
|
return metadata
|
|
|
|
|
|
def _anthropic_event(event: Mapping[str, JsonValue]) -> bytes:
|
|
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
|
|
|
|
|
|
def _anthropic_stream(model: JsonValue) -> tuple[bytes, ...]:
|
|
message: Final = {**_anthropic_message(model), "content": [], "stop_reason": None}
|
|
events: Final[tuple[Mapping[str, JsonValue], ...]] = (
|
|
{"type": "message_start", "message": message},
|
|
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
|
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
|
|
{"type": "content_block_stop", "index": 0},
|
|
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}},
|
|
{"type": "message_stop"},
|
|
)
|
|
return tuple(_anthropic_event(event) for event in events)
|
|
|
|
|
|
def _respond_streaming_or_not(request: Request) -> Reply:
|
|
if request.method == "GET":
|
|
return Reply(body=b'{"object":"list","data":[]}')
|
|
if request.target.startswith("/v1/messages"):
|
|
return _anthropic(request)
|
|
if JSON_OBJECT.validate_json(request.body).get("stream") is True:
|
|
return Reply(content_type="text/event-stream", chunks=text_stream(f"chatcmpl-{uuid4()}"))
|
|
return _respond(request)
|
|
|
|
|
|
def _stream(proxy: Gateway, key: str, path: str, body: Mapping[str, JsonValue]) -> None:
|
|
with proxy.client.stream("POST", path, json=dict(body), headers={"Authorization": f"Bearer {key}"}) as response:
|
|
assert response.status_code == 200, response.read().decode()
|
|
lines: Final = tuple(response.iter_lines())
|
|
assert any(line.startswith("data:") for line in lines), lines
|
|
|
|
|
|
def _call(proxy: Gateway, key: str, path: str, body: Mapping[str, JsonValue]) -> None:
|
|
response: Final = proxy.request("POST", path, body, key=key)
|
|
assert response.status_code == 200, response.text
|
|
|
|
|
|
def test_every_inference_surface_and_stream_stores_filtered_metadata(gateway: Gateway, tmp_path: Path) -> None:
|
|
with (
|
|
wire_server(_respond_streaming_or_not) as wire,
|
|
_isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated,
|
|
):
|
|
model: Final = isolated.scenario.model(
|
|
model="deepseek/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-key"
|
|
)
|
|
anthropic: Final = isolated.scenario.model(
|
|
model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key"
|
|
)
|
|
prompt: Final[list[JsonValue]] = [{"role": "user", "content": f"surfaces {uuid4()}"}]
|
|
calls: Final[tuple[Callable[[Gateway, str, str, Mapping[str, JsonValue]], None], ...]] = (
|
|
_stream,
|
|
_call,
|
|
_stream,
|
|
_call,
|
|
_stream,
|
|
)
|
|
requests: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = (
|
|
("/v1/chat/completions", {"model": model, "messages": prompt, "stream": True}),
|
|
("/v1/messages", {"model": anthropic, "max_tokens": 16, "messages": prompt}),
|
|
("/v1/messages", {"model": anthropic, "max_tokens": 16, "messages": prompt, "stream": True}),
|
|
("/v1/responses", {"model": model, "input": f"surfaces {uuid4()}"}),
|
|
("/v1/responses", {"model": model, "input": f"surfaces {uuid4()}", "stream": True}),
|
|
)
|
|
for send, (path, body) in zip(calls, requests, strict=True):
|
|
send(isolated.proxy, isolated.key, path, body)
|
|
|
|
rows: Final = isolated.rows(len(requests))
|
|
for row in rows:
|
|
assert row["status"] == "success", row
|
|
_assert_filtered(row, "usage_object")
|
|
|
|
|
|
def test_failed_request_keeps_its_error_and_status_but_drops_excluded_metadata(
|
|
gateway: Gateway, tmp_path: Path
|
|
) -> None:
|
|
def fail(request: Request) -> Reply:
|
|
if request.method == "GET":
|
|
return Reply(body=b'{"object":"list","data":[]}')
|
|
return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}')
|
|
|
|
with (
|
|
wire_server(fail) as wire,
|
|
_isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated,
|
|
):
|
|
model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic")
|
|
response: Final = isolated.proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": "fail"}]},
|
|
key=isolated.key,
|
|
)
|
|
assert response.status_code >= 500, response.text
|
|
|
|
row: Final = isolated.rows(1)[0]
|
|
assert row["status"] == "failure", row
|
|
metadata: Final = _assert_filtered(row, "error_information")
|
|
assert metadata["status"] == "failure", metadata
|
|
assert "scripted upstream failure" in json.dumps(metadata["error_information"]), metadata
|
|
|
|
|
|
def test_response_cache_hit_row_is_filtered_and_charged_nothing(gateway: Gateway, tmp_path: Path) -> None:
|
|
config: Final = _config(
|
|
tmp_path, {"exclude": list(EXCLUDED)}, litellm_settings={"cache": True, "cache_params": {"type": "local"}}
|
|
)
|
|
with wire_server(_respond) as wire, _isolated(gateway, config, tmp_path) as isolated:
|
|
model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic")
|
|
text: Final = f"cache {uuid4()}"
|
|
first: Final = string_value(isolated.proxy.chat(model, key=isolated.key, text=text)["id"])
|
|
isolated.proxy.chat(model, key=isolated.key, text=text)
|
|
|
|
rows: Final = isolated.rows(2)
|
|
assert len([request for request in wire.drain() if request.method == "POST"]) == 1
|
|
paid, hit = sorted(rows, key=lambda row: row["cache_hit"] == "True")
|
|
assert paid["request_id"] == first and float(str(paid["spend"])) > 0, rows
|
|
assert hit["cache_hit"] == "True" and string_value(hit["request_id"]).startswith(first + "_cache_hit"), rows
|
|
assert float(str(hit["spend"])) == 0, rows
|
|
for row in rows:
|
|
_assert_filtered(row, "usage_object")
|
|
|
|
|
|
def test_batch_cost_row_is_filtered_and_charged_once(gateway: Gateway, tmp_path: Path) -> None:
|
|
with _isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated:
|
|
handle: Final = register_scenario(f"metadata-fields-batch-{uuid4().hex[:12]}", batch_routes("gpt-4o-mini"))
|
|
isolated.scenario.cleanups.callback(delete_scenario, handle)
|
|
model: Final = isolated.scenario.model(api_base=handle.api_base())
|
|
uploaded: Final = isolated.proxy.request_multipart(
|
|
"/v1/files",
|
|
{"purpose": "batch", "model": model},
|
|
{"file": ("in.jsonl", batch_input_file(model), "application/jsonl")},
|
|
key=isolated.key,
|
|
)
|
|
assert uploaded.status_code == 200, uploaded.text
|
|
batch: Final = isolated.proxy.post(
|
|
"/v1/batches",
|
|
{
|
|
"input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]),
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"model": model,
|
|
},
|
|
key=isolated.key,
|
|
)
|
|
batch_path: Final = f"/v1/batches/{string_value(batch['id'])}"
|
|
retrievals: Final = tuple(isolated.proxy.request("GET", batch_path, key=isolated.key) for _ in range(2))
|
|
assert all(r.status_code == 200 and r.json()["status"] == "completed" for r in retrievals), retrievals
|
|
|
|
row: Final = isolated.rows(1, "call_type='aretrieve_batch'")[0]
|
|
assert float(str(row["spend"])) > 0, row
|
|
metadata: Final = _assert_filtered(row, "usage_object")
|
|
assert (metadata["batch_successful_requests"], metadata["batch_failed_requests"]) == (2, 3), metadata
|
|
|
|
|
|
def _update_retention(proxy: Gateway, value: JsonValue) -> httpx.Response:
|
|
return proxy.request(
|
|
"POST",
|
|
"/config/field/update",
|
|
{"field_name": "spend_logs_metadata_fields", "field_value": value, "config_type": "general_settings"},
|
|
)
|
|
|
|
|
|
def test_runtime_retention_update_rejects_invalid_values_and_applies_valid_ones_without_restart(
|
|
gateway: Gateway, tmp_path: Path
|
|
) -> None:
|
|
with (
|
|
scratch_database() as database_url,
|
|
wire_server(_respond) as wire,
|
|
_isolated(gateway, _config(tmp_path, None), tmp_path, database_url) as isolated,
|
|
):
|
|
model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic")
|
|
invalid: Final[tuple[JsonValue, ...]] = (
|
|
{"include": ["usage_object"], "exclude": ["model_map_information"]},
|
|
{},
|
|
{"exclude": ["not_a_metadata_field"]},
|
|
{"exclude": ["status"]},
|
|
)
|
|
|
|
isolated.proxy.chat(model, key=isolated.key)
|
|
assert set(EXCLUDED) <= set(object_value(isolated.rows(1)[0]["metadata"]))
|
|
for value in invalid:
|
|
assert _update_retention(isolated.proxy, value).status_code == 400, value
|
|
isolated.proxy.chat(model, key=isolated.key)
|
|
assert set(EXCLUDED) <= set(object_value(isolated.rows(2)[-1]["metadata"]))
|
|
|
|
assert _update_retention(isolated.proxy, {"exclude": list(EXCLUDED)}).status_code == 200
|
|
isolated.proxy.chat(model, key=isolated.key)
|
|
_assert_filtered(isolated.rows(3)[-1])
|
|
|
|
for value in invalid:
|
|
assert _update_retention(isolated.proxy, value).status_code == 400, value
|
|
isolated.proxy.chat(model, key=isolated.key)
|
|
_assert_filtered(isolated.rows(4)[-1])
|
|
|
|
|
|
SAVINGS_FIELDS: Final = ("autorouter_savings_estimate", "autorouter_savings")
|
|
|
|
|
|
def _anthropic_message(model: JsonValue) -> dict[str, JsonValue]:
|
|
return {
|
|
"id": f"msg_{uuid4().hex}",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"model": model,
|
|
"content": [{"type": "text", "text": "ok"}],
|
|
"stop_reason": "end_turn",
|
|
"stop_sequence": None,
|
|
"usage": {
|
|
"input_tokens": 12,
|
|
"output_tokens": 2,
|
|
"cache_creation_input_tokens": 0,
|
|
"cache_read_input_tokens": 0,
|
|
"cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0},
|
|
},
|
|
}
|
|
|
|
|
|
def _anthropic(request: Request) -> Reply:
|
|
if request.target.startswith("/v1/messages/count_tokens"):
|
|
return Reply(body=b'{"input_tokens":12}')
|
|
body: Final = JSON_OBJECT.validate_json(request.body)
|
|
if body.get("stream") is True:
|
|
return Reply(content_type="text/event-stream", chunks=_anthropic_stream(body["model"]))
|
|
return Reply(body=json.dumps(_anthropic_message(body["model"])).encode())
|
|
|
|
|
|
def _router_models(wire: Wire, router: str) -> tuple[Mapping[str, JsonValue], ...]:
|
|
tiers: Final = {"cheap": "anthropic/claude-sonnet-4-6", "frontier": "anthropic/claude-opus-4-8"}
|
|
return (
|
|
*(
|
|
{"model_name": name, "litellm_params": {"model": model, "api_base": wire.url, "api_key": "synthetic"}}
|
|
for name, model in tiers.items()
|
|
),
|
|
{
|
|
"model_name": router,
|
|
"litellm_params": {
|
|
"model": "auto_router/complexity_router",
|
|
"complexity_router_config": {
|
|
"tiers": {"SIMPLE": "cheap", "MEDIUM": "cheap", "COMPLEX": "frontier", "REASONING": "frontier"},
|
|
},
|
|
},
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("retention", "stored"),
|
|
[(None, SAVINGS_FIELDS), ({"exclude": list(SAVINGS_FIELDS)}, ())],
|
|
ids=["unset-keeps-savings", "excluded-savings-stay-out"],
|
|
)
|
|
@pytest.mark.timeout(240)
|
|
def test_delayed_autorouter_savings_publication_respects_retention(
|
|
gateway: Gateway, tmp_path: Path, retention: Mapping[str, JsonValue] | None, stored: tuple[str, ...]
|
|
) -> None:
|
|
router: Final = f"router-{uuid4().hex[:8]}"
|
|
with wire_server(_anthropic) as wire:
|
|
config: Final = _config(tmp_path, retention, model_list=_router_models(wire, router))
|
|
with _isolated(gateway, config, tmp_path) as isolated:
|
|
response: Final = isolated.proxy.request(
|
|
"POST",
|
|
"/v1/messages",
|
|
{"model": router, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]},
|
|
key=isolated.key,
|
|
headers={"x-litellm-session-id": f"session-{uuid4()}"},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
|
|
published: Final = isolated.rows(
|
|
1,
|
|
"metadata::jsonb ? 'routing_decision' AND NOT metadata::jsonb ? 'autorouter_baseline_observation' "
|
|
'AND request_id IN (SELECT request_id FROM "LiteLLM_AutoRouterBaselineObservation" '
|
|
"WHERE publication IS NOT NULL)",
|
|
seconds=150,
|
|
)[0]
|
|
metadata: Final = object_value(published["metadata"])
|
|
assert {name for name in SAVINGS_FIELDS if name in metadata} == set(stored), metadata
|