litellm/tests/e2e/models.py

"""Shared pydantic request/response models for the e2e gateway.

Only the fields the tests read are modelled; pydantic ignores the rest, so a
response validates without mirroring every proxy field. No untyped dicts.
"""

from __future__ import annotations

from collections.abc import Sequence
from datetime import datetime
from typing import Final, Literal

from e2e_http import PartialBody
from pydantic import (
    AliasChoices,
    BaseModel,
    ConfigDict,
    Field,
    JsonValue,
    RootModel,
    model_serializer,
    model_validator,
)

# ---------- keys ----------


class ModelBudgetEntry(BaseModel):
    budget_limit: float = Field(validation_alias=AliasChoices("budget_limit", "max_budget"))
    time_period: str = Field(validation_alias=AliasChoices("time_period", "budget_duration"))
    rpm_limit: int | None = None
    tpm_limit: int | None = None


class BudgetWindow(BaseModel):
    budget_duration: str
    max_budget: float


class BudgetWindowState(BudgetWindow):
    reset_at: datetime | None = None


class KeyLoggingCallbackVars(BaseModel):
    langfuse_public_key: str | None = Field(default=None, repr=False)
    langfuse_secret_key: str | None = Field(default=None, repr=False)
    langfuse_host: str | None = Field(default=None, repr=False)


class KeyLoggingCallback(BaseModel):
    callback_name: str
    callback_type: str = "success_and_failure"
    callback_vars: KeyLoggingCallbackVars


class KeyMetadata(BaseModel):
    logging: list[KeyLoggingCallback] | None = None
    priority: str | None = None
    batch_enqueued_token_limit: int | None = None
    tag: str | None = None
    guardrails: list[str] | None = None


class ObjectPermission(BaseModel):
    mcp_servers: list[str] | None = None
    mcp_access_groups: list[str] | None = None
    mcp_toolsets: list[str] | None = None
    mcp_tool_search_enabled: bool | None = None


class KeyGenerateBody(BaseModel):
    models: list[str] = []
    duration: str | None = None
    max_budget: float | None = None
    soft_budget: float | None = None
    budget_duration: str | None = None
    user_id: str | None = None
    team_id: str | None = None
    project_id: str | None = None
    organization_id: str | None = None
    budget_id: str | None = None
    key_alias: str | None = None
    model_max_budget: dict[str, ModelBudgetEntry] | None = None
    budget_fallbacks: dict[str, list[str]] | None = None
    budget_limits: list[BudgetWindow] | None = None
    tpm_limit: int | None = None
    rpm_limit: int | None = None
    allowed_routes: list[str] | None = None
    allowed_passthrough_routes: list[str] | None = None
    metadata: KeyMetadata | None = None
    object_permission: ObjectPermission | None = None
    router_settings: RouterSettingsOverride | None = None


class KeyGenerateResponse(BaseModel):
    key: str
    token: str | None = None
    key_alias: str | None = None
    models: list[str] = []
    max_budget: float | None = None
    tpm_limit: int | None = None
    rpm_limit: int | None = None
    budget_duration: str | None = None
    team_id: str | None = None
    metadata: KeyMetadata | None = None


class KeyRegenerateBody(BaseModel):
    key: str
    grace_period: str | None = None


class KeyResetSpendBody(BaseModel):
    reset_to: float


class KeyResetSpendResponse(BaseModel):
    spend: float
    previous_spend: float


class KeyDeleteBody(BaseModel):
    keys: list[str]


class KeyDeleteByAliasBody(BaseModel):
    key_aliases: list[str]


class AuditLogParams(BaseModel):
    object_id: str
    action: str
    table_name: str
    page_size: int


class AuditLogEntry(BaseModel):
    id: str
    changed_by: str | None = None
    changed_by_api_key: str | None = None
    action: str
    table_name: str
    object_id: str
    before_value: object | None = None


class AuditLogPage(BaseModel):
    audit_logs: list[AuditLogEntry]
    total: int


class KeyInfoParams(BaseModel):
    key: str


class LiteLLMBudgetTable(BaseModel):
    max_budget: float | None = None
    soft_budget: float | None = None
    budget_duration: str | None = None
    budget_reset_at: str | None = None


class KeyInfo(BaseModel):
    key_alias: str | None = None
    status: str | None = None
    metadata: KeyMetadata | None = None
    models: list[str] = []
    tpm_limit: int | None = None
    rpm_limit: int | None = None
    project_id: str | None = None
    organization_id: str | None = None
    team_id: str | None = None
    blocked: bool | None = None
    spend: float | None = None
    max_budget: float | None = None
    budget_duration: str | None = None
    budget_reset_at: str | None = None
    budget_id: str | None = None
    litellm_budget_table: LiteLLMBudgetTable | None = None
    budget_limits: list[BudgetWindowState] | None = None
    object_permission: ObjectPermission | None = None


class KeyInfoResponse(BaseModel):
    info: KeyInfo


# ---------- customers ----------


class CustomerNewBody(BaseModel):
    user_id: str


class CustomerResponse(BaseModel):
    user_id: str | None = None


class CustomerInfoParams(BaseModel):
    end_user_id: str


class CustomerDeleteBody(BaseModel):
    user_ids: list[str]


# ---------- chat / embeddings ----------


class ChatMetadata(BaseModel):
    tags: list[str] | None = None


class ImageUrl(BaseModel):
    url: str


class TextContentPart(BaseModel):
    type: str = "text"
    text: str
    cache_control: CacheControl | None = None


class ImageContentPart(BaseModel):
    type: str = "image_url"
    image_url: ImageUrl


ContentPart = TextContentPart | ImageContentPart


class ChatMessage(BaseModel):
    role: str
    content: str | list[ContentPart]
    cache_control: CacheControl | None = None


class CacheControl(BaseModel):
    type: str = "ephemeral"
    ttl: str | None = None


class TextBlock(BaseModel):
    type: str = "text"
    text: str
    cache_control: CacheControl | None = None


class RichMessage(BaseModel):
    role: str
    content: list[TextBlock]


