fix: apply code review comments

Author: yf-yang

pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_agent.py

@@ -115,6 +115,7 @@ async def wrapped_run_workflow(
             infer_name: bool = True,
             toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
             event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+            response_prefix: str | None = None,
             **_deprecated_kwargs: Never,
         ) -> AgentRunResult[Any]:
             with self._dbos_overrides():
@@ -131,6 +132,7 @@ async def wrapped_run_workflow(
                     infer_name=infer_name,
                     toolsets=toolsets,
                     event_stream_handler=event_stream_handler,
+                    response_prefix=response_prefix,
                     **_deprecated_kwargs,
                 )
 
@@ -152,6 +154,7 @@ def wrapped_run_sync_workflow(
             infer_name: bool = True,
             toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
             event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+            response_prefix: str | None = None,
             **_deprecated_kwargs: Never,
         ) -> AgentRunResult[Any]:
             with self._dbos_overrides():
@@ -168,6 +171,7 @@ def wrapped_run_sync_workflow(
                     infer_name=infer_name,
                     toolsets=toolsets,
                     event_stream_handler=event_stream_handler,
+                    response_prefix=response_prefix,
                     **_deprecated_kwargs,
                 )
 
@@ -240,6 +244,7 @@ async def run(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
     ) -> AgentRunResult[OutputDataT]: ...
 
     @overload
@@ -258,6 +263,7 @@ async def run(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
     ) -> AgentRunResult[RunOutputDataT]: ...
 
     async def run(
@@ -275,6 +281,7 @@ async def run(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
         **_deprecated_kwargs: Never,
     ) -> AgentRunResult[Any]:
         """Run the agent with a user prompt in async mode.
@@ -308,6 +315,7 @@ async def main():
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
             toolsets: Optional additional toolsets for this run.
             event_stream_handler: Optional event stream handler to use for this run.
+            response_prefix: Optional response prefix to use for this run.
 
         Returns:
             The result of the run.
@@ -325,6 +333,7 @@ async def main():
             infer_name=infer_name,
             toolsets=toolsets,
             event_stream_handler=event_stream_handler,
+            response_prefix=response_prefix,
             **_deprecated_kwargs,
         )
 
@@ -344,6 +353,7 @@ def run_sync(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
     ) -> AgentRunResult[OutputDataT]: ...
 
     @overload
@@ -362,6 +372,7 @@ def run_sync(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
     ) -> AgentRunResult[RunOutputDataT]: ...
 
     def run_sync(
@@ -379,6 +390,7 @@ def run_sync(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
         **_deprecated_kwargs: Never,
     ) -> AgentRunResult[Any]:
         """Synchronously run the agent with a user prompt.
@@ -411,6 +423,7 @@ def run_sync(
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
             toolsets: Optional additional toolsets for this run.
             event_stream_handler: Optional event stream handler to use for this run.
+            response_prefix: Optional response prefix to use for this run.
 
         Returns:
             The result of the run.
@@ -428,6 +441,7 @@ def run_sync(
             infer_name=infer_name,
             toolsets=toolsets,
             event_stream_handler=event_stream_handler,
+            response_prefix=response_prefix,
             **_deprecated_kwargs,
         )
 
@@ -447,6 +461,7 @@ def run_stream(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
     ) -> AbstractAsyncContextManager[StreamedRunResult[AgentDepsT, OutputDataT]]: ...
 
     @overload
@@ -465,6 +480,7 @@ def run_stream(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
     ) -> AbstractAsyncContextManager[StreamedRunResult[AgentDepsT, RunOutputDataT]]: ...
 
     @asynccontextmanager
@@ -483,6 +499,7 @@ async def run_stream(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        response_prefix: str | None = None,
         **_deprecated_kwargs: Never,
     ) -> AsyncIterator[StreamedRunResult[AgentDepsT, Any]]:
         """Run the agent with a user prompt in async mode, returning a streamed response.
@@ -513,6 +530,7 @@ async def main():
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
             toolsets: Optional additional toolsets for this run.
             event_stream_handler: Optional event stream handler to use for this run. It will receive all the events up until the final result is found, which you can then read or stream from inside the context manager.
+            response_prefix: Optional response prefix to use for this run.
 
         Returns:
             The result of the run.
@@ -537,6 +555,7 @@ async def main():
             infer_name=infer_name,
             toolsets=toolsets,
             event_stream_handler=event_stream_handler,
+            response_prefix=response_prefix,
             **_deprecated_kwargs,
         ) as result:
             yield result
@@ -556,6 +575,7 @@ def iter(
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        response_prefix: str | None = None,
         **_deprecated_kwargs: Never,
     ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, OutputDataT]]: ...
 
@@ -574,6 +594,7 @@ def iter(
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        response_prefix: str | None = None,
         **_deprecated_kwargs: Never,
     ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, RunOutputDataT]]: ...
 
@@ -592,6 +613,7 @@ async def iter(
         usage: _usage.RunUsage | None = None,
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        response_prefix: str | None = None,
         **_deprecated_kwargs: Never,
     ) -> AsyncIterator[AgentRun[AgentDepsT, Any]]:
         """A contextmanager which can be used to iterate over the agent graph's nodes as they are executed.
@@ -627,6 +649,7 @@ async def main():
                     system_prompts=(),
                     system_prompt_functions=[],
                     system_prompt_dynamic_functions={},
+                    response_prefix=None,
                 ),
                 ModelRequestNode(
                     request=ModelRequest(
@@ -637,6 +660,7 @@ async def main():
                             )
                         ]
                     )
