374 lines
12 KiB
Python
374 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Cohere Chat v2 API protocol.
|
|
|
|
The bulk of the wire types come straight from the official ``cohere``
|
|
Python SDK so we stay in sync with the upstream specification and
|
|
avoid re-declaring the schema. We only own three things locally:
|
|
|
|
1. The top-level request body model (the SDK doesn't ship one — its
|
|
``ClientV2.chat`` takes the body as kwargs), with vLLM-specific
|
|
extensions (``kv_transfer_params`` / ``chat_template_kwargs``).
|
|
2. The non-streaming response envelope (the SDK exposes the message
|
|
shape via :class:`AssistantMessageResponse` but no full response
|
|
wrapper).
|
|
3. The streaming discriminated union (the SDK exports each event type
|
|
individually but not as a combined ``Annotated[Union[...],
|
|
discriminator]``).
|
|
|
|
Importing this module pulls in the ``cohere`` package. The router that
|
|
mounts ``POST /cohere/v2/chat`` guards on that import succeeding so vLLM still
|
|
boots without the SDK installed.
|
|
|
|
See https://docs.cohere.com/reference/chat for the upstream spec.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Annotated, Any, Literal
|
|
|
|
from cohere import types as _sdk
|
|
from cohere.types import (
|
|
AssistantChatMessageV2,
|
|
AssistantMessageResponse,
|
|
ChatMessageV2,
|
|
ChatRequestSafetyMode,
|
|
Citation,
|
|
CitationOptions,
|
|
Document,
|
|
ResponseFormatV2,
|
|
SystemChatMessageV2,
|
|
Thinking,
|
|
ToolCallV2,
|
|
ToolChatMessageV2,
|
|
ToolV2,
|
|
UserChatMessageV2,
|
|
)
|
|
from pydantic import BaseModel, Field, field_validator
|
|
|
|
import vllm.envs as envs
|
|
from vllm.exceptions import VLLMValidationError
|
|
|
|
# Re-export the SDK wire-format types alongside our local extensions so
|
|
# ``vllm.entrypoints.cohere.serving`` and friends can import everything
|
|
# they need from this module.
|
|
__all__ = [
|
|
"AssistantChatMessageV2",
|
|
"AssistantMessageResponse",
|
|
"ChatMessageV2",
|
|
"Citation",
|
|
"CitationEndEvent",
|
|
"CitationOptions",
|
|
"CitationStartEvent",
|
|
"CohereChatV2Request",
|
|
"CohereChatV2Response",
|
|
"CohereError",
|
|
"CohereFinishReason",
|
|
"CohereLogprobItem",
|
|
"CohereUsage",
|
|
"CohereUsageBilledUnits",
|
|
"CohereUsageTokens",
|
|
"ContentDeltaEvent",
|
|
"ContentEndEvent",
|
|
"ContentStartEvent",
|
|
"Document",
|
|
"MessageEndEvent",
|
|
"MessageStartEvent",
|
|
"ResponseFormatV2",
|
|
"SystemChatMessageV2",
|
|
"Thinking",
|
|
"ToolCallDeltaEvent",
|
|
"ToolCallEndEvent",
|
|
"ToolCallStartEvent",
|
|
"ToolCallV2",
|
|
"ToolChatMessageV2",
|
|
"ToolPlanDeltaEvent",
|
|
"ToolV2",
|
|
"UserChatMessageV2",
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Errors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tool choice / finish reasons
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# These literals aren't first-class enums in the SDK but are documented at
|
|
# https://docs.cohere.com/reference/chat. We declare them here so the
|
|
# request/response models can validate them.
|
|
|
|
CohereToolChoice = Literal["REQUIRED", "NONE"]
|
|
|
|
CohereFinishReason = Literal[
|
|
"COMPLETE",
|
|
"STOP_SEQUENCE",
|
|
"MAX_TOKENS",
|
|
"TOOL_CALL",
|
|
"ERROR",
|
|
"TIMEOUT",
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
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: (
|
|
Annotated[list[str], Field(max_length=envs.VLLM_MAX_STOP_STRINGS)] | 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 VLLMValidationError("model is required", parameter="model")
|
|
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 VLLMValidationError(
|
|
"max_tokens must be non-negative",
|
|
parameter="max_tokens",
|
|
value=v,
|
|
)
|
|
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 VLLMValidationError(
|
|
"messages must contain at least one message", parameter="messages"
|
|
)
|
|
return v
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Usage / Logprobs
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# The Cohere SDK only exposes a v1 ``ApiMetaBilledUnits``. The v2 usage
|
|
# envelope is documented separately in the OpenAPI spec and we declare it
|
|
# here.
|
|
|
|
|
|
class CohereUsageBilledUnits(BaseModel):
|
|
input_tokens: float | None = None
|
|
output_tokens: float | None = None
|
|
search_units: float | None = None
|
|
classifications: float | None = None
|
|
|
|
|
|
class CohereUsageTokens(BaseModel):
|
|
input_tokens: float | None = None
|
|
output_tokens: float | None = None
|
|
|
|
|
|
class CohereUsage(BaseModel):
|
|
billed_units: CohereUsageBilledUnits | None = None
|
|
tokens: CohereUsageTokens | None = None
|
|
cached_tokens: float | None = None
|
|
|
|
|
|
class CohereLogprobItem(BaseModel):
|
|
text: str | None = None
|
|
token_ids: list[int]
|
|
logprobs: list[float] | None = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Non-streaming response
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
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."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Streaming events
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Cohere V2 streams a sequence of typed JSON events delivered as Server-
|
|
# Sent Events; each event's ``type`` field carries the discriminator. The
|
|
# SDK exposes a Pydantic model per event but none of them declare ``type``
|
|
# as a field (the SDK relies on its own deserializer for discrimination),
|
|
# so a naive ``model_dump_json()`` would silently drop the discriminator
|
|
# and break clients that demux on ``type``.
|
|
#
|
|
# We therefore subclass each SDK event and bake the wire-format ``type``
|
|
# string in as a ``Literal`` field with a default. ``model_dump()`` now
|
|
# emits ``type`` for free, and surfaces can simply construct the event
|
|
# class and serialize it -- no manual ``type`` parameter needed.
|
|
#
|
|
# See https://docs.cohere.com/v2/docs/streaming and the OpenAPI
|
|
# ``StreamedChatResponseV2`` schema for the wire-format reference.
|
|
|
|
|
|
class MessageStartEvent(_sdk.ChatMessageStartEvent):
|
|
type: Literal["message-start"] = "message-start"
|
|
|
|
|
|
class ContentStartEvent(_sdk.ChatContentStartEvent):
|
|
type: Literal["content-start"] = "content-start"
|
|
|
|
|
|
class ContentDeltaEvent(_sdk.ChatContentDeltaEvent):
|
|
type: Literal["content-delta"] = "content-delta"
|
|
|
|
|
|
class ContentEndEvent(_sdk.ChatContentEndEvent):
|
|
type: Literal["content-end"] = "content-end"
|
|
|
|
|
|
class ToolPlanDeltaEvent(_sdk.ChatToolPlanDeltaEvent):
|
|
type: Literal["tool-plan-delta"] = "tool-plan-delta"
|
|
|
|
|
|
class ToolCallStartEvent(_sdk.ChatToolCallStartEvent):
|
|
type: Literal["tool-call-start"] = "tool-call-start"
|
|
|
|
|
|
class ToolCallDeltaEvent(_sdk.ChatToolCallDeltaEvent):
|
|
type: Literal["tool-call-delta"] = "tool-call-delta"
|
|
|
|
|
|
class ToolCallEndEvent(_sdk.ChatToolCallEndEvent):
|
|
type: Literal["tool-call-end"] = "tool-call-end"
|
|
|
|
|
|
class CitationStartEvent(_sdk.CitationStartEvent):
|
|
type: Literal["citation-start"] = "citation-start"
|
|
|
|
|
|
class CitationEndEvent(_sdk.CitationEndEvent):
|
|
type: Literal["citation-end"] = "citation-end"
|
|
|
|
|
|
class MessageEndEvent(_sdk.ChatMessageEndEvent):
|
|
type: Literal["message-end"] = "message-end"
|