class ThinkingParam(BaseModel):
    """Extended-thinking control shared by Anthropic and DeepSeek reasoner models.
    DeepSeek accepts only ``type`` (enabled/disabled) and ignores budget_tokens;
    Anthropic also honors budget_tokens. Sending ``type="disabled"`` is the
    product-facing way a caller turns reasoning off (LIT-3686 / GH #27453)."""

    type: Literal["enabled", "disabled"]
    budget_tokens: int | None = None


class ChatToolFunction(BaseModel):
    name: str
    description: str | None = None
    parameters: dict[str, object] | None = None


class ChatTool(BaseModel):
    type: str = "function"
    function: ChatToolFunction


class AnthropicToolBody(BaseModel, frozen=True):
    type: Literal["computer_20250124", "web_fetch_20250910", "text_editor_20250728", "url", "mcp"]
    name: str | None = None
    function: dict[str, JsonValue] | None = None
    max_uses: int | None = None
    url: str | None = None
    authorization_token: str | None = None
    server_label: str | None = None
    server_url: str | None = None
    require_approval: str | None = None
    headers: dict[str, str] | None = None


class AnthropicExtraBody(BaseModel, frozen=True):
    tools: Sequence[AnthropicToolBody] | None = None
    web_search_options: dict[str, JsonValue] | None = None


class McpChatTool(BaseModel):
    """An MCP server attached to a chat completion (OpenAI `type: "mcp"` tool).
    `server_url` selects the gateway-registered server by its alias suffix; with
    `require_approval="never"` the gateway lists, calls, and feeds the server's
    tools back to the model in one agentic turn."""

    type: Literal["mcp"] = "mcp"
    server_url: str
    require_approval: str
    server_label: str | None = None
    allowed_tools: list[str] | None = None


class GeminiUrlContextTool(BaseModel, frozen=True):
    urlContext: dict[str, JsonValue] = Field(default_factory=dict)


class ToolCallFunction(BaseModel):
    name: str | None = None
    arguments: str | None = None


class ToolCall(BaseModel):
    id: str | None = None
    type: str | None = None
    function: ToolCallFunction = ToolCallFunction()
    cache_control: CacheControl | None = None


class ThinkingBlock(BaseModel):
    type: str
    thinking: str | None = None
    signature: str | None = None
    data: str | None = None


class ChatAssistantTurn(BaseModel):
    role: Literal["assistant"] = "assistant"
    content: str | None = None
    reasoning_content: str | None = None
    thinking_blocks: list[ThinkingBlock] | None = None
    tool_calls: list[ToolCall] | None = None


class ChatToolResultTurn(BaseModel):
    role: Literal["tool"] = "tool"
    tool_call_id: str
    content: str | list[ContentPart]


type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn


class ChatStreamOptions(BaseModel):
    include_usage: bool


class ChatBody(BaseModel):
    model: str
    messages: Sequence[ChatTurn]
    stream: bool = False
    stream_options: ChatStreamOptions | None = None
    max_tokens: int | None = None
    max_completion_tokens: int | None = None
    temperature: float | None = None
    user: str | None = None
    safety_identifier: str | None = None
    metadata: ChatMetadata | None = None
    reasoning_effort: str | None = None
    thinking: ThinkingParam | None = None
    service_tier: str | None = None
    prompt_cache_key: str | None = None
    prompt_cache_retention: str | None = None
    tools: Sequence[ChatTool | McpChatTool | GeminiUrlContextTool] | None = None
    tool_choice: str | None = None
    guardrails: list[str] | None = None
    include_guardrail_response: bool | None = None
    response_format: dict[str, object] | None = None
    chat_template_kwargs: dict[str, bool] | None = None
    cache: dict[str, bool] | None = {"no-cache": True}


RoutingStrategy = Literal[
    "simple-shuffle",
    "least-busy",
    "usage-based-routing-v2",
    "latency-based-routing",
    "cost-based-routing",
]


class RouterSettingsOverride(BaseModel):
    """Router settings a test scopes below the global config: sent per request as
    `router_settings_override` in a /chat/completions body (the reliability suite's
    fallback, retry, routing-strategy, and deadline knobs) or stored on a key as
    `router_settings` at /key/generate (the auto-router suite's tag filtering
    switch). Serialized exclude_none, so an override sets only the knobs a test
    exercises. Each fallbacks map is model_name -> the ordered fallback model_names
    to try."""

    fallbacks: list[dict[str, list[str]]] | None = None
    context_window_fallbacks: list[dict[str, list[str]]] | None = None
    content_policy_fallbacks: list[dict[str, list[str]]] | None = None
    num_retries: int | None = None
    routing_strategy: RoutingStrategy | None = None
    model_group_retry_policy: dict[str, dict[str, int]] | None = None
    enable_tag_filtering: bool | None = None


class DeploymentExtraBody(BaseModel):
    """`litellm_params.extra_body` of a deployment whose upstream is another LiteLLM
    proxy: forwarded verbatim in every request body, so the inner proxy honors the
    same per-request router knobs an end user could send it."""

    router_settings_override: RouterSettingsOverride | None = None


class ReliabilityChatBody(ChatBody):
    """A /chat/completions body carrying a per-request router_settings_override.
    Composes ChatBody (no attribute repetition) and adds the override; serialized
    exclude_none so an absent override never leaks into the request."""

    router_settings_override: RouterSettingsOverride | None = None


class McpToolFunctionRef(BaseModel):
    name: str


class McpListedTool(BaseModel):
    """One entry of `mcp_list_tools`: a tool the gateway listed from the
    attached MCP server and exposed to the model, in OpenAI function shape."""

    function: McpToolFunctionRef | None = None


class McpToolCall(BaseModel):
    """One entry of `mcp_tool_calls`: a tool the model asked the gateway to run."""

    function: McpToolFunctionRef | None = None


class McpCallResult(BaseModel):
    """One entry of `mcp_call_results`: what the gateway got back from executing
    a tool upstream on the caller's behalf."""

    name: str | None = None
    result: str | None = None


class McpResponseMetadata(BaseModel):
    """`choices[].message.provider_specific_fields` MCP section: which tools the
    gateway listed from the attached server, which the model called, and their
    results. Populated only when the completion drove an MCP server."""

    mcp_list_tools: list[McpListedTool] | None = None
    mcp_tool_calls: list[McpToolCall] | None = None
    mcp_call_results: list[McpCallResult] | None = None


