Skip to content

Commit a7da62e

Browse files
committed
Switch to just storing ModelRequestParameters and ModelSettings on ModelRequest
1 parent 3dc80a0 commit a7da62e

File tree

5 files changed

+86
-66
lines changed

5 files changed

+86
-66
lines changed

pydantic_ai_slim/pydantic_ai/_agent_graph.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -495,12 +495,11 @@ async def _prepare_request(
495495

496496
model_request_parameters = await _prepare_request_parameters(ctx)
497497

498-
# Populate tool tracking on the ModelRequest (the last request in the original history)
499-
self.request.function_tools = model_request_parameters.function_tools
500-
self.request.builtin_tools = model_request_parameters.builtin_tools
501-
self.request.output_tools = model_request_parameters.output_tools
502-
503498
model_settings = ctx.deps.model_settings
499+
# Record metadata on the ModelRequest (the last request in the original history)
500+
self.request.model_request_parameters = model_request_parameters
501+
self.request.model_settings = model_settings
502+
504503
usage = ctx.state.usage
505504
if ctx.deps.usage_limits.count_tokens_before_request:
506505
# Copy to avoid modifying the original usage object with the counted usage
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
from __future__ import annotations as _annotations
2+
3+
from dataclasses import dataclass, field
4+
from functools import cached_property
5+
from typing import TYPE_CHECKING
6+
7+
from . import _utils
8+
from ._tool_types import ToolDefinition
9+
from .builtin_tools import AbstractBuiltinTool
10+
11+
if TYPE_CHECKING:
12+
from ._output import OutputObjectDefinition
13+
from .output import OutputMode
14+
15+
__all__ = ('ModelRequestParameters',)
16+
17+
18+
@dataclass(repr=False, kw_only=True)
19+
class ModelRequestParameters:
20+
"""Configuration for an agent's request to a model, specifically related to tools and output handling."""
21+
22+
function_tools: list[ToolDefinition] = field(default_factory=list)
23+
builtin_tools: list[AbstractBuiltinTool] = field(default_factory=list)
24+
25+
output_mode: OutputMode = 'text'
26+
output_object: OutputObjectDefinition | None = None
27+
output_tools: list[ToolDefinition] = field(default_factory=list)
28+
prompted_output_template: str | None = None
29+
allow_text_output: bool = True
30+
allow_image_output: bool = False
31+
32+
@cached_property
33+
def tool_defs(self) -> dict[str, ToolDefinition]:
34+
return {tool_def.name: tool_def for tool_def in [*self.function_tools, *self.output_tools]}
35+
36+
@cached_property
37+
def prompted_output_instructions(self) -> str | None:
38+
if self.output_mode == 'prompted' and self.prompted_output_template and self.output_object:
39+
from ._output import PromptedOutputSchema
40+
41+
return PromptedOutputSchema.build_instructions(self.prompted_output_template, self.output_object)
42+
return None
43+
44+
__repr__ = _utils.dataclasses_no_defaults_repr

pydantic_ai_slim/pydantic_ai/messages.py

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,21 @@
1616
from typing_extensions import deprecated
1717

1818
from . import _otel_messages, _utils
19-
from ._tool_types import ToolDefinition
19+
from ._model_request_parameters import ModelRequestParameters
2020
from ._utils import generate_tool_call_id as _generate_tool_call_id, now_utc as _now_utc
21-
from .builtin_tools import AbstractBuiltinTool
2221
from .exceptions import UnexpectedModelBehavior
22+
from .settings import ModelSettings
2323
from .usage import RequestUsage
2424

2525
if TYPE_CHECKING:
2626
from .models.instrumented import InstrumentationSettings
2727

28+
ModelRequestParametersField = ModelRequestParameters | None
29+
ModelSettingsField = ModelSettings | None
30+
else: # pragma: no cover
31+
ModelRequestParametersField = Any
32+
ModelSettingsField = Any
33+
2834

2935
AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg', 'audio/ogg', 'audio/flac', 'audio/aiff', 'audio/aac']
3036
ImageMediaType: TypeAlias = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
@@ -947,26 +953,18 @@ class ModelRequest:
947953
instructions: str | None = None
948954
"""The instructions for the model."""
949955

950-
function_tools: Annotated[list[ToolDefinition] | None, pydantic.Field(exclude=True, repr=False)] = field(
951-
default=None, repr=False
952-
)
953-
"""Function tools that were available for this request.
954-
955-
Available for introspection during a run. This field is excluded from serialization.
956-
"""
957-
958-
builtin_tools: Annotated[list[AbstractBuiltinTool] | None, pydantic.Field(exclude=True, repr=False)] = field(
959-
default=None, repr=False
960-
)
961-
"""Builtin tools that were available for this request.
956+
model_request_parameters: Annotated[
957+
ModelRequestParametersField, pydantic.Field(exclude=True, repr=False)
958+
] = field(default=None, repr=False)
959+
"""Full request parameters captured for this request.
962960
963961
Available for introspection during a run. This field is excluded from serialization.
964962
"""
965963

966-
output_tools: Annotated[list[ToolDefinition] | None, pydantic.Field(exclude=True, repr=False)] = field(
964+
model_settings: Annotated[ModelSettingsField, pydantic.Field(exclude=True, repr=False)] = field(
967965
default=None, repr=False
968966
)
969-
"""Output tools that were available for this request.
967+
"""Effective model settings that were applied to this request.
970968
971969
Available for introspection during a run. This field is excluded from serialization.
972970
"""

pydantic_ai_slim/pydantic_ai/models/__init__.py

Lines changed: 2 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,11 @@
1919
import httpx
2020
from typing_extensions import TypeAliasType, TypedDict
2121

22-
from .. import _utils
2322
from .._json_schema import JsonSchemaTransformer
24-
from .._output import OutputObjectDefinition, PromptedOutputSchema
23+
from .._model_request_parameters import ModelRequestParameters
24+
from .._output import OutputObjectDefinition
2525
from .._parts_manager import ModelResponsePartsManager
2626
from .._run_context import RunContext
27-
from ..builtin_tools import AbstractBuiltinTool
2827
from ..exceptions import UserError
2928
from ..messages import (
3029
BaseToolCallPart,
@@ -45,7 +44,6 @@
4544
ToolCallPart,
4645
VideoUrl,
4746
)
48-
from ..output import OutputMode
4947
from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
5048
from ..providers import Provider, infer_provider
5149
from ..settings import ModelSettings, merge_model_settings
@@ -308,33 +306,6 @@
308306
"""
309307

310308

311-
@dataclass(repr=False, kw_only=True)
312-
class ModelRequestParameters:
313-
"""Configuration for an agent's request to a model, specifically related to tools and output handling."""
314-
315-
function_tools: list[ToolDefinition] = field(default_factory=list)
316-
builtin_tools: list[AbstractBuiltinTool] = field(default_factory=list)
317-
318-
output_mode: OutputMode = 'text'
319-
output_object: OutputObjectDefinition | None = None
320-
output_tools: list[ToolDefinition] = field(default_factory=list)
321-
prompted_output_template: str | None = None
322-
allow_text_output: bool = True
323-
allow_image_output: bool = False
324-
325-
@cached_property
326-
def tool_defs(self) -> dict[str, ToolDefinition]:
327-
return {tool_def.name: tool_def for tool_def in [*self.function_tools, *self.output_tools]}
328-
329-
@cached_property
330-
def prompted_output_instructions(self) -> str | None:
331-
if self.output_mode == 'prompted' and self.prompted_output_template and self.output_object:
332-
return PromptedOutputSchema.build_instructions(self.prompted_output_template, self.output_object)
333-
return None
334-
335-
__repr__ = _utils.dataclasses_no_defaults_repr
336-
337-
338309
class Model(ABC):
339310
"""Abstract class for a model."""
340311

tests/test_messages.py

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@
2727
VideoUrl,
2828
)
2929
from pydantic_ai.builtin_tools import ImageGenerationTool
30-
from pydantic_ai.models import ToolDefinition
30+
from pydantic_ai.models import ModelRequestParameters, ToolDefinition
31+
from pydantic_ai.settings import ModelSettings
3132

