litellm/tests/integration/cost_calculation/cost_tracking_case.py

from __future__ import annotations

import json
from collections.abc import Mapping
from pathlib import Path
from types import MappingProxyType
from typing import Annotated, Final, Literal, TypeAlias

from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator

CASES_PATH: Final = Path(__file__).resolve().parent / "cost_tracking_cases.json"
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
    input_cost_per_token: float | None = None
    input_cost_per_query: float | None = None
    output_cost_per_token: float | None = None
    input_cost_per_token_batches: float | None = None
    output_cost_per_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


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


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)


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


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
    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",
        ]
        | 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)

    @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
        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


class _CasesFile(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    cost_map: dict[str, CostMapEntry]
    cases: tuple[CostTrackingTestCase, ...]
    batch_cases: tuple[BatchCostCase, ...] = ()
    realtime_cases: tuple[RealtimeCostCase, ...] = ()


_PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType(
    {
        "anthropic": "anthropic",
        "bedrock": "bedrock",
        "bedrock_converse": "bedrock/converse",
        "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": "",
        "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",
            }
        ),
        "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({}),
        "perplexity": MappingProxyType({}),
        "deepseek": MappingProxyType({}),
        "xai": MappingProxyType({}),
        "azure_ai": MappingProxyType({}),
        "groq": MappingProxyType({}),
        "mistral": MappingProxyType({}),
        "cohere_chat": MappingProxyType({}),
    }
)

_LOADED: Final = _CasesFile.model_validate_json(CASES_PATH.read_bytes())
COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType(dict(_LOADED.cost_map))
CASES: Final[tuple[CostTrackingTestCase, ...]] = _LOADED.cases
BATCH_CASES: Final[tuple[BatchCostCase, ...]] = _LOADED.batch_cases
REALTIME_CASES: Final[tuple[RealtimeCostCase, ...]] = _LOADED.realtime_cases
_ALL_CASES: Final = CASES + BATCH_CASES + REALTIME_CASES
_LITELLM_MODELS: Final = tuple(case.litellm_model for case in _ALL_CASES)