+                    response_prefix=None,
                 ),
                 CallToolsNode(
                     model_response=ModelResponse(
@@ -666,6 +690,7 @@ async def main():
             usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
             infer_name: Whether to try to infer the agent name from the call frame if it's not set.
             toolsets: Optional additional toolsets for this run.
+            response_prefix: Optional response prefix to use for this run.
 
         Returns:
             The result of the run.
@@ -688,6 +713,7 @@ async def main():
                 usage=usage,
                 infer_name=infer_name,
                 toolsets=toolsets,
+                response_prefix=response_prefix,
                 **_deprecated_kwargs,
             ) as run:
                 yield run

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -626,7 +626,12 @@ class AnthropicStreamedResponse(StreamedResponse):
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         current_block: BetaContentBlock | None = None
-        first_text_delta = True
+
+        # Handle response prefix by emitting it as the first text event
+        if response_prefix := self.model_request_parameters.response_prefix:
+            maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=response_prefix)
+            if maybe_event is not None:
+                yield maybe_event
 
         async for event in self._response:
             if isinstance(event, BetaRawMessageStartEvent):
@@ -670,12 +675,9 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
 
             elif isinstance(event, BetaRawContentBlockDeltaEvent):
                 if isinstance(event.delta, BetaTextDelta):
-                    content = event.delta.text
-                    # Prepend response prefix to the first text delta if provided
-                    if first_text_delta and self.model_request_parameters.response_prefix:
-                        content = self.model_request_parameters.response_prefix + content
-                        first_text_delta = False
-                    maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=event.index, content=content)
+                    maybe_event = self._parts_manager.handle_text_delta(
+                        vendor_part_id=event.index, content=event.delta.text
+                    )
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
                 elif isinstance(event.delta, BetaThinkingDelta):

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -525,13 +525,9 @@ def _process_response(
             ]
 
         if choice.message.content is not None:
-            content = choice.message.content
-            # Prepend response prefix if provided
-            if model_request_parameters.response_prefix:
-                content = model_request_parameters.response_prefix + content
             items.extend(
                 (replace(part, id='content', provider_name=self.system) if isinstance(part, ThinkingPart) else part)
-                for part in split_content_into_text_and_thinking(content, self.profile.thinking_tags)
+                for part in split_content_into_text_and_thinking(choice.message.content, self.profile.thinking_tags)
             )
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
@@ -1243,7 +1239,12 @@ class OpenAIStreamedResponse(StreamedResponse):
     _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
-        first_text_delta = True
+        # Handle response prefix by emitting it as the first text event
+        if response_prefix := self.model_request_parameters.response_prefix:
+            maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=response_prefix)
+            if maybe_event is not None:
+                yield maybe_event
+
         async for chunk in self._response:
             self._usage += _map_usage(chunk)
 
@@ -1266,10 +1267,6 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             # Handle the text part of the response
             content = choice.delta.content
             if content is not None:
-                # Prepend response prefix to the first text delta if provided
-                if first_text_delta and self.model_request_parameters.response_prefix:
-                    content = self.model_request_parameters.response_prefix + content
-                    first_text_delta = False
                 maybe_event = self._parts_manager.handle_text_delta(
                     vendor_part_id='content',
                     content=content,

pydantic_ai_slim/pydantic_ai/profiles/deepseek.py

@@ -5,4 +5,7 @@
 
 def deepseek_model_profile(model_name: str) -> ModelProfile | None:
     """Get the model profile for a DeepSeek model."""
-    return ModelProfile(ignore_streamed_leading_whitespace='r1' in model_name)
+    return ModelProfile(
+        ignore_streamed_leading_whitespace='r1' in model_name,
+        supports_response_prefix=True,
+    )

pydantic_ai_slim/pydantic_ai/profiles/openai.py

@@ -77,16 +77,12 @@ def openai_model_profile(model_name: str) -> ModelProfile:
     # See https://github.com/pydantic/pydantic-ai/issues/974 for more details.
     openai_system_prompt_role = 'user' if model_name.startswith('o1-mini') else None
 
-    # Enable response prefix for DeepSeek and OpenRouter models
-    supports_response_prefix = 'deepseek' in model_name.lower() or 'openrouter' in model_name.lower()
-
     return OpenAIModelProfile(
         json_schema_transformer=OpenAIJsonSchemaTransformer,
         supports_json_schema_output=True,
         supports_json_object_output=True,
         openai_unsupported_model_settings=openai_unsupported_model_settings,
         openai_system_prompt_role=openai_system_prompt_role,
-        supports_response_prefix=supports_response_prefix,
         openai_chat_supports_web_search=supports_web_search,
     )
 

pydantic_ai_slim/pydantic_ai/providers/openrouter.py

@@ -70,7 +70,9 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
 
         # As OpenRouterProvider is always used with OpenAIChatModel, which used to unconditionally use OpenAIJsonSchemaTransformer,
         # we need to maintain that behavior unless json_schema_transformer is set explicitly
-        return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
+        return OpenAIModelProfile(
+            json_schema_transformer=OpenAIJsonSchemaTransformer, supports_response_prefix=True
+        ).update(profile)
 
     @overload
     def __init__(self) -> None: ...