Allow modifying the input sent to the model

Author: rm-openai

src/agents/run.py

@@ -4,7 +4,7 @@
 import copy
 import inspect
 from dataclasses import dataclass, field
-from typing import Any, Generic, cast
+from typing import Any, Callable, Generic, cast
 
 from openai.types.responses import ResponseCompletedEvent
 from openai.types.responses.response_prompt_param import (
@@ -56,6 +56,7 @@
 from .tracing.span_data import AgentSpanData
 from .usage import Usage
 from .util import _coro, _error_tracing
+from .util._types import MaybeAwaitable
 
 DEFAULT_MAX_TURNS = 10
 
@@ -81,6 +82,27 @@ def get_default_agent_runner() -> AgentRunner:
     return DEFAULT_AGENT_RUNNER
 
 
+@dataclass
+class ModelInputData:
+    """Container for the data that will be sent to the model."""
+
+    input: list[TResponseInputItem]
+    instructions: str | None
+
+
+@dataclass
+class CallModelData(Generic[TContext]):
+    """Data passed to `RunConfig.call_model_input_filter` prior to model call."""
+
+    model_data: ModelInputData
+    agent: Agent[TContext]
+    context: TContext | None
+
+
+# Type alias for the optional input filter callback
+CallModelInputFilter = Callable[[CallModelData[Any]], MaybeAwaitable[ModelInputData]]
+
+
 @dataclass
 class RunConfig:
     """Configures settings for the entire agent run."""
@@ -139,6 +161,16 @@ class RunConfig:
     An optional dictionary of additional metadata to include with the trace.
     """
 
+    call_model_input_filter: CallModelInputFilter | None = None
+    """
+    Optional callback that is invoked immediately before calling the model. It receives the current
+    agent, context and the model input (instructions and input items), and must return a possibly
+    modified `ModelInputData` to use for the model call.
+
+    This allows you to edit the input sent to the model e.g. to stay within a token limit.
+    For example, you can use this to add a system prompt to the input.
+    """
+
 
 class RunOptions(TypedDict, Generic[TContext]):
     """Arguments for ``AgentRunner`` methods."""
@@ -863,10 +895,40 @@ async def _run_single_turn_streamed(
         input = ItemHelpers.input_to_new_input_list(streamed_result.input)
         input.extend([item.to_input_item() for item in streamed_result.new_items])
 
+        # Allow user to modify model input right before the streaming call, if configured
+        effective_instructions = system_prompt
+        effective_input: list[TResponseInputItem] = input
+        if run_config.call_model_input_filter is not None:
+            try:
+                model_input = ModelInputData(
+                    input=copy.deepcopy(effective_input),
+                    instructions=effective_instructions,
+                )
+                filter_payload: CallModelData[TContext] = CallModelData(
+                    model_data=model_input,
+                    agent=agent,
+                    context=context_wrapper.context,
+                )
+                maybe_updated = run_config.call_model_input_filter(filter_payload)
+                updated = (
+                    await maybe_updated
+                    if inspect.isawaitable(maybe_updated)
+                    else maybe_updated
+                )
+                if not isinstance(updated, ModelInputData):
+                    raise UserError("call_model_input_filter must return a ModelInputData instance")
+                effective_input = updated.input
+                effective_instructions = updated.instructions
+            except Exception as e:
+                _error_tracing.attach_error_to_current_span(
+                    SpanError(message="Error in call_model_input_filter", data={"error": str(e)})
+                )
+                raise
+
         # 1. Stream the output events
         async for event in model.stream_response(
-            system_prompt,
-            input,
+            effective_instructions,
+            effective_input,
             model_settings,
             all_tools,
             output_schema,
@@ -1034,7 +1096,6 @@ async def _get_single_step_result_from_streamed_response(
         run_config: RunConfig,
         tool_use_tracker: AgentToolUseTracker,
     ) -> SingleStepResult:
-
         original_input = streamed_result.input
         pre_step_items = streamed_result.new_items
         event_queue = streamed_result._event_queue
@@ -1161,13 +1222,46 @@ async def _get_new_response(
         previous_response_id: str | None,
         prompt_config: ResponsePromptParam | None,
     ) -> ModelResponse:
+        # Allow user to modify model input right before the call, if configured
+        effective_instructions = system_prompt
+        effective_input: list[TResponseInputItem] = input
+
+        if run_config.call_model_input_filter is not None:
+            try:
+                model_input = ModelInputData(
+                    input=copy.deepcopy(effective_input),
+                    instructions=effective_instructions,
+                )
+                filter_payload: CallModelData[TContext] = CallModelData(
+                    model_data=model_input,
+                    agent=agent,
+                    context=context_wrapper.context,
+                )
+                maybe_updated = run_config.call_model_input_filter(filter_payload)
+                updated = (
+                    await maybe_updated
+                    if inspect.isawaitable(maybe_updated)
+                    else maybe_updated
+                )
+                # Basic validation to prevent accidental None returns
+                if not isinstance(updated, ModelInputData):
+                    raise UserError("call_model_input_filter must return a ModelInputData instance")
+                effective_input = updated.input
+                effective_instructions = updated.instructions
+            except Exception as e:
+                # Attach to trace and re-raise as user error so it's clear
+                _error_tracing.attach_error_to_current_span(
+                    SpanError(message="Error in call_model_input_filter", data={"error": str(e)})
+                )
+                raise
+
         model = cls._get_model(agent, run_config)
         model_settings = agent.model_settings.resolve(run_config.model_settings)
         model_settings = RunImpl.maybe_reset_tool_choice(agent, tool_use_tracker, model_settings)
 
         new_response = await model.get_response(
-            system_instructions=system_prompt,
-            input=input,
+            system_instructions=effective_instructions,
+            input=effective_input,
             model_settings=model_settings,
             tools=all_tools,
             output_schema=output_schema,