Adds dynamic to system_prompt decorator, allowing reevaluation

pydantic_ai_slim/pydantic_ai/_system_prompt.py

@@ -12,6 +12,7 @@
 @dataclass
 class SystemPromptRunner(Generic[AgentDeps]):
     function: SystemPromptFunc[AgentDeps]
+    dynamic: bool = False
     _takes_ctx: bool = field(init=False)
     _is_async: bool = field(init=False)
 

pydantic_ai_slim/pydantic_ai/agent.py

@@ -526,16 +526,37 @@ def system_prompt(self, func: Callable[[], str], /) -> Callable[[], str]: ...
     @overload
     def system_prompt(self, func: Callable[[], Awaitable[str]], /) -> Callable[[], Awaitable[str]]: ...
 
+    @overload
+    def system_prompt(
+        self, func: _system_prompt.SystemPromptFunc[AgentDeps] | None = None, /, *, dynamic: bool = False
+    ) -> Any:
+
+        if func is None:
+            def decorator(
+                func_: _system_prompt.SystemPromptFunc[AgentDeps],
+            ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
+                self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func_, dynamic=dynamic))
+                return func_
+            return decorator
+        else:
+            self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func, dynamic=dynamic))
+            return func
+
     def system_prompt(
-        self, func: _system_prompt.SystemPromptFunc[AgentDeps], /
-    ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
+        self,
+        func: _system_prompt.SystemPromptFunc[AgentDeps] | None = None,
+        /,
+        *,
+        dynamic: bool = False,
+    ) -> Any:
         """Decorator to register a system prompt function.
 
         Optionally takes [`RunContext`][pydantic_ai.tools.RunContext] as its only argument.
         Can decorate a sync or async functions.
 
-        Overloads for every possible signature of `system_prompt` are included so the decorator doesn't obscure
-        the type of the function, see `tests/typed_agent.py` for tests.
+        Args:
+            func: The function to decorate
+            dynamic: If True, the system prompt will be reevaluated when messages_history is present.
 
         Example:
         ```python
@@ -547,17 +568,21 @@ def system_prompt(
         def simple_system_prompt() -> str:
             return 'foobar'
 
-        @agent.system_prompt
+        @agent.system_prompt(dynamic=True)
         async def async_system_prompt(ctx: RunContext[str]) -> str:
             return f'{ctx.deps} is the best'
-
-        result = agent.run_sync('foobar', deps='spam')
-        print(result.data)
-        #> success (no tool calls)
         ```
         """
-        self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func))
-        return func
+        if func is None:
+            def decorator(
+                func_: _system_prompt.SystemPromptFunc[AgentDeps],
+            ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
+                self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func_, dynamic=dynamic))
+                return func_
+            return decorator
+        else:
+            self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func, dynamic=dynamic))
+            return func
 
     @overload
     def result_validator(
@@ -830,8 +855,8 @@ 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]:
+    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:
@@ -846,6 +871,22 @@ async def _prepare_messages(
         if message_history:
             # shallow copy messages
             messages.extend(message_history)
+
+            # If there are any dynamic system prompts, we need to reevaluate them
+            if any(runner.dynamic for runner in self._system_prompt_functions):
+                # Get fresh system prompts
+                new_sys_parts = await self._sys_parts(run_context)
+
+                # Replace the system prompts in the existing messages
+                for msg in messages:
+                    if isinstance(msg, _messages.ModelRequest):
+                        # Keep non-system parts and add new system parts
+                        non_system_parts = [
+                            part for part in msg.parts 
+                            if not isinstance(part, _messages.SystemPromptPart)
+                        ]
+                        msg.parts = new_sys_parts + non_system_parts
+
             messages.append(_messages.ModelRequest([_messages.UserPromptPart(user_prompt)]))
         else:
             parts = await self._sys_parts(run_context)

tests/test_agent.py

@@ -1230,3 +1230,233 @@ def test_double_capture_run_messages(set_event_loop: None) -> None:
             ModelResponse(parts=[TextPart(content='success (no tool calls)')], timestamp=IsNow(tz=timezone.utc)),
         ]
     )