class OutMessage(BaseModel):
    role: str | None = None
    content: str | None = None
    reasoning_content: str | None = None
    thinking_blocks: list[ThinkingBlock] | None = None
    tool_calls: list[ToolCall] | None = None
    provider_specific_fields: McpResponseMetadata | None = None


class ChatChoice(BaseModel):
    message: OutMessage | None = None
    finish_reason: str | None = None


class PromptTokensDetails(BaseModel):
    cached_tokens: int | None = None
    cache_write_tokens: int | None = None


class CompletionTokensDetails(BaseModel):
    reasoning_tokens: int | None = None


class Usage(BaseModel):
    prompt_tokens: int | None = None
    completion_tokens: int | None = None
    total_tokens: int | None = None
    cache_read_input_tokens: int | None = None
    cache_creation_input_tokens: int | None = None
    prompt_tokens_details: PromptTokensDetails | None = None
    completion_tokens_details: CompletionTokensDetails | None = None


class GuardrailInformationEntry(BaseModel):
    guardrail_name: str
    guardrail_status: str
    guardrail_mode: object | None = None
    guardrail_response: object | None = None
    duration: float | None = None


class GeminiUrlMetadata(BaseModel, frozen=True):
    retrievedUrl: str
    urlRetrievalStatus: str


class GeminiUrlContextMetadata(BaseModel, frozen=True):
    urlMetadata: Sequence[GeminiUrlMetadata]


class ChatResponse(BaseModel):
    id: str | None = None
    object: str | None = None
    model: str | None = None
    choices: list[ChatChoice] = []
    usage: Usage | None = None
    service_tier: str | None = None
    guardrail_information: list[GuardrailInformationEntry] | None = None
    vertex_ai_url_context_metadata: Sequence[GeminiUrlContextMetadata] | None = None


class ResponsesInputTokensDetails(BaseModel):
    cached_tokens: int | None = None


class ResponsesUsage(BaseModel):
    """`/v1/responses` usage shape: input/output tokens, not prompt/completion."""

    input_tokens: int
    output_tokens: int
    total_tokens: int | None = None
    input_tokens_details: ResponsesInputTokensDetails | None = None


class ResponsesApiResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")

    id: str | None = None
    usage: ResponsesUsage | None = None


class ResponsesStreamEvent(BaseModel):
    """One `/v1/responses` SSE event; `response.completed` carries the final response."""

    model_config = ConfigDict(extra="ignore")

    type: str
    response: ResponsesApiResponse | None = None


# ---------- anthropic /v1/messages + count_tokens ----------


class JsonSchemaProperty(BaseModel):
    """One property in a tool's JSON-Schema `input_schema`. Only `type` is
    modelled; the endpoints under test read no further into the schema."""

    type: str


class ToolInputSchema(BaseModel):
    type: str = "object"
    properties: dict[str, JsonSchemaProperty] = {}
    required: list[str] = []


class AnthropicServerTool(BaseModel):
    """An Anthropic-managed tool the upstream executes itself. It carries no
    `input_schema`; `type` is the SDK-version-pinned identifier LiteLLM keys its
    per-provider translation on, and `name` is the unsuffixed canonical name the
    upstream accepts."""

    type: str
    name: str


class AnthropicToolSearchTool(AnthropicServerTool):
    """The tool_search discovery tool, e.g. ``tool_search_tool_regex_20251119``."""


class AnthropicWebSearchTool(AnthropicServerTool):
    """The web_search server tool, e.g. ``web_search_20250305``. Distinct from
    Claude Code's client-side ``WebSearch`` tool, which is an ordinary custom
    tool the CLI executes and feeds back as a tool_result."""

    max_uses: int | None = None


class AnthropicCustomTool(BaseModel):
    name: str
    description: str
    input_schema: ToolInputSchema


type AnthropicTool = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool


class AnthropicContentBlock(BaseModel):
    """One block of a `content` array. Only the fields a test reads are
    declared; `extra="allow"` keeps the rest (a `server_tool_use` block's
    `input`, a `tool_search_tool_result` block's nested `content`) so an
    assistant turn read off the wire can be replayed into history verbatim
    instead of being silently flattened to its text."""

    model_config = ConfigDict(extra="allow")
    type: str | None = None
    text: str | None = None
    id: str | None = None
    name: str | None = None
    input: dict[str, object] | None = None


class AnthropicToolResultBlock(BaseModel):
    """The user-turn answer to a client-side `tool_use`. `tool_use_id` must be
    the id the model actually emitted; an invented one is rejected by
    Anthropic's own schema validator, which Bedrock inherits."""

    type: Literal["tool_result"] = "tool_result"
    tool_use_id: str
    content: str


class AnthropicAssistantTurn(BaseModel):
    role: Literal["assistant"] = "assistant"
    content: list[AnthropicContentBlock]


class AnthropicToolResultTurn(BaseModel):
    role: Literal["user"] = "user"
    content: list[AnthropicToolResultBlock]


type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn


class AnthropicToolChoice(BaseModel):
    type: Literal["auto", "any", "tool", "none"]
    name: str | None = None


class AnthropicMessagesBody(BaseModel):
    model: str
    messages: list[AnthropicMessage]
    max_tokens: int
    stream: bool | None = None
    tools: list[AnthropicTool] | None = None
    tool_choice: AnthropicToolChoice | None = None
    thinking: ThinkingParam | None = None
    guardrails: list[str] | None = None
    cache: dict[str, bool] | None = {"no-cache": True}


class ResponsesStreamBody(BaseModel):
    """POST /v1/responses body in the subset the spend tests stream with.
    `input` stays a plain string: the tests only drive single-turn prompts."""

    model: str
    input: str
    stream: bool = True
    max_output_tokens: int | None = None
    cache: dict[str, bool] | None = {"no-cache": True}


class ResponsesFunctionTool(BaseModel):
    model_config = ConfigDict(frozen=True)
    type: Literal["function"] = "function"
    name: str
    description: str | None = None
    parameters: ToolInputSchema


