Expose the "store" parameter through ModelSettings (#357)

Closes https://github.com/openai/openai-agents-python/issues/173

This will also set stored completions to True by default, encouraging a
best practice.
This commit is contained in:
Steven Heidel 2025-03-26 16:01:28 -07:00 committed by GitHub
parent 4854a745e8
commit aad8accc86
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 52 additions and 13 deletions

View file

@ -1,6 +1,6 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass, fields, replace
from typing import Literal from typing import Literal
@ -30,8 +30,9 @@ class ModelSettings:
tool_choice: Literal["auto", "required", "none"] | str | None = None tool_choice: Literal["auto", "required", "none"] | str | None = None
"""The tool choice to use when calling the model.""" """The tool choice to use when calling the model."""
parallel_tool_calls: bool | None = False parallel_tool_calls: bool | None = None
"""Whether to use parallel tool calls when calling the model.""" """Whether to use parallel tool calls when calling the model.
Defaults to False if not provided."""
truncation: Literal["auto", "disabled"] | None = None truncation: Literal["auto", "disabled"] | None = None
"""The truncation strategy to use when calling the model.""" """The truncation strategy to use when calling the model."""
@ -39,18 +40,19 @@ class ModelSettings:
max_tokens: int | None = None max_tokens: int | None = None
"""The maximum number of output tokens to generate.""" """The maximum number of output tokens to generate."""
store: bool | None = None
"""Whether to store the generated model response for later retrieval.
Defaults to True if not provided."""
def resolve(self, override: ModelSettings | None) -> ModelSettings: def resolve(self, override: ModelSettings | None) -> ModelSettings:
"""Produce a new ModelSettings by overlaying any non-None values from the """Produce a new ModelSettings by overlaying any non-None values from the
override on top of this instance.""" override on top of this instance."""
if override is None: if override is None:
return self return self
return ModelSettings(
temperature=override.temperature or self.temperature, changes = {
top_p=override.top_p or self.top_p, field.name: getattr(override, field.name)
frequency_penalty=override.frequency_penalty or self.frequency_penalty, for field in fields(self)
presence_penalty=override.presence_penalty or self.presence_penalty, if getattr(override, field.name) is not None
tool_choice=override.tool_choice or self.tool_choice, }
parallel_tool_calls=override.parallel_tool_calls or self.parallel_tool_calls, return replace(self, **changes)
truncation=override.truncation or self.truncation,
max_tokens=override.max_tokens or self.max_tokens,
)

View file

@ -518,6 +518,9 @@ class OpenAIChatCompletionsModel(Model):
f"Response format: {response_format}\n" f"Response format: {response_format}\n"
) )
# Match the behavior of Responses where store is True when not given
store = model_settings.store if model_settings.store is not None else True
ret = await self._get_client().chat.completions.create( ret = await self._get_client().chat.completions.create(
model=self.model, model=self.model,
messages=converted_messages, messages=converted_messages,
@ -532,6 +535,7 @@ class OpenAIChatCompletionsModel(Model):
parallel_tool_calls=parallel_tool_calls, parallel_tool_calls=parallel_tool_calls,
stream=stream, stream=stream,
stream_options={"include_usage": True} if stream else NOT_GIVEN, stream_options={"include_usage": True} if stream else NOT_GIVEN,
store=store,
extra_headers=_HEADERS, extra_headers=_HEADERS,
) )

View file

@ -246,6 +246,7 @@ class OpenAIResponsesModel(Model):
stream=stream, stream=stream,
extra_headers=_HEADERS, extra_headers=_HEADERS,
text=response_format, text=response_format,
store=self._non_null_or_not_given(model_settings.store),
) )
def _get_client(self) -> AsyncOpenAI: def _get_client(self) -> AsyncOpenAI:

View file

@ -14,8 +14,10 @@ from agents import (
InputGuardrail, InputGuardrail,
InputGuardrailTripwireTriggered, InputGuardrailTripwireTriggered,
ModelBehaviorError, ModelBehaviorError,
ModelSettings,
OutputGuardrail, OutputGuardrail,
OutputGuardrailTripwireTriggered, OutputGuardrailTripwireTriggered,
RunConfig,
RunContextWrapper, RunContextWrapper,
Runner, Runner,
UserError, UserError,
@ -634,3 +636,31 @@ async def test_tool_use_behavior_custom_function():
assert len(result.raw_responses) == 2, "should have two model responses" assert len(result.raw_responses) == 2, "should have two model responses"
assert result.final_output == "the_final_output", "should have used the custom function" assert result.final_output == "the_final_output", "should have used the custom function"
@pytest.mark.asyncio
async def test_model_settings_override():
model = FakeModel()
agent = Agent(
name="test",
model=model,
model_settings=ModelSettings(temperature=1.0, max_tokens=1000)
)
model.add_multiple_turn_outputs(
[
[
get_text_message("a_message"),
],
]
)
await Runner.run(
agent,
input="user_message",
run_config=RunConfig(model_settings=ModelSettings(0.5)),
)
# temperature is overridden by Runner.run, but max_tokens is not
assert model.last_turn_args["model_settings"].temperature == 0.5
assert model.last_turn_args["model_settings"].max_tokens == 1000

View file

@ -226,6 +226,7 @@ async def test_fetch_response_non_stream(monkeypatch) -> None:
# Ensure expected args were passed through to OpenAI client. # Ensure expected args were passed through to OpenAI client.
kwargs = completions.kwargs kwargs = completions.kwargs
assert kwargs["stream"] is False assert kwargs["stream"] is False
assert kwargs["store"] is True
assert kwargs["model"] == "gpt-4" assert kwargs["model"] == "gpt-4"
assert kwargs["messages"][0]["role"] == "system" assert kwargs["messages"][0]["role"] == "system"
assert kwargs["messages"][0]["content"] == "sys" assert kwargs["messages"][0]["content"] == "sys"
@ -279,6 +280,7 @@ async def test_fetch_response_stream(monkeypatch) -> None:
) )
# Check OpenAI client was called for streaming # Check OpenAI client was called for streaming
assert completions.kwargs["stream"] is True assert completions.kwargs["stream"] is True
assert completions.kwargs["store"] is True
assert completions.kwargs["stream_options"] == {"include_usage": True} assert completions.kwargs["stream_options"] == {"include_usage": True}
# Response is a proper openai Response # Response is a proper openai Response
assert isinstance(response, Response) assert isinstance(response, Response)