main
py 245 lines 7.86 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import json
5
6 import anthropic
7 import pytest
8 import pytest_asyncio
9
10 from tests.utils import RemoteOpenAIServer
11
12 MODEL_NAME = "Qwen/Qwen3-0.6B"
13
14
15 @pytest.fixture(scope="module")
16 def server():
17 args = [
18 "--max-model-len",
19 "2048",
20 "--enforce-eager",
21 "--enable-prompt-tokens-details",
22 "--enable-auto-tool-choice",
23 "--tool-call-parser",
24 "hermes",
25 "--served-model-name",
26 "claude-3-7-sonnet-latest",
27 ]
28
29 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
30 yield remote_server
31
32
33 @pytest_asyncio.fixture
34 async def client(server):
35 async with server.get_async_client_anthropic() as async_client:
36 yield async_client
37
38
39 @pytest.mark.asyncio
40 async def test_simple_messages(client: anthropic.AsyncAnthropic):
41 resp = await client.messages.create(
42 model="claude-3-7-sonnet-latest",
43 max_tokens=1024,
44 messages=[{"role": "user", "content": "how are you!"}],
45 )
46 assert resp.stop_reason == "end_turn"
47 assert resp.role == "assistant"
48
49 print(f"Anthropic response: {resp.model_dump_json()}")
50
51
52 @pytest.mark.asyncio
53 async def test_system_message(client: anthropic.AsyncAnthropic):
54 resp = await client.messages.create(
55 model="claude-3-7-sonnet-latest",
56 max_tokens=1024,
57 system="you are a helpful assistant",
58 messages=[{"role": "user", "content": "how are you!"}],
59 )
60 assert resp.stop_reason == "end_turn"
61 assert resp.role == "assistant"
62
63 print(f"Anthropic response: {resp.model_dump_json()}")
64
65
66 @pytest.mark.asyncio
67 async def test_anthropic_streaming(client: anthropic.AsyncAnthropic):
68 resp = await client.messages.create(
69 model="claude-3-7-sonnet-latest",
70 max_tokens=1024,
71 messages=[{"role": "user", "content": "how are you!"}],
72 stream=True,
73 )
74
75 first_chunk = None
76 chunk_count = 0
77 async for chunk in resp:
78 chunk_count += 1
79 if first_chunk is None and chunk.type == "message_start":
80 first_chunk = chunk
81 print(chunk.model_dump_json())
82
83 assert chunk_count > 0
84 assert first_chunk is not None, "message_start chunk was never observed"
85 assert first_chunk.message is not None, "first chunk should include message"
86 assert first_chunk.message.usage is not None, (
87 "first chunk should include usage stats"
88 )
89 assert first_chunk.message.usage.output_tokens == 0
90 assert first_chunk.message.usage.input_tokens > 5
91
92
93 @pytest.mark.asyncio
94 async def test_anthropic_tool_call(client: anthropic.AsyncAnthropic):
95 resp = await client.messages.create(
96 model="claude-3-7-sonnet-latest",
97 max_tokens=1024,
98 messages=[
99 {"role": "user", "content": "What's the weather like in New York today?"}
100 ],
101 tools=[
102 {
103 "name": "get_current_weather",
104 "description": "Useful for querying the weather in a specified city.",
105 "input_schema": {
106 "type": "object",
107 "properties": {
108 "location": {
109 "type": "string",
110 "description": "City or region, for example: "
111 "New York, London, Tokyo, etc.",
112 }
113 },
114 "required": ["location"],
115 },
116 }
117 ],
118 stream=False,
119 )
120 assert resp.stop_reason == "tool_use"
121 assert resp.role == "assistant"
122
123 print(f"Anthropic response: {resp.model_dump_json()}")
124
125
126 @pytest.mark.asyncio
127 async def test_anthropic_tool_call_streaming(client: anthropic.AsyncAnthropic):
128 resp = await client.messages.create(
129 model="claude-3-7-sonnet-latest",
130 max_tokens=1024,
131 messages=[
132 {
133 "role": "user",
134 "content": "What's the weather like in New York today?",
135 }
136 ],
137 tools=[
138 {
139 "name": "get_current_weather",
140 "description": "Useful for querying the weather in a specified city.",
141 "input_schema": {
142 "type": "object",
143 "properties": {
144 "location": {
145 "type": "string",
146 "description": "City or region, for example: "
147 "New York, London, Tokyo, etc.",
148 }
149 },
150 "required": ["location"],
151 },
152 }
153 ],
154 stream=True,
155 )
156
157 async for chunk in resp:
158 print(chunk.model_dump_json())
159
160
161 @pytest.mark.asyncio
162 async def test_anthropic_structured_output(client: anthropic.AsyncAnthropic):
163 response = await client.messages.create(
164 model="claude-3-7-sonnet-latest",
165 max_tokens=1024,
166 messages=[
167 {
168 "role": "user",
169 "content": "Extract the key information from this email:"
170 "John Smith (john@example.com) is interested in our "
171 "Enterprise plan and wants to schedule a demo for next Tuesday at 2pm.",
172 }
173 ],
174 output_config={
175 "format": {
176 "type": "json_schema",
177 "schema": {
178 "type": "object",
179 "properties": {
180 "name": {"type": "string"},
181 "email": {"type": "string"},
182 "plan_interest": {"type": "string"},
183 "demo_requested": {"type": "boolean"},
184 },
185 "required": ["name", "email", "plan_interest", "demo_requested"],
186 "additionalProperties": False,
187 },
188 }
189 },
190 )
191 print(response.content[0].text)
192 json_obj = json.loads(response.content[0].text)
193 for key in ["name", "email", "plan_interest", "demo_requested"]:
194 assert key in json_obj, f"Missing key in output: {key}"
195
196
197 @pytest.mark.asyncio
198 async def test_anthropic_streaming_cache_usage(client: anthropic.AsyncAnthropic):
199 async def get_stream_usage(resp):
200 prompt_tokens = None
201 usage = None
202 async for chunk in resp:
203 if (
204 chunk.type == "message_start"
205 and chunk.message is not None
206 and chunk.message.usage is not None
207 ):
208 prompt_tokens = chunk.message.usage.input_tokens
209 elif chunk.type == "message_delta" and chunk.usage is not None:
210 usage = chunk.usage
211
212 assert usage is not None
213 assert usage.input_tokens >= 0
214 assert usage.output_tokens >= 0
215 cache_created = usage.cache_creation_input_tokens
216 cache_read = usage.cache_read_input_tokens
217 assert cache_read is not None
218 assert cache_created is not None
219 assert cache_created >= 0
220 assert cache_read >= 0
221 assert prompt_tokens == usage.input_tokens + cache_created + cache_read
222 return usage
223
224 request = dict(
225 model="claude-3-7-sonnet-latest",
226 max_tokens=1,
227 extra_body={"temperature": 0.0},
228 messages=[
229 {
230 "role": "user",
231 "content": "Cache coverage sentinel. " * 256
232 + "Answer with exactly one word: ok.",
233 }
234 ],
235 stream=True,
236 )
237
238 cold_usage = await get_stream_usage(await client.messages.create(**request))
239 assert cold_usage.cache_read_input_tokens == 0
240 assert cold_usage.cache_creation_input_tokens is not None
241 assert cold_usage.cache_creation_input_tokens > 0
242
243 warm_usage = await get_stream_usage(await client.messages.create(**request))
244 assert warm_usage.cache_read_input_tokens is not None
245 assert warm_usage.cache_read_input_tokens > 0