main
py 78 lines 2.46 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import json
4
5 import openai
6 import pytest
7 from pydantic import BaseModel
8
9
10 @pytest.mark.asyncio
11 async def test_structured_output(client: openai.AsyncOpenAI):
12 response = await client.responses.create(
13 input=[
14 {"role": "system", "content": "Extract the event information."},
15 {
16 "role": "user",
17 "content": "Alice and Bob are going to a science fair on Friday.",
18 },
19 ],
20 text={
21 "format": {
22 "type": "json_schema",
23 "name": "calendar_event",
24 "schema": {
25 "type": "object",
26 "properties": {
27 "event_name": {"type": "string"},
28 "date": {"type": "string"},
29 "participants": {"type": "array", "items": {"type": "string"}},
30 },
31 "required": ["event_name", "date", "participants"],
32 "additionalProperties": False,
33 },
34 "description": "A calendar event.",
35 "strict": True,
36 }
37 },
38 )
39 print(response)
40
41 # NOTE: The JSON schema is applied to the output text, not reasoning.
42 output_text = response.output[-1].content[0].text
43 event = json.loads(output_text)
44
45 assert event["event_name"].lower() == "science fair"
46 assert event["date"] == "Friday"
47 participants = event["participants"]
48 assert len(participants) == 2
49 assert participants[0] == "Alice"
50 assert participants[1] == "Bob"
51
52
53 @pytest.mark.asyncio
54 async def test_structured_output_with_parse(client: openai.AsyncOpenAI):
55 class CalendarEvent(BaseModel):
56 event_name: str
57 date: str
58 participants: list[str]
59
60 response = await client.responses.parse(
61 model=None,
62 instructions="Extract the event information.",
63 input="Alice and Bob are going to a science fair on Friday.",
64 text_format=CalendarEvent,
65 )
66 print(response)
67
68 # The output is successfully parsed.
69 event = response.output_parsed
70 assert event is not None
71
72 # The output is correct.
73 assert event.event_name.lower() == "science fair"
74 assert event.date == "Friday"
75 participants = event.participants
76 assert len(participants) == 2
77 assert participants[0] == "Alice"
78 assert participants[1] == "Bob"