[WIP] Allow customization of prompts(ModelRetry,ToolReturn), tool_args, tool_descriptions sent to the Model.

docs/deferred-tools.md

@@ -150,6 +150,7 @@ print(result.all_messages())
                 content="File 'README.md' updated: 'Hello, world!'",
                 tool_call_id='update_file_readme',
                 timestamp=datetime.datetime(...),
+                return_kind='tool-executed',
             )
         ],
         run_id='...',
@@ -161,12 +162,14 @@ print(result.all_messages())
                 content="File '.env' updated: ''",
                 tool_call_id='update_file_dotenv',
                 timestamp=datetime.datetime(...),
+                return_kind='tool-executed',
             ),
             ToolReturnPart(
                 tool_name='delete_file',
                 content='Deleting files is not allowed',
                 tool_call_id='delete_file',
                 timestamp=datetime.datetime(...),
+                return_kind='tool-denied',
             ),
             UserPromptPart(
                 content='Now create a backup of README.md',
@@ -195,6 +198,7 @@ print(result.all_messages())
                 content="File 'README.md.bak' updated: 'Hello, world!'",
                 tool_call_id='update_file_backup',
                 timestamp=datetime.datetime(...),
+                return_kind='tool-executed',
             )
         ],
         run_id='...',
@@ -348,6 +352,7 @@ async def main():
                     content=42,
                     tool_call_id='pyd_ai_tool_call_id',
                     timestamp=datetime.datetime(...),
+                    return_kind='tool-executed',
                 )
             ],
             run_id='...',

docs/testing.md

@@ -156,6 +156,7 @@ async def test_forecast():
                     content='Sunny with a chance of rain',
                     tool_call_id=IsStr(),
                     timestamp=IsNow(tz=timezone.utc),
+                    return_kind='tool-executed',
                 ),
             ],
             run_id=IsStr(),

docs/tools.md

@@ -108,6 +108,7 @@ print(dice_result.all_messages())
                 content='4',
                 tool_call_id='pyd_ai_tool_call_id',
                 timestamp=datetime.datetime(...),
+                return_kind='tool-executed',
             )
         ],
         run_id='...',
@@ -130,6 +131,7 @@ print(dice_result.all_messages())
                 content='Anne',
                 tool_call_id='pyd_ai_tool_call_id',
                 timestamp=datetime.datetime(...),
+                return_kind='tool-executed',
             )
         ],
         run_id='...',

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -138,6 +138,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
 
     model: models.Model
     model_settings: ModelSettings | None
+    prompt_templates: _messages.PromptTemplates | None
     usage_limits: _usage.UsageLimits
     max_result_retries: int
     end_strategy: EndStrategy
@@ -509,6 +510,10 @@ async def _prepare_request(
         # Update the new message index to ensure `result.new_messages()` returns the correct messages
         ctx.deps.new_message_index -= len(original_history) - len(message_history)
 
+        prompt_templates = ctx.deps.prompt_templates
+        if prompt_templates:
+            _apply_prompt_templates(message_history, prompt_templates, run_context)
+
         # Merge possible consecutive trailing `ModelRequest`s into one, with tool call parts before user parts,
         # but don't store it in the message history on state. This is just for the benefit of model classes that want clear user/assistant boundaries.
         # See `tests/test_tools.py::test_parallel_tool_return_with_deferred` for an example where this is necessary
@@ -783,6 +788,11 @@ def _handle_final_result(
     ) -> End[result.FinalResult[NodeRunEndT]]:
         messages = ctx.state.message_history
 
+        if tool_responses and ctx.deps.prompt_templates:
+            run_ctx = build_run_context(ctx)
+            for part in tool_responses:
+                ctx.deps.prompt_templates.apply_template(part, run_ctx)
+
         # For backwards compatibility, append a new ModelRequest using the tool returns and retries
         if tool_responses:
             messages.append(_messages.ModelRequest(parts=tool_responses, run_id=ctx.state.run_id))
@@ -871,13 +881,15 @@ async def process_tool_calls(  # noqa: C901
                     tool_name=call.tool_name,
                     content='Final result processed.',
                     tool_call_id=call.tool_call_id,
+                    return_kind='final-result-processed',
                 )
             else:
                 yield _messages.FunctionToolCallEvent(call)
                 part = _messages.ToolReturnPart(
                     tool_name=call.tool_name,
                     content='Output tool not used - a final result was already processed.',
                     tool_call_id=call.tool_call_id,
+                    return_kind='output-tool-not-executed',
                 )
                 yield _messages.FunctionToolResultEvent(part)
 
@@ -902,6 +914,7 @@ async def process_tool_calls(  # noqa: C901
                     tool_name=call.tool_name,
                     content='Final result processed.',
                     tool_call_id=call.tool_call_id,
+                    return_kind='final-result-processed',
                 )
                 output_parts.append(part)
                 final_result = result.FinalResult(result_data, call.tool_name, call.tool_call_id)
@@ -915,6 +928,7 @@ async def process_tool_calls(  # noqa: C901
                     tool_name=call.tool_name,
                     content='Tool not executed - a final result was already processed.',
                     tool_call_id=call.tool_call_id,
+                    return_kind='function-tool-not-executed',
                 )
             )
     else:
