Fix streaming gpt-oss using Ollama (#3035)

Author: DouweM

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -526,16 +526,20 @@ def _process_response(self, response: chat.ChatCompletion | str) -> ModelRespons
 
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
+
         # The `reasoning_content` field is only present in DeepSeek models.
         # https://api-docs.deepseek.com/guides/reasoning_model
         if reasoning_content := getattr(choice.message, 'reasoning_content', None):
             items.append(ThinkingPart(id='reasoning_content', content=reasoning_content, provider_name=self.system))
 
-        # NOTE: We don't currently handle OpenRouter `reasoning_details`:
-        # - https://openrouter.ai/docs/use-cases/reasoning-tokens#preserving-reasoning-blocks
-        # NOTE: We don't currently handle OpenRouter/gpt-oss `reasoning`:
+        # The `reasoning` field is only present in gpt-oss via Ollama and OpenRouter.
         # - https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot#chat-completions-api
         # - https://openrouter.ai/docs/use-cases/reasoning-tokens#basic-usage-with-reasoning-tokens
+        if reasoning := getattr(choice.message, 'reasoning', None):
+            items.append(ThinkingPart(id='reasoning', content=reasoning, provider_name=self.system))
+
+        # NOTE: We don't currently handle OpenRouter `reasoning_details`:
+        # - https://openrouter.ai/docs/use-cases/reasoning-tokens#preserving-reasoning-blocks
         # If you need this, please file an issue.
 
         vendor_details: dict[str, Any] = {}
@@ -1492,6 +1496,17 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     provider_name=self.provider_name,
                 )
 
+            # The `reasoning` field is only present in gpt-oss via Ollama and OpenRouter.
+            # - https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot#chat-completions-api
+            # - https://openrouter.ai/docs/use-cases/reasoning-tokens#basic-usage-with-reasoning-tokens
+            if reasoning := getattr(choice.delta, 'reasoning', None):  # pragma: no cover
+                yield self._parts_manager.handle_thinking_delta(
+                    vendor_part_id='reasoning',
+                    id='reasoning',
+                    content=reasoning,
+                    provider_name=self.provider_name,
+                )
+
             for dtc in choice.delta.tool_calls or []:
                 maybe_event = self._parts_manager.handle_tool_call_delta(
                     vendor_part_id=dtc.index,

pydantic_ai_slim/pydantic_ai/profiles/harmony.py

@@ -10,4 +10,6 @@ def harmony_model_profile(model_name: str) -> ModelProfile | None:
     See <https://cookbook.openai.com/articles/openai-harmony> for more details.
     """
     profile = openai_model_profile(model_name)
-    return OpenAIModelProfile(openai_supports_tool_choice_required=False).update(profile)
+    return OpenAIModelProfile(
+        openai_supports_tool_choice_required=False, ignore_streamed_leading_whitespace=True
+    ).update(profile)

pydantic_ai_slim/pydantic_ai/providers/ollama.py

@@ -11,6 +11,7 @@
 from pydantic_ai.profiles.cohere import cohere_model_profile
 from pydantic_ai.profiles.deepseek import deepseek_model_profile
 from pydantic_ai.profiles.google import google_model_profile
+from pydantic_ai.profiles.harmony import harmony_model_profile
 from pydantic_ai.profiles.meta import meta_model_profile
 from pydantic_ai.profiles.mistral import mistral_model_profile
 from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
@@ -50,6 +51,7 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
             'deepseek': deepseek_model_profile,
             'mistral': mistral_model_profile,
             'command': cohere_model_profile,
+            'gpt-oss': harmony_model_profile,
         }
 
         profile = None

tests/providers/test_ollama.py

@@ -9,6 +9,7 @@
 from pydantic_ai.profiles.cohere import cohere_model_profile
 from pydantic_ai.profiles.deepseek import deepseek_model_profile
 from pydantic_ai.profiles.google import GoogleJsonSchemaTransformer, google_model_profile
+from pydantic_ai.profiles.harmony import harmony_model_profile
 from pydantic_ai.profiles.meta import meta_model_profile
 from pydantic_ai.profiles.mistral import mistral_model_profile
 from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer
@@ -77,6 +78,7 @@ def test_ollama_provider_model_profile(mocker: MockerFixture):
     mistral_model_profile_mock = mocker.patch(f'{ns}.mistral_model_profile', wraps=mistral_model_profile)
     qwen_model_profile_mock = mocker.patch(f'{ns}.qwen_model_profile', wraps=qwen_model_profile)
     cohere_model_profile_mock = mocker.patch(f'{ns}.cohere_model_profile', wraps=cohere_model_profile)
+    harmony_model_profile_mock = mocker.patch(f'{ns}.harmony_model_profile', wraps=harmony_model_profile)
 
     meta_profile = provider.model_profile('llama3.2')
     meta_model_profile_mock.assert_called_with('llama3.2')
@@ -115,6 +117,12 @@ def test_ollama_provider_model_profile(mocker: MockerFixture):
     assert cohere_profile is not None
     assert cohere_profile.json_schema_transformer == OpenAIJsonSchemaTransformer
 
+    harmony_profile = provider.model_profile('gpt-oss')
+    harmony_model_profile_mock.assert_called_with('gpt-oss')
+    assert harmony_profile is not None
+    assert harmony_profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+    assert harmony_profile.ignore_streamed_leading_whitespace is True
+
     unknown_profile = provider.model_profile('unknown-model')
     assert unknown_profile is not None
     assert unknown_profile.json_schema_transformer == OpenAIJsonSchemaTransformer