mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
test(e2e): every spend row a virtual key writes joins its token across all write paths
One aliased key owned by a user with an email drives chat, queued chat, messages, responses, embeddings, the Gemini passthrough, a batch file upload, and a batch create against a live proxy. Each row must carry api_key equal to the key's LiteLLM_VerificationToken.token and the alias in metadata, and /spend/logs?api_key= and /user/daily/activity must report the key with its alias and email. Health-check rows must keep the literal service-account key, and the batch cost row for a completed marker batch must join the key that created it. A re-hashed api_key (the v1.99.0 regression fixed by #39568 and #39572) now fails the Buildkite e2e stage naming the write path Resolves MAT-180
This commit is contained in:
mateo-berri
parent
096984bfc2
commit
3020a13e24
6 changed files with 642 additions and 3 deletions
|
|
@ -181,13 +181,14 @@ quota_management.<behavior>.<variant>.<assertion>
|
|||
| team_multi_window | fallback | spend_counter
|
||||
<spend_tracking> chat_completions | stream | messages_bridge | embeddings
|
||||
| cache_hit | key_rollup | concurrent_burst | tags | end_user
|
||||
| per_model | failure | spend_calculate | pagination
|
||||
| per_model | failure | spend_calculate | pagination | key_attribution
|
||||
assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm
|
||||
| blocks_then_resets | resets_windows_independently | alerts_without_blocking
|
||||
| isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys
|
||||
| routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost
|
||||
| matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows
|
||||
| writes_failure_row | returns_cost | keeps_total
|
||||
| writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email
|
||||
| health_rows_keep_service_account | batch_cost_joins_key
|
||||
e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages]
|
||||
quota_management.budget.key.blocks_over_limit exercised_on=[chat_completions]
|
||||
```
|
||||
|
|
|
|||
|
|
@ -58,3 +58,7 @@
|
|||
- {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"}
|
||||
- {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"}
|
||||
- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"}
|
||||
- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, and batch create carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"}
|
||||
- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"}
|
||||
- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"}
|
||||
- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The _batch_cost row the batch cost poller writes once a batch completes carries the submitting key's token hash; it rides a cross-run marker batch whose metadata records the expected hash, since completion can lag by up to the polling interval. Cold start with no completed marker is a documented vacuous pass"}
|
||||
|
|
|
|||
|
|
@ -84,6 +84,7 @@ class KeyGenerateBody(BaseModel):
|
|||
|
||||
class KeyGenerateResponse(BaseModel):
|
||||
key: str
|
||||
token: str | None = None
|
||||
key_alias: str | None = None
|
||||
models: list[str] = []
|
||||
max_budget: float | None = None
|
||||
|
|
@ -672,6 +673,7 @@ class GuardrailRunRecord(BaseModel):
|
|||
|
||||
|
||||
class SpendLogMetadata(BaseModel):
|
||||
user_api_key_alias: str | None = None
|
||||
applied_guardrails: list[str] | None = None
|
||||
guardrail_information: list[GuardrailRunRecord] | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = (
|
|||
("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"),
|
||||
("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"),
|
||||
("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"),
|
||||
("openai-gpt-4o-mini", "openai/gpt-4o-mini", "OPENAI_API_KEY"),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -15,9 +15,14 @@ import time
|
|||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import (
|
||||
FileUploadForm,
|
||||
Headers,
|
||||
NoBody,
|
||||
ProbeResult,
|
||||
Result,
|
||||
|
|
@ -35,6 +40,8 @@ from models import (
|
|||
DateRangeParams,
|
||||
EmbedBody,
|
||||
EmbedResponse,
|
||||
KeyGenerateBody,
|
||||
KeyGenerateResponse,
|
||||
OpenAPISchema,
|
||||
SpendCalculateBody,
|
||||
SpendCalculateResponse,
|
||||
|
|
@ -43,13 +50,24 @@ from models import (
|
|||
SpendLogsPageParams,
|
||||
SpendTagsResponse,
|
||||
TagSpend,
|
||||
UserDeleteBody,
|
||||
UserDeleteResponse,
|
||||
UserNewBody,
|
||||
UserNewResponse,
|
||||
UserRole,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
from proxy_client import Converged, ProxyClient, await_converged
|
||||
|
||||
__all__ = [
|
||||
"BatchCreateBody",
|
||||
"BatchObject",
|
||||
"DailyActivityKeyBreakdown",
|
||||
"FileObject",
|
||||
"ProbeResult",
|
||||
"ResponseIdentity",
|
||||
"SpendClient",
|
||||
"SpendLogRow",
|
||||
"StreamingResponse",
|
||||
"build_client",
|
||||
"is_ok",
|
||||
"unique_marker",
|
||||
|
|
@ -57,6 +75,112 @@ __all__ = [
|
|||
]
|
||||
|
||||
|
||||
class GeminiApiKeyHeaders(Headers):
|
||||
x_goog_api_key: str = Field(serialization_alias="x-goog-api-key")
|
||||
content_type: str = Field(default="application/json", serialization_alias="Content-Type")
|
||||
|
||||
|
||||
class GeminiPart(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class GeminiContent(BaseModel):
|
||||
parts: list[GeminiPart]
|
||||
|
||||
|
||||
class GeminiGenerationConfig(BaseModel):
|
||||
maxOutputTokens: int
|
||||
|
||||
|
||||
class GeminiGenerateBody(BaseModel):
|
||||
contents: list[GeminiContent]
|
||||
generationConfig: GeminiGenerationConfig
|
||||
|
||||
|
||||
class ResponsesBody(BaseModel):
|
||||
model: str
|
||||
input: str
|
||||
cache: dict[str, bool] | None = {"no-cache": True}
|
||||
|
||||
|
||||
class QueuedChatBody(ChatBody):
|
||||
priority: int = 0
|
||||
|
||||
|
||||
class ResponseIdentity(BaseModel):
|
||||
id: str | None = None
|
||||
|
||||
|
||||
class HealthParams(BaseModel):
|
||||
model: str
|
||||
|
||||
|
||||
class ModelQuery(BaseModel):
|
||||
model: str
|
||||
|
||||
|
||||
class FileObject(BaseModel):
|
||||
id: str
|
||||
|
||||
|
||||
class BatchCreateBody(BaseModel):
|
||||
input_file_id: str
|
||||
endpoint: str = "/v1/chat/completions"
|
||||
completion_window: str = "24h"
|
||||
model: str
|
||||
metadata: dict[str, str]
|
||||
|
||||
|
||||
class BatchObject(BaseModel):
|
||||
id: str
|
||||
status: str
|
||||
metadata: dict[str, str] | None = None
|
||||
created_at: int | None = None
|
||||
|
||||
|
||||
class BatchList(BaseModel):
|
||||
data: list[BatchObject] = []
|
||||
|
||||
|
||||
class BatchListQuery(BaseModel):
|
||||
model: str
|
||||
limit: int
|
||||
|
||||
|
||||
class DailyActivityParams(BaseModel):
|
||||
start_date: str
|
||||
end_date: str
|
||||
api_key: str
|
||||
|
||||
|
||||
class DailyActivityKeyMetadata(BaseModel):
|
||||
key_alias: str | None = None
|
||||
team_id: str | None = None
|
||||
user_email: str | None = None
|
||||
|
||||
|
||||
class DailyActivityKeyMetrics(BaseModel):
|
||||
api_requests: int = 0
|
||||
|
||||
|
||||
class DailyActivityKeyBreakdown(BaseModel):
|
||||
metrics: DailyActivityKeyMetrics
|
||||
metadata: DailyActivityKeyMetadata
|
||||
|
||||
|
||||
class DailyActivityBreakdown(BaseModel):
|
||||
api_keys: dict[str, DailyActivityKeyBreakdown] = {}
|
||||
|
||||
|
||||
class DailyActivityRow(BaseModel):
|
||||
date: str
|
||||
breakdown: DailyActivityBreakdown
|
||||
|
||||
|
||||
class DailyActivityResponse(BaseModel):
|
||||
results: list[DailyActivityRow] = []
|
||||
|
||||
|
||||
def _chat_body(
|
||||
model: str,
|
||||
content: str,
|
||||
|
|
@ -207,6 +331,164 @@ class SpendClient:
|
|||
def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult:
|
||||
return self.proxy.transport.probe(path, params=params)
|
||||
|
||||
def create_user(self, *, email: str, role: UserRole, user_id: str) -> str:
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/user/new",
|
||||
headers=self.proxy.transport.master,
|
||||
json=UserNewBody(user_email=email, user_role=role, user_id=user_id),
|
||||
response_type=UserNewResponse,
|
||||
)
|
||||
).user_id
|
||||
|
||||
def delete_user(self, user_id: str) -> None:
|
||||
_ = self.proxy.transport.post(
|
||||
"/user/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
json=UserDeleteBody(user_ids=[user_id]),
|
||||
response_type=UserDeleteResponse,
|
||||
)
|
||||
|
||||
def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse:
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/key/generate",
|
||||
headers=self.proxy.transport.master,
|
||||
json=body,
|
||||
response_type=KeyGenerateResponse,
|
||||
)
|
||||
)
|
||||
|
||||
def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=_chat_body(model, content, max_tokens=max_tokens),
|
||||
)
|
||||
|
||||
def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/queue/chat/completions",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=QueuedChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
max_tokens=max_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/v1/messages",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
max_tokens=max_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
def send_responses(self, key: str, model: str, content: str) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/v1/responses",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ResponsesBody(model=model, input=content),
|
||||
)
|
||||
|
||||
def send_embed(self, key: str, model: str, content: str) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/embeddings",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=EmbedBody(model=model, input=content),
|
||||
)
|
||||
|
||||
def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
f"/gemini/v1beta/models/{model}:generateContent",
|
||||
headers=GeminiApiKeyHeaders(x_goog_api_key=key),
|
||||
json=GeminiGenerateBody(
|
||||
contents=[GeminiContent(parts=[GeminiPart(text=content)])],
|
||||
generationConfig=GeminiGenerationConfig(maxOutputTokens=max_tokens),
|
||||
),
|
||||
)
|
||||
|
||||
def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject:
|
||||
return unwrap(
|
||||
self.proxy.transport.upload(
|
||||
"/v1/files",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
form=FileUploadForm(purpose="batch"),
|
||||
filename="key_attribution.jsonl",
|
||||
content=content,
|
||||
params=ModelQuery(model=model),
|
||||
response_type=FileObject,
|
||||
)
|
||||
)
|
||||
|
||||
def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject:
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/v1/batches",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=body,
|
||||
response_type=BatchObject,
|
||||
)
|
||||
)
|
||||
|
||||
def list_batches(self, key: str, model: str, *, limit: int) -> list[BatchObject]:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/v1/batches",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=BatchListQuery(model=model, limit=limit),
|
||||
response_type=BatchList,
|
||||
)
|
||||
).data
|
||||
|
||||
def retrieve_batch(self, key: str, batch_id: str) -> BatchObject:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
f"/v1/batches/{batch_id}",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=NoBody(),
|
||||
response_type=BatchObject,
|
||||
)
|
||||
)
|
||||
|
||||
def health(self, model: str) -> ProbeResult:
|
||||
return self.proxy.transport.probe("/health", params=HealthParams(model=model))
|
||||
|
||||
def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None:
|
||||
response: Final = unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/user/daily/activity",
|
||||
headers=self.proxy.transport.master,
|
||||
params=DailyActivityParams(
|
||||
start_date=start.strftime("%Y-%m-%d"),
|
||||
end_date=end.strftime("%Y-%m-%d"),
|
||||
api_key=token,
|
||||
),
|
||||
response_type=DailyActivityResponse,
|
||||
)
|
||||
)
|
||||
return next(
|
||||
(row.breakdown.api_keys[token] for row in response.results if token in row.breakdown.api_keys),
|
||||
None,
|
||||
)
|
||||
|
||||
def poll_daily_activity_for_key(
|
||||
self, token: str, *, start: datetime, end: datetime, min_requests: int
|
||||
) -> DailyActivityKeyBreakdown | None:
|
||||
outcome: Final = await_converged(
|
||||
lambda: self.daily_activity_for_key(token, start=start, end=end),
|
||||
converged=lambda found: found is not None and found.metrics.api_requests >= min_requests,
|
||||
timeout=self.proxy.poll_timeout,
|
||||
interval=self.proxy.poll_interval,
|
||||
now=time.monotonic,
|
||||
sleep=time.sleep,
|
||||
)
|
||||
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
|
||||
|
||||
def openapi(self) -> OpenAPISchema:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,349 @@
|
|||
"""Every spend row a live proxy writes joins its virtual key (MAT-180).
|
||||
|
||||
One virtual key with an alias, owned by a user with an email, drives every spend
|
||||
write path a key can reach: /chat/completions, /queue/chat/completions,
|
||||
/v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch
|
||||
input file upload, and a batch create. Each row those calls write must carry
|
||||
`api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash
|
||||
/key/generate returns as `token`), which is the join /spend/logs?api_key= and
|
||||
/user/daily/activity rely on to report key_alias and user_email. A row keyed by a
|
||||
re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a
|
||||
key-hash-* row with no alias and no email in the customer's usage exports.
|
||||
|
||||
The health-check service account writes rows too; those must stay keyed by the
|
||||
literal service-account name, never by a hash of it. The batch cost row is
|
||||
written by the CheckBatchCost poller once the batch completes, up to an hour
|
||||
later, so it rides a cross-run baton like the batches suite: each run submits a
|
||||
one-line marker batch whose metadata records the token it expects on the cost
|
||||
row, and asserts on the newest completed marker from any run (a cold start with
|
||||
no completed marker is a documented vacuous pass, never a skip).
|
||||
|
||||
/spend/logs carries no email field, so the email assertion lives on
|
||||
/user/daily/activity alone; /spend/logs is held to the alias in metadata.
|
||||
"""
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final, Iterator
|
||||
|
||||
import pytest
|
||||
|
||||
from models import ChatMessage, KeyGenerateBody
|
||||
from proxy_client import Converged, await_converged
|
||||
from spend_e2e_client import (
|
||||
BatchCreateBody,
|
||||
BatchObject,
|
||||
DailyActivityKeyBreakdown,
|
||||
ResponseIdentity,
|
||||
SpendClient,
|
||||
SpendLogRow,
|
||||
StreamingResponse,
|
||||
unique_marker,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
CHAT_MODEL: Final = "gemini-2.5-flash"
|
||||
MESSAGES_MODEL: Final = "claude-haiku-4-5"
|
||||
RESPONSES_MODEL: Final = "openai-responses-codex"
|
||||
EMBED_MODEL: Final = "openai-text-embedding-3-small"
|
||||
BATCH_MODEL: Final = "openai-gpt-4o-mini"
|
||||
BATCH_BACKEND_MODEL: Final = "gpt-4o-mini"
|
||||
HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check"
|
||||
BATON_MARKER_KEY: Final = "litellm_e2e_suite"
|
||||
BATON_MARKER_VALUE: Final = "key-attribution-baton"
|
||||
BATON_EXPECTED_TOKEN_KEY: Final = "expected_api_key"
|
||||
BATON_POLL_SECONDS: Final = 300.0
|
||||
BATON_POLL_INTERVAL_SECONDS: Final = 10.0
|
||||
BATON_LIST_LIMIT: Final = 100
|
||||
MAX_TOKENS: Final = 8
|
||||
WRITE_PATHS: Final = (
|
||||
"chat_completions",
|
||||
"queue_chat_completions",
|
||||
"messages",
|
||||
"responses",
|
||||
"embeddings",
|
||||
"gemini_passthrough",
|
||||
"batch_file_upload",
|
||||
"batch_create",
|
||||
)
|
||||
|
||||
|
||||
class BatchLineBody(BaseModel):
|
||||
model: str
|
||||
messages: list[ChatMessage]
|
||||
max_tokens: int
|
||||
|
||||
|
||||
class BatchLine(BaseModel):
|
||||
custom_id: str
|
||||
method: str = "POST"
|
||||
url: str = "/v1/chat/completions"
|
||||
body: BatchLineBody
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AttributedKey:
|
||||
key: str
|
||||
token: str
|
||||
alias: str
|
||||
email: str
|
||||
user_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WritePath:
|
||||
name: str
|
||||
request_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DrivenKey:
|
||||
identity: AttributedKey
|
||||
paths: tuple[WritePath, ...]
|
||||
started_at: datetime
|
||||
|
||||
|
||||
def _body_id(name: str, sent: StreamingResponse) -> WritePath:
|
||||
assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}"
|
||||
response_id: Final = ResponseIdentity.model_validate_json(sent.body).id
|
||||
assert response_id, f"{name} answered without a response id: {sent.body[:300]}"
|
||||
return WritePath(name=name, request_id=response_id)
|
||||
|
||||
|
||||
def _call_id(name: str, sent: StreamingResponse) -> WritePath:
|
||||
assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}"
|
||||
assert sent.call_id, f"{name} answered without an x-litellm-call-id header"
|
||||
return WritePath(name=name, request_id=sent.call_id)
|
||||
|
||||
|
||||
def _batch_jsonl(marker: str) -> bytes:
|
||||
line: Final = BatchLine(
|
||||
custom_id=marker,
|
||||
body=BatchLineBody(
|
||||
model=BATCH_BACKEND_MODEL,
|
||||
messages=[ChatMessage(role="user", content=f"Reply with the word ok. {marker}")],
|
||||
max_tokens=MAX_TOKENS,
|
||||
),
|
||||
)
|
||||
return f"{line.model_dump_json()}\n".encode()
|
||||
|
||||
|
||||
def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]:
|
||||
uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _batch_jsonl(marker))
|
||||
created: Final = client.create_batch(
|
||||
identity.key,
|
||||
BatchCreateBody(
|
||||
input_file_id=uploaded.id,
|
||||
model=BATCH_MODEL,
|
||||
metadata={
|
||||
BATON_MARKER_KEY: BATON_MARKER_VALUE,
|
||||
BATON_EXPECTED_TOKEN_KEY: identity.token,
|
||||
"run": marker,
|
||||
},
|
||||
),
|
||||
)
|
||||
return (
|
||||
WritePath(name="batch_file_upload", request_id=uploaded.id),
|
||||
WritePath(name="batch_create", request_id=created.id),
|
||||
)
|
||||
|
||||
|
||||
def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]:
|
||||
marker: Final = unique_marker()
|
||||
prompt: Final = f"Reply with the word ok. {marker}"
|
||||
key: Final = identity.key
|
||||
return (
|
||||
_body_id("chat_completions", client.send_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
|
||||
_body_id("queue_chat_completions", client.send_queued_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
|
||||
_body_id("messages", client.send_messages(key, MESSAGES_MODEL, prompt, max_tokens=MAX_TOKENS)),
|
||||
_body_id("responses", client.send_responses(key, RESPONSES_MODEL, prompt)),
|
||||
_call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)),
|
||||
_call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)),
|
||||
*_drive_batch(client, identity, marker),
|
||||
)
|
||||
|
||||
|
||||
def _completed_baton(listed: list[BatchObject]) -> BatchObject | None:
|
||||
return max(
|
||||
(
|
||||
batch
|
||||
for batch in listed
|
||||
if batch.status == "completed" and (batch.metadata or {}).get(BATON_MARKER_KEY) == BATON_MARKER_VALUE
|
||||
),
|
||||
key=lambda batch: batch.created_at or 0,
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def _newest_completed_baton(client: SpendClient, key: str) -> BatchObject | None:
|
||||
outcome: Final = await_converged(
|
||||
lambda: _completed_baton(client.list_batches(key, BATCH_MODEL, limit=BATON_LIST_LIMIT)),
|
||||
converged=lambda completed: completed is not None,
|
||||
timeout=BATON_POLL_SECONDS,
|
||||
interval=BATON_POLL_INTERVAL_SECONDS,
|
||||
now=time.monotonic,
|
||||
sleep=time.sleep,
|
||||
)
|
||||
return outcome.result if isinstance(outcome, Converged) else None
|
||||
|
||||
|
||||
def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]:
|
||||
return [
|
||||
row
|
||||
for row in client.proxy.spend_logs_window(
|
||||
start=started_at - timedelta(minutes=1), end=datetime.now(timezone.utc) + timedelta(minutes=1)
|
||||
)
|
||||
if HEALTH_SERVICE_ACCOUNT in (row.request_tags or [])
|
||||
]
|
||||
|
||||
|
||||
def _health_rows_since(client: SpendClient, started_at: datetime) -> list[SpendLogRow]:
|
||||
outcome: Final = await_converged(
|
||||
lambda: _health_rows_between(client, started_at),
|
||||
converged=lambda rows: bool(rows),
|
||||
timeout=client.proxy.poll_timeout,
|
||||
interval=client.proxy.poll_interval,
|
||||
now=time.monotonic,
|
||||
sleep=time.sleep,
|
||||
)
|
||||
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
|
||||
|
||||
|
||||
class TestKeyAttribution:
|
||||
@pytest.fixture(scope="class")
|
||||
def driven(self, client: SpendClient) -> Iterator[DrivenKey]:
|
||||
marker: Final = unique_marker()
|
||||
user_id: Final = client.create_user(
|
||||
email=f"key-attribution-{marker}@example.com",
|
||||
role="proxy_admin",
|
||||
user_id=f"key-attribution-{marker}",
|
||||
)
|
||||
record: Final = client.generate_key_record(
|
||||
KeyGenerateBody(models=[], user_id=user_id, key_alias=f"key-attribution-{marker}")
|
||||
)
|
||||
assert record.token, "/key/generate answered without the key's token hash"
|
||||
assert record.key_alias, "/key/generate dropped the key alias"
|
||||
identity: Final = AttributedKey(
|
||||
key=record.key,
|
||||
token=record.token,
|
||||
alias=record.key_alias,
|
||||
email=f"key-attribution-{marker}@example.com",
|
||||
user_id=user_id,
|
||||
)
|
||||
started_at: Final = datetime.now(timezone.utc)
|
||||
try:
|
||||
yield DrivenKey(
|
||||
identity=identity,
|
||||
paths=_drive_every_write_path(client, identity),
|
||||
started_at=started_at,
|
||||
)
|
||||
finally:
|
||||
client.proxy.delete_key(identity.key)
|
||||
client.delete_user(identity.user_id)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.key_attribution.joins_key",
|
||||
exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"],
|
||||
)
|
||||
def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None:
|
||||
assert tuple(path.name for path in driven.paths) == WRITE_PATHS
|
||||
found: Final = tuple(
|
||||
(path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths
|
||||
)
|
||||
unwritten: Final = [path.name for path, rows in found if not rows]
|
||||
assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}"
|
||||
unjoined: Final = [
|
||||
(path.name, row.call_type, row.api_key)
|
||||
for path, rows in found
|
||||
for row in rows
|
||||
if row.api_key != driven.identity.token
|
||||
]
|
||||
assert not unjoined, (
|
||||
"spend rows whose api_key does not join LiteLLM_VerificationToken.token "
|
||||
f"{driven.identity.token}: {unjoined}"
|
||||
)
|
||||
unaliased: Final = [
|
||||
(path.name, row.call_type, row.metadata.user_api_key_alias if row.metadata else None)
|
||||
for path, rows in found
|
||||
for row in rows
|
||||
if row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias
|
||||
]
|
||||
assert not unaliased, f"spend rows written without key alias {driven.identity.alias!r}: {unaliased}"
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.key_attribution.reports_alias_and_email",
|
||||
exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"],
|
||||
)
|
||||
def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None:
|
||||
expected_ids: Final = frozenset(path.request_id for path in driven.paths)
|
||||
rows: Final = client.poll_logs_for_key(
|
||||
driven.identity.key,
|
||||
min_rows=len(driven.paths),
|
||||
predicate=lambda found: expected_ids <= frozenset(row.request_id or "" for row in found),
|
||||
)
|
||||
missing: Final = expected_ids - frozenset(row.request_id or "" for row in rows)
|
||||
assert not missing, (
|
||||
f"/spend/logs?api_key= does not return {len(missing)} of {len(expected_ids)} rows for the key: "
|
||||
f"{sorted(path.name for path in driven.paths if path.request_id in missing)}"
|
||||
)
|
||||
aliases: Final = frozenset(row.metadata.user_api_key_alias if row.metadata else None for row in rows)
|
||||
assert aliases == {driven.identity.alias}, f"/spend/logs rows carry aliases {sorted(map(str, aliases))}"
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.key_attribution.reports_alias_and_email",
|
||||
exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"],
|
||||
)
|
||||
def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None:
|
||||
breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key(
|
||||
driven.identity.token,
|
||||
start=driven.started_at - timedelta(days=1),
|
||||
end=datetime.now(timezone.utc) + timedelta(days=1),
|
||||
min_requests=len(driven.paths),
|
||||
)
|
||||
assert breakdown is not None, (
|
||||
f"/user/daily/activity?api_key={driven.identity.token} has no api_keys breakdown: "
|
||||
"the key's rows did not aggregate under its token"
|
||||
)
|
||||
assert breakdown.metrics.api_requests >= len(driven.paths), (
|
||||
f"/user/daily/activity counts {breakdown.metrics.api_requests} requests for the key, "
|
||||
f"expected at least {len(driven.paths)}"
|
||||
)
|
||||
assert breakdown.metadata.key_alias == driven.identity.alias, f"key_alias={breakdown.metadata.key_alias!r}"
|
||||
assert breakdown.metadata.user_email == driven.identity.email, f"user_email={breakdown.metadata.user_email!r}"
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.key_attribution.health_rows_keep_service_account",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_health_check_rows_keep_the_service_account_key(self, client: SpendClient) -> None:
|
||||
started_at: Final = datetime.now(timezone.utc)
|
||||
probe: Final = client.health(CHAT_MODEL)
|
||||
assert probe.healthy, f"/health?model={CHAT_MODEL} answered {probe.status_code}: {probe.body[:300]}"
|
||||
rows: Final = _health_rows_since(client, started_at)
|
||||
assert rows, f"/health?model={CHAT_MODEL} wrote no {HEALTH_SERVICE_ACCOUNT}-tagged spend row"
|
||||
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.covers(
|
||||
"quota_management.spend_tracking.key_attribution.batch_cost_joins_key",
|
||||
exercised_on=["batches"],
|
||||
)
|
||||
def test_completed_batch_cost_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None:
|
||||
completed: Final = _newest_completed_baton(client, driven.identity.key)
|
||||
if completed is None:
|
||||
return
|
||||
expected_token: Final = (completed.metadata or {}).get(BATON_EXPECTED_TOKEN_KEY)
|
||||
assert expected_token, f"marker batch {completed.id} lost its expected token metadata: {completed.metadata}"
|
||||
fetched: Final = client.retrieve_batch(driven.identity.key, completed.id)
|
||||
assert fetched.status == "completed", f"listed-completed marker retrieved as {fetched.status!r}"
|
||||
rows: Final = client.proxy.poll_logs_for_request_id(
|
||||
f"{fetched.id}_batch_cost",
|
||||
predicate=lambda found: any((row.spend or 0) > 0 for row in found),
|
||||
)
|
||||
priced: Final = [row for row in rows if (row.spend or 0) > 0]
|
||||
assert priced, f"completed batch {fetched.id} has no positive-cost spend row under {fetched.id}_batch_cost"
|
||||
unjoined: Final = [(row.call_type, row.api_key) for row in priced if row.api_key != expected_token]
|
||||
assert not unjoined, f"batch cost rows whose api_key does not join the key's token {expected_token}: {unjoined}"
|
||||
Loading…
Add table
Reference in a new issue