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" ALIAS_MARKER: Final = "$MARKER" 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 use_openai_responses_path: bool | 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 input_cost_per_token_ultrafast: float | None = None output_cost_per_token_ultrafast: float | None = None input_cost_per_token_above_272k_tokens_ultrafast: float | None = None output_cost_per_token_above_272k_tokens_ultrafast: float | None = None input_cost_per_token_above_200k_tokens_priority: float | None = None output_cost_per_token_above_200k_tokens_priority: float | None = None cache_read_input_token_cost_flex: float | None = None cache_read_input_token_cost_priority: float | None = None cache_read_input_token_cost_ultrafast: float | None = None cache_read_input_token_cost_above_272k_tokens_flex: float | None = None cache_read_input_token_cost_above_272k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_ultrafast: float | None = None cache_read_input_token_cost_above_200k_tokens_priority: float | None = None cache_creation_input_token_cost_flex: float | None = None cache_creation_input_token_cost_priority: float | None = None cache_creation_input_token_cost_ultrafast: float | None = None cache_creation_input_token_cost_above_272k_tokens_flex: float | None = None cache_creation_input_token_cost_above_272k_tokens_priority: float | None = None cache_creation_input_token_cost_above_272k_tokens_ultrafast: 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 OffPeakPricing(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") hours_utc: str input_cost_per_token: float output_cost_per_token: float ALWAYS_OPEN_OFF_PEAK_WINDOW: Final = "00:00-00:00" 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 off_peak_pricing: OffPeakPricing | 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 created_after_first_frame: bool = False dialect: Literal["openai", "xai"] = "openai" 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 alias: str | None = None 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 @field_validator("alias") @classmethod def validate_alias(cls, value: str | None) -> str | None: if value is not None and ALIAS_MARKER not in value: raise ValueError(f"alias must contain {ALIAS_MARKER} so every run registers a unique model_name") return value @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": "vertex_ai", "vertex_ai-language-models": "vertex_ai", "vertex_ai-image-models": "vertex_ai", "vertex_ai-anthropic_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": MappingProxyType( {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} ), "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-anthropic_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, } ) def with_response_usage( base: CostTrackingTestCase, name: str, usage: Mapping[str, JsonValue], expected: Expected, choice_updates: Mapping[str, JsonValue] | None = None, ) -> CostTrackingTestCase: if not isinstance(base.response, JsonResponse): raise TypeError(f"{base.name}: with_response_usage needs a JsonResponse base") body: Final = {**base.response.body, "usage": dict(usage)} response_body: Final = _with_choice_updates(body, choice_updates) if choice_updates is not None else body return base.model_copy( update={ "name": name, "response": base.response.model_copy(update={"body": response_body}), "expected": expected, } ) def _with_choice_updates( body: Mapping[str, JsonValue], choice_updates: Mapping[str, JsonValue] ) -> dict[str, JsonValue]: choices: Final = body.get("choices") if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict): raise TypeError("stored body has no choice object") return {**body, "choices": [{**choices[0], **choice_updates}, *choices[1:]]}