Support overriding agent instructions #2926

mwildehahn · 2025-09-17T00:47:36Z

I just started working on a library: https://github.com/mwildehahn/pydantic-ai-gepa to integrate https://github.com/gepa-ai/gepa into pydantic-ai. This would provide similar functionality to https://dspy.ai/ where you can provide a signature and then let an LLM handle constructing the prompt.

This is very experimental and I just started on this yesterday, but I wanted to propose a small extension to pydantic-ai that would allow us to override system_prompt and instructions in the same way we can override toolsets etc. With that minimal surface area, I can hook in something like pydantic-ai-gepa to optimize the prompts and then run those optimized prompts.

LMK if there are better ways to override the system prompts on demand.

pydantic_ai_slim/pydantic_ai/agent/__init__.py

DouweM · 2025-09-18T15:35:20Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

+        else:
+            system_prompts = self._system_prompts
+
+        if override_instructions := self._override_instructions.get():


Could we move this to some top-level method, like the other override context var usages?

I didn't follow this one -- I think we're doing the same as the existing ones?

What I mean is that we're currently only calling self._override_<foo>.get() from private getters on Agent that return either the overridden value, or the original, and then those getters are used in the place where we actually need the value.

So I suggest making _get_instructions_literal_and_functions a method that doesn't take an argument, but itself checks whether to use the overridden or original value. We can make that easier by merging the self._instructions and self._instructions_functions variables into just self._instructions, so that we don't need to call _get_instructions_literal_and_functions from __init__ anymore, just here.

And actually, if we move this up above the get_instructions function, we can reference the same instructions and instructions_functions variables inside there, so we only call _instructions_literal_and_functions once.

pydantic_ai_slim/pydantic_ai/agent/__init__.py

mwildehahn · 2025-09-18T19:14:12Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

+            if isinstance(instruction, str):
+                literal_parts.append(instruction)
+            elif callable(instruction):
+                func = cast(_system_prompt.SystemPromptFunc[AgentDepsT], instruction)


I don't love this, but not sure how to appease the type checker. This was marked as "unknown" otherwise.

Does it work if we change elif callable(instruction): to just else:, like we had in the original code?

Nope, i think because we have an explicit:

functions: list[_system_prompt.SystemPromptRunner[AgentDepsT]] = []

@mwildehahn Hmm ok, once this PR is otherwise ready I'll have a look to see if I can clean up the typing here a bit.

DouweM · 2025-09-18T20:06:51Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

+            if isinstance(instruction, str):
+                literal_parts.append(instruction)
+            elif callable(instruction):
+                func = cast(_system_prompt.SystemPromptFunc[AgentDepsT], instruction)


Does it work if we change elif callable(instruction): to just else:, like we had in the original code?

DouweM · 2025-09-18T20:11:37Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

+        else:
+            system_prompts = self._system_prompts
+
+        if override_instructions := self._override_instructions.get():


What I mean is that we're currently only calling self._override_<foo>.get() from private getters on Agent that return either the overridden value, or the original, and then those getters are used in the place where we actually need the value.

So I suggest making _get_instructions_literal_and_functions a method that doesn't take an argument, but itself checks whether to use the overridden or original value. We can make that easier by merging the self._instructions and self._instructions_functions variables into just self._instructions, so that we don't need to call _get_instructions_literal_and_functions from __init__ anymore, just here.

And actually, if we move this up above the get_instructions function, we can reference the same instructions and instructions_functions variables inside there, so we only call _instructions_literal_and_functions once.

pydantic_ai_slim/pydantic_ai/agent/__init__.py

pydantic_ai_slim/pydantic_ai/agent/abstract.py

DouweM · 2025-09-19T17:43:52Z

dbostest.sqlite

Please remove this file, it should've gotten cleaned up automatically 🤔

DouweM · 2025-09-19T17:46:33Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

-            else:
-                self._instructions_functions.append(_system_prompt.SystemPromptRunner(instruction))
-        self._instructions = self._instructions.strip() or None
+        self._instructions, self._instructions_functions = self._instructions_literal_and_functions(instructions)


I was thinking we could store self._instructions = instructions, and then call _get_instructions_literal_and_functions where we need it. I don't think we need the 2 private vars

DouweM · 2025-09-19T17:46:54Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

        self._entered_count = 0
        self._exit_stack = None

