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 Literal

from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, 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 = None
    langfuse_secret_key: str | None = None
    langfuse_host: str | None = None
    wandb_api_key: str | None = None
    weave_project_id: str | None = None


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


class ObjectPermission(BaseModel):
    mcp_servers: list[str] | None = None
    mcp_access_groups: list[str] | 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
    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


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 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
    metadata: KeyMetadata | None = None
    models: list[str] = []
    tpm_limit: int | None = None
    rpm_limit: int | None = None
    team_id: str | None = None
    blocked: bool | None = None
    spend: float | None = None
    max_budget: float | 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


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


ContentPart = TextContentPart | ImageContentPart


class ChatMessage(BaseModel):
    role: str
    content: str | list[ContentPart]


class CacheControl(BaseModel):
    type: str = "ephemeral"


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 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 ToolCallFunction(BaseModel):
    name: str | None = None
    arguments: str | None = None


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


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


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


type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn


class ChatBody(BaseModel):
    model: str
    messages: Sequence[ChatTurn]
    stream: bool = False
    max_tokens: int | None = None
    max_completion_tokens: int | None = None
    temperature: float | None = None
    user: 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
    tools: Sequence[ChatTool | McpChatTool] | None = None
    tool_choice: str | None = None
    guardrails: list[str] | 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}


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 and retry 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
    enable_tag_filtering: bool | 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
    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


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


# ---------- 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 AnthropicMessagesBody(BaseModel):
    model: str
    messages: list[AnthropicMessage]
    max_tokens: int
    stream: bool | None = None
    tools: list[AnthropicTool] | None = None
    guardrails: list[str] | None = None
    cache: dict[str, bool] | None = {"no-cache": True}


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


class CountTokensResponse(BaseModel):
    """`/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


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


class McpServerCreateBody(BaseModel):
    """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
    `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
    are discovered and registered via DCR when left unset. `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
    auth_type: str | None = None
    oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None
    authorization_url: str | None = None
    token_url: 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 EmbedBody(BaseModel):
    model: str
    input: str
    cache: dict[str, bool] | None = {"no-cache": True}


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


# ---------- 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 OcrPage(BaseModel):
    index: int
    markdown: str


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


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


class SpendLogRow(BaseModel):
    request_id: str | None = None
    api_key: str | None = None
    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
    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


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


class SpendLogsParams(BaseModel):
    request_id: str | None = None
    api_key: str | None = None

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


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 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 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 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 `_priority` 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 = None
    litellm_credential_name: str | None = None
    api_base: str | None = None
    api_version: str | None = None
    realtime_protocol: str | None = None
    aws_access_key_id: str | None = None
    aws_secret_access_key: str | None = None
    aws_region_name: str | None = None
    vertex_project: str | None = None
    vertex_location: str | None = None
    vertex_credentials: str | None = None
    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 = None
    s3_secret_access_key: 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
    extra_headers: dict[str, str] | None = None
    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 | None = None
    timeout: float | None = None
    tpm: int | None = None
    weight: int | None = None


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


class ModelInfoBody(BaseModel):
    # 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
    access_groups: list[str] | None = None
    team_id: str | None = None
    allowed_fails_policy: dict[str, int] | 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]
    credential_info: dict[str, str] = {}


class CredentialCreateResponse(BaseModel):
    success: bool


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


class KeyUpdateBody(BaseModel):
    key: str
    models: list[str] | None = None
    key_alias: str | None = None
    tpm_limit: int | None = None
    rpm_limit: int | 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
    redirect_url: str


class UiSessionClaims(BaseModel):
    user_id: str
    key: str
    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


class TeamNewResponse(BaseModel):
    team_id: str


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


class TeamInfoParams(BaseModel):
    team_id: str


class TeamData(BaseModel):
    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


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


class TeamDeleteBody(BaseModel):
    team_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


class UserNewResponse(BaseModel):
    user_id: str


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


class UserInfoParams(BaseModel):
    user_id: str


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 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 OrgInfoResponse(BaseModel):
    organization_id: str
    organization_alias: str | None = None
    models: list[str] = []


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


# ---------- 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 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] = []