fix: add instructions on further requests besides UserPromptNode (#1503)

Author: Kludex

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -3,11 +3,11 @@
 import asyncio
 import dataclasses
 import json
-from collections.abc import AsyncIterator, Iterator, Sequence
+from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from dataclasses import field
-from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast
 
 from opentelemetry.trace import Span, Tracer
 from typing_extensions import TypeGuard, TypeVar, assert_never
@@ -87,6 +87,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     usage_limits: _usage.UsageLimits
     max_result_retries: int
     end_strategy: EndStrategy
+    get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]]
 
     output_schema: _output.OutputSchema[OutputDataT] | None
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
@@ -141,7 +142,9 @@ async def _get_first_message(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> _messages.ModelRequest:
         run_context = build_run_context(ctx)
-        history, next_message = await self._prepare_messages(self.user_prompt, ctx.state.message_history, run_context)
+        history, next_message = await self._prepare_messages(
+            self.user_prompt, ctx.state.message_history, ctx.deps.get_instructions, run_context
+        )
         ctx.state.message_history = history
         run_context.messages = history
 
@@ -155,6 +158,7 @@ async def _prepare_messages(
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None,
         message_history: list[_messages.ModelMessage] | None,
+        get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]],
         run_context: RunContext[DepsT],
     ) -> tuple[list[_messages.ModelMessage], _messages.ModelRequest]:
         try:
