feat: add RunContext support to history processors (#2002)

Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: Wh1isper

docs/message-history.md

@@ -386,6 +386,34 @@ long_conversation_history: list[ModelMessage] = []  # Your long conversation his
 # result = agent.run_sync('What did we discuss?', message_history=long_conversation_history)
 ```
 
+#### `RunContext` parameter
+
+History processors can optionally accept a [`RunContext`][pydantic_ai.tools.RunContext] parameter to access
+additional information about the current run, such as dependencies, model information, and usage statistics:
+
+```python {title="context_aware_processor.py"}
+from pydantic_ai import Agent
+from pydantic_ai.messages import ModelMessage
+from pydantic_ai.tools import RunContext
+
+
+def context_aware_processor(
+    ctx: RunContext[None],
+    messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    # Access current usage
+    current_tokens = ctx.usage.total_tokens
+
+    # Filter messages based on context
+    if current_tokens > 1000:
+        return messages[-3:]  # Keep only recent messages when token usage is high
+    return messages
+
+agent = Agent('openai:gpt-4o', history_processors=[context_aware_processor])
+```
+
+This allows for more sophisticated message processing based on the current state of the agent run.
+
 #### Summarize Old Messages
 
 Use an LLM to summarize older messages to preserve context while reducing tokens.

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -12,6 +12,7 @@
 from opentelemetry.trace import Tracer
 from typing_extensions import TypeGuard, TypeVar, assert_never
 
+from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx  # type: ignore
 from pydantic_ai._utils import is_async_callable, run_in_executor
 from pydantic_graph import BaseNode, Graph, GraphRunContext
 from pydantic_graph.nodes import End, NodeRunEndT
@@ -50,8 +51,20 @@
 
 _HistoryProcessorSync = Callable[[list[_messages.ModelMessage]], list[_messages.ModelMessage]]
 _HistoryProcessorAsync = Callable[[list[_messages.ModelMessage]], Awaitable[list[_messages.ModelMessage]]]
-HistoryProcessor = Union[_HistoryProcessorSync, _HistoryProcessorAsync]
-"""A function that processes a list of model messages and returns a list of model messages."""
+_HistoryProcessorSyncWithCtx = Callable[[RunContext[DepsT], list[_messages.ModelMessage]], list[_messages.ModelMessage]]
+_HistoryProcessorAsyncWithCtx = Callable[
+    [RunContext[DepsT], list[_messages.ModelMessage]], Awaitable[list[_messages.ModelMessage]]
+]
+HistoryProcessor = Union[
+    _HistoryProcessorSync,
+    _HistoryProcessorAsync,
+    _HistoryProcessorSyncWithCtx[DepsT],
+    _HistoryProcessorAsyncWithCtx[DepsT],
+]
+"""A function that processes a list of model messages and returns a list of model messages.
+
+Can optionally accept a `RunContext` as a parameter.
+"""
 
 
 @dataclasses.dataclass
@@ -92,7 +105,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     output_schema: _output.OutputSchema[OutputDataT] | None
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
 
-    history_processors: Sequence[HistoryProcessor]
+    history_processors: Sequence[HistoryProcessor[DepsT]]
 
     function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False)
     mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
@@ -328,7 +341,9 @@ async def _stream(
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
-        message_history = await _process_message_history(ctx.state.message_history, ctx.deps.history_processors)
+        message_history = await _process_message_history(
+            ctx.state.message_history, ctx.deps.history_processors, build_run_context(ctx)
+        )
         async with ctx.deps.model.request_stream(
             message_history, model_settings, model_request_parameters
         ) as streamed_response:
@@ -352,7 +367,9 @@ async def _make_request(
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
-        message_history = await _process_message_history(ctx.state.message_history, ctx.deps.history_processors)
+        message_history = await _process_message_history(
+            ctx.state.message_history, ctx.deps.history_processors, build_run_context(ctx)
+        )
         model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
         ctx.state.usage.incr(_usage.Usage())
 
@@ -881,12 +898,24 @@ def build_agent_graph(
 
 async def _process_message_history(
     messages: list[_messages.ModelMessage],
-    processors: Sequence[HistoryProcessor],
+    processors: Sequence[HistoryProcessor[DepsT]],
+    run_context: RunContext[DepsT],
 ) -> list[_messages.ModelMessage]:
     """Process message history through a sequence of processors."""
     for processor in processors:
+        takes_ctx = is_takes_ctx(processor)
+
         if is_async_callable(processor):
-            messages = await processor(messages)
+            if takes_ctx:
+                messages = await processor(run_context, messages)
+            else:
+                async_processor = cast(_HistoryProcessorAsync, processor)
+                messages = await async_processor(messages)
         else:
-            messages = await run_in_executor(processor, messages)
+            if takes_ctx:
+                sync_processor_with_ctx = cast(_HistoryProcessorSyncWithCtx[DepsT], processor)
+                messages = await run_in_executor(sync_processor_with_ctx, run_context, messages)
+            else:
+                sync_processor = cast(_HistoryProcessorSync, processor)
+                messages = await run_in_executor(sync_processor, messages)
     return messages

pydantic_ai_slim/pydantic_ai/_function_schema.py

@@ -8,7 +8,7 @@
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Callable, cast
+from typing import TYPE_CHECKING, Any, Callable, Union, cast
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -17,7 +17,7 @@
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import get_origin
+from typing_extensions import Concatenate, ParamSpec, TypeIs, TypeVar, get_origin
 
 from pydantic_ai.tools import RunContext
 
@@ -218,7 +218,16 @@ def function_schema(  # noqa: C901
     )
 
 
-def _takes_ctx(function: Callable[..., Any]) -> bool:
+P = ParamSpec('P')
+R = TypeVar('R')
+
+
+WithCtx = Callable[Concatenate[RunContext[Any], P], R]
+WithoutCtx = Callable[P, R]
+TargetFunc = Union[WithCtx[P, R], WithoutCtx[P, R]]
+
+
+def _takes_ctx(function: TargetFunc[P, R]) -> TypeIs[WithCtx[P, R]]:
     """Check if a function takes a `RunContext` first argument.
 
     Args:

pydantic_ai_slim/pydantic_ai/agent.py

@@ -180,7 +180,7 @@ def __init__(
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
-        history_processors: Sequence[HistoryProcessor] | None = None,
+        history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
     ) -> None: ...
 
     @overload
@@ -210,7 +210,7 @@ def __init__(
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
-        history_processors: Sequence[HistoryProcessor] | None = None,
+        history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
     ) -> None: ...
 
     def __init__(
@@ -235,7 +235,7 @@ def __init__(
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
-        history_processors: Sequence[HistoryProcessor] | None = None,
+        history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
         **_deprecated_kwargs: Any,
     ):
         """Create an agent.

tests/test_history_processor.py

@@ -1,11 +1,13 @@
 from collections.abc import AsyncIterator
+from typing import Any, cast
 
 import pytest
 from inline_snapshot import snapshot
 
 from pydantic_ai import Agent
 from pydantic_ai.messages import ModelMessage, ModelRequest, ModelRequestPart, ModelResponse, TextPart, UserPromptPart
 from pydantic_ai.models.function import AgentInfo, FunctionModel
+from pydantic_ai.tools import RunContext
 from pydantic_ai.usage import Usage
 
 from .conftest import IsDatetime
@@ -201,3 +203,101 @@ async def async_processor(messages: list[ModelMessage]) -> list[ModelMessage]:
             ModelRequest(parts=[UserPromptPart(content='Question 2', timestamp=IsDatetime())]),
         ]
     )
+
+
+async def test_history_processor_with_context(function_model: FunctionModel, received_messages: list[ModelMessage]):
+    """Test history processor that takes RunContext."""
+
+    def context_processor(ctx: RunContext[str], messages: list[ModelMessage]) -> list[ModelMessage]:
+        # Access deps from context
+        prefix = ctx.deps
+        processed: list[ModelMessage] = []
+        for msg in messages:
+            if isinstance(msg, ModelRequest):
+                new_parts: list[ModelRequestPart] = []
+                for part in msg.parts:
+                    if isinstance(part, UserPromptPart):
+                        new_parts.append(UserPromptPart(content=f'{prefix}: {part.content}'))
+                    else:
+                        new_parts.append(part)  # pragma: no cover
+                processed.append(ModelRequest(parts=new_parts))
+            else:
+                processed.append(msg)  # pragma: no cover
+        return processed
+
+    agent = Agent(function_model, history_processors=[context_processor], deps_type=str)
+    await agent.run('test', deps='PREFIX')
+
+    # Verify the prefix was added
+    assert len(received_messages) == 1
+    assert isinstance(received_messages[0], ModelRequest)
+    user_part = received_messages[0].parts[0]
+    assert isinstance(user_part, UserPromptPart)
+    assert user_part.content == 'PREFIX: test'
+
+
+async def test_history_processor_with_context_async(
+    function_model: FunctionModel, received_messages: list[ModelMessage]
+):
+    """Test async history processor that takes RunContext."""
+
+    async def async_context_processor(ctx: RunContext[Any], messages: list[ModelMessage]) -> list[ModelMessage]:
+        return messages[-1:]  # Keep only the last message
+
+    message_history = [
+        ModelRequest(parts=[UserPromptPart(content='Question 1')]),
+        ModelResponse(parts=[TextPart(content='Answer 1')]),
+        ModelRequest(parts=[UserPromptPart(content='Question 2')]),
+        ModelResponse(parts=[TextPart(content='Answer 2')]),
+    ]
+
+    agent = Agent(function_model, history_processors=[async_context_processor])
+    await agent.run('Question 3', message_history=message_history)
+
+    # Should have filtered to only recent messages
+    assert len(received_messages) <= 2  # Last message from history + new message
+
+
+async def test_history_processor_mixed_signatures(function_model: FunctionModel, received_messages: list[ModelMessage]):
+    """Test mixing processors with and without context."""
+
+    def simple_processor(messages: list[ModelMessage]) -> list[ModelMessage]:
+        # Filter out responses
+        return [msg for msg in messages if isinstance(msg, ModelRequest)]
+
+    def context_processor(ctx: RunContext[Any], messages: list[ModelMessage]) -> list[ModelMessage]:
+        # Add prefix based on deps
+        prefix = getattr(ctx.deps, 'prefix', 'DEFAULT')
+        processed: list[ModelMessage] = []
+        for msg in messages:
+            if isinstance(msg, ModelRequest):
+                new_parts: list[ModelRequestPart] = []
+                for part in msg.parts:
+                    if isinstance(part, UserPromptPart):
+                        new_parts.append(UserPromptPart(content=f'{prefix}: {part.content}'))
+                    else:
+                        new_parts.append(part)  # pragma: no cover
+                processed.append(ModelRequest(parts=new_parts))
+            else:
+                processed.append(msg)  # pragma: no cover
+        return processed
+
+    message_history = [
+        ModelRequest(parts=[UserPromptPart(content='Question 1')]),
+        ModelResponse(parts=[TextPart(content='Answer 1')]),
+    ]
+
+    # Create deps with prefix attribute
+    class Deps:
+        prefix = 'TEST'
+
+    agent = Agent(function_model, history_processors=[simple_processor, context_processor], deps_type=Deps)
+    await agent.run('Question 2', message_history=message_history, deps=Deps())
+
+    # Should have filtered responses and added prefix
+    assert len(received_messages) == 2
+    for msg in received_messages:
+        assert isinstance(msg, ModelRequest)
+        user_part = msg.parts[0]
+        assert isinstance(user_part, UserPromptPart)
+        assert cast(str, user_part.content).startswith('TEST: ')