+
+
+def test_messages_history_reevaluate_system_prompt_clean(set_event_loop: None):
+    agent = Agent('test', system_prompt='Foobar')
+
+    dynamic_value = "A"
+
+    @agent.system_prompt
+    async def func(ctx):
+        return dynamic_value
+
+    res = agent.run_sync('Hello')
+
+    assert res.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content='Foobar', 
+                        part_kind='system-prompt'),
+                    SystemPromptPart(
+                        content=dynamic_value,
+                        part_kind='system-prompt'),
+                    UserPromptPart(
+                        content='Hello',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response')
+        ]
+    )
+
+    dynamic_value = "B"
+
+    @agent.system_prompt
+    async def func_two(ctx):
+        return dynamic_value + "!"
+
+    res_two = agent.run_sync('World', message_history=res.all_messages())
+
+    assert res_two.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content='Foobar', 
+                        part_kind='system-prompt'),
+                    SystemPromptPart(
+                        content="A", #Remains the same
+                        part_kind='system-prompt'),
+                    UserPromptPart(
+                        content='Hello',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response'),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='World',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response')
+        ]
+    )
+
+def test_messages_history_reevaluate_system_prompt_dirty(set_event_loop: None):
+    agent = Agent('test', system_prompt='Foobar')
+
+    dynamic_value = "A"
+
+    @agent.system_prompt(dynamic=True)
+    async def func(ctx):
+        return dynamic_value
+
+    res = agent.run_sync('Hello')
+
+    assert res.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content='Foobar', 
+                        part_kind='system-prompt'),
+                    SystemPromptPart(
+                        content=dynamic_value,
+                        part_kind='system-prompt'),
+                    UserPromptPart(
+                        content='Hello',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response')
+        ]
+    )
+
+    dynamic_value = "B"
+
+    @agent.system_prompt
+    async def func_two(ctx):
+        return "This is a new prompt, but it wont reach the model"
+
+    res_two = agent.run_sync('World', message_history=res.all_messages())
+
+    assert res_two.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content='Foobar', 
+                        part_kind='system-prompt'),
+                    SystemPromptPart(
+                        content="B", # Updated value since dirty
+                        part_kind='system-prompt'),
+                    UserPromptPart(
+                        content='Hello',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response'),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='World',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response')
+        ]
+    )
+
+
+
+def test_messages_history_reevaluate_system_prompt_dirty(set_event_loop: None):
+    agent = Agent('test', system_prompt='Foobar')
+
+    dynamic_value = "A"
+
+    @agent.system_prompt(dynamic=True)
+    async def func(ctx):
+        return dynamic_value
+
+    res = agent.run_sync('Hello')
+
+    assert res.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content='Foobar', 
+                        part_kind='system-prompt'),
+                    SystemPromptPart(
+                        content=dynamic_value,
+                        part_kind='system-prompt'),
+                    UserPromptPart(
+                        content='Hello',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response')
+        ]
+    )
+
+    dynamic_value = "B"
+
+    @agent.system_prompt(dynamic=True)
+    async def func_two(ctx):
+        return "This is a new prompt, and model will know"
+
+    res_two = agent.run_sync('World', message_history=res.all_messages())
+
+    assert res_two.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content='Foobar', 
+                        part_kind='system-prompt'),
+                    SystemPromptPart(
+                        content="B", # Updated value since dirty
+                        part_kind='system-prompt'),
+                    SystemPromptPart(
+                        content="This is a new prompt, and model will know",
+                        part_kind='system-prompt'),
+                    UserPromptPart(
+                        content='Hello',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response'),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='World',
+                        timestamp=IsNow(tz=timezone.utc), part_kind='user-prompt')
+                ], kind='request'), 
+            ModelResponse(
+                parts=[
+                    TextPart(content='success (no tool calls)', part_kind='text')
+                ], 
+                timestamp=IsNow(tz=timezone.utc), kind='response')
+        ]
+    )
+
+