fix: move reasoning wrapper types to llama-stack-api and fix mypy errors

src/llama_stack/core/routers/inference.py

@@ -18,10 +18,6 @@
 from llama_stack.core.datatypes import ModelWithOwner
 from llama_stack.core.request_headers import get_authenticated_user
 from llama_stack.log import get_logger
-from llama_stack.providers.inline.responses.builtin.responses.types import (
-    OpenAIChatCompletionChunkWithReasoning,
-    OpenAIChatCompletionWithReasoning,
-)
 from llama_stack.providers.utils.inference.inference_store import InferenceStore
 from llama_stack.telemetry.inference_metrics import (
     create_inference_metric_attributes,
@@ -41,10 +37,12 @@
     ModelTypeError,
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAIChatCompletionResponseMessage,
     OpenAIChatCompletionToolCall,
     OpenAIChatCompletionToolCallFunction,
+    OpenAIChatCompletionWithReasoning,
     OpenAIChoice,
     OpenAIChoiceLogprobs,
     OpenAICompletion,

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

@@ -28,6 +28,7 @@
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionCustomToolCall,
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAIChatCompletionResponseMessage,
@@ -36,6 +37,7 @@
     OpenAIChatCompletionToolChoiceAllowedTools,
     OpenAIChatCompletionToolChoiceCustomTool,
     OpenAIChatCompletionToolChoiceFunctionTool,
+    OpenAIChatCompletionWithReasoning,
     OpenAIChoice,
     OpenAIChoiceLogprobs,
     OpenAIFinishReason,
@@ -108,7 +110,6 @@
     AssistantMessageWithReasoning,
     ChatCompletionContext,
     ChatCompletionResult,
-    OpenAIChatCompletionChunkWithReasoning,
 )
 from .utils import (
     convert_chat_choice_to_response_message,
@@ -543,6 +544,12 @@ async def create_response(self) -> AsyncIterator[OpenAIResponseObjectStream]:
                     **(self.extra_body or {}),
                 )
                 # Use reasoning-aware method when reasoning is explicitly requested
+                completion_result: (
+                    OpenAIChatCompletion
+                    | AsyncIterator[OpenAIChatCompletionChunk]
+                    | OpenAIChatCompletionWithReasoning
+                    | AsyncIterator[OpenAIChatCompletionChunkWithReasoning]
+                )
                 if self.reasoning and self.reasoning.effort and self.reasoning.effort != "none":
                     try:
                         # Pass a copy — the router mutates params.model (strips provider prefix).

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

