@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 250 lines 8.47 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for ``vllm/entrypoints/cohere/protocol.py``.
4
5 The module is mostly a thin wrapper around the official ``cohere`` SDK
6 types plus a few local additions:
7
8 * :class:`CohereError` envelope.
9 * :class:`CohereChatV2Request` (model and messages required,
10 ``max_tokens`` non-negative).
11 * :class:`CohereChatV2Response` plus the usage / logprob helpers.
12 * The streaming event subclasses that bake a wire-format ``type``
13 discriminator into ``model_dump()`` so SSE consumers can demux on it.
14 """
15
16 import pytest
17 from pydantic import ValidationError
18
19 import vllm.envs as envs
20 from vllm.entrypoints.cohere.protocol import (
21 AssistantMessageResponse,
22 CitationEndEvent,
23 CitationStartEvent,
24 CohereChatV2Request,
25 CohereChatV2Response,
26 CohereError,
27 CohereUsageTokens,
28 ContentDeltaEvent,
29 ContentEndEvent,
30 ContentStartEvent,
31 MessageEndEvent,
32 MessageStartEvent,
33 ToolCallDeltaEvent,
34 ToolCallEndEvent,
35 ToolCallStartEvent,
36 ToolPlanDeltaEvent,
37 )
38 from vllm.exceptions import VLLMValidationError
39
40 # ======================================================================
41 # CohereError
42 # ======================================================================
43
44
45 class TestCohereError:
46 def test_model_dump_excludes_none(self):
47 err = CohereError(message="boom")
48 assert err.model_dump(exclude_none=True) == {"message": "boom"}
49
50
51 # ======================================================================
52 # CohereChatV2Request
53 # ======================================================================
54
55
56 class TestCohereChatV2Request:
57 def test_empty_model_rejected(self):
58 with pytest.raises(VLLMValidationError, match="model is required") as exc_info:
59 CohereChatV2Request(model="", messages=[{"role": "user", "content": "hi"}])
60 assert exc_info.value.parameter == "model"
61 assert exc_info.value.value is None
62
63 def test_negative_max_tokens_rejected(self):
64 with pytest.raises(VLLMValidationError, match="non-negative") as exc_info:
65 CohereChatV2Request(
66 model="m",
67 messages=[{"role": "user", "content": "hi"}],
68 max_tokens=-1,
69 )
70 assert exc_info.value.parameter == "max_tokens"
71 assert exc_info.value.value == -1
72
73 def test_empty_messages_rejected(self):
74 with pytest.raises(
75 VLLMValidationError, match="messages must contain at least one message"
76 ) as exc_info:
77 CohereChatV2Request(model="m", messages=[])
78 assert exc_info.value.parameter == "messages"
79 assert exc_info.value.value is None
80
81 def test_zero_max_tokens_allowed(self):
82 # Zero is allowed (the docs allow 0 -> return prompt only).
83 req = CohereChatV2Request(
84 model="m",
85 messages=[{"role": "user", "content": "hi"}],
86 max_tokens=0,
87 )
88 assert req.max_tokens == 0
89
90 def test_invalid_tool_choice_rejected(self):
91 with pytest.raises(ValidationError):
92 CohereChatV2Request(
93 model="m",
94 messages=[{"role": "user", "content": "hi"}],
95 tool_choice="ANY", # not REQUIRED/NONE
96 )
97
98 def test_stop_sequences_bounded_by_env_limit(self):
99 messages = [{"role": "user", "content": "hi"}]
100 limit = envs.VLLM_MAX_STOP_STRINGS
101
102 req = CohereChatV2Request(
103 model="m",
104 messages=messages,
105 stop_sequences=[f"s{i}" for i in range(limit)],
106 )
107 assert len(req.stop_sequences) == limit
108
109 with pytest.raises(ValidationError):
110 CohereChatV2Request(
111 model="m",
112 messages=messages,
113 stop_sequences=[f"s{i}" for i in range(limit + 1)],
114 )
115
116
117 # ======================================================================
118 # Usage / Logprobs
119 # ======================================================================
120
121
122 class TestUsage:
123 def test_tokens_serialization(self):
124 u = CohereUsageTokens(input_tokens=10, output_tokens=5)
125 assert u.model_dump() == {"input_tokens": 10.0, "output_tokens": 5.0}
126
127
128 # ======================================================================
129 # CohereChatV2Response
130 # ======================================================================
131
132
133 class TestCohereChatV2Response:
134 def test_invalid_finish_reason_rejected(self):
135 msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}])
136 with pytest.raises(ValidationError):
137 CohereChatV2Response(
138 id="r1",
139 finish_reason="NOT_A_REASON",
140 message=msg,
141 )
142
143
144 # ======================================================================
145 # Streaming event ``type`` discriminator baked into model_dump()
146 # ======================================================================
147
148
149 class TestStreamingEventTypeField:
150 """Each event subclass adds a ``type: Literal[...]`` field with a
151 default so ``model_dump()`` always emits the wire-format discriminator
152 (the parent SDK classes don't declare ``type`` as a Pydantic field).
153 """
154
155 @pytest.mark.parametrize(
156 "cls, expected_type, kwargs",
157 [
158 (
159 MessageStartEvent,
160 "message-start",
161 {"id": "a", "delta": {"message": {"role": "assistant"}}},
162 ),
163 (
164 ContentStartEvent,
165 "content-start",
166 {
167 "index": 0,
168 "delta": {"message": {"content": {"type": "text", "text": ""}}},
169 },
170 ),
171 (
172 ContentDeltaEvent,
173 "content-delta",
174 {
175 "index": 0,
176 "delta": {"message": {"content": {"text": "hi"}}},
177 },
178 ),
179 (ContentEndEvent, "content-end", {"index": 0}),
180 (
181 ToolPlanDeltaEvent,
182 "tool-plan-delta",
183 {"delta": {"message": {"tool_plan": "thinking"}}},
184 ),
185 (
186 ToolCallStartEvent,
187 "tool-call-start",
188 {
189 "index": 0,
190 "delta": {
191 "message": {
192 "tool_calls": {
193 "id": "c1",
194 "type": "function",
195 "function": {"name": "f", "arguments": ""},
196 }
197 }
198 },
199 },
200 ),
201 (
202 ToolCallDeltaEvent,
203 "tool-call-delta",
204 {
205 "index": 0,
206 "delta": {
207 "message": {"tool_calls": {"function": {"arguments": "{}"}}}
208 },
209 },
210 ),
211 (ToolCallEndEvent, "tool-call-end", {"index": 0}),
212 (
213 CitationStartEvent,
214 "citation-start",
215 {
216 "index": 0,
217 "delta": {
218 "message": {
219 "citations": {"start": 0, "end": 5, "text": "hello"}
220 }
221 },
222 },
223 ),
224 (CitationEndEvent, "citation-end", {"index": 0}),
225 (
226 MessageEndEvent,
227 "message-end",
228 {"id": "a", "delta": {"finish_reason": "COMPLETE"}},
229 ),
230 ],
231 )
232 def test_type_field_default(self, cls, expected_type, kwargs):
233 ev = cls(**kwargs)
234 # type field is auto-populated from the Literal default.
235 assert ev.type == expected_type
236 # The discriminator must be present in the serialized payload so
237 # clients reading the stream can demux on it.
238 dumped = ev.model_dump(exclude_none=True)
239 assert dumped["type"] == expected_type
240 # Same in JSON form (what ``_emit`` serializes).
241 assert f'"type":"{expected_type}"' in ev.model_dump_json(exclude_none=True)
242
243 def test_type_field_cannot_be_overridden_to_wrong_value(self):
244 # Literal types reject any value other than the bake-in default.
245 with pytest.raises(ValidationError):
246 MessageStartEvent(
247 id="a",
248 delta={"message": {"role": "assistant"}},
249 type="other",
250 )