Skip to content

Commit acced35

Browse files
committed
Return ModelRequestParameters and ModelSettings on the ModelRequest object
1 parent 365b67b commit acced35

File tree

5 files changed

+124
-32
lines changed

5 files changed

+124
-32
lines changed

pydantic_ai_slim/pydantic_ai/_agent_graph.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,10 @@ async def _prepare_request(
496496
model_request_parameters = await _prepare_request_parameters(ctx)
497497

498498
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+
499503
usage = ctx.state.usage
500504
if ctx.deps.usage_limits.count_tokens_before_request:
501505
# Copy to avoid modifying the original usage object with the counted usage
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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, Any
6+
7+
from . import _utils
8+
from .builtin_tools import AbstractBuiltinTool
9+
10+
if TYPE_CHECKING:
11+
from .tools import ToolDefinition
12+
else: # pragma: no cover
13+
ToolDefinition = Any
14+
15+
if TYPE_CHECKING:
16+
from ._output import OutputObjectDefinition
17+
from .output import OutputMode
18+
19+
__all__ = ('ModelRequestParameters',)
20+
21+
22+
@dataclass(repr=False, kw_only=True)
23+
class ModelRequestParameters:
24+
"""Configuration for an agent's request to a model, specifically related to tools and output handling."""
25+
26+
function_tools: list[ToolDefinition] = field(default_factory=list)
27+
builtin_tools: list[AbstractBuiltinTool] = field(default_factory=list)
28+
29+
output_mode: OutputMode = 'text'
30+
output_object: OutputObjectDefinition | None = None
31+
output_tools: list[ToolDefinition] = field(default_factory=list)
32+
prompted_output_template: str | None = None
33+
allow_text_output: bool = True
34+
allow_image_output: bool = False
35+
36+
@cached_property
37+
def tool_defs(self) -> dict[str, ToolDefinition]:
38+
return {tool_def.name: tool_def for tool_def in [*self.function_tools, *self.output_tools]}
39+
40+
@cached_property
41+
def prompted_output_instructions(self) -> str | None:
42+
if self.output_mode == 'prompted' and self.prompted_output_template and self.output_object:
43+
from ._output import PromptedOutputSchema
44+
45+
return PromptedOutputSchema.build_instructions(self.prompted_output_template, self.output_object)
46+
return None
47+
48+
__repr__ = _utils.dataclasses_no_defaults_repr

pydantic_ai_slim/pydantic_ai/messages.py

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

1818
from . import _otel_messages, _utils
19+
from ._model_request_parameters import ModelRequestParameters
1920
from ._utils import generate_tool_call_id as _generate_tool_call_id, now_utc as _now_utc
2021
from .exceptions import UnexpectedModelBehavior
22+
from .settings import ModelSettings
2123
from .usage import RequestUsage
2224

2325
if TYPE_CHECKING:
2426
from .models.instrumented import InstrumentationSettings
2527

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

2735
AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg', 'audio/ogg', 'audio/flac', 'audio/aiff', 'audio/aac']
2836
ImageMediaType: TypeAlias = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
@@ -945,6 +953,22 @@ class ModelRequest:
945953
instructions: str | None = None
946954
"""The instructions for the model."""
947955

956+
model_request_parameters: Annotated[ModelRequestParametersField, pydantic.Field(exclude=True, repr=False)] = field(
957+
default=None, repr=False
958+
)
959+
"""Full request parameters captured for this request.
960+
961+
Available for introspection during a run. This field is excluded from serialization.
962+
"""
963+
964+
model_settings: Annotated[ModelSettingsField, pydantic.Field(exclude=True, repr=False)] = field(
965+
default=None, repr=False
966+
)
967+
"""Effective model settings that were applied to this request.
968+
969+
Available for introspection during a run. This field is excluded from serialization.
970+
"""
971+
948972
kind: Literal['request'] = 'request'
949973
"""Message type identifier, this is available on all parts as a discriminator."""
950974

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: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@
2626
UserPromptPart,
2727
VideoUrl,
2828
)
29+
from pydantic_ai.builtin_tools import ImageGenerationTool
30+
from pydantic_ai.models import ModelRequestParameters, ToolDefinition
31+
from pydantic_ai.settings import ModelSettings
2932

3033
from .conftest import IsDatetime, IsNow, IsStr
3134

@@ -404,7 +407,7 @@ def test_pre_usage_refactor_messages_deserializable():
404407
content='What is the capital of Mexico?',
405408
timestamp=IsNow(tz=timezone.utc),
406409
)
407-
]
410+
],
408411
),
409412
ModelResponse(
410413
parts=[TextPart(content='Mexico City.')],
@@ -605,3 +608,45 @@ def test_binary_content_validation_with_optional_identifier():
605608
'identifier': 'foo',
606609
}
607610
)
611+
612+
613+
def test_model_request_tool_tracking_excluded_from_serialization():
614+
"""Test that request metadata is accessible but not serialized."""
615+
tool_def = ToolDefinition(
616+
name='test_tool',
617+
description='A test tool',
618+
parameters_json_schema={'type': 'object', 'properties': {}},
619+
)
620+
output_tool_def = ToolDefinition(
621+
name='request_output',
622+
description='An output tool',
623+
parameters_json_schema={'type': 'object', 'properties': {}},
624+
)
625+
626+
model_request_parameters = ModelRequestParameters(
627+
function_tools=[tool_def],
628+
builtin_tools=[ImageGenerationTool()],
629+
output_tools=[output_tool_def],
630+
)
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+
)
639+
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]
648+
649+
# Serialize - fields ARE excluded
650+
serialized = ModelMessagesTypeAdapter.dump_python([request], mode='json')
651+
assert 'model_request_parameters' not in serialized[0]
652+
assert 'model_settings' not in serialized[0]

0 commit comments

Comments
 (0)