Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 40 additions & 4 deletions src/agents/model_settings.py
Original file line number Diff line number Diff line change
@@ -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]]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tconley1428 can you confirm that requests to the openAI API still work when headers are set like this?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can confirm that our normal requests work, but I don't I know if they actually have additional headers (I think a user agent is added by default?), and I'm pretty sure they don't actually have an instance of Omit


@dataclass
class ModelSettings:
Expand Down
31 changes: 31 additions & 0 deletions tests/model_settings/test_serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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