Merge pull request #42052 from BerriAI/litellm_cost_shard_provider_wires

test(integration): provider wire cost cases
This commit is contained in:
kerry-berri 2026-09-21 14:24:35 -07:00 • committed by GitHub
commit 34d5f9d41b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 1285 additions and 79 deletions

View file

@ -1,6 +1,7 @@
from __future__ import annotations
import argparse
import base64
from collections import deque
from collections.abc import Mapping
import json
@ -23,6 +24,7 @@ from starlette.routing import Route
from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations
from integration.cost_calculation.cost_tracking_case import (
BinaryResponse,
EventStreamEvent,
EventStreamResponse,
JsonResponse,
SseResponse,
@ -226,8 +228,25 @@ class Provider:
)
return Response(content=stream_body.encode(), media_type=response.content_type)
case EventStreamResponse():
events: Final = (
tuple(
EventStreamEvent(
event_type="chunk",
payload={
"bytes": base64.b64encode(
json.dumps(event.payload, separators=(",", ":"))
.replace("$REQUEST_ID", scenario_id)
.encode()
).decode(),
},
)
for event in response.events
)
if response.framing == "invoke"
else response.events
)
event_body: Final = b"".join(
_aws_event_frame(event.event_type, event.payload, scenario_id) for event in response.events
_aws_event_frame(event.event_type, event.payload, scenario_id) for event in events
)
return Response(content=event_body, media_type=response.content_type)

View file

@ -1340,52 +1340,52 @@
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[whisper-next-transcriptions-per-second]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[whisper-verbose-next-transcriptions-duration]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-4o-transcribe-next-transcriptions-tokens]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[nova-next-transcriptions-per-second]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-whisper-next-transcriptions-deployment]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[tts-next-speech-per-character]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[tts-next-hd-speech-per-character]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-tts-next-speech-deployment]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-standard]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-hd]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-wide]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-two]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-image-next-images-low]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[imagen-next-images-one]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[amazon-nova-canvas-next-images-one]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-image-next-images-edit]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-text-embeddings-4-large-deployment]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
@ -1579,6 +1579,51 @@
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[xai-grok-5-provider_reported_cost]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-invoke-haiku-json]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-invoke-haiku-stream]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-profile-base-model]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-eu-regional-key]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-apac-bare-fallback]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-nova-2-pro]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[bedrock-converse-mistral-large-3-stream]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-ai-gpt-5.4-mini-latest]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-ai-gpt-5.4-mini-latest-stream]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-pinned-gpt-5.4-mini-stream]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[groq-qwen-3.8-json]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[groq-qwen-3.8-stream_x_groq_recount]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[cohere-command-a-v2-tokens]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[mistral-medium-2604-json]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[openai-deployment-pricing-override]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/mcp/test_mcp_lifecycle.py::test_health_intersects_route_restricted_key_grants_in_both_management_modes": [
"other.mcp.health.restricted_keys_intersect_grants_in_both_modes"
],

View file

@ -163,6 +163,18 @@ def register_scenario_deployment(
"api_key": case.api_key,
"api_base": handle.api_base(),
**case.litellm_params,
**(
{
key: value
for key, value in (
("input_cost_per_token", case.deployment.input_cost_per_token),
("output_cost_per_token", case.deployment.output_cost_per_token),
)
if value is not None
}
if case.deployment is not None
else {}
),
**(
{"vertex_credentials": _vertex_service_account_json(control_url)}
if case.rates.litellm_provider.startswith("vertex_ai")

View file

@ -86,6 +86,8 @@ class Deployment(BaseModel):
model: str | None = None
base_model: str | None = None
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
class WavUpload(BaseModel):
@ -131,6 +133,7 @@ class EventStreamResponse(BaseModel):
content_type: Literal["application/vnd.amazon.eventstream"]
events: tuple[EventStreamEvent, ...]
framing: Literal["converse", "invoke"] = "converse"
class BinaryResponse(BaseModel):
@ -239,9 +242,11 @@ class CostTrackingTestCase(BaseModel):
)
if prefix is None:
raise ValueError(f"unsupported cost-map provider {provider} for {self.model}")
return self.deployment.model if self.deployment and self.deployment.model is not None else (
self.model if prefix == "" else f"{prefix}/{self.model}"
)
if self.deployment and self.deployment.model is not None:
return self.deployment.model
if prefix == "" or self.model.startswith(f"{prefix}/"):
return self.model
return f"{prefix}/{self.model}"
@property
def litellm_params(self) -> Mapping[str, str]:
@ -301,6 +306,10 @@ _PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType(
"perplexity": "",
"deepseek": "",
"xai": "",
"azure_ai": "azure_ai",
"groq": "groq",
"mistral": "mistral",
"cohere_chat": "cohere_chat",
}
)
_LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
@ -342,6 +351,10 @@ _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
"perplexity": MappingProxyType({}),
"deepseek": MappingProxyType({}),
"xai": MappingProxyType({}),
"azure_ai": MappingProxyType({}),
"groq": MappingProxyType({}),
"mistral": MappingProxyType({}),
"cohere_chat": MappingProxyType({}),
}
)

File diff suppressed because it is too large Load diff