class ResponsesToolBody(BaseModel):
    model_config = ConfigDict(frozen=True)
    model: str
    input: str
    tools: list[ResponsesFunctionTool]
    tool_choice: Literal["required"] | None = None
    guardrails: list[str]
    max_output_tokens: int = 128


class ResponsesOutputItem(BaseModel):
    model_config = ConfigDict(frozen=True)
    type: str
    name: str | None = None


class ResponsesToolOutput(BaseModel):
    model_config = ConfigDict(frozen=True)
    output: list[ResponsesOutputItem]


class CountTokensBody(BaseModel):
    """POST /v1/messages/count_tokens body: the /v1/messages shape minus
    max_tokens (the endpoint only counts the prompt)."""

    model: str
    messages: list[ChatMessage]


class AnthropicMessagesResponse(BaseModel):
    """A /v1/messages answer. `content` is the Anthropic-native passthrough
    shape; `choices` is the OpenAI-normalized shape LiteLLM emits for some
    providers (e.g. Bedrock Converse). Presence of either proves the proxy
    accepted and round-tripped the request. `extra="allow"` keeps the other
    top-level keys so a shape-check failure can report the actual response keys
    for triage."""

    model_config = ConfigDict(extra="allow")
    model: str | None = None
    content: list[AnthropicContentBlock] | None = None
    choices: list[ChatChoice] | None = None
    usage: Usage | None = None


TokenCountSource = Literal["provider", "token_count_model", "local_estimate"]


class LitellmCountSource(BaseModel, frozen=True):
    litellm_count_source: TokenCountSource | None = None
    litellm_count_model: str | None = None


class CountTokensResponse(LitellmCountSource, frozen=True):
    """`/v1/messages/count_tokens` answer. `input_tokens` is required so a 200
    whose body lacks it fails validation instead of passing vacuously."""

    input_tokens: int


class ResponsesInputTokensResponse(LitellmCountSource, frozen=True):
    object: str
    input_tokens: int


class TokenCounterBody(BaseModel, frozen=True):
    model: str
    messages: list[ChatMessage]


class TokenCounterResponse(BaseModel, frozen=True):
    total_tokens: int
    request_model: str
    model_used: str
    tokenizer_type: str
    count_source: TokenCountSource | None = None


class AnthropicErrorBody(BaseModel):
    type: str
    message: str


class AnthropicErrorEvent(BaseModel):
    type: Literal["error"]
    error: AnthropicErrorBody


# ---------- mcp servers ----------


class McpInfo(BaseModel):
    """The `mcp_info` display block stored on an MCP server; only the fields the
    lifecycle test writes and reads back."""

    server_name: str | None = None
    description: str | None = None
    logo_url: str | None = None


class McpServerCreateBody(BaseModel):
    """POST /v1/mcp/server. `allow_all_keys` false scopes the server to keys
    granted it through object_permission."""

    alias: str
    url: str
    transport: str = "http"
    allow_all_keys: bool = True
    server_name: str | None = None
    description: str | None = None
    mcp_info: McpInfo | None = None


class McpServerUpdateBody(PartialBody):
    """PUT /v1/mcp/server: a field left unset keeps its stored value, a field set
    to None is cleared."""

    server_id: str
    alias: str | None = None
    description: str | None = None


class McpServerInfo(BaseModel):
    """Response of POST /v1/mcp/server and GET /v1/mcp/server/{server_id}."""

    server_id: str
    alias: str | None = None
    url: str | None = None
    auth_type: str | None = None
    oauth2_flow: str | None = None
    allow_all_keys: bool | None = None


class McpServerRow(McpServerInfo):
    """A stored MCP server as the create, get, and list routes return it: the
    fields the lifecycle test asserts survive the round trip."""

    server_name: str | None = None
    transport: str | None = None
    description: str | None = None
    mcp_info: McpInfo | None = None


class McpServerListResponse(RootModel[list[McpServerRow]]):
    """GET /v1/mcp/server answers with a bare array of servers."""


class ToolsetTool(BaseModel):
    server_id: str
    tool_name: str


class ToolsetCreateBody(BaseModel):
    toolset_name: str
    description: str | None = None
    tools: list[ToolsetTool]


class ToolsetUpdateBody(PartialBody):
    """PUT /v1/mcp/toolset: a field left unset keeps its stored value, a field set
    to None is cleared."""

    toolset_id: str
    description: str | None = None
    tools: list[ToolsetTool] | None = None


class ToolsetRow(BaseModel):
    """A stored toolset as POST /v1/mcp/toolset, GET /v1/mcp/toolset/{toolset_id},
    and each row of GET /v1/mcp/toolset return it."""

    toolset_id: str
    toolset_name: str
    description: str | None = None
    tools: list[ToolsetTool] = Field(default_factory=list)


class ToolsetListResponse(RootModel[list[ToolsetRow]]):
    """GET /v1/mcp/toolset answers with a bare array of toolsets."""


class EmbedBody(BaseModel):
    model: str
    input: str
    cache: dict[str, bool] | None = {"no-cache": True}


class EmbedResponse(BaseModel):
    model: str | None = None


# ---------- videos ----------


class VideoCreateBody(BaseModel):
    model: str
    prompt: str
    seconds: str | None = None


class VideoCreateResponse(BaseModel):
    id: str
    status: str | None = None


# ---------- rerank ----------


class RerankBody(BaseModel):
    model: str
    query: str
    documents: list[str]
    top_n: int
    cache: dict[str, bool] | None = {"no-cache": True}


class RerankItem(BaseModel):
    index: int | None = None
    relevance_score: float | None = None


class RerankResponse(BaseModel):
    results: list[RerankItem] = []


# ---------- ocr ----------


class OcrDocument(BaseModel):
    """A document for /v1/ocr in Mistral OCR format: a document_url for PDFs/docs
    or an image_url for images. exclude_none on serialize drops the unset one."""

    type: str
    document_url: str | None = None
    image_url: str | None = None


class OcrBody(BaseModel):
    model: str
    document: OcrDocument


class OcrForm(BaseModel):
    """Multipart /v1/ocr form fields; the document travels as the `file` part."""

    model: str