@@ -13,8 +13,6 @@
 
 from llama_stack_api import (
     OpenAIAssistantMessageParam,
-    OpenAIChatCompletion,
-    OpenAIChatCompletionChunk,
     OpenAIChatCompletionToolCall,
     OpenAIFinishReason,
     OpenAIMessageParam,
@@ -50,30 +48,6 @@ class AssistantMessageWithReasoning(OpenAIAssistantMessageParam):
     reasoning_content: str | None = None
 
 
-@dataclass
-class OpenAIChatCompletionWithReasoning:
-    """Internal wrapper: a CC response with extracted reasoning content.
-
-    Returned by openai_chat_completions_with_reasoning for non-streaming.
-    The Responses layer unwraps .completion and reads .reasoning_content.
-    """
-
-    completion: OpenAIChatCompletion
-    reasoning_content: str | None = None
-
-
-@dataclass
-class OpenAIChatCompletionChunkWithReasoning:
-    """Internal wrapper: a CC streaming chunk with extracted reasoning content.
-
-    Yielded by openai_chat_completions_with_reasoning for streaming.
-    The Responses layer unwraps .chunk and reads .reasoning_content.
-    """
-
-    chunk: OpenAIChatCompletionChunk
-    reasoning_content: str | None = None
-
-
 def _json_equal(a: str, b: str) -> bool:
     """Compare two JSON strings by value, falling back to string comparison."""
     try:

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

@@ -11,14 +11,14 @@
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.responses.builtin.responses.types import (
     AssistantMessageWithReasoning,
-    OpenAIChatCompletionChunkWithReasoning,
-    OpenAIChatCompletionWithReasoning,
 )
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 from llama_stack_api import (
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionRequestWithExtraBody,
+    OpenAIChatCompletionWithReasoning,
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAIEmbeddingsRequestWithExtraBody,

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

@@ -13,16 +13,16 @@
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.responses.builtin.responses.types import (
     AssistantMessageWithReasoning,
-    OpenAIChatCompletionChunkWithReasoning,
-    OpenAIChatCompletionWithReasoning,
 )
 from llama_stack.providers.remote.inference.ollama.config import OllamaImplConfig
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 from llama_stack_api import (
     HealthResponse,
     HealthStatus,
     Model,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionRequestWithExtraBody,
+    OpenAIChatCompletionWithReasoning,
     UnsupportedModelError,
 )
 

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

@@ -12,8 +12,6 @@
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.responses.builtin.responses.types import (
     AssistantMessageWithReasoning,
-    OpenAIChatCompletionChunkWithReasoning,
-    OpenAIChatCompletionWithReasoning,
 )
 from llama_stack.providers.utils.inference.http_client import _build_network_client_kwargs
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
@@ -24,9 +22,11 @@
     ModelType,
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionContentPartImageParam,
     OpenAIChatCompletionContentPartTextParam,
     OpenAIChatCompletionRequestWithExtraBody,
+    OpenAIChatCompletionWithReasoning,
     RerankData,
     RerankResponse,
 )

src/llama_stack_api/__init__.py

@@ -256,6 +256,8 @@
     ModelStore,
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,
+    OpenAIChatCompletionChunkWithReasoning,
+    OpenAIChatCompletionWithReasoning,
     OpenAIChatCompletionChunk,
     OpenAIChatCompletionContentPartImageParam,
     OpenAIChatCompletionContentPartParam,
@@ -803,6 +805,8 @@
     "OpenAIAssistantMessageParam",
     "OpenAIChatCompletion",
     "OpenAIChatCompletionChunk",
+    "OpenAIChatCompletionChunkWithReasoning",
+    "OpenAIChatCompletionWithReasoning",
     "OpenAIChatCompletionContentPartImageParam",
     "OpenAIChatCompletionContentPartParam",
     "OpenAIChatCompletionContentPartTextParam",

src/llama_stack_api/inference/__init__.py

@@ -43,6 +43,7 @@
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionContentPartImageParam,
     OpenAIChatCompletionContentPartParam,
     OpenAIChatCompletionContentPartTextParam,
@@ -61,6 +62,7 @@
     OpenAIChatCompletionUsage,
     OpenAIChatCompletionUsageCompletionTokensDetails,
     OpenAIChatCompletionUsagePromptTokensDetails,
+    OpenAIChatCompletionWithReasoning,
     OpenAIChoice,
     OpenAIChoiceDelta,
     OpenAIChoiceLogprobs,
@@ -191,6 +193,8 @@
     "OpenAIChatCompletionUsage",
     "OpenAIChatCompletion",
     "OpenAIChatCompletionChunk",
+    "OpenAIChatCompletionChunkWithReasoning",
+    "OpenAIChatCompletionWithReasoning",
     "OpenAICompletionLogprobs",
     "OpenAICompletionChoice",
     "OpenAICompletion",

src/llama_stack_api/inference/api.py

@@ -15,7 +15,9 @@
     ListOpenAIChatCompletionResponse,
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionRequestWithExtraBody,
+    OpenAIChatCompletionWithReasoning,
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAICompletionWithInputMessages,
@@ -67,15 +69,15 @@ async def openai_chat_completion(
     async def openai_chat_completions_with_reasoning(
         self,
         params: OpenAIChatCompletionRequestWithExtraBody,
-    ) -> object:
+    ) -> OpenAIChatCompletionWithReasoning | AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
         """Chat completion with reasoning token extraction.
 
         Internal method used by the Responses implementation when reasoning
         is requested. Returns internal wrapper types that carry reasoning
         alongside the CC response:
         - OpenAIChatCompletionWithReasoning (non-streaming)
         - AsyncIterator[OpenAIChatCompletionChunkWithReasoning] (streaming)
-        These are defined in llama_stack.providers.inline.responses.builtin.responses.types.
+        These are defined in llama_stack_api.inference.models.
 
         Default raises NotImplementedError so unsupported providers fail
         loudly instead of silently returning no reasoning.

src/llama_stack_api/inference/models.py

@@ -1038,6 +1038,32 @@ class RerankRequest(BaseModel):
     )
 
 
+class OpenAIChatCompletionWithReasoning(BaseModel):
+    """Internal wrapper: a CC response with extracted reasoning content.
+
+    Returned by openai_chat_completions_with_reasoning for non-streaming.
+    The Responses layer unwraps .completion and reads .reasoning_content.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    completion: Any = Field(..., description="The chat completion response.")
+    reasoning_content: str | None = Field(None, description="Extracted reasoning content, if any.")
+
+
+class OpenAIChatCompletionChunkWithReasoning(BaseModel):
+    """Internal wrapper: a CC streaming chunk with extracted reasoning content.
+
+    Yielded by openai_chat_completions_with_reasoning for streaming.
+    The Responses layer unwraps .chunk and reads .reasoning_content.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    chunk: Any = Field(..., description="The chat completion chunk.")
+    reasoning_content: str | None = Field(None, description="Extracted reasoning content, if any.")
+
+
 __all__ = [
     # Sampling
     "GreedySamplingStrategy",
@@ -1071,6 +1097,8 @@ class RerankRequest(BaseModel):
     "RerankResponse",
     # OpenAI Compatibility
     "OpenAIChatCompletionContentPartTextParam",
+    "OpenAIChatCompletionWithReasoning",
+    "OpenAIChatCompletionChunkWithReasoning",
     "OpenAIImageURL",
     "OpenAIChatCompletionContentPartImageParam",
     "OpenAIFileFile",