@@ -169,7 +173,7 @@ async def _prepare_messages(
                 ctx_messages.used = True
 
         parts: list[_messages.ModelRequestPart] = []
-        instructions = await self._instructions(run_context)
+        instructions = await get_instructions(run_context)
         if message_history:
             # Shallow copy messages
             messages.extend(message_history)
@@ -210,15 +214,6 @@ async def _sys_parts(self, run_context: RunContext[DepsT]) -> list[_messages.Mod
                 messages.append(_messages.SystemPromptPart(prompt))
         return messages
 
-    async def _instructions(self, run_context: RunContext[DepsT]) -> str | None:
-        if self.instructions is None and not self.instructions_functions:
-            return None
-
-        instructions = self.instructions or ''
-        for instructions_runner in self.instructions_functions:
-            instructions += await instructions_runner.run(run_context)
-        return instructions
-
 
 async def _prepare_request_parameters(
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
@@ -479,7 +474,11 @@ async def _handle_tool_calls(
         else:
             if tool_responses:
                 parts.extend(tool_responses)
-            self._next_node = ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=parts))
+            run_context = build_run_context(ctx)
+            instructions = await ctx.deps.get_instructions(run_context)
+            self._next_node = ModelRequestNode[DepsT, NodeRunEndT](
+                _messages.ModelRequest(parts=parts, instructions=instructions)
+            )
 
     def _handle_final_result(
         self,

pydantic_ai_slim/pydantic_ai/agent.py

@@ -620,6 +620,15 @@ async def main():
             },
         )
 
+        async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
+            if self._instructions is None and not self._instructions_functions:
+                return None
+
+            instructions = self._instructions or ''
+            for instructions_runner in self._instructions_functions:
+                instructions += await instructions_runner.run(run_context)
+            return instructions
+
         graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunOutputDataT](
             user_deps=deps,
             prompt=user_prompt,
@@ -635,6 +644,7 @@ async def main():
             mcp_servers=self._mcp_servers,
             run_span=run_span,
             tracer=tracer,
+            get_instructions=get_instructions,
         )
         start_node = _agent_graph.UserPromptNode[AgentDepsT](
             user_prompt=user_prompt,

tests/models/cassettes/test_openai/test_openai_instructions_with_tool_calls_keep_instructions.yaml

@@ -0,0 +1,207 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '419'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+    method: POST
+    parsed_body:
+      messages:
+      - content: You are a helpful assistant.
+        role: system
+      - content: What is the temperature in Tokyo?
+        role: user
+      model: gpt-4.1-mini
+      n: 1
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: ''
+          name: get_temperature
+          parameters:
+            additionalProperties: false
+            properties:
+              city:
+                type: string
+            required:
+            - city
+            type: object
+          strict: true
+        type: function
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    headers:
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '1089'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '490'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: tool_calls
+        index: 0
+        logprobs: null
+        message:
+          annotations: []
+          content: null
+          refusal: null
+          role: assistant
+          tool_calls:
+          - function:
+              arguments: '{"city":"Tokyo"}'
+              name: get_temperature
+            id: call_bhZkmIKKItNGJ41whHUHB7p9
+            type: function
+      created: 1744810634
+      id: chatcmpl-BMxEwRA0p0gJ52oKS7806KAlfMhqq
+      model: gpt-4.1-mini-2025-04-14
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: fp_38647f5e19
+      usage:
+        completion_tokens: 15
+        completion_tokens_details:
+          accepted_prediction_tokens: 0
+          audio_tokens: 0
+          reasoning_tokens: 0
+          rejected_prediction_tokens: 0
+        prompt_tokens: 50
+        prompt_tokens_details:
+          audio_tokens: 0
+          cached_tokens: 0
+        total_tokens: 65
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '665'
+      content-type:
+      - application/json
+      cookie:
+      - __cf_bm=x.H2GlMeh.t_Q.gVlCXrh3.ggn9lKjhmUeG_ToNThLs-1744810635-1.0.1.1-tiHwqGvBw3eEy_y9_q5nx7B.7YCbLb9cXdDj6DklLmtFllOFe708mKwYvGd8fY2y5bO2NOagULipA7MxfwW9P0hlnRSiJZbZBO9tjrUweFc;
+        _cfuvid=VlHcJdsIsxGEt2lddKu_5Am_lfyYndl9JB2Ezy.aygo-1744810635187-0.0.1.1-604800000
+      host:
+      - api.openai.com
+    method: POST
+    parsed_body:
+      messages:
+      - content: You are a helpful assistant.
+        role: system
+      - content: What is the temperature in Tokyo?
+        role: user
+      - role: assistant
+        tool_calls:
+        - function:
+            arguments: '{"city":"Tokyo"}'
+            name: get_temperature
+          id: call_bhZkmIKKItNGJ41whHUHB7p9
+          type: function
+      - content: '20.0'
+        role: tool
+        tool_call_id: call_bhZkmIKKItNGJ41whHUHB7p9
+      model: gpt-4.1-mini
+      n: 1
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: ''
+          name: get_temperature
+          parameters:
+            additionalProperties: false
+            properties:
+              city:
+                type: string
+            required:
+            - city
+            type: object
+          strict: true
+        type: function
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    headers:
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '867'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '949'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        logprobs: null
+        message:
+          annotations: []
+          content: The temperature in Tokyo is currently 20.0 degrees Celsius.
+          refusal: null
+          role: assistant
+      created: 1744810635
+      id: chatcmpl-BMxEx6B8JEj6oDC45MOWKp0phg8UP
+      model: gpt-4.1-mini-2025-04-14
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: fp_38647f5e19
+      usage:
+        completion_tokens: 15
+        completion_tokens_details:
+          accepted_prediction_tokens: 0
+          audio_tokens: 0
+          reasoning_tokens: 0
+          rejected_prediction_tokens: 0
+        prompt_tokens: 75
+        prompt_tokens_details:
+          audio_tokens: 0
+          cached_tokens: 0
+        total_tokens: 90
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/test_openai.py

@@ -1215,3 +1215,41 @@ async def test_openai_model_without_system_prompt(allow_model_requests: None, op
     assert result.output == snapshot(
         "That's right—I am a potato! A spud of many talents, here to help you out. How can this humble potato be of service today?"
     )
+
+
+@pytest.mark.vcr()
+async def test_openai_instructions_with_tool_calls_keep_instructions(allow_model_requests: None, openai_api_key: str):
+    m = OpenAIModel('gpt-4.1-mini', provider=OpenAIProvider(api_key=openai_api_key))
+    agent = Agent(m, instructions='You are a helpful assistant.')
+
+    @agent.tool_plain
+    async def get_temperature(city: str) -> float:
+        return 20.0
+
+    result = await agent.run('What is the temperature in Tokyo?')
+    assert result.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[UserPromptPart(content='What is the temperature in Tokyo?', timestamp=IsDatetime())],
+                instructions='You are a helpful assistant.',
+            ),
+            ModelResponse(
+                parts=[ToolCallPart(tool_name='get_temperature', args='{"city":"Tokyo"}', tool_call_id=IsStr())],
+                model_name='gpt-4.1-mini-2025-04-14',
+                timestamp=IsDatetime(),
+            ),
+            ModelRequest(
+                parts=[
+                    ToolReturnPart(
+                        tool_name='get_temperature', content=20.0, tool_call_id=IsStr(), timestamp=IsDatetime()
+                    )
+                ],
+                instructions='You are a helpful assistant.',
+            ),
+            ModelResponse(
+                parts=[TextPart(content='The temperature in Tokyo is currently 20.0 degrees Celsius.')],
+                model_name='gpt-4.1-mini-2025-04-14',
+                timestamp=IsDatetime(),
+            ),
+        ]
+    )