class OcrPage(BaseModel):
    index: int
    markdown: str


class OcrUsageInfo(BaseModel):
    pages_processed: int | None = None


class OcrResponse(BaseModel):
    object: str | None = None
    model: str | None = None
    pages: list[OcrPage] = []
    usage_info: OcrUsageInfo | None = None


# ---------- completions ----------


class CompletionBody(BaseModel):
    model: str
    prompt: str
    max_tokens: int = 8
    n: int = 1


class CompletionChoice(BaseModel):
    text: str = ""


class CompletionResponse(BaseModel):
    choices: list[CompletionChoice] = []


# ---------- images ----------


class ImageGenerationBody(BaseModel):
    model: str
    prompt: str
    n: int = 1
    size: str = "1024x1024"
    quality: str = "low"


class ImageDatum(BaseModel):
    url: str | None = None
    b64_json: str | None = None


class ImageGenerationResponse(BaseModel):
    data: list[ImageDatum] = []


class ImageEditForm(BaseModel):
    """POST /v1/images/edits form fields; the image travels as the `image` multipart part."""

    model: str
    prompt: str
    size: str = "1024x1024"
    quality: str = "low"


class SearchBody(BaseModel):
    """POST /v1/search/{search_tool_name} body (Perplexity-compatible)."""

    query: str
    max_results: int = 2


class SearchResultItem(BaseModel):
    title: str = ""
    url: str = ""
    snippet: str = ""


class SearchResponse(BaseModel):
    results: list[SearchResultItem] = []


class SearchToolLiteLLMParamsBody(BaseModel):
    search_provider: str


class SearchToolBody(BaseModel):
    search_tool_name: str
    litellm_params: SearchToolLiteLLMParamsBody


class SearchToolCreateBody(BaseModel):
    """POST /search_tools body: the tool as it would sit under `search_tools:` in the config."""

    search_tool: SearchToolBody


class SearchToolCreateResponse(BaseModel):
    search_tool_id: str


# ---------- audio ----------


class SpeechBody(BaseModel):
    model: str
    input: str
    voice: str = "alloy"


class TranscriptionForm(BaseModel):
    model: str


class TranscriptionResponse(BaseModel):
    text: str = ""


# ---------- moderations ----------


class ModerationBody(BaseModel):
    model: str
    input: str


class ModerationResult(BaseModel):
    flagged: bool


class ModerationResponse(BaseModel):
    results: list[ModerationResult] = []


# ---------- spend logs ----------


class GuardrailEntityMatch(BaseModel):
    entity_type: str
    score: float
    start: int
    end: int


class GuardrailModeRecord(BaseModel):
    tags: dict[str, str | list[str]] | None = None
    default: str | list[str] | None = None


class GuardrailRunRecord(BaseModel):
    guardrail_name: str | None = None
    guardrail_mode: str | list[str] | GuardrailModeRecord | None = None
    guardrail_status: str | None = None
    guardrail_provider: str | None = None
    masked_entity_count: dict[str, int] | None = None
    guardrail_response: object | None = None


class SpendLogErrorInformation(BaseModel):
    error_code: str | None = None
    error_class: str | None = None
    error_message: str | None = None
    normalized_error: str | None = None


class SpendLogMetadata(BaseModel):
    user_api_key_alias: str | None = None
    applied_guardrails: list[str] | None = None
    guardrail_information: list[GuardrailRunRecord] | None = None
    error_information: SpendLogErrorInformation | None = None


class SpendLogRow(BaseModel):
    request_id: str | None = None
    api_key: str | None = Field(default=None, repr=False)
    model: str | None = None
    spend: float | None = None
    status: str | None = None
    cache_hit: str | None = None
    call_type: str | None = None
    custom_llm_provider: str | None = None
    model_id: str | None = None
    team_id: str | None = None
    user: str | None = None
    end_user: str | None = None
    prompt_tokens: int | None = None
    completion_tokens: int | None = None
    total_tokens: int | None = None
    request_tags: list[str] | None = None
    session_id: str | None = None
    metadata: SpendLogMetadata | None = None
    proxy_server_request: JsonValue = None
    response: JsonValue = None
    litellm_call_id: str | None = None


class SpendLogs(RootModel[list[SpendLogRow]]):
    pass


class SpendLogsParams(BaseModel):
    request_id: str | None = None
    api_key: str | None = Field(default=None, repr=False)

    @model_validator(mode="after")
    def require_filter(self) -> SpendLogsParams:
        if self.request_id is None and self.api_key is None:
            raise ValueError(
                "unfiltered /spend/logs returns the entire spend table and OOMs the "
                "runner on long-lived environments; filter by request_id or api_key, "
                "or use ProxyClient.spend_logs_window for a bounded /spend/logs/v2 read"
            )
        return self


class SpendLogsPageParams(BaseModel):
    """Query for /spend/logs/v2, which requires an explicit date window and
    serves pages of at most 100 rows."""

    start_date: str
    end_date: str
    page: int
    page_size: int
    api_key: str | None = Field(default=None, repr=False)


class SessionSpendLogsParams(BaseModel):
    """Query for /spend/logs/session/ui, the session view the Admin UI logs page
    opens: every row whose session_id equals the given one, newest first."""

    session_id: str
    page: int = 1
    page_size: int = 100


class SpendLogsPage(BaseModel):
    data: list[SpendLogRow] = []
    total: int
    page: int
    page_size: int
    total_pages: int


# ---------- spend calculate ----------


class SpendCalculateBody(BaseModel):
    model: str
    messages: list[ChatMessage]


class SpendCalculateResponse(BaseModel):
    cost: float


# ---------- spend tags ----------


class TagSpend(BaseModel):
    individual_request_tag: str | None = None
    log_count: int | None = None
    total_spend: float | None = None


class SpendTagsResponse(RootModel[list[TagSpend]]):
    """GET /spend/tags answers with a bare array of per-tag aggregates, not an
    object wrapping them (that's /global/spend/tags). Read the rows off .root."""


# ---------- route probing ----------


class DateRangeParams(BaseModel):
    start_date: str
    end_date: str