3233
from .conftest import IsDatetime, IsNow, IsStr
3334

@@ -407,8 +408,6 @@ def test_pre_usage_refactor_messages_deserializable():
407408
timestamp=IsNow(tz=timezone.utc),
408409
)
409410
],
410-
function_tools=None,
411-
builtin_tools=None,
412411
),
413412
ModelResponse(
414413
parts=[TextPart(content='Mexico City.')],
@@ -612,7 +611,7 @@ def test_binary_content_validation_with_optional_identifier():
612611

613612

614613
def test_model_request_tool_tracking_excluded_from_serialization():
615-
"""Test that function, builtin, and output tools are not serialized in the request."""
614+
"""Test that request metadata is accessible but not serialized."""
616615
tool_def = ToolDefinition(
617616
name='test_tool',
618617
description='A test tool',
@@ -624,21 +623,30 @@ def test_model_request_tool_tracking_excluded_from_serialization():
624623
parameters_json_schema={'type': 'object', 'properties': {}},
625624
)
626625

627-
request = ModelRequest(
628-
parts=[UserPromptPart('test prompt')],
629-
instructions='test instructions',
626+
model_request_parameters = ModelRequestParameters(
630627
function_tools=[tool_def],
631628
builtin_tools=[ImageGenerationTool()],
632629
output_tools=[output_tool_def],
633630
)
631+
model_settings = ModelSettings(max_tokens=256)
632+
633+
request = ModelRequest(
634+
parts=[UserPromptPart('test prompt')],
635+
instructions='test instructions',
636+
model_request_parameters=model_request_parameters,
637+
model_settings=model_settings,
638+
)
634639

635-
# Verify the fields are accessible
636-
assert request.function_tools == [tool_def]
637-
assert request.builtin_tools == [ImageGenerationTool()]
638-
assert request.output_tools == [output_tool_def]
640+
# Verify the metadata is accessible
641+
assert request.model_request_parameters is model_request_parameters
642+
assert request.model_settings == model_settings
643+
params = request.model_request_parameters
644+
assert params is not None
645+
assert params.function_tools == [tool_def]
646+
assert params.builtin_tools == [ImageGenerationTool()]
647+
assert params.output_tools == [output_tool_def]
639648

640649
# Serialize - fields ARE excluded
641650
serialized = ModelMessagesTypeAdapter.dump_python([request], mode='json')
642-
assert 'function_tools' not in serialized[0]
643-
assert 'builtin_tools' not in serialized[0]
644-
assert 'output_tools' not in serialized[0]
651+
assert 'model_request_parameters' not in serialized[0]
652+
assert 'model_settings' not in serialized[0]

0 commit comments

Comments
 (0)