+    def _get_instructions_literal_and_functions(


Let's just call this _get_instructions

DouweM · 2025-09-19T17:47:35Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

+            instructions, instructions_functions = self._instructions_literal_and_functions(override_instructions.value)
+        return instructions, instructions_functions
+
+    def _instructions_literal_and_functions(


With my suggestion above, we should not need this as a separate method anymore, so we can move its contents into the get method.

DouweM · 2025-09-19T17:50:40Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

+            if isinstance(instruction, str):
+                literal_parts.append(instruction)
+            elif callable(instruction):
+                func = cast(_system_prompt.SystemPromptFunc[AgentDepsT], instruction)


@mwildehahn Hmm ok, once this PR is otherwise ready I'll have a look to see if I can clean up the typing here a bit.

DouweM · 2025-09-19T17:51:20Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

        usage_limits = usage_limits or _usage.UsageLimits()

        async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
+            literal, functions = self._get_instructions_literal_and_functions()


With the changes I suggested above, this can be:

Suggested change

literal, functions = self._get_instructions_literal_and_functions()

instructions, instructions_functions = self._get_instructions()

And can we move this out of the function, and then use the same instructions and instructions_functions when we build the UserPromptNode below, to save calling this method twice?

mwildehahn · 2025-09-26T20:13:49Z

@DouweM i think i've addressed everything and fixed the weird type errors -- issue was I wasn't passing generic with Instruction alias

DouweM · 2025-09-29T23:27:15Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

        self._entered_count = 0
        self._exit_stack = None

+    def _get_instructions(


Minor thing, but please move this to where the other _get_ methods are -- it's not important enough to be right at the top of the class :)

DouweM · 2025-09-29T23:29:03Z

tests/test_override_instructions.py

Please move the relevant tests to tests/test_agent.py, where the other overrides are tested

DouweM · 2025-09-29T23:29:25Z

tests/test_override_instructions.py

+def test_first_request_skips_non_requests():
+    """Helper ignores non-request messages until it finds a request."""
+    response = ModelResponse(parts=())
+    request = ModelRequest(parts=())
+    assert _first_request([response, request]) is request
+
+
+def test_first_request_raises_without_model_request():
+    """Helper raises when no model request is present."""
+    response = ModelResponse(parts=())
+    with pytest.raises(AssertionError, match='no ModelRequest found'):
+        _first_request([response])


No need to test the helpers :)

DouweM · 2025-09-29T23:29:59Z

tests/test_override_instructions.py

+def _first_request(messages: list[ModelMessage]) -> ModelRequest:
+    """Helper to extract the first ModelRequest from captured messages."""
+    assert messages, 'no messages captured'
+    for m in messages:
+        if isinstance(m, ModelRequest):
+            return m
+    raise AssertionError('no ModelRequest found in captured messages')
+
+
+def _system_prompt_texts(parts: Sequence[ModelRequestPart]) -> list[str]:
+    """Helper to extract system prompt text content from message parts."""
+    return [p.content for p in parts if isinstance(p, SystemPromptPart)]


Do we really need these helpers? I'd rather just repeat this line, and use assert isinstance(messages[0], ModelRequest) to ensure the first message is a request.

DouweM · 2025-09-29T23:36:11Z

pydantic_ai_slim/pydantic_ai/agent/__init__.py

+        elif callable(self._instructions):
+            instructions_list = [self._instructions, instruction]
+        else:
+            instructions_list = [*self._instructions, instruction]  # pragma: no cover


Sorry for continuing to come up with more ways to refactor this, but what do you think about making self._instructions and self._override_instructions always hold a list[str | _system_prompt.SystemPromptFunc[AgentDepsT]], and change __init__ and override to store list(instructions) so that any single items automatically get wrapped in a list, and this method and the one above can be much simpler?

DouweM self-assigned this Sep 18, 2025

DouweM requested changes Sep 18, 2025

View reviewed changes

DouweM added the awaiting author revision label Sep 18, 2025

mwildehahn force-pushed the mh/gepa branch from c84806a to 8e036eb Compare September 18, 2025 19:10

mwildehahn commented Sep 18, 2025

View reviewed changes

DouweM requested changes Sep 18, 2025

View reviewed changes

mwildehahn added 12 commits September 18, 2025 14:16

Support overriding agent instructions

28b8e8b

Review feedback

4b36d26

Fix type error

b1c0892

Rename test file

fd49ab2

Fix lint errors

a1d25cd

Formatting tweaks

2ad2a99

Add coverage

a36c1bf

Fix lint errors

e7d531c

Fix lint error

fa10e67

Add helper function

d51624b

Change variable name

696c4ba

Add more tests

01c6f47

mwildehahn force-pushed the mh/gepa branch from 3cd7853 to 01c6f47 Compare September 18, 2025 21:26

DouweM requested changes Sep 19, 2025

View reviewed changes

mwildehahn added 6 commits September 23, 2025 20:40

Address review comments

c7a1cf1

Lint

e9de7f2

Fix instruction typing

3293846

Fix typing

88577f7

Ignore coverage

3048f31

...cool cool

ce9eab5

DouweM requested changes Sep 29, 2025

View reviewed changes

	literal, functions = self._get_instructions_literal_and_functions()
	instructions, instructions_functions = self._get_instructions()

Support overriding agent instructions #2926

Are you sure you want to change the base?

Support overriding agent instructions #2926

Conversation

mwildehahn commented Sep 17, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mwildehahn commented Sep 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!