Annotating the openai.Omit type so that ModelSettings can be serialized by pydantic (openai#938)

Because `openai.Omit` is not a type which can be serialized by `pydantic`, it produces difficulty in consistently serializing/deserializing types across the agents SDK, many of which are `pydantic` types. This adds an annotation to enable `pydantic` to serialize `Omit`, in particular in the case of `ModelSettings` which contains `Omit` in its `extra_headers`

Author: tconley1428

src/agents/model_settings.py

@@ -1,14 +1,50 @@
 from __future__ import annotations
 
 import dataclasses
+from collections.abc import Mapping
 from dataclasses import dataclass, fields, replace
-from typing import Any, Literal
+from typing import Annotated, Any, Literal, Union
 
-from openai._types import Body, Headers, Query
+from openai import Omit as _Omit
+from openai._types import Body, Query
 from openai.types.responses import ResponseIncludable
 from openai.types.shared import Reasoning
-from pydantic import BaseModel
-
+from pydantic import BaseModel, GetCoreSchemaHandler
+from pydantic_core import core_schema
+from typing_extensions import TypeAlias
+
+
+class _OmitTypeAnnotation:
+ @classmethod
+ def __get_pydantic_core_schema__(
+ cls,
+ _source_type: Any,
+ _handler: GetCoreSchemaHandler,
+ ) -> core_schema.CoreSchema:
+ def validate_from_none(value: None) -> _Omit:
+ return _Omit()
+
+ from_none_schema = core_schema.chain_schema(
+ [
+ core_schema.none_schema(),
+ core_schema.no_info_plain_validator_function(validate_from_none),
+ ]
+ )
+ return core_schema.json_or_python_schema(
+ json_schema=from_none_schema,
+ python_schema=core_schema.union_schema(
+ [
+ # check if it's an instance first before doing any further work
+ core_schema.is_instance_schema(_Omit),
+ from_none_schema,
+ ]
+ ),
+ serialization=core_schema.plain_serializer_function_ser_schema(
+ lambda instance: None
+ ),
+ )
+Omit = Annotated[_Omit, _OmitTypeAnnotation]
+Headers: TypeAlias = Mapping[str, Union[str, Omit]]
 
 @dataclass
 class ModelSettings:

tests/model_settings/test_serialization.py

@@ -2,6 +2,8 @@
 from dataclasses import fields
 
 from openai.types.shared import Reasoning
+from pydantic import TypeAdapter
+from pydantic_core import to_json
 
 from agents.model_settings import ModelSettings
 
@@ -132,3 +134,32 @@ def test_extra_args_resolve_both_none() -> None:
  assert resolved.extra_args is None
  assert resolved.temperature == 0.5
  assert resolved.top_p == 0.9
+
+def test_pydantic_serialization() -> None:
+
+ """Tests whether ModelSettings can be serialized with Pydantic."""
+
+ # First, lets create a ModelSettings instance
+ model_settings = ModelSettings(
+ temperature=0.5,
+ top_p=0.9,
+ frequency_penalty=0.0,
+ presence_penalty=0.0,
+ tool_choice="auto",
+ parallel_tool_calls=True,
+ truncation="auto",
+ max_tokens=100,
+ reasoning=Reasoning(),
+ metadata={"foo": "bar"},
+ store=False,
+ include_usage=False,
+ extra_query={"foo": "bar"},
+ extra_body={"foo": "bar"},
+ extra_headers={"foo": "bar"},
+ extra_args={"custom_param": "value", "another_param": 42},
+ )
+
+ json = to_json(model_settings)
+ deserialized = TypeAdapter(ModelSettings).validate_json(json)
+
+ assert model_settings == deserialized