Function test_parse_pydantic_model_enum

tests/lib/chat/test_completions.py:239–288

(client: OpenAI, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch)

@pytest.mark.respx(base_url=base_url)
def test_parse_pydantic_model_enum(client: OpenAI, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None:
    class Color(Enum):
        """The detected color"""

        RED = "red"
        BLUE = "blue"
        GREEN = "green"

    class ColorDetection(BaseModel):
        color: Color
        hex_color_code: str = Field(description="The hex color code of the detected color")

    if PYDANTIC_V1:
        ColorDetection.update_forward_refs(**locals())  # type: ignore

    completion = make_snapshot_request(
        lambda c: c.chat.completions.parse(
            model="gpt-4o-2024-08-06",
            messages=[
                {"role": "user", "content": "What color is a Coke can?"},
            ],
            response_format=ColorDetection,
        ),
        content_snapshot=snapshot(
            '{"id": "chatcmpl-ABfvjIatz0zrZu50gRbMtlp0asZpz", "object": "chat.completion", "created": 1727346151, "model": "gpt-4o-2024-08-06", "choices": [{"index": 0, "message": {"role": "assistant", "content": "{\\"color\\":\\"red\\",\\"hex_color_code\\":\\"#FF0000\\"}", "refusal": null}, "logprobs": null, "finish_reason": "stop"}], "usage": {"prompt_tokens": 109, "completion_tokens": 14, "total_tokens": 123, "completion_tokens_details": {"reasoning_tokens": 0}}, "system_fingerprint": "fp_5050236cbd"}'
        ),
        path="/chat/completions",
        mock_client=client,
        respx_mock=respx_mock,
    )

    assert print_obj(completion.choices[0], monkeypatch) == snapshot(
        """\
ParsedChoice(
    finish_reason='stop',
    index=0,
    logprobs=None,
    message=ParsedChatCompletionMessage(
        annotations=None,
        audio=None,
        content='{"color":"red","hex_color_code":"#FF0000"}',
        function_call=None,
        parsed=ColorDetection(color=<Color.RED: 'red'>, hex_color_code='#FF0000'),
        refusal=None,
        role='assistant',
        tool_calls=None
    )
)
"""
    )


@pytest.mark.respx(base_url=base_url)