diff --git a/src/openai/_legacy_response.py b/src/openai/_legacy_response.py index 1a0814448..4585cd742 100644 --- a/src/openai/_legacy_response.py +++ b/src/openai/_legacy_response.py @@ -13,7 +13,7 @@ import pydantic from ._types import NoneType -from ._utils import is_given +from ._utils import is_given, extract_type_arg, is_annotated_type from ._models import BaseModel, is_basemodel from ._constants import RAW_RESPONSE_HEADER from ._streaming import Stream, AsyncStream, is_stream_class_type, extract_stream_chunk_type @@ -174,6 +174,10 @@ def elapsed(self) -> datetime.timedelta: return self.http_response.elapsed def _parse(self, *, to: type[_T] | None = None) -> R | _T: + # unwrap `Annotated[T, ...]` -> `T` + if to and is_annotated_type(to): + to = extract_type_arg(to, 0) + if self._stream: if to: if not is_stream_class_type(to): @@ -215,6 +219,11 @@ def _parse(self, *, to: type[_T] | None = None) -> R | _T: ) cast_to = to if to is not None else self._cast_to + + # unwrap `Annotated[T, ...]` -> `T` + if is_annotated_type(cast_to): + cast_to = extract_type_arg(cast_to, 0) + if cast_to is NoneType: return cast(R, None) diff --git a/src/openai/_models.py b/src/openai/_models.py index 810891497..af68e6181 100644 --- a/src/openai/_models.py +++ b/src/openai/_models.py @@ -30,7 +30,16 @@ AnyMapping, HttpxRequestFiles, ) -from ._utils import is_list, is_given, is_mapping, parse_date, parse_datetime, strip_not_given +from ._utils import ( + is_list, + is_given, + is_mapping, + parse_date, + parse_datetime, + strip_not_given, + extract_type_arg, + is_annotated_type, +) from ._compat import ( PYDANTIC_V2, ConfigDict, @@ -275,6 +284,9 @@ def construct_type(*, value: object, type_: type) -> object: If the given value does not match the expected type then it is returned as-is. """ + # unwrap `Annotated[T, ...]` -> `T` + if is_annotated_type(type_): + type_ = extract_type_arg(type_, 0) # we need to use the origin class for any types that are subscripted generics # e.g. Dict[str, object] diff --git a/src/openai/_response.py b/src/openai/_response.py index f684b08c2..47f484ef7 100644 --- a/src/openai/_response.py +++ b/src/openai/_response.py @@ -25,7 +25,7 @@ import pydantic from ._types import NoneType -from ._utils import is_given, extract_type_var_from_base +from ._utils import is_given, extract_type_arg, is_annotated_type, extract_type_var_from_base from ._models import BaseModel, is_basemodel from ._constants import RAW_RESPONSE_HEADER, OVERRIDE_CAST_TO_HEADER from ._streaming import Stream, AsyncStream, is_stream_class_type, extract_stream_chunk_type @@ -121,6 +121,10 @@ def __repr__(self) -> str: ) def _parse(self, *, to: type[_T] | None = None) -> R | _T: + # unwrap `Annotated[T, ...]` -> `T` + if to and is_annotated_type(to): + to = extract_type_arg(to, 0) + if self._is_sse_stream: if to: if not is_stream_class_type(to): @@ -162,6 +166,11 @@ def _parse(self, *, to: type[_T] | None = None) -> R | _T: ) cast_to = to if to is not None else self._cast_to + + # unwrap `Annotated[T, ...]` -> `T` + if is_annotated_type(cast_to): + cast_to = extract_type_arg(cast_to, 0) + if cast_to is NoneType: return cast(R, None) diff --git a/tests/test_legacy_response.py b/tests/test_legacy_response.py index 995250a58..45025f81d 100644 --- a/tests/test_legacy_response.py +++ b/tests/test_legacy_response.py @@ -1,4 +1,6 @@ import json +from typing import cast +from typing_extensions import Annotated import httpx import pytest @@ -63,3 +65,20 @@ def test_response_parse_custom_model(client: OpenAI) -> None: obj = response.parse(to=CustomModel) assert obj.foo == "hello!" assert obj.bar == 2 + + +def test_response_parse_annotated_type(client: OpenAI) -> None: + response = LegacyAPIResponse( + raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})), + client=client, + stream=False, + stream_cls=None, + cast_to=str, + options=FinalRequestOptions.construct(method="get", url="/foo"), + ) + + obj = response.parse( + to=cast("type[CustomModel]", Annotated[CustomModel, "random metadata"]), + ) + assert obj.foo == "hello!" + assert obj.bar == 2 diff --git a/tests/test_models.py b/tests/test_models.py index 713bd2cb1..d8a3c9ca5 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,14 +1,14 @@ import json from typing import Any, Dict, List, Union, Optional, cast from datetime import datetime, timezone -from typing_extensions import Literal +from typing_extensions import Literal, Annotated import pytest import pydantic from pydantic import Field from openai._compat import PYDANTIC_V2, parse_obj, model_dump, model_json -from openai._models import BaseModel +from openai._models import BaseModel, construct_type class BasicModel(BaseModel): @@ -571,3 +571,15 @@ class OurModel(BaseModel): foo: Optional[str] = None takes_pydantic(OurModel()) + + +def test_annotated_types() -> None: + class Model(BaseModel): + value: str + + m = construct_type( + value={"value": "foo"}, + type_=cast(Any, Annotated[Model, "random metadata"]), + ) + assert isinstance(m, Model) + assert m.value == "foo" diff --git a/tests/test_response.py b/tests/test_response.py index 7c99327b4..af153b67c 100644 --- a/tests/test_response.py +++ b/tests/test_response.py @@ -1,5 +1,6 @@ import json -from typing import List +from typing import List, cast +from typing_extensions import Annotated import httpx import pytest @@ -157,3 +158,37 @@ async def test_async_response_parse_custom_model(async_client: AsyncOpenAI) -> N obj = await response.parse(to=CustomModel) assert obj.foo == "hello!" assert obj.bar == 2 + + +def test_response_parse_annotated_type(client: OpenAI) -> None: + response = APIResponse( + raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})), + client=client, + stream=False, + stream_cls=None, + cast_to=str, + options=FinalRequestOptions.construct(method="get", url="/foo"), + ) + + obj = response.parse( + to=cast("type[CustomModel]", Annotated[CustomModel, "random metadata"]), + ) + assert obj.foo == "hello!" + assert obj.bar == 2 + + +async def test_async_response_parse_annotated_type(async_client: AsyncOpenAI) -> None: + response = AsyncAPIResponse( + raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})), + client=async_client, + stream=False, + stream_cls=None, + cast_to=str, + options=FinalRequestOptions.construct(method="get", url="/foo"), + ) + + obj = await response.parse( + to=cast("type[CustomModel]", Annotated[CustomModel, "random metadata"]), + ) + assert obj.foo == "hello!" + assert obj.bar == 2 diff --git a/tests/utils.py b/tests/utils.py index 216b33355..43c3cb5cf 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -14,6 +14,8 @@ is_list, is_list_type, is_union_type, + extract_type_arg, + is_annotated_type, ) from openai._compat import PYDANTIC_V2, field_outer_type, get_model_fields from openai._models import BaseModel @@ -49,6 +51,10 @@ def assert_matches_type( path: list[str], allow_none: bool = False, ) -> None: + # unwrap `Annotated[T, ...]` -> `T` + if is_annotated_type(type_): + type_ = extract_type_arg(type_, 0) + if allow_none and value is None: return