vllm.entrypoints.cohere.protocol ¶

class CohereChatV2Request(BaseModel):
    """Cohere Chat v2 request body.

    Mirrors the schema documented at https://docs.cohere.com/reference/chat.
    All structured fields delegate to the official SDK types so the body
    schema stays in sync with the upstream spec.
    """

    model: str
    messages: list[ChatMessageV2]
    stream: bool | None = False

    # Tooling
    tools: list[ToolV2] | None = None
    strict_tools: bool | None = None
    tool_choice: CohereToolChoice | None = None

    # Grounding
    documents: list[str | Document] | None = None
    citation_options: CitationOptions | None = None

    # Output
    response_format: ResponseFormatV2 | None = None
    safety_mode: ChatRequestSafetyMode | None = None
    max_tokens: int | None = None
    stop_sequences: list[str] | None = None

    # Sampling
    temperature: float | None = None
    seed: int | None = None
    frequency_penalty: float | None = None
    presence_penalty: float | None = None
    k: int | None = None
    p: float | None = None
    logprobs: bool | None = None

    # Reasoning
    thinking: Thinking | None = None

    # Scheduling
    priority: int | None = None

    # vLLM-specific extensions (not in Cohere spec). These mirror what the
    # Anthropic and OpenAI surfaces already expose so V2 callers can reach
    # the same engine knobs when needed.
    kv_transfer_params: dict[str, Any] | None = Field(
        default=None,
        description="KVTransfer parameters used for disaggregated serving.",
    )
    chat_template_kwargs: dict[str, Any] | None = Field(
        default=None,
        description=(
            "Additional keyword args to pass to the chat template renderer. "
            "Will be accessible by the template."
        ),
    )

    @field_validator("model")
    @classmethod
    def _validate_model(cls, v: str) -> str:
        if not v:
            raise ValueError("model is required")
        return v

    @field_validator("max_tokens")
    @classmethod
    def _validate_max_tokens(cls, v: int | None) -> int | None:
        if v is not None and v < 0:
            raise ValueError("max_tokens must be non-negative")
        return v

    @field_validator("messages", mode="before")
    @classmethod
    def _normalize_message_roles(cls, v: Any) -> Any:
        """Rewrite OpenAI-style ``developer`` roles to ``system``.

        Cohere's v2 ``ChatMessageV2`` discriminated union only admits
        the four literal roles ``user`` / ``assistant`` / ``system`` /
        ``tool``. OpenAI's ``developer`` role is documented as
        high-priority system instructions, so we alias it onto
        ``system`` *before* the SDK discriminator runs (otherwise it
        rejects the message with a ``literal_error`` against each union
        member). Mirrors ``_role_to_melody`` in the renderer so the
        same alias is honoured no matter which surface the message
        arrives through.

        ``mode="before"`` is required so the rewrite happens before the
        ``list[ChatMessageV2]`` coercion runs the SDK's discriminated
        union; a default-mode validator would never see ``developer``
        because validation would have already failed. On any
        structural malformation (non-iterable input, items without a
        dict-shaped ``role`` field, etc.) we hand ``v`` back unchanged
        and let Pydantic's normal coercion surface a precise error.
        """
        try:
            return [
                {**msg, "role": "system"}
                if msg.get("role", "").lower() == "developer"
                else msg
                for msg in v
            ]
        except (AttributeError, TypeError):
            return v

    @field_validator("messages")
    @classmethod
    def _validate_messages(cls, v: list[ChatMessageV2]) -> list[ChatMessageV2]:
        if not v:
            raise ValueError("messages must contain at least one message")
        return v

class CohereChatV2Response(BaseModel):
    """Cohere Chat v2 non-streaming response body.

    Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in
    the documented v2 response envelope (``id``, ``finish_reason``,
    ``usage``, ``logprobs``). The single constructor in
    :class:`CohereServingChatV2._chat_completion_to_v2` is responsible for
    supplying a non-empty ``id`` (falling back to a synthesized one if
    the upstream response is missing it) to this model.
    """

    id: str
    finish_reason: CohereFinishReason
    message: AssistantMessageResponse
    usage: CohereUsage | None = None
    logprobs: list[CohereLogprobItem] | None = None

    # vLLM-specific extension.
    kv_transfer_params: dict[str, Any] | None = Field(
        default=None, description="KVTransfer parameters."
    )

class CohereError(BaseModel):
    """Top-level error body returned by ``/cohere/v2/chat`` error responses.

    Cohere's documented error schemas are uniform: ``{message, id}``.
    """

    message: str
    id: str | None = None