Remove last run messages (#536)

Author: samuelcolvin

docs/agents.md

@@ -408,10 +408,10 @@ user_id=123 message='Hello John, would you be free for coffee sometime next week
 
 If models behave unexpectedly (e.g., the retry limit is exceeded, or their API returns `503`), agent runs will raise [`UnexpectedModelBehavior`][pydantic_ai.exceptions.UnexpectedModelBehavior].
 
-In these cases, [`agent.last_run_messages`][pydantic_ai.Agent.last_run_messages] can be used to access the messages exchanged during the run to help diagnose the issue.
+In these cases, [`capture_run_messages`][pydantic_ai.capture_run_messages] can be used to access the messages exchanged during the run to help diagnose the issue.
 
 ```python
-from pydantic_ai import Agent, ModelRetry, UnexpectedModelBehavior
+from pydantic_ai import Agent, ModelRetry, UnexpectedModelBehavior, capture_run_messages
 
 agent = Agent('openai:gpt-4o')
 
@@ -424,68 +424,76 @@ def calc_volume(size: int) -> int:  # (1)!
         raise ModelRetry('Please try again.')
 
 
-try:
-    result = agent.run_sync('Please get me the volume of a box with size 6.')
-except UnexpectedModelBehavior as e:
-    print('An error occurred:', e)
-    #> An error occurred: Tool exceeded max retries count of 1
-    print('cause:', repr(e.__cause__))
-    #> cause: ModelRetry('Please try again.')
-    print('messages:', agent.last_run_messages)
-    """
-    messages:
-    [
-        ModelRequest(
-            parts=[
-                UserPromptPart(
-                    content='Please get me the volume of a box with size 6.',
-                    timestamp=datetime.datetime(...),
-                    part_kind='user-prompt',
-                )
-            ],
-            kind='request',
-        ),
-        ModelResponse(
-            parts=[
-                ToolCallPart(
-                    tool_name='calc_volume',
-                    args=ArgsDict(args_dict={'size': 6}),
-                    tool_call_id=None,
-                    part_kind='tool-call',
-                )
-            ],
-            timestamp=datetime.datetime(...),
-            kind='response',
-        ),
-        ModelRequest(
-            parts=[
-                RetryPromptPart(
-                    content='Please try again.',
-                    tool_name='calc_volume',
-                    tool_call_id=None,
-                    timestamp=datetime.datetime(...),
-                    part_kind='retry-prompt',
-                )
-            ],
-            kind='request',
-        ),
-        ModelResponse(
-            parts=[
-                ToolCallPart(
-                    tool_name='calc_volume',
-                    args=ArgsDict(args_dict={'size': 6}),
-                    tool_call_id=None,
-                    part_kind='tool-call',
-                )
-            ],
-            timestamp=datetime.datetime(...),
-            kind='response',
-        ),
-    ]
-    """
-else:
-    print(result.data)
+with capture_run_messages() as messages:  # (2)!
+    try:
+        result = agent.run_sync('Please get me the volume of a box with size 6.')
+    except UnexpectedModelBehavior as e:
+        print('An error occurred:', e)
+        #> An error occurred: Tool exceeded max retries count of 1
+        print('cause:', repr(e.__cause__))
+        #> cause: ModelRetry('Please try again.')
+        print('messages:', messages)
+        """
+        messages:
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Please get me the volume of a box with size 6.',
+                        timestamp=datetime.datetime(...),
+                        part_kind='user-prompt',
+                    )
+                ],
+                kind='request',
+            ),
+            ModelResponse(
+                parts=[
+                    ToolCallPart(
+                        tool_name='calc_volume',
+                        args=ArgsDict(args_dict={'size': 6}),
+                        tool_call_id=None,
+                        part_kind='tool-call',
+                    )
+                ],
+                timestamp=datetime.datetime(...),
+                kind='response',
+            ),
+            ModelRequest(
+                parts=[
+                    RetryPromptPart(
+                        content='Please try again.',
+                        tool_name='calc_volume',
+                        tool_call_id=None,
+                        timestamp=datetime.datetime(...),
+                        part_kind='retry-prompt',
+                    )
+                ],
+                kind='request',
+            ),
+            ModelResponse(
+                parts=[
+                    ToolCallPart(
+                        tool_name='calc_volume',
+                        args=ArgsDict(args_dict={'size': 6}),
+                        tool_call_id=None,
+                        part_kind='tool-call',
+                    )
+                ],
+                timestamp=datetime.datetime(...),
+                kind='response',
+            ),
+        ]
+        """
+    else:
+        print(result.data)
 ```
+
 1. Define a tool that will raise `ModelRetry` repeatedly in this case.
+2. [`capture_run_messages`][pydantic_ai.capture_run_messages] is used to capture the messages exchanged during the run.
 
 _(This example is complete, it can be run "as is")_
+
+!!! note
+    You may not call [`run`][pydantic_ai.Agent.run], [`run_sync`][pydantic_ai.Agent.run_sync], or [`run_stream`][pydantic_ai.Agent.run_stream] more than once within a single `capture_run_messages` context.
+
+    If you try to do so, a [`UserError`][pydantic_ai.exceptions.UserError] will be raised.

docs/api/agent.md

@@ -1,21 +1,3 @@
-# `pydantic_ai.Agent`
+# `pydantic_ai.agent`
 
-::: pydantic_ai.Agent
-    options:
-      members:
-        - __init__
-        - name
-        - run
-        - run_sync
-        - run_stream
-        - model
-        - override
-        - last_run_messages
-        - system_prompt
-        - tool
-        - tool_plain
-        - result_validator
-
-::: pydantic_ai.agent.EndStrategy
-    options:
-      show_root_heading: true
+::: pydantic_ai.agent

docs/testing-evals.md

@@ -95,7 +95,7 @@ import pytest
 
 from dirty_equals import IsNow
 
-from pydantic_ai import models
+from pydantic_ai import models, capture_run_messages
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.messages import (
     ArgsDict,
@@ -118,14 +118,15 @@ models.ALLOW_MODEL_REQUESTS = False  # (2)!
 async def test_forecast():
     conn = DatabaseConn()
     user_id = 1
-    with weather_agent.override(model=TestModel()):  # (3)!
-        prompt = 'What will the weather be like in London on 2024-11-28?'
-        await run_weather_forecast([(prompt, user_id)], conn)  # (4)!
+    with capture_run_messages() as messages:
+        with weather_agent.override(model=TestModel()):  # (3)!
+            prompt = 'What will the weather be like in London on 2024-11-28?'
+            await run_weather_forecast([(prompt, user_id)], conn)  # (4)!
 
     forecast = await conn.get_forecast(user_id)
     assert forecast == '{"weather_forecast":"Sunny with a chance of rain"}'  # (5)!
 
-    assert weather_agent.last_run_messages == [  # (6)!
+    assert messages == [  # (6)!
         ModelRequest(
             parts=[
                 SystemPromptPart(
@@ -178,7 +179,7 @@ async def test_forecast():
 3. We're using [`Agent.override`][pydantic_ai.agent.Agent.override] to replace the agent's model with [`TestModel`][pydantic_ai.models.test.TestModel], the nice thing about `override` is that we can replace the model inside agent without needing access to the agent `run*` methods call site.
 4. Now we call the function we want to test inside the `override` context manager.
 5. But default, `TestModel` will return a JSON string summarising the tools calls made, and what was returned. If you wanted to customise the response to something more closely aligned with the domain, you could add [`custom_result_text='Sunny'`][pydantic_ai.models.test.TestModel.custom_result_text] when defining `TestModel`.
-6. So far we don't actually know which tools were called and with which values, we can use the [`last_run_messages`][pydantic_ai.agent.Agent.last_run_messages] attribute to inspect messages from the most recent run and assert the exchange between the agent and the model occurred as expected.
+6. So far we don't actually know which tools were called and with which values, we can use [`capture_run_messages`][pydantic_ai.capture_run_messages] to inspect messages from the most recent run and assert the exchange between the agent and the model occurred as expected.
 7. The [`IsNow`][dirty_equals.IsNow] helper allows us to use declarative asserts even with data which will contain timestamps that change over time.
 8. `TestModel` isn't doing anything clever to extract values from the prompt, so these values are hardcoded.
 

pydantic_ai_slim/pydantic_ai/__init__.py

@@ -1,11 +1,12 @@
 from importlib.metadata import version
 
-from .agent import Agent
+from .agent import Agent, capture_run_messages
 from .exceptions import AgentRunError, ModelRetry, UnexpectedModelBehavior, UsageLimitExceeded, UserError
 from .tools import RunContext, Tool
 
 __all__ = (
     'Agent',
+    'capture_run_messages',
     'RunContext',
     'Tool',
     'AgentRunError',

pydantic_ai_slim/pydantic_ai/agent.py

@@ -5,12 +5,13 @@
 import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass, field
 from types import FrameType
 from typing import Any, Callable, Generic, Literal, cast, final, overload
 
 import logfire_api
-from typing_extensions import assert_never
+from typing_extensions import assert_never, deprecated
 
 from . import (
     _result,
@@ -35,7 +36,7 @@
     ToolPrepareFunc,
 )
 
-__all__ = ('Agent',)
+__all__ = 'Agent', 'capture_run_messages', 'EndStrategy'
 
 _logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
 
@@ -89,12 +90,6 @@ class Agent(Generic[AgentDeps, ResultData]):
     be merged with this value, with the runtime argument taking priority.
     """
 
-    last_run_messages: list[_messages.ModelMessage] | None
-    """The messages from the last run, useful when a run raised an exception.
-
-    Note: these are not used by the agent, e.g. in future runs, they are just stored for developers' convenience.
-    """
-
     _result_schema: _result.ResultSchema[ResultData] | None = field(repr=False)
     _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]] = field(repr=False)
     _allow_text_result: bool = field(repr=False)
@@ -161,7 +156,6 @@ def __init__(
         self.end_strategy = end_strategy
         self.name = name
         self.model_settings = model_settings
-        self.last_run_messages = None
         self._result_schema = _result.ResultSchema[result_type].build(
             result_type, result_tool_name, result_tool_description
         )
@@ -234,7 +228,7 @@ async def run(
         ) as run_span:
             run_context = RunContext(deps, 0, [], None, model_used)
             messages = await self._prepare_messages(user_prompt, message_history, run_context)
-            self.last_run_messages = run_context.messages = messages
+            run_context.messages = messages
 
             for tool in self._function_tools.values():
                 tool.current_retry = 0
@@ -393,7 +387,7 @@ async def main():
         ) as run_span:
             run_context = RunContext(deps, 0, [], None, model_used)
             messages = await self._prepare_messages(user_prompt, message_history, run_context)
-            self.last_run_messages = run_context.messages = messages
+            run_context.messages = messages
 
             for tool in self._function_tools.values():
                 tool.current_retry = 0
@@ -614,7 +608,7 @@ async def result_validator_deps(ctx: RunContext[str], data: str) -> str:
         #> success (no tool calls)
         ```
         """
-        self._result_validators.append(_result.ResultValidator(func))
+        self._result_validators.append(_result.ResultValidator[AgentDeps, Any](func))
         return func
 
     @overload
@@ -835,14 +829,25 @@ async def add_tool(tool: Tool[AgentDeps]) -> None:
     async def _prepare_messages(
         self, user_prompt: str, message_history: list[_messages.ModelMessage] | None, run_context: RunContext[AgentDeps]
     ) -> list[_messages.ModelMessage]:
+        try:
+            messages = _messages_ctx_var.get()
+        except LookupError:
+            messages = []
+        else:
+            if messages:
+                raise exceptions.UserError(
+                    'The capture_run_messages() context manager may only be used to wrap '
+                    'one call to run(), run_sync(), or run_stream().'
+                )
+
         if message_history:
             # shallow copy messages
-            messages = message_history.copy()
+            messages.extend(message_history)
             messages.append(_messages.ModelRequest([_messages.UserPromptPart(user_prompt)]))
         else:
             parts = await self._sys_parts(run_context)
             parts.append(_messages.UserPromptPart(user_prompt))
-            messages: list[_messages.ModelMessage] = [_messages.ModelRequest(parts)]
+            messages.append(_messages.ModelRequest(parts))
 
         return messages
 
@@ -1119,6 +1124,51 @@ def _infer_name(self, function_frame: FrameType | None) -> None:
                             self.name = name
                             return
 
+    @property
+    @deprecated(
+        'The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.', category=None
+    )
+    def last_run_messages(self) -> list[_messages.ModelMessage]:
+        raise AttributeError('The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.')
+
+
+_messages_ctx_var: ContextVar[list[_messages.ModelMessage]] = ContextVar('var')
+
+
+@contextmanager
+def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
+    """Context manager to access the messages used in a [`run`][pydantic_ai.Agent.run], [`run_sync`][pydantic_ai.Agent.run_sync], or [`run_stream`][pydantic_ai.Agent.run_stream] call.
+
+    Useful when a run may raise an exception, see [model errors](../agents.md#model-errors) for more information.
+
+    Examples:
+    ```python
+    from pydantic_ai import Agent, capture_run_messages
+
+    agent = Agent('test')
+
+    with capture_run_messages() as messages:
+        try:
+            result = agent.run_sync('foobar')
+        except Exception:
+            print(messages)
+            raise
+    ```
+
+    !!! note
+        You may not call `run`, `run_sync`, or `run_stream` more than once within a single `capture_run_messages` context.
+        If you try to do so, a [`UserError`][pydantic_ai.exceptions.UserError] will be raised.
+    """
+    try:
+        yield _messages_ctx_var.get()
+    except LookupError:
+        messages: list[_messages.ModelMessage] = []
+        token = _messages_ctx_var.set(messages)
+        try:
+            yield messages
+        finally:
+            _messages_ctx_var.reset(token)
+
 
 @dataclass
 class _MarkFinalResult(Generic[ResultData]):