1import json
  2from typing import Any, Union, cast
  3from typing_extensions import Annotated
  4
  5import httpx
  6import pytest
  7import pydantic
  8
  9from openai import OpenAI, BaseModel
 10from openai._streaming import Stream
 11from openai._base_client import FinalRequestOptions
 12from openai._legacy_response import LegacyAPIResponse
 13
 14from .utils import rich_print_str
 15
 16
 17class PydanticModel(pydantic.BaseModel): ...
 18
 19
 20def test_response_parse_mismatched_basemodel(client: OpenAI) -> None:
 21    response = LegacyAPIResponse(
 22        raw=httpx.Response(200, content=b"foo"),
 23        client=client,
 24        stream=False,
 25        stream_cls=None,
 26        cast_to=str,
 27        options=FinalRequestOptions.construct(method="get", url="/foo"),
 28    )
 29
 30    with pytest.raises(
 31        TypeError,
 32        match="Pydantic models must subclass our base model type, e.g. `from openai import BaseModel`",
 33    ):
 34        response.parse(to=PydanticModel)
 35
 36
 37@pytest.mark.parametrize(
 38    "content, expected",
 39    [
 40        ("false", False),
 41        ("true", True),
 42        ("False", False),
 43        ("True", True),
 44        ("TrUe", True),
 45        ("FalSe", False),
 46    ],
 47)
 48def test_response_parse_bool(client: OpenAI, content: str, expected: bool) -> None:
 49    response = LegacyAPIResponse(
 50        raw=httpx.Response(200, content=content),
 51        client=client,
 52        stream=False,
 53        stream_cls=None,
 54        cast_to=str,
 55        options=FinalRequestOptions.construct(method="get", url="/foo"),
 56    )
 57
 58    result = response.parse(to=bool)
 59    assert result is expected
 60
 61
 62def test_response_parse_custom_stream(client: OpenAI) -> None:
 63    response = LegacyAPIResponse(
 64        raw=httpx.Response(200, content=b"foo"),
 65        client=client,
 66        stream=True,
 67        stream_cls=None,
 68        cast_to=str,
 69        options=FinalRequestOptions.construct(method="get", url="/foo"),
 70    )
 71
 72    stream = response.parse(to=Stream[int])
 73    assert stream._cast_to == int
 74
 75
 76class CustomModel(BaseModel):
 77    foo: str
 78    bar: int
 79
 80
 81def test_response_parse_custom_model(client: OpenAI) -> None:
 82    response = LegacyAPIResponse(
 83        raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})),
 84        client=client,
 85        stream=False,
 86        stream_cls=None,
 87        cast_to=str,
 88        options=FinalRequestOptions.construct(method="get", url="/foo"),
 89    )
 90
 91    obj = response.parse(to=CustomModel)
 92    assert obj.foo == "hello!"
 93    assert obj.bar == 2
 94
 95
 96def test_response_basemodel_request_id(client: OpenAI) -> None:
 97    response = LegacyAPIResponse(
 98        raw=httpx.Response(
 99            200,
100            headers={"x-request-id": "my-req-id"},
101            content=json.dumps({"foo": "hello!", "bar": 2}),
102        ),
103        client=client,
104        stream=False,
105        stream_cls=None,
106        cast_to=str,
107        options=FinalRequestOptions.construct(method="get", url="/foo"),
108    )
109
110    obj = response.parse(to=CustomModel)
111    assert obj._request_id == "my-req-id"
112    assert obj.foo == "hello!"
113    assert obj.bar == 2
114    assert obj.to_dict() == {"foo": "hello!", "bar": 2}
115    assert "_request_id" not in rich_print_str(obj)
116    assert "__exclude_fields__" not in rich_print_str(obj)
117
118
119def test_response_parse_annotated_type(client: OpenAI) -> None:
120    response = LegacyAPIResponse(
121        raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})),
122        client=client,
123        stream=False,
124        stream_cls=None,
125        cast_to=str,
126        options=FinalRequestOptions.construct(method="get", url="/foo"),
127    )
128
129    obj = response.parse(
130        to=cast("type[CustomModel]", Annotated[CustomModel, "random metadata"]),
131    )
132    assert obj.foo == "hello!"
133    assert obj.bar == 2
134
135
136class OtherModel(pydantic.BaseModel):
137    a: str
138
139
140@pytest.mark.parametrize("client", [False], indirect=True)  # loose validation
141def test_response_parse_expect_model_union_non_json_content(client: OpenAI) -> None:
142    response = LegacyAPIResponse(
143        raw=httpx.Response(200, content=b"foo", headers={"Content-Type": "application/text"}),
144        client=client,
145        stream=False,
146        stream_cls=None,
147        cast_to=str,
148        options=FinalRequestOptions.construct(method="get", url="/foo"),
149    )
150
151    obj = response.parse(to=cast(Any, Union[CustomModel, OtherModel]))
152    assert isinstance(obj, str)
153    assert obj == "foo"