@@ -973,6 +987,7 @@ async def process_tool_calls(  # noqa: C901
                             tool_name=call.tool_name,
                             content='Tool not executed - a final result was already processed.',
                             tool_call_id=call.tool_call_id,
+                            return_kind='function-tool-not-executed',
                         )
                     )
         elif calls:
@@ -1129,6 +1144,7 @@ async def _call_tool(
                 tool_name=tool_call.tool_name,
                 content=tool_call_result.message,
                 tool_call_id=tool_call.tool_call_id,
+                return_kind='tool-denied',
             ), None
         elif isinstance(tool_call_result, exceptions.ModelRetry):
             m = _messages.RetryPromptPart(
@@ -1191,6 +1207,7 @@ async def _call_tool(
         tool_call_id=tool_call.tool_call_id,
         content=tool_return.return_value,  # type: ignore
         metadata=tool_return.metadata,
+        return_kind='tool-executed',
     )
 
     return return_part, tool_return.content or None
@@ -1361,3 +1378,11 @@ def _clean_message_history(messages: list[_messages.ModelMessage]) -> list[_mess
             else:
                 clean_messages.append(message)
     return clean_messages
+
+
+def _apply_prompt_templates(
+    messages: list[_messages.ModelMessage], prompt_templates: _messages.PromptTemplates, ctx: RunContext[Any]
+):
+    for msg in messages:
+        for msg_part in msg.parts:
+            prompt_templates.apply_template(msg_part, ctx)

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -125,6 +125,9 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     be merged with this value, with the runtime argument taking priority.
     """
 
+    prompt_templates: _messages.PromptTemplates | None
+    """Optional prompt templates used to customize the system-injected messages for this agent."""
+
     _output_type: OutputSpec[OutputDataT]
 
     instrument: InstrumentationSettings | bool | None
@@ -167,6 +170,7 @@ def __init__(
         deps_type: type[AgentDepsT] = NoneType,
         name: str | None = None,
         model_settings: ModelSettings | None = None,
+        prompt_templates: _messages.PromptTemplates | None = None,
         retries: int = 1,
         validation_context: Any | Callable[[RunContext[AgentDepsT]], Any] = None,
         output_retries: int | None = None,
@@ -219,6 +223,7 @@ def __init__(
         deps_type: type[AgentDepsT] = NoneType,
         name: str | None = None,
         model_settings: ModelSettings | None = None,
+        prompt_templates: _messages.PromptTemplates | None = None,
         retries: int = 1,
         validation_context: Any | Callable[[RunContext[AgentDepsT]], Any] = None,
         output_retries: int | None = None,
@@ -252,6 +257,8 @@ def __init__(
             name: The name of the agent, used for logging. If `None`, we try to infer the agent name from the call frame
                 when the agent is first run.
             model_settings: Optional model request settings to use for this agent's runs, by default.
+            prompt_templates: Optional prompt templates to customize how system-injected messages
+                (like retry prompts or tool return wrappers) are rendered for this agent.
             retries: The default number of retries to allow for tool calls and output validation, before raising an error.
                 For model request retries, see the [HTTP Request Retries](../retries.md) documentation.
             validation_context: Pydantic [validation context](https://docs.pydantic.dev/latest/concepts/validators/#validation-context) used to validate tool arguments and outputs.
@@ -295,6 +302,7 @@ def __init__(
         self._name = name
         self.end_strategy = end_strategy
         self.model_settings = model_settings
+        self.prompt_templates = prompt_templates
 
         self._output_type = output_type
         self.instrument = instrument
@@ -357,6 +365,9 @@ def __init__(
         self._override_instructions: ContextVar[
             _utils.Option[list[str | _system_prompt.SystemPromptFunc[AgentDepsT]]]
         ] = ContextVar('_override_instructions', default=None)
+        self._override_prompt_templates: ContextVar[_utils.Option[_messages.PromptTemplates]] = ContextVar(
+            '_override_prompt_templates', default=None
+        )
 
         self._enter_lock = Lock()
         self._entered_count = 0
@@ -410,7 +421,7 @@ def event_stream_handler(self) -> EventStreamHandler[AgentDepsT] | None:
         return self._event_stream_handler
 
     def __repr__(self) -> str:
-        return f'{type(self).__name__}(model={self.model!r}, name={self.name!r}, end_strategy={self.end_strategy!r}, model_settings={self.model_settings!r}, output_type={self.output_type!r}, instrument={self.instrument!r})'
+        return f'{type(self).__name__}(model={self.model!r}, name={self.name!r}, end_strategy={self.end_strategy!r}, model_settings={self.model_settings!r}, prompt_templates={self.prompt_templates!r},output_type={self.output_type!r}, instrument={self.instrument!r})'
 
     @overload
     def iter(
@@ -424,6 +435,7 @@ def iter(
         instructions: Instructions[AgentDepsT] = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
+        prompt_templates: _messages.PromptTemplates | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
@@ -443,6 +455,7 @@ def iter(
         instructions: Instructions[AgentDepsT] = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
+        prompt_templates: _messages.PromptTemplates | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
@@ -462,6 +475,7 @@ async def iter(
         instructions: Instructions[AgentDepsT] = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
+        prompt_templates: _messages.PromptTemplates | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
@@ -538,6 +552,8 @@ async def main():
             instructions: Optional additional instructions to use for this run.
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
+            prompt_templates: Optional prompt templates to override how system-generated parts are
+                phrased for this specific run, falling back to the agent's defaults if omitted.
             usage_limits: Optional limits on model request count or token usage.
             usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
@@ -588,6 +604,7 @@ async def main():
         merged_settings = merge_model_settings(model_used.settings, self.model_settings)
         model_settings = merge_model_settings(merged_settings, model_settings)
         usage_limits = usage_limits or _usage.UsageLimits()
+        prompt_templates = self._get_prompt_templates(prompt_templates)
 
         instructions_literal, instructions_functions = self._get_instructions(additional_instructions=instructions)
 
@@ -615,6 +632,7 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
             new_message_index=len(message_history) if message_history else 0,
             model=model_used,
             model_settings=model_settings,
+            prompt_templates=prompt_templates,
             usage_limits=usage_limits,
             max_result_retries=self._max_result_retries,
             end_strategy=self.end_strategy,
@@ -749,6 +767,7 @@ def override(
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET,
         instructions: Instructions[AgentDepsT] | _utils.Unset = _utils.UNSET,
+        prompt_templates: _messages.PromptTemplates | _utils.Unset = _utils.UNSET,
     ) -> Iterator[None]:
         """Context manager to temporarily override agent name, dependencies, model, toolsets, tools, or instructions.
 
@@ -762,6 +781,7 @@ def override(
             toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run.
             tools: The tools to use instead of the tools registered with the agent.
             instructions: The instructions to use instead of the instructions registered with the agent.
+            prompt_templates: The prompt templates to use instead of the prompt templates passed to the agent constructor and agent run.
         """
         if _utils.is_set(name):
             name_token = self._override_name.set(_utils.Some(name))
@@ -794,6 +814,11 @@ def override(
         else:
             instructions_token = None
 
+        if _utils.is_set(prompt_templates):
+            prompt_templates_token = self._override_prompt_templates.set(_utils.Some(prompt_templates))
+        else:
+            prompt_templates_token = None
+
         try:
             yield
         finally:
@@ -809,6 +834,8 @@ def override(
                 self._override_tools.reset(tools_token)
             if instructions_token is not None:
                 self._override_instructions.reset(instructions_token)
+            if prompt_templates_token is not None:
+                self._override_prompt_templates.reset(prompt_templates_token)
 
     @overload
     def instructions(
@@ -1323,6 +1350,19 @@ def _get_deps(self: Agent[T, OutputDataT], deps: T) -> T:
         else:
             return deps
 
+    def _get_prompt_templates(
+        self, prompt_templates: _messages.PromptTemplates | None
+    ) -> _messages.PromptTemplates | None:
+        """Get prompt_templates for a run.
+
+        If we've overridden prompt_templates via `_override_prompt_templates`, use that,
+        otherwise use the prompt_templates passed to the call, falling back to the agent default.
+        """
+        if some_prompt_templates := self._override_prompt_templates.get():
+            return some_prompt_templates.value
+        else:
+            return prompt_templates or self.prompt_templates
+
     def _normalize_instructions(
         self,
         instructions: Instructions[AgentDepsT],