refactor: extract shared reasoning util and clean up provider implementations

Remove duplicated reasoning logic from vLLM, Ollama, and Bedrock providers
into a shared reasoning.py utility. Remove dead _prepare_reasoning_params
stubs, simplify return types to streaming-only, eliminate type ignores in
streaming.py, and add multi-turn reasoning + tool call integration test.

Signed-off-by: robinnarsinghranabhat <robinnarsingha123@gmail.com>

Author: robinnarsinghranabhat

src/llama_stack/core/routers/inference.py

@@ -42,7 +42,6 @@
     OpenAIChatCompletionResponseMessage,
     OpenAIChatCompletionToolCall,
     OpenAIChatCompletionToolCallFunction,
-    OpenAIChatCompletionWithReasoning,
     OpenAIChoice,
     OpenAIChoiceLogprobs,
     OpenAICompletion,
@@ -272,7 +271,7 @@ async def openai_chat_completion(
     async def openai_chat_completions_with_reasoning(
         self,
         params: OpenAIChatCompletionRequestWithExtraBody,
-    ) -> OpenAIChatCompletionWithReasoning | AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
+    ) -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
         """Called by the Responses layer when a user requests reasoning.
 
         Routes to the provider's reasoning-aware CC implementation, which

src/llama_stack/providers/inline/responses/builtin/responses/streaming.py

@@ -37,7 +37,6 @@
     OpenAIChatCompletionToolChoiceAllowedTools,
     OpenAIChatCompletionToolChoiceCustomTool,
     OpenAIChatCompletionToolChoiceFunctionTool,
-    OpenAIChatCompletionWithReasoning,
     OpenAIChoice,
     OpenAIChoiceLogprobs,
     OpenAIFinishReason,
@@ -547,7 +546,6 @@ async def create_response(self) -> AsyncIterator[OpenAIResponseObjectStream]:
                 completion_result: (
                     OpenAIChatCompletion
                     | AsyncIterator[OpenAIChatCompletionChunk]
-                    | OpenAIChatCompletionWithReasoning
                     | AsyncIterator[OpenAIChatCompletionChunkWithReasoning]
                 )
                 if self.reasoning and self.reasoning.effort and self.reasoning.effort != "none":
@@ -730,20 +728,18 @@ def _separate_tool_calls(
 
         for choice in current_response.choices:
             # Convert response message to input message format for multi-turn.
-            # Use AssistantMessageWithReasoning if reasoning was present in the
-            # CC response. Providers will be check for this AssistantMessageWithReasoning
-            # message
+            # Assign base type first, then narrow to AssistantMessageWithReasoning
+            # if reasoning was present in the CC response.
+            message: OpenAIAssistantMessageParam | AssistantMessageWithReasoning = OpenAIAssistantMessageParam(
+                content=choice.message.content,
+                tool_calls=choice.message.tool_calls,
+            )
             if reasoning_content:
                 message = AssistantMessageWithReasoning(
                     content=choice.message.content,
                     tool_calls=choice.message.tool_calls,
                     reasoning_content=reasoning_content,
                 )
-            else:
-                message = OpenAIAssistantMessageParam(  # type: ignore[assignment]
-                    content=choice.message.content,
-                    tool_calls=choice.message.tool_calls,
-                )
             next_turn_messages.append(message)
             logger.debug("Choice message content", content=choice.message.content)
             logger.debug("Choice message tool_calls", tool_calls=choice.message.tool_calls)
@@ -1123,7 +1119,7 @@ async def _process_streaming_chunks(
                 # chunk: OpenAIChatCompletionChunk annotation above.
                 if chunk_choice.delta.tool_calls:
                     for tool_call in chunk_choice.delta.tool_calls:
-                        response_tool_call = chat_response_tool_calls.get(tool_call.index, None)  # type: ignore[arg-type]
+                        response_tool_call = chat_response_tool_calls.get(tool_call.index, None)
                         # Create new tool call entry if this is the first chunk for this index
                         is_new_tool_call = response_tool_call is None
                         if is_new_tool_call:
@@ -1133,16 +1129,16 @@ async def _process_streaming_chunks(
                             if tool_call_dict.get("function") and tool_call_dict["function"].get("arguments") is None:
                                 tool_call_dict["function"]["arguments"] = "{}"
                             response_tool_call = OpenAIChatCompletionToolCall(**tool_call_dict)
-                            chat_response_tool_calls[tool_call.index] = response_tool_call  # type: ignore[index]
+                            chat_response_tool_calls[tool_call.index] = response_tool_call
 
                             # Create item ID for this tool call for streaming events
                             tool_call_item_id = f"fc_{uuid.uuid4()}"
-                            tool_call_item_ids[tool_call.index] = tool_call_item_id  # type: ignore[index]
+                            tool_call_item_ids[tool_call.index] = tool_call_item_id
 
                             # Emit output_item.added event for the new function call
                             self.sequence_number += 1
-                            is_mcp_tool = tool_call.function.name and tool_call.function.name in self.mcp_tool_to_server  # type: ignore[union-attr]
-                            if not is_mcp_tool and tool_call.function.name not in _SERVER_SIDE_BUILTIN_TOOL_NAMES:  # type: ignore[union-attr]
+                            is_mcp_tool = tool_call.function.name and tool_call.function.name in self.mcp_tool_to_server
+                            if not is_mcp_tool and tool_call.function.name not in _SERVER_SIDE_BUILTIN_TOOL_NAMES:
                                 # for MCP tools (and even other non-function tools) we emit an output message item later
                                 function_call_item = OpenAIResponseOutputMessageFunctionToolCall(
                                     arguments="",  # Will be filled incrementally via delta events

src/llama_stack/providers/remote/inference/bedrock/bedrock.py

@@ -5,20 +5,18 @@
 # the root directory of this source tree.
 
 from collections.abc import AsyncIterator
+from typing import cast
 
 from openai import AuthenticationError
 
 from llama_stack.log import get_logger
-from llama_stack.providers.inline.responses.builtin.responses.types import (
-    AssistantMessageWithReasoning,
-)
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack.providers.utils.inference.reasoning import map_reasoning_messages, wrap_chunks_with_reasoning
 from llama_stack_api import (
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
     OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionRequestWithExtraBody,
-    OpenAIChatCompletionWithReasoning,
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAIEmbeddingsRequestWithExtraBody,
@@ -65,56 +63,24 @@ async def openai_completion(
             "See https://docs.aws.amazon.com/bedrock/latest/userguide/inference-chat-completions.html"
         )
 
-    def _prepare_reasoning_params(self, params: OpenAIChatCompletionRequestWithExtraBody) -> None:
-        """Adapt CC request params to match what Bedrock expects for reasoning.
-
-        No-op for now. Override if Bedrock needs specific param adjustments.
-        """
-        pass
-
     async def openai_chat_completions_with_reasoning(
         self,
         params: OpenAIChatCompletionRequestWithExtraBody,
-    ) -> OpenAIChatCompletionWithReasoning | AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
+    ) -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
         """Chat completion with reasoning support for Bedrock.
 
         Extracts reasoning from Bedrock's response and wraps it in internal
