openai-python

1
import json
2

3
import httpx
4
import pytest
5
import pydantic
6

7
from openai import OpenAI, BaseModel
8
from openai._streaming import Stream
9
from openai._base_client import FinalRequestOptions
10
from openai._legacy_response import LegacyAPIResponse
11

12

13
class PydanticModel(pydantic.BaseModel):
14
    ...
15

16

17
def test_response_parse_mismatched_basemodel(client: OpenAI) -> None:
18
    response = LegacyAPIResponse(
19
        raw=httpx.Response(200, content=b"foo"),
20
        client=client,
21
        stream=False,
22
        stream_cls=None,
23
        cast_to=str,
24
        options=FinalRequestOptions.construct(method="get", url="/foo"),
25
    )
26

27
    with pytest.raises(
28
        TypeError,
29
        match="Pydantic models must subclass our base model type, e.g. `from openai import BaseModel`",
30
    ):
31
        response.parse(to=PydanticModel)
32

33

34
def test_response_parse_custom_stream(client: OpenAI) -> None:
35
    response = LegacyAPIResponse(
36
        raw=httpx.Response(200, content=b"foo"),
37
        client=client,
38
        stream=True,
39
        stream_cls=None,
40
        cast_to=str,
41
        options=FinalRequestOptions.construct(method="get", url="/foo"),
42
    )
43

44
    stream = response.parse(to=Stream[int])
45
    assert stream._cast_to == int
46

47

48
class CustomModel(BaseModel):
49
    foo: str
50
    bar: int
51

52

53
def test_response_parse_custom_model(client: OpenAI) -> None:
54
    response = LegacyAPIResponse(
55
        raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})),
56
        client=client,
57
        stream=False,
58
        stream_cls=None,
59
        cast_to=str,
60
        options=FinalRequestOptions.construct(method="get", url="/foo"),
61
    )
62

63
    obj = response.parse(to=CustomModel)
64
    assert obj.foo == "hello!"
65
    assert obj.bar == 2
66