class RouteSpec(RootModel[dict[str, object]]):
    """One /openapi.json path entry: a map of HTTP method -> operation. Only the
    method names are read, so the operation specs stay opaque."""

    @property
    def methods(self) -> frozenset[str]:
        return frozenset(method.lower() for method in self.root)


class OpenAPISchema(BaseModel):
    paths: dict[str, RouteSpec] = {}


# ---------- model info / custom pricing ----------


class CustomPricing(BaseModel):
    """The per-token custom-pricing fields a deployment can override in
    litellm_params - the token-cost subset of litellm's CustomPricingLiteLLMParams
    the proxy applies to chat spend. All optional: a config sets only what it
    overrides, and /model/info echoes the rates the proxy resolved."""

    model_config = ConfigDict(extra="ignore")
    mode: str | None = None
    input_cost_per_token: float | None = None
    output_cost_per_token: float | None = None
    cache_read_input_token_cost: float | None = None
    cache_creation_input_token_cost: float | None = None

    def overrides(self) -> dict[str, float]:
        """The rates actually declared (non-null) - e.g. those a config.yml sets."""
        declared = {
            "input_cost_per_token": self.input_cost_per_token,
            "output_cost_per_token": self.output_cost_per_token,
            "cache_read_input_token_cost": self.cache_read_input_token_cost,
            "cache_creation_input_token_cost": self.cache_creation_input_token_cost,
        }
        return {field: rate for field, rate in declared.items() if rate is not None}

    def token_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """Spend for a fresh (uncached) call under these rates: the proxy's
        custom-pricing formula (prompt * input + completion * output)."""
        assert self.input_cost_per_token is not None and self.output_cost_per_token is not None, (
            "custom pricing has no per-token rates"
        )
        return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token


class ModelInfoEntry(BaseModel):
    """One /model/info row. `litellm_params` is the configured deployment (carries
    any custom-pricing override); `model_info` is the price the proxy resolved for
    it - the override merged over the cost-map defaults."""

    model_config = ConfigDict(protected_namespaces=())
    model_name: str
    litellm_params: CustomPricing = CustomPricing()
    model_info: CustomPricing = CustomPricing()


class ModelInfoResponse(BaseModel):
    data: list[ModelInfoEntry] = []


class RouterCurrentValues(BaseModel):
    """The `current_values` block of GET /router/settings: the router knobs the
    proxy is actually running with (only the ones a test preconditions on)."""

    optional_pre_call_checks: tuple[str, ...] = ()


class RouterSettingsResponse(BaseModel):
    current_values: RouterCurrentValues


class ConfigListParams(BaseModel):
    config_type: Literal["general_settings"]


class ConfigField(BaseModel):
    """One row of GET /config/list: a general_settings field and the value the
    proxy is running with, the two fields a test preconditions on."""

    model_config = ConfigDict(extra="ignore")
    field_name: str
    field_value: JsonValue = None


class ConfigFieldList(RootModel[tuple[ConfigField, ...]]):
    """GET /config/list answers with a bare array of general_settings fields."""


class CostMapEntry(BaseModel):
    model_config = ConfigDict(extra="ignore")
    litellm_provider: str | None = None
    mode: str | None = None
    deprecation_date: str | None = None
    input_cost_per_token: float | None = None
    output_cost_per_token: float | None = None
    cache_read_input_token_cost: float | None = None
    supports_function_calling: bool | None = None
    supports_reasoning: bool | None = None
    supports_response_schema: bool | None = None


class CostMap(RootModel[dict[str, CostMapEntry]]):
    pass


class BedrockGuardrailParams(BaseModel):
    guardrail: Literal["bedrock"] = "bedrock"
    mode: Literal["pre_call"] = "pre_call"
    default_on: bool = False
    guardrailIdentifier: str
    guardrailVersion: str


class BedrockGuardrailSpec(BaseModel):
    guardrail_name: str
    litellm_params: BedrockGuardrailParams


class GuardrailCreateBody(BaseModel):
    guardrail: BedrockGuardrailSpec


class GuardrailCreateResponse(BaseModel):
    guardrail_id: str


class FileEntry(BaseModel):
    id: str


class FileListResponse(BaseModel):
    """GET /files answer. `data` is required on purpose: a 200 whose body lacks
    the OpenAI-format file list must fail validation, not pass vacuously."""

    data: list[FileEntry]


class FineTuningJobsParams(BaseModel):
    custom_llm_provider: Literal["openai", "azure"]


class FineTuningJobEntry(BaseModel):
    id: str


class FineTuningJobsResponse(BaseModel):
    """GET /fine_tuning/jobs answer; `data` required for the same reason as
    FileListResponse."""

    data: list[FineTuningJobEntry]


# ---------- model management ----------


class CacheControlInjectionPoint(BaseModel):
    location: Literal["message"]
    role: str | None = None
    index: int | None = None


class PromptCacheOptions(BaseModel):
    mode: Literal["implicit", "explicit"]