-        types so the Responses layer can read reasoning as a typed field.
+        OpenAIChatCompletionChunkWithReasoning so the Responses layer can
+        read reasoning as a typed field.
         """
         if not params.stream:
             raise NotImplementedError("Non-streaming reasoning is not yet supported for Bedrock")
 
         params = params.model_copy()
-        self._prepare_reasoning_params(params)
-
-        # Bedrock's CC endpoint expects 'reasoning' on assistant messages, but
-        # that field isn't part of the official CC spec. Convert to dicts so we
-        # can rename reasoning_content → reasoning.
-        mapped_messages: list = []
-        for msg in params.messages:
-            if isinstance(msg, AssistantMessageWithReasoning) and msg.reasoning_content:
-                msg_dict = msg.model_dump(exclude_none=True)
-                msg_dict["reasoning"] = msg_dict.pop("reasoning_content")
-                mapped_messages.append(msg_dict)
-            else:
-                mapped_messages.append(msg)
-        params.messages = mapped_messages
-
+        params.messages = map_reasoning_messages(params.messages, reasoning_field="reasoning")
         result = await self.openai_chat_completion(params)
-
-        async def _wrap_chunks() -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
-            async for chunk in result:
-                reasoning = None
-                for choice in chunk.choices or []:
-                    reasoning = getattr(choice.delta, "reasoning", None) or getattr(
-                        choice.delta, "reasoning_content", None
-                    )
-                yield OpenAIChatCompletionChunkWithReasoning(
-                    chunk=chunk,
-                    reasoning_content=reasoning,
-                )
-
-        return _wrap_chunks()
+        result = cast(AsyncIterator[OpenAIChatCompletionChunk], result)
+        return wrap_chunks_with_reasoning(result, reasoning_fields=("reasoning", "reasoning_content"))
 
     async def openai_chat_completion(
         self,

src/llama_stack/providers/remote/inference/ollama/ollama.py

@@ -7,22 +7,21 @@
 
 import asyncio
 from collections.abc import AsyncIterator
+from typing import cast
 
 from ollama import AsyncClient as AsyncOllamaClient
 
 from llama_stack.log import get_logger
-from llama_stack.providers.inline.responses.builtin.responses.types import (
-    AssistantMessageWithReasoning,
-)
 from llama_stack.providers.remote.inference.ollama.config import OllamaImplConfig
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack.providers.utils.inference.reasoning import map_reasoning_messages, wrap_chunks_with_reasoning
 from llama_stack_api import (
     HealthResponse,
     HealthStatus,
     Model,
+    OpenAIChatCompletionChunk,
     OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionRequestWithExtraBody,
-    OpenAIChatCompletionWithReasoning,
     UnsupportedModelError,
 )
 
@@ -79,62 +78,26 @@ def get_api_key(self):
     def get_base_url(self):
         return str(self.config.base_url)
 
-    def _prepare_reasoning_params(self, params: OpenAIChatCompletionRequestWithExtraBody) -> None:
-        """Adapt CC request params to match what Ollama expects for reasoning.
-
-        Each provider may need different param adjustments. For Ollama:
-        - If reasoning_effort is not set, default to "none" so Ollama
-          doesn't apply its own default (medium).
-
-        Override this in other providers if they need different mapping,
-        e.g. converting effort levels to boolean flags.
-        """
-        if params.reasoning_effort is None:
-            params.reasoning_effort = "none"
-
     async def openai_chat_completions_with_reasoning(
         self,
         params: OpenAIChatCompletionRequestWithExtraBody,
-    ) -> OpenAIChatCompletionWithReasoning | AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
+    ) -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
         """Chat completion with reasoning support for Ollama.
 
         Extracts reasoning from Ollama's response and wraps it in internal
