litellm/tests/integration/cost_calculation/cost_tracking_case.py
Kerry Lu 5f74a79c75
test(cost): add xai grok-4.7, Bedrock runtime chat completions and Bedrock Mantle cost cases (#46154)
* test(cost): add xai grok-4.7, Bedrock runtime chat completions and Bedrock Mantle cost cases

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(cost): bring back known_bug and carry the xai grok-4.7 stream breakdown bug as a strict xfail case

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-13 09:06:19 -07:00

631 lines
22 KiB
Python

from __future__ import annotations
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import Annotated, Final, Literal, TypeAlias
from integration.cost_calculation.cost_map import COST_MAP_ENTRIES
from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator
PRIOR_RESPONSE_ID_MARKER: Final = "$PRIOR_RESPONSE_ID"
class SearchContextCostPerQuery(BaseModel):
model_config = ConfigDict(frozen=True)
search_context_size_low: float | None = None
search_context_size_medium: float | None = None
search_context_size_high: float | None = None
class ProviderSpecificEntry(BaseModel):
model_config = ConfigDict(frozen=True)
fast: float | None = None
us: float | None = None
class TieredPrice(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
range: tuple[float, float]
input_cost_per_token: float
output_cost_per_token: float
class CostMapEntry(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
litellm_provider: str
mode: str
max_tokens: int | None = None
max_input_tokens: int | None = None
max_output_tokens: int | None = None
supports_function_calling: bool | None = None
supports_none_reasoning_effort: bool | None = None
supports_reasoning: bool | None = None
supports_prompt_caching: bool | None = None
supported_endpoints: tuple[str, ...] | None = None
supports_bedrock_runtime_chat_completions_tools_with_reasoning: bool | None = None
supports_bedrock_runtime_chat_completions_response_format: bool | None = None
cache_creation_input_token_cost_above_100k_tokens: float | None = None
cache_creation_input_token_cost_above_1hr_above_100k_tokens: float | None = None
cache_read_input_token_cost_above_100k_tokens: float | None = None
input_cost_per_token_above_100k_tokens: float | None = None
input_cost_per_token_above_272k_tokens_flex: float | None = None
input_cost_per_token_above_272k_tokens_priority: float | None = None
output_cost_per_token_above_100k_tokens: float | None = None
output_cost_per_token_above_272k_tokens_flex: float | None = None
output_cost_per_token_above_272k_tokens_priority: float | None = None
input_cost_per_token: float | None = None
input_cost_per_token_decisions: float | None = None
input_cost_per_query: float | None = None
output_cost_per_token: float | None = None
output_cost_per_token_decisions: float | None = None
cache_read_input_token_cost_decisions: float | None = None
cache_creation_input_token_cost_decisions: float | None = None
input_cost_per_token_above_272k_tokens: float | None = None
input_cost_per_token_above_272k_tokens_decisions: float | None = None
output_cost_per_token_above_272k_tokens: float | None = None
output_cost_per_token_above_272k_tokens_decisions: float | None = None
cache_read_input_token_cost_above_272k_tokens: float | None = None
cache_read_input_token_cost_above_272k_tokens_decisions: float | None = None
cache_creation_input_token_cost_above_272k_tokens: float | None = None
cache_creation_input_token_cost_above_272k_tokens_decisions: float | None = None
input_cost_per_token_batches: float | None = None
output_cost_per_token_batches: float | None = None
output_cost_per_image_token_batches: float | None = None
input_cost_per_token_above_128k_tokens: float | None = None
output_cost_per_token_above_128k_tokens: float | None = None
output_vector_size: int | None = None
input_cost_per_token_batches: float | None = None
cache_read_input_token_cost: float | None = None
cache_creation_input_token_cost: float | None = None
cache_creation_input_token_cost_above_1hr: float | None = None
cache_creation_input_token_cost_above_1hr_above_200k_tokens: float | None = None
cache_read_input_token_cost_above_200k_tokens: float | None = None
cache_creation_input_token_cost_above_200k_tokens: float | None = None
input_cost_per_token_above_200k_tokens: float | None = None
output_cost_per_token_above_200k_tokens: float | None = None
cache_read_input_audio_token_cost: float | None = None
tiered_pricing: tuple[TieredPrice, ...] | None = None
output_cost_per_reasoning_token: float | None = None
input_cost_per_audio_token: float | None = None
input_cost_per_second: float | None = None
output_cost_per_second: float | None = None
input_cost_per_character: float | None = None
output_cost_per_character: float | None = None
input_cost_per_image: float | None = None
output_cost_per_image: float | None = None
output_cost_per_audio_token: float | None = None
input_cost_per_image_token: float | None = None
output_cost_per_image_token: float | None = None
input_cost_per_video_token: float | None = None
input_cost_per_token_flex: float | None = None
output_cost_per_token_flex: float | None = None
input_cost_per_token_priority: float | None = None
output_cost_per_token_priority: float | None = None
search_context_cost_per_query: SearchContextCostPerQuery | None = None
web_search_billing_unit: str | None = None
google_maps_grounding_cost_per_query: float | None = None
file_search_cost_per_1k_calls: float | None = None
provider_specific_entry: ProviderSpecificEntry | None = None
class Deployment(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
model: str | None = None
base_model: str | None = None
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
search_cost_per_query: float | None = None
class WavUpload(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
kind: Literal["wav"]
seconds: float
class PngUpload(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
kind: Literal["png"]
Upload: TypeAlias = Annotated[WavUpload | PngUpload, Field(discriminator="kind")]
class JsonResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
content_type: Literal["application/json"]
body: dict[str, JsonValue]
status: int = 200
headers: dict[str, str] = Field(default_factory=dict)
class SseResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
content_type: Literal["text/event-stream"]
frames: tuple[str, ...]
frame_delay_ms: int = Field(default=0, ge=0)
headers: dict[str, str] = Field(default_factory=dict)
class EventStreamEvent(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
event_type: str
payload: dict[str, JsonValue]
class EventStreamResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
content_type: Literal["application/vnd.amazon.eventstream"]
events: tuple[EventStreamEvent, ...]
framing: Literal["converse", "invoke"] = "converse"
class BinaryResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
content_type: Literal["audio/mpeg"]
length: int
class TextResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
content_type: Literal["application/jsonl"]
body: str
status: int = 200
class RoutedResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
content_type: Literal["application/x-routed"]
routes: dict[str, JsonResponse | TextResponse]
class RealtimeResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
content_type: Literal["application/x-realtime"]
events: tuple[dict[str, JsonValue], ...]
session_model: str | None = None
session_type: str | None = None
created_event: Literal["session.created", "transcription_session.created"] = "session.created"
created_repeats: int = 1
StoredResponse: TypeAlias = Annotated[
JsonResponse | SseResponse | EventStreamResponse | BinaryResponse | RoutedResponse | RealtimeResponse,
Field(discriminator="content_type"),
]
class ExactExpected(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
spend: float
input_cost: float
output_cost: float
prompt_tokens: int
completion_tokens: int
cache_read_cost: float | None = None
cache_creation_cost: float | None = None
reasoning_cost: float | None = None
tool_usage_cost: float | None = None
breakdown_persisted: bool = True
cost_header: bool = True
rollups: bool = False
class RecountRates(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
input_cost_per_token: float
output_cost_per_token: float
class RecountExpected(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
recount: RecountRates
prompt_tokens: int | None = None
completion_tokens: int | None = None
min_completion_tokens: int | None = None
max_completion_tokens: int | None = None
class FailureDetails(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
status: int
class FailureExpected(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
failure: FailureDetails
Expected: TypeAlias = ExactExpected | RecountExpected | FailureExpected
class CostTrackingTestCase(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
name: str
covers: str
known_bug: str | None = None
model: str
endpoint: (
Literal[
"/v1/chat/completions",
"/v1/responses",
"/v1/messages",
"/v1/embeddings",
"/v1/rerank",
"/v1/completions",
"/v1/moderations",
"/v1/audio/transcriptions",
"/v1/audio/speech",
"/v1/images/generations",
"/v1/images/edits",
"/v1/systemone",
"/v1/decisions",
]
| Annotated[str, Field(pattern=r"^/(gemini|anthropic|bedrock)/")]
) = "/v1/chat/completions"
deployment: Deployment | None = None
upload: Upload | None = None
request: dict[str, JsonValue]
response: StoredResponse
expected: Expected
fallback_from: StoredResponse | None = None
disconnect_after_frames: int | None = Field(default=None, ge=1)
upstream_body: dict[str, JsonValue] | None = None
@property
def rates(self) -> CostMapEntry:
return COST_MAP[self.model]
@property
def litellm_model(self) -> str:
provider: Final = self.rates.litellm_provider
prefix: Final = (
"openai"
if provider == "openai"
and (
self.endpoint == "/v1/responses"
or self.rates.mode
in {"chat", "embedding", "moderation", "audio_transcription", "audio_speech", "image_generation"}
)
else "openai/responses"
if provider == "openai"
else _PROVIDER_PREFIXES.get(provider)
)
if prefix is None:
raise ValueError(f"unsupported cost-map provider {provider} for {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]:
return _LITELLM_PARAMS[self.rates.litellm_provider]
@property
def api_key(self) -> str:
return "sk-scripted-provider"
@property
def base_model(self) -> str | None:
return self.deployment.base_model if self.deployment else None
@property
def passthrough_provider(self) -> Literal["gemini", "anthropic", "bedrock"] | None:
provider: Final = self.endpoint.removeprefix("/").split("/", 1)[0]
if provider == "gemini":
return "gemini"
if provider == "anthropic":
return "anthropic"
if provider == "bedrock":
return "bedrock"
return None
@property
def reports_provider_cost(self) -> bool:
if not isinstance(self.response, JsonResponse):
return False
if "x-litellm-response-cost" in self.response.headers:
return True
usage: Final = self.response.body.get("usage")
return isinstance(usage, dict) and isinstance(usage.get("cost"), (int, float))
@property
def chains_prior_response(self) -> bool:
return self.request.get("previous_response_id") == PRIOR_RESPONSE_ID_MARKER
@property
def can_chain_prior_response(self) -> bool:
return (
self.chains_prior_response
and self.endpoint == "/v1/responses"
and isinstance(self.response, JsonResponse)
and isinstance(self.response.body.get("id"), str)
and not isinstance(self.expected, FailureExpected)
and not (isinstance(self.expected, ExactExpected) and self.expected.rollups)
)
class BatchOutputLine(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
status_code: int
prompt_tokens: int | None = None
completion_tokens: int | None = None
cached_tokens: int | None = None
@field_validator("status_code")
@classmethod
def validate_status_code(cls, value: int) -> int:
if value != 200 and not 400 <= value <= 499:
raise ValueError("status_code must be 200 or a 4xx status")
return value
@model_validator(mode="after")
def validate_success_tokens(self) -> BatchOutputLine:
if self.status_code == 200 and (self.prompt_tokens is None or self.completion_tokens is None):
raise ValueError("successful batch output lines require prompt and completion tokens")
return self
def render(self, index: int, model: str, request_id: str) -> dict[str, JsonValue]:
if self.status_code != 200:
return {
"id": f"batch_req_{index}",
"custom_id": f"r{index}",
"response": None,
"error": {"code": "bad_request", "message": "failed"},
}
if self.prompt_tokens is None or self.completion_tokens is None:
raise ValueError("successful batch output lines require prompt and completion tokens")
usage: Final = {
"prompt_tokens": self.prompt_tokens,
"completion_tokens": self.completion_tokens,
"total_tokens": self.prompt_tokens + self.completion_tokens,
**(
{"prompt_tokens_details": {"cached_tokens": self.cached_tokens}}
if self.cached_tokens is not None
else {}
),
}
return {
"id": f"batch_req_{index}",
"custom_id": f"r{index}",
"response": {
"status_code": 200,
"request_id": f"{request_id}-{index}",
"body": {
"id": f"chatcmpl-{request_id}-{index}",
"object": "chat.completion",
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": usage,
},
},
"error": None,
}
class BatchCostCase(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
name: str
covers: str
model: str
litellm_model: str
output_lines: tuple[BatchOutputLine, ...]
expected: ExactExpected
@property
def request_count(self) -> int:
return len(self.output_lines) or 2
@property
def completed_count(self) -> int:
return sum(line.status_code == 200 for line in self.output_lines)
@property
def failed_count(self) -> int:
return self.request_count - self.completed_count
class RealtimeTurn(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
input_tokens: int
output_tokens: int
input_text_tokens: int
input_audio_tokens: int
input_cached_tokens: int
output_text_tokens: int
output_audio_tokens: int
@model_validator(mode="after")
def validate_token_totals(self) -> RealtimeTurn:
if self.input_text_tokens + self.input_audio_tokens != self.input_tokens:
raise ValueError("input text and audio tokens must equal input_tokens")
if self.output_text_tokens + self.output_audio_tokens != self.output_tokens:
raise ValueError("output text and audio tokens must equal output_tokens")
if self.input_cached_tokens > self.input_text_tokens:
raise ValueError("input_cached_tokens must not exceed input_text_tokens")
return self
def render(self, index: int, request_id: str) -> dict[str, JsonValue]:
return {
"type": "response.done",
"event_id": f"evt_{request_id}_{index}",
"response": {
"id": f"resp_{request_id}_{index}",
"object": "realtime.response",
"status": "completed",
"output": [],
"usage": {
"total_tokens": self.input_tokens + self.output_tokens,
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"input_token_details": {
"text_tokens": self.input_text_tokens,
"audio_tokens": self.input_audio_tokens,
"cached_tokens": self.input_cached_tokens,
"cached_tokens_details": {
"text_tokens": self.input_cached_tokens,
"audio_tokens": 0,
},
},
"output_token_details": {
"text_tokens": self.output_text_tokens,
"audio_tokens": self.output_audio_tokens,
},
},
},
}
class RealtimeCostCase(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
name: str
covers: str
model: str
litellm_model: str
turns: tuple[RealtimeTurn, ...] = Field(min_length=0)
session_model: str | None = None
expected: ExactExpected
_PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType(
{
"anthropic": "anthropic",
"bedrock": "bedrock",
"bedrock_converse": "bedrock/converse",
"bedrock_mantle": "bedrock_mantle",
"deepgram": "deepgram",
"text-completion-openai": "text-completion-openai",
"cohere": "cohere",
"vertex_ai-language-models": "vertex_ai",
"vertex_ai-image-models": "vertex_ai",
"vertex_ai-embedding-models": "vertex_ai",
"gemini": "",
"together_ai": "",
"fireworks_ai": "",
"azure": "",
"dashscope": "",
"openrouter": "",
"perplexity": "",
"typesafe": "typesafe",
"cloudflare": "cloudflare",
"databricks": "databricks",
"deepseek": "",
"xai": "",
"azure_ai": "azure_ai",
"groq": "groq",
"mistral": "mistral",
"cohere_chat": "cohere_chat",
}
)
_LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
{
"anthropic": MappingProxyType({}),
"bedrock": MappingProxyType(
{
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
"aws_secret_access_key": "scripted-secret",
"aws_region_name": "us-east-1",
}
),
"bedrock_converse": MappingProxyType(
{
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
"aws_secret_access_key": "scripted-secret",
"aws_region_name": "us-east-1",
}
),
"bedrock_mantle": MappingProxyType(
{
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
"aws_secret_access_key": "scripted-secret",
"aws_region_name": "us-east-1",
}
),
"deepgram": MappingProxyType({}),
"text-completion-openai": MappingProxyType({}),
"cohere": MappingProxyType({}),
"vertex_ai-language-models": MappingProxyType(
{"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"}
),
"vertex_ai-image-models": MappingProxyType(
{"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"}
),
"vertex_ai-embedding-models": MappingProxyType(
{"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"}
),
"gemini": MappingProxyType({}),
"together_ai": MappingProxyType({}),
"fireworks_ai": MappingProxyType({}),
"azure": MappingProxyType({"api_version": "2025-04-01-preview"}),
"openai": MappingProxyType({}),
"dashscope": MappingProxyType({}),
"openrouter": MappingProxyType({}),
"typesafe": MappingProxyType({}),
"cloudflare": MappingProxyType({}),
"databricks": MappingProxyType({}),
"perplexity": MappingProxyType({}),
"deepseek": MappingProxyType({}),
"xai": MappingProxyType({}),
"azure_ai": MappingProxyType({}),
"groq": MappingProxyType({}),
"mistral": MappingProxyType({}),
"cohere_chat": MappingProxyType({}),
}
)
COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType(
{model: CostMapEntry.model_validate(entry) for model, entry in COST_MAP_ENTRIES.items()}
)
def with_usage(
base: CostTrackingTestCase, name: str, usage: Mapping[str, JsonValue], expected: Expected
) -> CostTrackingTestCase:
"""Same request and stored body as ``base``, with the provider usage block and the expected row swapped."""
if not isinstance(base.response, JsonResponse):
raise TypeError(f"{base.name}: with_usage needs a JsonResponse base")
body: Final = base.response.body
current: Final = body["usage"]
if not isinstance(current, dict):
raise TypeError(f"{base.name}: stored body has no usage object")
return base.model_copy(
update={
"name": name,
"response": base.response.model_copy(update={"body": {**body, "usage": {**current, **usage}}}),
"expected": expected,
}
)