class LiteLLMParamsBody(BaseModel):
    """POST /model/new litellm_params: `model` is the only required field; `api_key`
    et al may be an `os.environ/FOO` reference the proxy resolves at call time.
    The `*_cost_per_token` / `*_token_cost` fields register a per-deployment custom
    pricing override (the cache and service-tier rates only apply when both base
    rates are set, which is what makes the proxy register the deployment's full
    pricing entry); left None (and dropped from the body) the deployment keeps the
    backend's canonical rate."""

    model: str
    api_key: str | None = Field(default=None, repr=False)
    litellm_credential_name: str | None = None
    api_base: str | None = Field(default=None, repr=False)
    api_version: str | None = None
    responses_api_version: str | None = None
    realtime_protocol: str | None = None
    allowed_openai_params: list[str] | None = None
    drop_params: bool | None = None
    thought_signature_in_tool_call_id: bool | None = None
    aws_access_key_id: str | None = Field(default=None, repr=False)
    aws_secret_access_key: str | None = Field(default=None, repr=False)
    aws_region_name: str | None = None
    aws_bedrock_runtime_endpoint: str | None = Field(default=None, repr=False)
    vertex_project: str | None = None
    vertex_location: str | None = None
    vertex_credentials: str | None = Field(default=None, repr=False)
    gcs_bucket_name: str | None = None
    bucket_name: str | None = None
    s3_bucket_name: str | None = None
    s3_region_name: str | None = None
    s3_access_key_id: str | None = Field(default=None, repr=False)
    s3_secret_access_key: str | None = Field(default=None, repr=False)
    s3_encryption_key_id: str | None = None
    aws_batch_role_arn: str | None = None
    aws_role_name: str | None = None
    aws_session_name: str | None = None
    aws_external_id: str | None = None
    input_cost_per_token: float | None = None
    output_cost_per_token: float | None = None
    cache_read_input_token_cost: float | None = None
    cache_creation_input_token_cost: float | None = None
    input_cost_per_token_priority: float | None = None
    output_cost_per_token_priority: float | None = None
    input_cost_per_token_balanced: float | None = None
    output_cost_per_token_balanced: float | None = None
    cache_read_input_token_cost_balanced: float | None = None
    input_cost_per_token_flex: float | None = None
    output_cost_per_token_flex: float | None = None
    cache_read_input_token_cost_flex: float | None = None
    max_tokens: int | None = None
    max_completion_tokens: int | None = None
    custom_llm_provider: str | None = None
    reasoning_effort: str | None = None
    extra_headers: dict[str, str] | None = Field(default=None, repr=False)
    use_in_pass_through: bool | None = None
    complexity_router_config: dict[str, object] | None = None
    auto_router_config: str | None = None
    auto_router_default_model: str | None = None
    auto_router_embedding_model: str | None = None
    tags: list[str] | None = None
    mock_response: str | list[float] | None = None
    timeout: float | None = None
    max_retries: int | None = None
    cooldown_time: float | None = None
    extra_body: DeploymentExtraBody | None = None
    cache_control_injection_points: list[CacheControlInjectionPoint] | None = None
    prompt_cache_options: PromptCacheOptions | None = None
    tpm: int | None = None
    weight: int | None = None
    order: int | None = None


ModelMode = Literal["chat", "embedding", "batch", "realtime", "image_generation"]


class TokenCountModelBody(BaseModel, frozen=True):
    model: str
    region: str | None = None


class ModelInfoBody(BaseModel):
    enforce_supported_endpoints: bool | None = None
    # id is left unset so the proxy assigns a unique model_id per deployment.
    # Pinning it to the model_name made re-registrations of a fixed-name model
    # (e.g. the batch suite's openai-batch) collide on the model_id unique
    # constraint when a prior run's teardown had not removed the row.
    id: str | None = None
    mode: ModelMode | None = None
    supported_endpoints: list[str] | None = None
    access_groups: list[str] | None = None
    team_id: str | None = None
    allowed_fails: int | None = None
    allowed_fails_policy: dict[str, int] | None = None
    token_count_model: TokenCountModelBody | None = None


class ModelNewBody(BaseModel):
    model_config = ConfigDict(protected_namespaces=())
    model_name: str
    litellm_params: LiteLLMParamsBody
    model_info: ModelInfoBody


class ModelNewResponse(BaseModel):
    model_config = ConfigDict(protected_namespaces=())
    model_id: str


class ModelUpdateBody(BaseModel):
    """POST /model/update body: the target deployment (`model_info.id`) plus the
    `litellm_params` to merge over its stored params. The handler overlays only the
    non-null fields, so a body carrying `input_cost_per_token` re-prices the
    deployment while leaving its other params intact."""

    model_config = ConfigDict(protected_namespaces=())
    litellm_params: LiteLLMParamsBody
    model_info: ModelInfoBody


class ModelListEntry(BaseModel):
    id: str


class ModelsListParams(BaseModel):
    """Query for GET /v1/models. A wildcard route such as ``openai/gpt-5.4*`` is
    listed only under ``return_wildcard_routes``; without it the route is dropped
    and only its expansions remain, so a readiness poll for the pattern itself
    never resolves."""

    return_wildcard_routes: bool = True


class ModelsListResponse(BaseModel):
    """GET /v1/models on the data plane: the deployments the gateway can actually
    serve right now. Used to confirm a freshly created model has propagated from
    the control plane before a test calls it."""

    data: tuple[ModelListEntry, ...] = ()


class ModelDeleteBody(BaseModel):
    id: str


class ConnectionTestBody(BaseModel):
    """POST /health/test_connection body, the API behind the Admin UI's Test
    Connection button: the deployment params as typed into the add-model form and
    the health-check mode picking which endpoint the probe calls. The endpoint
    rejects `os.environ/` references, so credentials are either literal values or
    omitted to fall through to the proxy's own environment."""

    litellm_params: LiteLLMParamsBody
    mode: Literal["chat", "completion", "embedding", "responses"]


class ConnectionTestResult(BaseModel):
    error: str | None = None


class ConnectionTestResponse(BaseModel):
    status: Literal["success", "error"]
    result: ConnectionTestResult | None = None


class CredentialCreateBody(BaseModel):
    credential_name: str
    credential_values: dict[str, str] = Field(repr=False)
    credential_info: dict[str, str] = {}


class CredentialCreateResponse(BaseModel):
    success: bool


# ---------- key / team / user / organization management ----------


class Cleared(BaseModel):
    """An explicit JSON null in a merge-patch body. The transport drops `None` fields
    before sending (`exclude_none`), so `None` means "leave the stored value alone"; a
    field set to `CLEAR` reaches the wire as `null`, which tells the proxy to clear it."""

    model_config = ConfigDict(frozen=True)

    @model_serializer
    def _as_null(self) -> None:
        return None


CLEAR: Final = Cleared()


class KeyUpdateBody(BaseModel):
    """POST /key/update is a merge patch: a field left `None` is dropped from the body and
    keeps its stored value, `CLEAR` sends an explicit null that clears it (`budget_duration`
    clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale."""

    key: str = Field(repr=False)
    project_id: str | Cleared | None = None
    models: list[str] | None = None
    key_alias: str | None = None
    tpm_limit: int | None = None
    rpm_limit: int | None = None
    max_budget: float | Cleared | None = None
    budget_duration: str | Cleared | None = None
    metadata: KeyMetadata | None = None