-        types so the Responses layer can read reasoning as a typed field.
+        OpenAIChatCompletionChunkWithReasoning so the Responses layer can
+        read reasoning as a typed field.
         """
         if not params.stream:
             raise NotImplementedError("Non-streaming reasoning is not yet supported for Ollama")
 
         params = params.model_copy()
-        self._prepare_reasoning_params(params)
-
-        # Ollama's CC endpoint expects 'reasoning' on assistant messages, but
-        # that field isn't part of the official CC spec. Convert to dicts so we
-        # can rename reasoning_content → reasoning.
-        mapped_messages: list = []
-        for msg in params.messages:
-            if isinstance(msg, AssistantMessageWithReasoning) and msg.reasoning_content:
-                msg_dict = msg.model_dump(exclude_none=True)
-                msg_dict["reasoning"] = msg_dict.pop("reasoning_content")
-                mapped_messages.append(msg_dict)
-            else:
-                mapped_messages.append(msg)
-        params.messages = mapped_messages
-
+        if params.reasoning_effort is None:
+            params.reasoning_effort = "none"
+        params.messages = map_reasoning_messages(params.messages, reasoning_field="reasoning")
         result = await self.openai_chat_completion(params)
-
-        async def _wrap_chunks() -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
-            async for chunk in result:
-                reasoning = None
-                for choice in chunk.choices or []:
-                    reasoning = getattr(choice.delta, "reasoning", None) or getattr(
-                        choice.delta, "reasoning_content", None
-                    )
-                yield OpenAIChatCompletionChunkWithReasoning(
-                    chunk=chunk,
-                    reasoning_content=reasoning,
-                )
-
-        return _wrap_chunks()
+        result = cast(AsyncIterator[OpenAIChatCompletionChunk], result)
+        return wrap_chunks_with_reasoning(result, reasoning_fields=("reasoning", "reasoning_content"))
 
     async def initialize(self) -> None:
         logger.info("checking connectivity to Ollama", base_url=self.config.base_url)

src/llama_stack/providers/remote/inference/vllm/vllm.py

@@ -4,17 +4,16 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 from collections.abc import AsyncIterator
+from typing import cast
 from urllib.parse import urljoin
 
 import httpx
 from pydantic import ConfigDict
 
 from llama_stack.log import get_logger
-from llama_stack.providers.inline.responses.builtin.responses.types import (
-    AssistantMessageWithReasoning,
-)
 from llama_stack.providers.utils.inference.http_client import _build_network_client_kwargs
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack.providers.utils.inference.reasoning import map_reasoning_messages, wrap_chunks_with_reasoning
 from llama_stack_api import (
     HealthResponse,
     HealthStatus,
@@ -26,7 +25,6 @@
     OpenAIChatCompletionContentPartImageParam,
     OpenAIChatCompletionContentPartTextParam,
     OpenAIChatCompletionRequestWithExtraBody,
-    OpenAIChatCompletionWithReasoning,
     RerankData,
     RerankResponse,
 )
@@ -117,58 +115,24 @@ async def openai_chat_completion(
 
         return await super().openai_chat_completion(params)
 
-    def _prepare_reasoning_params(self, params: OpenAIChatCompletionRequestWithExtraBody) -> None:
-        """Adapt CC request params to match what vLLM expects for reasoning.
-
-        No-op for now. Override if vLLM needs specific param adjustments,
-        e.g. mapping effort levels or moving params to extra_body.
-        """
-        pass
-
     async def openai_chat_completions_with_reasoning(
         self,
         params: OpenAIChatCompletionRequestWithExtraBody,
-    ) -> OpenAIChatCompletionWithReasoning | AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
+    ) -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
         """Chat completion with reasoning support for vLLM.
 
         Extracts reasoning from vLLM's response and wraps it in internal