def data_errors() -> tuple[str, ...]:
    case_models: Final = frozenset(case.model for case in _ALL_CASES) | frozenset(
        case.session_model for case in REALTIME_CASES if case.session_model is not None
    )
    unknown_models: Final = sorted(model for model in case_models if model not in COST_MAP)
    missing_cases: Final = sorted(model for model in COST_MAP if model not in case_models)
    duplicate_names: Final = sorted(
        name for name in {case.name for case in _ALL_CASES} if sum(case.name == name for case in _ALL_CASES) > 1
    )
    input_rates: Final = tuple(
        (entry.input_cost_per_token, model)
        for model, entry in COST_MAP.items()
        if entry.mode != "realtime"
    )
    shared_input_rates: Final = sorted(
        f"{rate}: {tuple(model for value, model in input_rates if value == rate)}"
        for rate in {value for value, _ in input_rates if value is not None}
        if sum(value == rate for value, _ in input_rates) > 1
    )
    recount_mismatches: Final = sorted(
        case.name
        for case in CASES
        if isinstance(case.expected, RecountExpected)
        and case.model in COST_MAP
        and (
            case.expected.recount.input_cost_per_token != (COST_MAP[case.model].input_cost_per_token or 0.0)
            or case.expected.recount.output_cost_per_token != (COST_MAP[case.model].output_cost_per_token or 0.0)
        )
    )
    component_mismatches: Final = sorted(
        case.name
        for case in CASES
        if isinstance(case.expected, ExactExpected)
        and any(
            component is not None
            for component in (
                case.expected.cache_read_cost,
                case.expected.cache_creation_cost,
                case.expected.reasoning_cost,
                case.expected.tool_usage_cost,
            )
        )
        and (
            (case.expected.cache_read_cost or 0.0) + (case.expected.cache_creation_cost or 0.0)
            > case.expected.input_cost
            or (case.expected.reasoning_cost or 0.0) > case.expected.output_cost
            or not _approx_equal(
                case.expected.input_cost
                + case.expected.output_cost
                + (case.expected.tool_usage_cost or 0.0),
                case.expected.spend,
            )
        )
    )
    failure_response_mismatches: Final = sorted(
        case.name
        for case in CASES
        if (
            isinstance(case.expected, FailureExpected)
            and (
                not isinstance(case.response, JsonResponse)
                or not 400 <= case.response.status <= 599
                or not 400 <= case.expected.failure.status <= 599
            )
        )
        or (
            not isinstance(case.expected, FailureExpected)
            and isinstance(case.response, JsonResponse)
            and case.response.status != 200
        )
    )
    invalid_opt_outs: Final = sorted(
        case.name
        for case in CASES
        if isinstance(case.expected, ExactExpected)
        and (
            (
                not case.expected.breakdown_persisted
                and case.passthrough_provider is None
                and case.rates.mode != "image_generation"
                and not case.reports_provider_cost
            )
            or (
                not case.expected.cost_header
                and case.passthrough_provider is None
                and not isinstance(case.response, SseResponse)
                and case.expected.spend != 0.0
            )
        )
    )
    invalid_fallbacks: Final = sorted(
        case.name
        for case in CASES
        if case.fallback_from is not None
        and (
            not isinstance(case.fallback_from, JsonResponse)
            or not 400 <= case.fallback_from.status <= 599
        )
    )
    invalid_disconnects: Final = sorted(
        case.name
        for case in CASES
        if case.disconnect_after_frames is not None
        and (
            not isinstance(case.response, SseResponse)
            or case.response.frame_delay_ms <= 0
            or not isinstance(case.expected, RecountExpected)
        )
    )
    invalid_rollup_ids: Final = sorted(
        case.name
        for case in CASES
        if isinstance(case.expected, ExactExpected)
        and case.expected.rollups
        and "$UNIQUE_ID" not in case.response.model_dump_json()
    )
    invalid_pinned_tool_ids: Final = sorted(
        case.name
        for case in CASES
        if isinstance(case.expected, RecountExpected)
        and (case.expected.prompt_tokens is not None or case.expected.completion_tokens is not None)
        and any(
            marker in case.response.model_dump_json()
            for marker in ('"id": "call_$REQUEST_ID"', '"id": "toolu_$REQUEST_ID"')
        )
    )
    invalid_prior_response_chains: Final = sorted(
        case.name
        for case in CASES
        if PRIOR_RESPONSE_ID_MARKER in json.dumps(case.request) and not case.can_chain_prior_response
    )
    return tuple(
        message
        for message in (
            f"case models absent from cost_map: {unknown_models}" if unknown_models else None,
            f"cost-map entries without cases: {missing_cases}" if missing_cases else None,
            f"duplicate case names: {duplicate_names}" if duplicate_names else None,
            f"cost-map entries share input_cost_per_token: {shared_input_rates}" if shared_input_rates else None,
            f"recount rates differ from cost-map rates: {recount_mismatches}" if recount_mismatches else None,
            f"breakdown components are inconsistent: {component_mismatches}" if component_mismatches else None,
            f"failure response statuses are inconsistent: {failure_response_mismatches}"
            if failure_response_mismatches
            else None,
            f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None,
            f"invalid fallback responses: {invalid_fallbacks}" if invalid_fallbacks else None,
            f"invalid disconnect cases: {invalid_disconnects}" if invalid_disconnects else None,
            f"rollup responses lack $UNIQUE_ID: {invalid_rollup_ids}" if invalid_rollup_ids else None,
            f"pinned tool IDs contain $REQUEST_ID: {invalid_pinned_tool_ids}"
            if invalid_pinned_tool_ids
            else None,
            f"{PRIOR_RESPONSE_ID_MARKER} needs a non-rollup, non-failure /v1/responses JSON response with a string id"
            f" as previous_response_id: {invalid_prior_response_chains}"
            if invalid_prior_response_chains
            else None,
        )
        if message is not None
    )


def _approx_equal(actual: float, expected: float) -> bool:
    return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)