Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 0189e28

Browse files
RobertCraigieRobertCraigie
authoredSep 26, 2024··
feat(structured outputs): add support for accessing raw responses (#1748)
1 parent af535ce commit 0189e28

File tree

2 files changed

+355
-68
lines changed

2 files changed

+355
-68
lines changed
 

‎src/openai/resources/beta/chat/completions.py

+179-67
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,21 @@
22

33
from __future__ import annotations
44

5-
from typing import Dict, List, Union, Iterable, Optional
5+
from typing import Dict, List, Type, Union, Iterable, Optional, cast
66
from functools import partial
77
from typing_extensions import Literal
88

99
import httpx
1010

11+
from .... import _legacy_response
1112
from ...._types import NOT_GIVEN, Body, Query, Headers, NotGiven
13+
from ...._utils import maybe_transform, async_maybe_transform
14+
from ...._compat import cached_property
1215
from ...._resource import SyncAPIResource, AsyncAPIResource
16+
from ...._response import to_streamed_response_wrapper, async_to_streamed_response_wrapper
1317
from ...._streaming import Stream
1418
from ....types.chat import completion_create_params
19+
from ...._base_client import make_request_options
1520
from ....lib._parsing import (
1621
ResponseFormatT,
1722
validate_input_tools as _validate_input_tools,
@@ -20,6 +25,7 @@
2025
)
2126
from ....types.chat_model import ChatModel
2227
from ....lib.streaming.chat import ChatCompletionStreamManager, AsyncChatCompletionStreamManager
28+
from ....types.chat.chat_completion import ChatCompletion
2329
from ....types.chat.chat_completion_chunk import ChatCompletionChunk
2430
from ....types.chat.parsed_chat_completion import ParsedChatCompletion
2531
from ....types.chat.chat_completion_tool_param import ChatCompletionToolParam
@@ -31,6 +37,25 @@
3137

3238

3339
class Completions(SyncAPIResource):
40+
@cached_property
41+
def with_raw_response(self) -> CompletionsWithRawResponse:
42+
"""
43+
This property can be used as a prefix for any HTTP method call to return the
44+
the raw response object instead of the parsed content.
45+
46+
For more information, see https://www.github.com/openai/openai-python#accessing-raw-response-data-eg-headers
47+
"""
48+
return CompletionsWithRawResponse(self)
49+
50+
@cached_property
51+
def with_streaming_response(self) -> CompletionsWithStreamingResponse:
52+
"""
53+
An alternative to `.with_raw_response` that doesn't eagerly read the response body.
54+
55+
For more information, see https://www.github.com/openai/openai-python#with_streaming_response
56+
"""
57+
return CompletionsWithStreamingResponse(self)
58+
3459
def parse(
3560
self,
3661
*,
@@ -113,39 +138,55 @@ class MathResponse(BaseModel):
113138
**(extra_headers or {}),
114139
}
115140

116-
raw_completion = self._client.chat.completions.create(
117-
messages=messages,
118-
model=model,
119-
response_format=_type_to_response_format(response_format),
120-
frequency_penalty=frequency_penalty,
121-
function_call=function_call,
122-
functions=functions,
123-
logit_bias=logit_bias,
124-
logprobs=logprobs,
125-
max_completion_tokens=max_completion_tokens,
126-
max_tokens=max_tokens,
127-
n=n,
128-
parallel_tool_calls=parallel_tool_calls,
129-
presence_penalty=presence_penalty,
130-
seed=seed,
131-
service_tier=service_tier,
132-
stop=stop,
133-
stream_options=stream_options,
134-
temperature=temperature,
135-
tool_choice=tool_choice,
136-
tools=tools,
137-
top_logprobs=top_logprobs,
138-
top_p=top_p,
139-
user=user,
140-
extra_headers=extra_headers,
141-
extra_query=extra_query,
142-
extra_body=extra_body,
143-
timeout=timeout,
144-
)
145-
return _parse_chat_completion(
146-
response_format=response_format,
147-
chat_completion=raw_completion,
148-
input_tools=tools,
141+
def parser(raw_completion: ChatCompletion) -> ParsedChatCompletion[ResponseFormatT]:
142+
return _parse_chat_completion(
143+
response_format=response_format,
144+
chat_completion=raw_completion,
145+
input_tools=tools,
146+
)
147+
148+
return self._post(
149+
"/chat/completions",
150+
body=maybe_transform(
151+
{
152+
"messages": messages,
153+
"model": model,
154+
"frequency_penalty": frequency_penalty,
155+
"function_call": function_call,
156+
"functions": functions,
157+
"logit_bias": logit_bias,
158+
"logprobs": logprobs,
159+
"max_completion_tokens": max_completion_tokens,
160+
"max_tokens": max_tokens,
161+
"n": n,
162+
"parallel_tool_calls": parallel_tool_calls,
163+
"presence_penalty": presence_penalty,
164+
"response_format": _type_to_response_format(response_format),
165+
"seed": seed,
166+
"service_tier": service_tier,
167+
"stop": stop,
168+
"stream": False,
169+
"stream_options": stream_options,
170+
"temperature": temperature,
171+
"tool_choice": tool_choice,
172+
"tools": tools,
173+
"top_logprobs": top_logprobs,
174+
"top_p": top_p,
175+
"user": user,
176+
},
177+
completion_create_params.CompletionCreateParams,
178+
),
179+
options=make_request_options(
180+
extra_headers=extra_headers,
181+
extra_query=extra_query,
182+
extra_body=extra_body,
183+
timeout=timeout,
184+
post_parser=parser,
185+
),
186+
# we turn the `ChatCompletion` instance into a `ParsedChatCompletion`
187+
# in the `parser` function above
188+
cast_to=cast(Type[ParsedChatCompletion[ResponseFormatT]], ChatCompletion),
189+
stream=False,
149190
)
150191

151192
def stream(
@@ -247,6 +288,25 @@ def stream(
247288

248289

249290
class AsyncCompletions(AsyncAPIResource):
291+
@cached_property
292+
def with_raw_response(self) -> AsyncCompletionsWithRawResponse:
293+
"""
294+
This property can be used as a prefix for any HTTP method call to return the
295+
the raw response object instead of the parsed content.
296+
297+
For more information, see https://www.github.com/openai/openai-python#accessing-raw-response-data-eg-headers
298+
"""
299+
return AsyncCompletionsWithRawResponse(self)
300+
301+
@cached_property
302+
def with_streaming_response(self) -> AsyncCompletionsWithStreamingResponse:
303+
"""
304+
An alternative to `.with_raw_response` that doesn't eagerly read the response body.
305+
306+
For more information, see https://www.github.com/openai/openai-python#with_streaming_response
307+
"""
308+
return AsyncCompletionsWithStreamingResponse(self)
309+
250310
async def parse(
251311
self,
252312
*,
@@ -329,39 +389,55 @@ class MathResponse(BaseModel):
329389
**(extra_headers or {}),
330390
}
331391

332-
raw_completion = await self._client.chat.completions.create(
333-
messages=messages,
334-
model=model,
335-
response_format=_type_to_response_format(response_format),
336-
frequency_penalty=frequency_penalty,
337-
function_call=function_call,
338-
functions=functions,
339-
logit_bias=logit_bias,
340-
logprobs=logprobs,
341-
max_completion_tokens=max_completion_tokens,
342-
max_tokens=max_tokens,
343-
n=n,
344-
parallel_tool_calls=parallel_tool_calls,
345-
presence_penalty=presence_penalty,
346-
seed=seed,
347-
service_tier=service_tier,
348-
stop=stop,
349-
stream_options=stream_options,
350-
temperature=temperature,
351-
tool_choice=tool_choice,
352-
tools=tools,
353-
top_logprobs=top_logprobs,
354-
top_p=top_p,
355-
user=user,
356-
extra_headers=extra_headers,
357-
extra_query=extra_query,
358-
extra_body=extra_body,
359-
timeout=timeout,
360-
)
361-
return _parse_chat_completion(
362-
response_format=response_format,
363-
chat_completion=raw_completion,
364-
input_tools=tools,
392+
def parser(raw_completion: ChatCompletion) -> ParsedChatCompletion[ResponseFormatT]:
393+
return _parse_chat_completion(
394+
response_format=response_format,
395+
chat_completion=raw_completion,
396+
input_tools=tools,
397+
)
398+
399+
return await self._post(
400+
"/chat/completions",
401+
body=await async_maybe_transform(
402+
{
403+
"messages": messages,
404+
"model": model,
405+
"frequency_penalty": frequency_penalty,
406+
"function_call": function_call,
407+
"functions": functions,
408+
"logit_bias": logit_bias,
409+
"logprobs": logprobs,
410+
"max_completion_tokens": max_completion_tokens,
411+
"max_tokens": max_tokens,
412+
"n": n,
413+
"parallel_tool_calls": parallel_tool_calls,
414+
"presence_penalty": presence_penalty,
415+
"response_format": _type_to_response_format(response_format),
416+
"seed": seed,
417+
"service_tier": service_tier,
418+
"stop": stop,
419+
"stream": False,
420+
"stream_options": stream_options,
421+
"temperature": temperature,
422+
"tool_choice": tool_choice,
423+
"tools": tools,
424+
"top_logprobs": top_logprobs,
425+
"top_p": top_p,
426+
"user": user,
427+
},
428+
completion_create_params.CompletionCreateParams,
429+
),
430+
options=make_request_options(
431+
extra_headers=extra_headers,
432+
extra_query=extra_query,
433+
extra_body=extra_body,
434+
timeout=timeout,
435+
post_parser=parser,
436+
),
437+
# we turn the `ChatCompletion` instance into a `ParsedChatCompletion`
438+
# in the `parser` function above
439+
cast_to=cast(Type[ParsedChatCompletion[ResponseFormatT]], ChatCompletion),
440+
stream=False,
365441
)
366442

367443
def stream(
@@ -461,3 +537,39 @@ def stream(
461537
response_format=response_format,
462538
input_tools=tools,
463539
)
540+
541+
542+
class CompletionsWithRawResponse:
543+
def __init__(self, completions: Completions) -> None:
544+
self._completions = completions
545+
546+
self.parse = _legacy_response.to_raw_response_wrapper(
547+
completions.parse,
548+
)
549+
550+
551+
class AsyncCompletionsWithRawResponse:
552+
def __init__(self, completions: AsyncCompletions) -> None:
553+
self._completions = completions
554+
555+
self.parse = _legacy_response.async_to_raw_response_wrapper(
556+
completions.parse,
557+
)
558+
559+
560+
class CompletionsWithStreamingResponse:
561+
def __init__(self, completions: Completions) -> None:
562+
self._completions = completions
563+
564+
self.parse = to_streamed_response_wrapper(
565+
completions.parse,
566+
)
567+
568+
569+
class AsyncCompletionsWithStreamingResponse:
570+
def __init__(self, completions: AsyncCompletions) -> None:
571+
self._completions = completions
572+
573+
self.parse = async_to_streamed_response_wrapper(
574+
completions.parse,
575+
)

‎tests/lib/chat/test_completions.py

+176-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import os
44
import json
55
from enum import Enum
6-
from typing import Any, List, Callable, Optional
6+
from typing import Any, List, Callable, Optional, Awaitable
77
from typing_extensions import Literal, TypeVar
88

99
import httpx
@@ -773,6 +773,139 @@ def test_parse_non_strict_tools(client: OpenAI) -> None:
773773
)
774774

775775

776+
@pytest.mark.respx(base_url=base_url)
777+
def test_parse_pydantic_raw_response(client: OpenAI, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None:
778+
class Location(BaseModel):
779+
city: str
780+
temperature: float
781+
units: Literal["c", "f"]
782+
783+
response = _make_snapshot_request(
784+
lambda c: c.beta.chat.completions.with_raw_response.parse(
785+
model="gpt-4o-2024-08-06",
786+
messages=[
787+
{
788+
"role": "user",
789+
"content": "What's the weather like in SF?",
790+
},
791+
],
792+
response_format=Location,
793+
),
794+
content_snapshot=snapshot(
795+
'{"id": "chatcmpl-ABrDYCa8W1w66eUxKDO8TQF1m6trT", "object": "chat.completion", "created": 1727389540, "model": "gpt-4o-2024-08-06", "choices": [{"index": 0, "message": {"role": "assistant", "content": "{\\"city\\":\\"San Francisco\\",\\"temperature\\":58,\\"units\\":\\"f\\"}", "refusal": null}, "logprobs": null, "finish_reason": "stop"}], "usage": {"prompt_tokens": 79, "completion_tokens": 14, "total_tokens": 93, "completion_tokens_details": {"reasoning_tokens": 0}}, "system_fingerprint": "fp_5050236cbd"}'
796+
),
797+
mock_client=client,
798+
respx_mock=respx_mock,
799+
)
800+
assert response.http_request.headers.get("x-stainless-helper-method") == "beta.chat.completions.parse"
801+
802+
completion = response.parse()
803+
message = completion.choices[0].message
804+
assert message.parsed is not None
805+
assert isinstance(message.parsed.city, str)
806+
assert print_obj(completion, monkeypatch) == snapshot(
807+
"""\
808+
ParsedChatCompletion[Location](
809+
choices=[
810+
ParsedChoice[Location](
811+
finish_reason='stop',
812+
index=0,
813+
logprobs=None,
814+
message=ParsedChatCompletionMessage[Location](
815+
content='{"city":"San Francisco","temperature":58,"units":"f"}',
816+
function_call=None,
817+
parsed=Location(city='San Francisco', temperature=58.0, units='f'),
818+
refusal=None,
819+
role='assistant',
820+
tool_calls=[]
821+
)
822+
)
823+
],
824+
created=1727389540,
825+
id='chatcmpl-ABrDYCa8W1w66eUxKDO8TQF1m6trT',
826+
model='gpt-4o-2024-08-06',
827+
object='chat.completion',
828+
service_tier=None,
829+
system_fingerprint='fp_5050236cbd',
830+
usage=CompletionUsage(
831+
completion_tokens=14,
832+
completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
833+
prompt_tokens=79,
834+
total_tokens=93
835+
)
836+
)
837+
"""
838+
)
839+
840+
841+
@pytest.mark.respx(base_url=base_url)
842+
@pytest.mark.asyncio
843+
async def test_async_parse_pydantic_raw_response(
844+
async_client: AsyncOpenAI, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch
845+
) -> None:
846+
class Location(BaseModel):
847+
city: str
848+
temperature: float
849+
units: Literal["c", "f"]
850+
851+
response = await _make_async_snapshot_request(
852+
lambda c: c.beta.chat.completions.with_raw_response.parse(
853+
model="gpt-4o-2024-08-06",
854+
messages=[
855+
{
856+
"role": "user",
857+
"content": "What's the weather like in SF?",
858+
},
859+
],
860+
response_format=Location,
861+
),
862+
content_snapshot=snapshot(
863+
'{"id": "chatcmpl-ABrDQWOiw0PK5JOsxl1D9ooeQgznq", "object": "chat.completion", "created": 1727389532, "model": "gpt-4o-2024-08-06", "choices": [{"index": 0, "message": {"role": "assistant", "content": "{\\"city\\":\\"San Francisco\\",\\"temperature\\":65,\\"units\\":\\"f\\"}", "refusal": null}, "logprobs": null, "finish_reason": "stop"}], "usage": {"prompt_tokens": 79, "completion_tokens": 14, "total_tokens": 93, "completion_tokens_details": {"reasoning_tokens": 0}}, "system_fingerprint": "fp_5050236cbd"}'
864+
),
865+
mock_client=async_client,
866+
respx_mock=respx_mock,
867+
)
868+
assert response.http_request.headers.get("x-stainless-helper-method") == "beta.chat.completions.parse"
869+
870+
completion = response.parse()
871+
message = completion.choices[0].message
872+
assert message.parsed is not None
873+
assert isinstance(message.parsed.city, str)
874+
assert print_obj(completion, monkeypatch) == snapshot(
875+
"""\
876+
ParsedChatCompletion[Location](
877+
choices=[
878+
ParsedChoice[Location](
879+
finish_reason='stop',
880+
index=0,
881+
logprobs=None,
882+
message=ParsedChatCompletionMessage[Location](
883+
content='{"city":"San Francisco","temperature":65,"units":"f"}',
884+
function_call=None,
885+
parsed=Location(city='San Francisco', temperature=65.0, units='f'),
886+
refusal=None,
887+
role='assistant',
888+
tool_calls=[]
889+
)
890+
)
891+
],
892+
created=1727389532,
893+
id='chatcmpl-ABrDQWOiw0PK5JOsxl1D9ooeQgznq',
894+
model='gpt-4o-2024-08-06',
895+
object='chat.completion',
896+
service_tier=None,
897+
system_fingerprint='fp_5050236cbd',
898+
usage=CompletionUsage(
899+
completion_tokens=14,
900+
completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
901+
prompt_tokens=79,
902+
total_tokens=93
903+
)
904+
)
905+
"""
906+
)
907+
908+
776909
@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
777910
def test_parse_method_in_sync(sync: bool, client: OpenAI, async_client: AsyncOpenAI) -> None:
778911
checking_client: OpenAI | AsyncOpenAI = client if sync else async_client
@@ -824,3 +957,45 @@ def _on_response(response: httpx.Response) -> None:
824957
client.close()
825958

826959
return result
960+
961+
962+
async def _make_async_snapshot_request(
963+
func: Callable[[AsyncOpenAI], Awaitable[_T]],
964+
*,
965+
content_snapshot: Any,
966+
respx_mock: MockRouter,
967+
mock_client: AsyncOpenAI,
968+
) -> _T:
969+
live = os.environ.get("OPENAI_LIVE") == "1"
970+
if live:
971+
972+
async def _on_response(response: httpx.Response) -> None:
973+
# update the content snapshot
974+
assert json.dumps(json.loads(await response.aread())) == content_snapshot
975+
976+
respx_mock.stop()
977+
978+
client = AsyncOpenAI(
979+
http_client=httpx.AsyncClient(
980+
event_hooks={
981+
"response": [_on_response],
982+
}
983+
)
984+
)
985+
else:
986+
respx_mock.post("/chat/completions").mock(
987+
return_value=httpx.Response(
988+
200,
989+
content=content_snapshot._old_value,
990+
headers={"content-type": "application/json"},
991+
)
992+
)
993+
994+
client = mock_client
995+
996+
result = await func(client)
997+
998+
if live:
999+
await client.close()
1000+
1001+
return result

0 commit comments

Comments
 (0)
Please sign in to comment.