add messages to RunContext (#257)

Author: samuelcolvin

pydantic_ai_slim/pydantic_ai/_result.py

@@ -10,9 +10,8 @@
 from pydantic import TypeAdapter, ValidationError
 from typing_extensions import Self, TypeAliasType, TypedDict
 
-from . import _utils, messages
+from . import _utils, messages as _messages
 from .exceptions import ModelRetry
-from .messages import ModelResponse, ToolCallPart
 from .result import ResultData
 from .tools import AgentDeps, ResultValidatorFunc, RunContext, ToolDefinition
 
@@ -28,7 +27,12 @@ def __post_init__(self):
         self._is_async = inspect.iscoroutinefunction(self.function)
 
     async def validate(
-        self, result: ResultData, deps: AgentDeps, retry: int, tool_call: messages.ToolCallPart | None
+        self,
+        result: ResultData,
+        deps: AgentDeps,
+        retry: int,
+        tool_call: _messages.ToolCallPart | None,
+        messages: list[_messages.Message],
     ) -> ResultData:
         """Validate a result but calling the function.
 
@@ -37,12 +41,13 @@ async def validate(
             deps: The agent dependencies.
             retry: The current retry number.
             tool_call: The original tool call message, `None` if there was no tool call.
+            messages: The messages exchanged so far in the conversation.
 
         Returns:
             Result of either the validated result data (ok) or a retry message (Err).
         """
         if self._takes_ctx:
-            args = RunContext(deps, retry, tool_call.tool_name if tool_call else None), result
+            args = RunContext(deps, retry, messages, tool_call.tool_name if tool_call else None), result
         else:
             args = (result,)
 
@@ -54,7 +59,7 @@ async def validate(
                 function = cast(Callable[[Any], ResultData], self.function)
                 result_data = await _utils.run_in_executor(function, *args)
         except ModelRetry as r:
-            m = messages.RetryPrompt(content=r.message)
+            m = _messages.RetryPrompt(content=r.message)
             if tool_call is not None:
                 m.tool_name = tool_call.tool_name
                 m.tool_call_id = tool_call.tool_call_id
@@ -66,7 +71,7 @@ async def validate(
 class ToolRetryError(Exception):
     """Internal exception used to signal a `ToolRetry` message should be returned to the LLM."""
 
-    def __init__(self, tool_retry: messages.RetryPrompt):
+    def __init__(self, tool_retry: _messages.RetryPrompt):
         self.tool_retry = tool_retry
         super().__init__()
 
@@ -108,10 +113,12 @@ def _build_tool(a: Any, tool_name_: str, multiple: bool) -> ResultTool[ResultDat
 
         return cls(tools=tools, allow_text_result=allow_text_result)
 
-    def find_tool(self, message: ModelResponse) -> tuple[ToolCallPart, ResultTool[ResultData]] | None:
+    def find_tool(
+        self, message: _messages.ModelResponse
+    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultData]] | None:
         """Find a tool that matches one of the calls."""
         for item in message.parts:
-            if isinstance(item, ToolCallPart):
+            if isinstance(item, _messages.ToolCallPart):
                 if result := self.tools.get(item.tool_name):
                     return item, result
 
@@ -168,7 +175,7 @@ def __init__(self, response_type: type[ResultData], name: str, description: str
         )
 
     def validate(
-        self, tool_call: messages.ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+        self, tool_call: _messages.ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
     ) -> ResultData:
         """Validate a result message.
 
@@ -182,7 +189,7 @@ def validate(
         """
         try:
             pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
-            if isinstance(tool_call.args, messages.ArgsJson):
+            if isinstance(tool_call.args, _messages.ArgsJson):
                 result = self.type_adapter.validate_json(
                     tool_call.args.args_json or '', experimental_allow_partial=pyd_allow_partial
                 )
@@ -192,7 +199,7 @@ def validate(
                 )
         except ValidationError as e:
             if wrap_validation_errors:
-                m = messages.RetryPrompt(
+                m = _messages.RetryPrompt(
                     tool_name=tool_call.tool_name,
                     content=e.errors(include_url=False),
                     tool_call_id=tool_call.tool_call_id,

pydantic_ai_slim/pydantic_ai/_system_prompt.py

@@ -21,7 +21,7 @@ def __post_init__(self):
 
     async def run(self, deps: AgentDeps) -> str:
         if self._takes_ctx:
-            args = (RunContext(deps, 0),)
+            args = (RunContext(deps, 0, [], None),)
         else:
             args = ()
 

pydantic_ai_slim/pydantic_ai/agent.py

@@ -243,7 +243,7 @@ async def run(
             while True:
                 run_step += 1
                 with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
-                    agent_model = await self._prepare_model(model_used, deps)
+                    agent_model = await self._prepare_model(model_used, deps, messages)
 
                 with _logfire.span('model request', run_step=run_step) as model_req_span:
                     model_response, request_cost = await agent_model.request(messages, model_settings)
@@ -255,7 +255,7 @@ async def run(
                 cost += request_cost
 
                 with _logfire.span('handle model response', run_step=run_step) as handle_span:
-                    final_result, response_messages = await self._handle_model_response(model_response, deps)
+                    final_result, response_messages = await self._handle_model_response(model_response, deps, messages)
 
                     # Add all messages to the conversation
                     messages.extend(response_messages)
@@ -391,7 +391,7 @@ async def main():
                 run_step += 1
 
                 with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
-                    agent_model = await self._prepare_model(model_used, deps)
+                    agent_model = await self._prepare_model(model_used, deps, messages)
 
                 with _logfire.span('model request {run_step=}', run_step=run_step) as model_req_span:
                     async with agent_model.request_stream(messages, model_settings) as model_response:
@@ -402,7 +402,7 @@ async def main():
 
                         with _logfire.span('handle model response') as handle_span:
                             final_result, response_messages = await self._handle_streamed_model_response(
-                                model_response, deps
+                                model_response, deps, messages
                             )
 
                             # Add all messages to the conversation
@@ -773,12 +773,14 @@ async def _get_model(self, model: models.Model | models.KnownModelName | None) -
 
         return model_, mode_selection
 
-    async def _prepare_model(self, model: models.Model, deps: AgentDeps) -> models.AgentModel:
+    async def _prepare_model(
+        self, model: models.Model, deps: AgentDeps, messages: list[_messages.Message]
+    ) -> models.AgentModel:
         """Create building tools and create an agent model."""
         function_tools: list[ToolDefinition] = []
 
         async def add_tool(tool: Tool[AgentDeps]) -> None:
-            ctx = RunContext(deps, tool.current_retry, tool.name)
+            ctx = RunContext(deps, tool.current_retry, messages, tool.name)
             if tool_def := await tool.prepare_tool_def(ctx):
                 function_tools.append(tool_def)
 
@@ -807,7 +809,7 @@ async def _prepare_messages(
         return new_message_index, messages
 
     async def _handle_model_response(
-        self, model_response: _messages.ModelResponse, deps: AgentDeps
+        self, model_response: _messages.ModelResponse, deps: AgentDeps, conv_messages: list[_messages.Message]
     ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.Message]]:
         """Process a non-streamed response from the model.
 
@@ -824,20 +826,20 @@ async def _handle_model_response(
 
         if texts:
             text = '\n\n'.join(texts)
-            return await self._handle_text_response(text, deps)
+            return await self._handle_text_response(text, deps, conv_messages)
         elif tool_calls:
-            return await self._handle_structured_response(tool_calls, deps)
+            return await self._handle_structured_response(tool_calls, deps, conv_messages)
         else:
             raise exceptions.UnexpectedModelBehavior('Received empty model response')
 
     async def _handle_text_response(
-        self, text: str, deps: AgentDeps
+        self, text: str, deps: AgentDeps, conv_messages: list[_messages.Message]
     ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.Message]]:
         """Handle a plain text response from the model for non-streaming responses."""
         if self._allow_text_result:
             result_data_input = cast(ResultData, text)
             try:
-                result_data = await self._validate_result(result_data_input, deps, None)
+                result_data = await self._validate_result(result_data_input, deps, None, conv_messages)
             except _result.ToolRetryError as e:
                 self._incr_result_retry()
                 return None, [e.tool_retry]
@@ -851,26 +853,24 @@ async def _handle_text_response(
             return None, [response]
 
     async def _handle_structured_response(
-        self, tool_calls: list[_messages.ToolCallPart], deps: AgentDeps
+        self, tool_calls: list[_messages.ToolCallPart], deps: AgentDeps, conv_messages: list[_messages.Message]
     ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.Message]]:
         """Handle a structured response containing tool calls from the model for non-streaming responses."""
         assert tool_calls, 'Expected at least one tool call'
 
         # First process any final result tool calls
-        final_result, final_messages = await self._process_final_tool_calls(tool_calls, deps)
+        final_result, final_messages = await self._process_final_tool_calls(tool_calls, deps, conv_messages)
 
         # Then process regular tools based on end strategy
         if self.end_strategy == 'early' and final_result:
             tool_messages = self._mark_skipped_function_tools(tool_calls)
         else:
-            tool_messages = await self._process_function_tools(tool_calls, deps)
+            tool_messages = await self._process_function_tools(tool_calls, deps, conv_messages)
 
         return final_result, [*final_messages, *tool_messages]
 
     async def _process_final_tool_calls(
-        self,
-        tool_calls: list[_messages.ToolCallPart],
-        deps: AgentDeps,
+        self, tool_calls: list[_messages.ToolCallPart], deps: AgentDeps, conv_messages: list[_messages.Message]
     ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.Message]]:
         """Process any final result tool calls and return the first valid result."""
         if not self._result_schema:
@@ -888,7 +888,7 @@ async def _process_final_tool_calls(
                 # This is the first result tool - try to use it
                 try:
                     result_data = result_tool.validate(call)
-                    result_data = await self._validate_result(result_data, deps, call)
+                    result_data = await self._validate_result(result_data, deps, call, conv_messages)
                 except _result.ToolRetryError as e:
                     self._incr_result_retry()
                     messages.append(e.tool_retry)
@@ -914,17 +914,15 @@ async def _process_final_tool_calls(
         return final_result, messages
 
     async def _process_function_tools(
-        self,
-        tool_calls: list[_messages.ToolCallPart],
-        deps: AgentDeps,
+        self, tool_calls: list[_messages.ToolCallPart], deps: AgentDeps, conv_messages: list[_messages.Message]
     ) -> list[_messages.Message]:
         """Process function (non-final) tool calls in parallel."""
         messages: list[_messages.Message] = []
         tasks: list[asyncio.Task[_messages.Message]] = []
 
         for call in tool_calls:
             if tool := self._function_tools.get(call.tool_name):
-                tasks.append(asyncio.create_task(tool.run(deps, call), name=call.tool_name))
+                tasks.append(asyncio.create_task(tool.run(deps, call, conv_messages), name=call.tool_name))
             elif self._result_schema is None or call.tool_name not in self._result_schema.tools:
                 messages.append(self._unknown_tool(call.tool_name))
 
@@ -958,7 +956,7 @@ def _mark_skipped_function_tools(
         return messages
 
     async def _handle_streamed_model_response(
-        self, model_response: models.EitherStreamedResponse, deps: AgentDeps
+        self, model_response: models.EitherStreamedResponse, deps: AgentDeps, conv_messages: list[_messages.Message]
     ) -> tuple[_MarkFinalResult[models.EitherStreamedResponse] | None, list[_messages.Message]]:
         """Process a streamed response from the model.
 
@@ -1015,7 +1013,7 @@ async def _handle_streamed_model_response(
                 if isinstance(item, _messages.ToolCallPart):
                     call = item
                     if tool := self._function_tools.get(call.tool_name):
-                        tasks.append(asyncio.create_task(tool.run(deps, call), name=call.tool_name))
+                        tasks.append(asyncio.create_task(tool.run(deps, call, conv_messages), name=call.tool_name))
                     else:
                         messages.append(self._unknown_tool(call.tool_name))
 
@@ -1025,10 +1023,16 @@ async def _handle_streamed_model_response(
             return None, messages
 
     async def _validate_result(
-        self, result_data: ResultData, deps: AgentDeps, tool_call: _messages.ToolCallPart | None
+        self,
+        result_data: ResultData,
+        deps: AgentDeps,
+        tool_call: _messages.ToolCallPart | None,
+        conv_messages: list[_messages.Message],
     ) -> ResultData:
         for validator in self._result_validators:
-            result_data = await validator.validate(result_data, deps, self._current_result_retry, tool_call)
+            result_data = await validator.validate(
+                result_data, deps, self._current_result_retry, tool_call, conv_messages
+            )
         return result_data
 
     def _incr_result_retry(self) -> None:

pydantic_ai_slim/pydantic_ai/messages.py

@@ -220,9 +220,13 @@ class ModelResponse:
     message_kind: Literal['model-response'] = 'model-response'
     """Message source identifier, this type is available on all messages as a discriminator."""
 
-    @staticmethod
-    def from_text(content: str, timestamp: datetime | None = None) -> ModelResponse:
-        return ModelResponse([TextPart(content)], timestamp=timestamp or _now_utc())
+    @classmethod
+    def from_text(cls, content: str, timestamp: datetime | None = None) -> ModelResponse:
+        return cls([TextPart(content)], timestamp=timestamp or _now_utc())
+
+    @classmethod
+    def from_tool_call(cls, tool_call: ToolCallPart) -> ModelResponse:
+        return cls([tool_call])
 
 
 Message = Union[SystemPrompt, UserPrompt, ToolReturn, RetryPrompt, ModelResponse]