|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import dataclasses |
| 4 | +from collections.abc import Mapping |
4 | 5 | from dataclasses import dataclass, fields, replace |
5 | | -from typing import Any, Literal |
| 6 | +from typing import Annotated, Any, Literal, Union |
6 | 7 |
|
7 | | -from openai._types import Body, Headers, Query |
| 8 | +from openai import Omit as _Omit |
| 9 | +from openai._types import Body, Query |
8 | 10 | from openai.types.responses import ResponseIncludable |
9 | 11 | from openai.types.shared import Reasoning |
10 | | -from pydantic import BaseModel |
11 | | - |
| 12 | +from pydantic import BaseModel, GetCoreSchemaHandler |
| 13 | +from pydantic_core import core_schema |
| 14 | +from typing_extensions import TypeAlias |
| 15 | + |
| 16 | + |
| 17 | +class _OmitTypeAnnotation: |
| 18 | + @classmethod |
| 19 | + def __get_pydantic_core_schema__( |
| 20 | + cls, |
| 21 | + _source_type: Any, |
| 22 | + _handler: GetCoreSchemaHandler, |
| 23 | + ) -> core_schema.CoreSchema: |
| 24 | + def validate_from_none(value: None) -> _Omit: |
| 25 | + return _Omit() |
| 26 | + |
| 27 | + from_none_schema = core_schema.chain_schema( |
| 28 | + [ |
| 29 | + core_schema.none_schema(), |
| 30 | + core_schema.no_info_plain_validator_function(validate_from_none), |
| 31 | + ] |
| 32 | + ) |
| 33 | + return core_schema.json_or_python_schema( |
| 34 | + json_schema=from_none_schema, |
| 35 | + python_schema=core_schema.union_schema( |
| 36 | + [ |
| 37 | + # check if it's an instance first before doing any further work |
| 38 | + core_schema.is_instance_schema(_Omit), |
| 39 | + from_none_schema, |
| 40 | + ] |
| 41 | + ), |
| 42 | + serialization=core_schema.plain_serializer_function_ser_schema( |
| 43 | + lambda instance: None |
| 44 | + ), |
| 45 | + ) |
| 46 | +Omit = Annotated[_Omit, _OmitTypeAnnotation] |
| 47 | +Headers: TypeAlias = Mapping[str, Union[str, Omit]] |
12 | 48 |
|
13 | 49 | @dataclass |
14 | 50 | class ModelSettings: |
|
0 commit comments