-        types (OpenAIChatCompletionChunkWithReasoning / OpenAIChatCompletionWithReasoning)
-        so the Responses layer can read reasoning as a typed field.
+        OpenAIChatCompletionChunkWithReasoning so the Responses layer can
+        read reasoning as a typed field.
         """
         if not params.stream:
             raise NotImplementedError("Non-streaming reasoning is not yet supported for vLLM")
 
         params = params.model_copy()
-        self._prepare_reasoning_params(params)
-
-        # vLLM's CC endpoint expects 'reasoning' on assistant messages, but
-        # that field isn't part of the official CC spec. Convert to dicts so we
-        # can rename reasoning_content → reasoning.
-        mapped_messages: list = []
-        for msg in params.messages:
-            if isinstance(msg, AssistantMessageWithReasoning) and msg.reasoning_content:
-                msg_dict = msg.model_dump(exclude_none=True)
-                msg_dict["reasoning"] = msg_dict.pop("reasoning_content")
-                mapped_messages.append(msg_dict)
-            else:
-                mapped_messages.append(msg)
-        params.messages = mapped_messages
-
+        params.messages = map_reasoning_messages(params.messages, reasoning_field="reasoning")
         result = await self.openai_chat_completion(params)
-
-        async def _wrap_chunks() -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
-            async for chunk in result:  # type: ignore[union-attr]
-                reasoning = None
-                for choice in chunk.choices or []:
-                    reasoning = getattr(choice.delta, "reasoning", None) or getattr(
-                        choice.delta, "reasoning_content", None
-                    )
-                yield OpenAIChatCompletionChunkWithReasoning(
-                    chunk=chunk,
-                    reasoning_content=reasoning,
-                )
-
-        return _wrap_chunks()
+        result = cast(AsyncIterator[OpenAIChatCompletionChunk], result)
+        return wrap_chunks_with_reasoning(result, reasoning_fields=("reasoning", "reasoning_content"))
 
     def construct_model_from_identifier(self, identifier: str) -> Model:
         # vLLM's /v1/models response does not expose a model task/type field, so classify by name.