class KeyBlockBody(BaseModel):
    key: str


class KeyListParams(BaseModel):
    key_alias: str


class KeyListResponse(BaseModel):
    total_count: int


# ---------- admin UI session ----------


class UiLoginBody(BaseModel):
    username: str
    password: str


class UiLoginResponse(BaseModel):
    token: str = Field(repr=False)
    redirect_url: str


class UiSessionClaims(BaseModel):
    user_id: str
    key: str = Field(repr=False)
    user_role: str
    login_method: Literal["sso", "username_password"]
    exp: int


class TeamMemberEntry(BaseModel):
    role: Literal["admin", "user"]
    user_id: str


class TeamMetadata(BaseModel):
    disable_global_guardrails: bool | None = None


class TeamNewBody(BaseModel):
    team_alias: str
    models: list[str] = []
    team_id: str | None = None
    organization_id: str | None = None
    metadata: TeamMetadata | None = None
    model_aliases: dict[str, str] | None = None


class TeamNewResponse(BaseModel):
    team_id: str


class TeamUpdateBody(BaseModel):
    team_id: str
    team_alias: str | None = None
    models: list[str] | None = None
    object_permission: ObjectPermission | None = None


class TeamInfoParams(BaseModel):
    team_id: str


class TeamData(BaseModel):
    organization_id: str | None = None
    team_alias: str | None = None
    models: list[str] = []
    members_with_roles: list[TeamMemberEntry] = []


class TeamInfoResponse(BaseModel):
    team_id: str
    team_info: TeamData


class TeamMemberAddBody(BaseModel):
    team_id: str
    member: TeamMemberEntry | list[TeamMemberEntry]


class TeamMemberDeleteBody(BaseModel):
    team_id: str
    user_id: str


class TeamDeleteBody(BaseModel):
    team_ids: list[str]


class ProjectCreateBody(BaseModel):
    team_id: str
    project_alias: str
    models: list[str]


class ProjectIdentity(BaseModel):
    project_id: str


class ProjectDeleteBody(BaseModel):
    project_ids: list[str]


class TeamListEntry(BaseModel):
    team_id: str


class TeamListResponse(RootModel[list[TeamListEntry]]):
    """GET /team/list answers with a bare array of team objects (not an object
    wrapping them). Only team_id is read; pydantic ignores the rest."""


UserRole = Literal["proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer"]


class UserNewBody(BaseModel):
    user_email: str
    user_role: UserRole
    user_id: str | None = None
    auto_create_key: bool | None = None
    max_budget: float | None = None


class UserNewResponse(BaseModel):
    user_id: str
    key: str | None = None


class UserUpdateBody(BaseModel):
    user_id: str
    user_role: UserRole


class UserInfoParams(BaseModel):
    user_id: str | None = None


class UserData(BaseModel):
    user_id: str | None = None
    user_email: str | None = None
    user_role: str | None = None


class UserInfoResponse(BaseModel):
    user_id: str
    user_info: UserData


class UserDeleteBody(BaseModel):
    user_ids: list[str]


class UserDeleteResponse(RootModel[int]):
    pass


class UserListParams(BaseModel):
    user_ids: str


class UserListRow(BaseModel):
    user_id: str


class UserListResponse(BaseModel):
    users: list[UserListRow]
    total: int


class UserKeyRow(BaseModel):
    token: str
    key_alias: str | None = None


class UserInfoWithKeysResponse(BaseModel):
    user_id: str | None = None
    keys: list[UserKeyRow] = []


class JwtKeyMappingRow(BaseModel):
    id: str
    jwt_claim_name: str
    jwt_claim_value: str
    created_by: str | None = None


class JwtKeyMappingListParams(BaseModel):
    size: int = 100


class JwtKeyMappingListResponse(BaseModel):
    mappings: list[JwtKeyMappingRow]
    total_count: int


class JwtKeyMappingDeleteBody(BaseModel):
    id: str


class JwtKeyMappingDeleteResponse(BaseModel):
    status: str


class OrgNewBody(BaseModel):
    organization_alias: str
    models: list[str] = []


class OrgNewResponse(BaseModel):
    organization_id: str


class OrgUpdateBody(BaseModel):
    organization_id: str
    organization_alias: str


class OrgInfoParams(BaseModel):
    organization_id: str


class OrgMembership(BaseModel):
    user_id: str
    user_role: str


class OrgInfoResponse(BaseModel):
    organization_id: str
    organization_alias: str | None = None
    models: list[str] = []
    members: tuple[OrgMembership, ...] = ()


class OrgMemberEntry(BaseModel):
    user_id: str
    role: Literal["org_admin", "internal_user"]


class OrgMemberAddBody(BaseModel):
    organization_id: str
    member: OrgMemberEntry


class OrgDeleteBody(BaseModel):
    organization_ids: list[str]


class OrgDeleteResponse(RootModel[tuple[OrgInfoResponse, ...]]):
    pass


# ---------- tags (management) ----------


class TagNewBody(BaseModel):
    name: str
    description: str | None = None


class TagDeleteBody(BaseModel):
    name: str


class TagListEntry(BaseModel):
    name: str
    description: str | None = None


class TagListResponse(RootModel[list[TagListEntry]]):
    """GET /tag/list answers with a bare array of tag configs (the stored tags plus
    any dynamically-seen spend tags), not an object wrapping them. Read the rows off
    .root."""


# ---------- health / lifecycle ----------


class ProcessMemory(BaseModel):
    ram_usage_mb: float | None = None
    system_memory_percent: float | None = None
    error: str | None = None


class MemorySummaryResponse(BaseModel):
    worker_pid: int
    hostname: str | None = None
    status: str
    memory: ProcessMemory


class ReadinessResponse(BaseModel):
    """GET /health/readiness (public probe). The low-detail payload a load
    balancer sees: `status` plus the resolved DB state (`connected`,
    `disconnected`, or `Not connected`)."""

    status: str
    db: str | None = None


class ReadinessDetailsResponse(ReadinessResponse):
    """GET /health/readiness/details (authenticated). Extends the public payload
    with the diagnostics only an authenticated caller may read."""

    litellm_version: str | None = None
    success_callbacks: list[str] = []