| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """vLLM chat-protocol extensions for grounded (citation-carrying) models. |
| 4 | |
| 5 | This module keeps the OpenAI chat completion protocol classes |
| 6 | (:class:`ChatMessage` / :class:`DeltaMessage`) free of Cohere-specific fields. |
| 7 | It provides two things: |
| 8 | |
| 9 | * :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation |
| 10 | representation produced by the Cohere parser |
| 11 | (:mod:`vllm.parser.cohere_command`) and consumed by |
| 12 | :mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere |
| 13 | SDK shape -- that conversion happens in the serving layer. |
| 14 | * :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage` |
| 15 | / :class:`DeltaMessage` subclasses that add a ``citations`` field. They |
| 16 | are only instantiated by the citation-aware serving handler |
| 17 | (:class:`vllm.entrypoints.cohere.serving.CohereServingChatV2`) and by the |
| 18 | Cohere reasoning parser's streaming path. Non-Cohere code paths continue |
| 19 | to construct plain :class:`ChatMessage` / :class:`DeltaMessage`. |
| 20 | |
| 21 | The module intentionally does *not* import the ``cohere`` Python SDK so the |
| 22 | reasoning parser (which is Cohere-model-specific but ships without the SDK |
| 23 | dependency) can import from here freely. |
| 24 | """ |
| 25 | |
| 26 | from __future__ import annotations |
| 27 | |
| 28 | from typing import Any, Literal |
| 29 | |
| 30 | from pydantic import Field, model_serializer |
| 31 | |
| 32 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 33 | from vllm.entrypoints.openai.chat_completion.protocol import ChatMessage |
| 34 | from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel |
| 35 | |
| 36 | |
| 37 | class CitationSource(OpenAIBaseModel): |
| 38 | """Source attribution for a :class:`Citation`. |
| 39 | |
| 40 | Mirrors the shape used by Cohere's Chat v2 API. ``type`` is the source |
| 41 | discriminator (``document`` or ``tool``); ``id`` is the citing |
| 42 | document/tool-output identifier; ``document`` and ``tool_output`` carry |
| 43 | the original payload that produced the citation. |
| 44 | |
| 45 | Sources are fully resolved by the parser at emit time |
| 46 | (see :func:`_melody_sources_to_vllm` in |
| 47 | :mod:`vllm.parser.cohere_command`) using the |
| 48 | ``POSITION_TO_SOURCE_KEY`` map forwarded through |
| 49 | ``chat_template_kwargs`` by |
| 50 | :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`. |
| 51 | That means every instance that reaches the serving / wire layer |
| 52 | already has ``type`` / ``id`` populated; the serving layer only has |
| 53 | to coerce to :class:`cohere.types.Citation` and rewrite |
| 54 | ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning models. |
| 55 | """ |
| 56 | |
| 57 | type: Literal["document", "tool"] | None = None |
| 58 | id: str | None = None |
| 59 | document: dict[str, Any] | None = None |
| 60 | tool_output: dict[str, Any] | None = None |
| 61 | |
| 62 | |
| 63 | class Citation(OpenAIBaseModel): |
| 64 | """A citation grounding a span of generated text in source material. |
| 65 | |
| 66 | vLLM-internal representation used by the Cohere reasoning parser and |
| 67 | the Cohere v2 serving layer. This is not the on-wire Cohere SDK shape: |
| 68 | conversion to the SDK's ``cohere.types.Citation`` happens in |
| 69 | :mod:`vllm.entrypoints.cohere.serving`. |
| 70 | """ |
| 71 | |
| 72 | start: int | None = None |
| 73 | """Start character offset in the surrounding text content.""" |
| 74 | end: int | None = None |
| 75 | """End character offset (exclusive) in the surrounding text content.""" |
| 76 | text: str | None = None |
| 77 | """The cited text snippet.""" |
| 78 | sources: list[CitationSource] = Field(default_factory=list) |
| 79 | """Source documents / tool outputs that ground this citation.""" |
| 80 | content_index: int | None = None |
| 81 | """Index of the content block this citation refers to (when the message |
| 82 | has multiple content blocks).""" |
| 83 | type: Literal["TEXT_CONTENT", "THINKING_CONTENT", "PLAN"] | None = None |
| 84 | """Which kind of content block this citation grounds: the user-visible |
| 85 | text (``TEXT_CONTENT``), a thinking block (``THINKING_CONTENT``), or a |
| 86 | tool-plan block (``PLAN``). ``None`` means unspecified.""" |
| 87 | |
| 88 | |
| 89 | class CohereChatMessage(ChatMessage): |
| 90 | """:class:`ChatMessage` extension carrying grounding citations. |
| 91 | |
| 92 | Only instantiated by :class:`CohereServingChatV2` (and any other |
| 93 | :class:`OpenAIServingChat` subclass that opts in by overriding |
| 94 | ``_create_chat_message``). Regular OpenAI-compatible handlers keep |
| 95 | emitting plain :class:`ChatMessage` so their response schema is |
| 96 | unchanged. |
| 97 | |
| 98 | The response envelope declares ``message: SerializeAsAny[ChatMessage]``, |
| 99 | so pydantic serializes this subclass with its own schema (including |
| 100 | ``citations``) when it flows through |
| 101 | :class:`ChatCompletionResponseChoice`. |
| 102 | """ |
| 103 | |
| 104 | citations: list[Citation] | None = None |
| 105 | |
| 106 | @model_serializer(mode="wrap") |
| 107 | def _serialize(self, handler): |
| 108 | # ``mode="wrap"`` fully overrides (rather than chains) the |
| 109 | # parent's ``@model_serializer``, so we explicitly delegate via |
| 110 | # ``super()._serialize(handler)`` to preserve the parent's |
| 111 | # cleanup (e.g. stripping empty ``tool_calls``). Then we drop an |
| 112 | # unset ``citations`` field so the wire matches the OpenAI-style |
| 113 | # contract that optional vLLM extensions are omitted rather than |
| 114 | # serialized as ``null``. |
| 115 | data = super()._serialize(handler) |
| 116 | if not data.get("citations"): |
| 117 | data.pop("citations", None) |
| 118 | return data |
| 119 | |
| 120 | |
| 121 | class CohereDeltaMessage(DeltaMessage): |
| 122 | """:class:`DeltaMessage` extension carrying grounding citations for streaming. |
| 123 | |
| 124 | Emitted by :class:`vllm.parser.cohere_command.CohereCommandParser` |
| 125 | on delta events whose payload includes citations. Non-Cohere parsers |
| 126 | return plain :class:`DeltaMessage`, so their streamed shape is |
| 127 | unchanged. |
| 128 | |
| 129 | The response envelope declares |
| 130 | ``delta: SerializeAsAny[DeltaMessage]``, so pydantic serializes this |
| 131 | subclass with its own schema when it flows through |
| 132 | :class:`ChatCompletionResponseStreamChoice`. |
| 133 | """ |
| 134 | |
| 135 | citations: list[Citation] | None = None |
| 136 | |
| 137 | @model_serializer(mode="wrap") |
| 138 | def _serialize(self, handler): |
| 139 | # See ``CohereChatMessage._serialize`` for why we delegate via |
| 140 | # ``super()`` instead of calling ``handler(self)`` directly. |
| 141 | data = super()._serialize(handler) |
| 142 | if not data.get("citations"): |
| 143 | data.pop("citations", None) |
| 144 | return data |