diff --git a/src/llama_stack/core/routers/inference.py b/src/llama_stack/core/routers/inference.py index c6afebf41ba..4259c558bed 100644 --- a/src/llama_stack/core/routers/inference.py +++ b/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, diff --git a/src/llama_stack/providers/inline/responses/builtin/responses/streaming.py b/src/llama_stack/providers/inline/responses/builtin/responses/streaming.py index 9a599ade205..e2c80d5f746 100644 --- a/src/llama_stack/providers/inline/responses/builtin/responses/streaming.py +++ b/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). diff --git a/src/llama_stack/providers/inline/responses/builtin/responses/types.py b/src/llama_stack/providers/inline/responses/builtin/responses/types.py index 4b4e7fc00b2..3eb563d8c65 100644 --- a/src/llama_stack/providers/inline/responses/builtin/responses/types.py +++ b/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: diff --git a/src/llama_stack/providers/remote/inference/bedrock/bedrock.py b/src/llama_stack/providers/remote/inference/bedrock/bedrock.py index 300ef9dd0ef..02800d21ed2 100644 --- a/src/llama_stack/providers/remote/inference/bedrock/bedrock.py +++ b/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, diff --git a/src/llama_stack/providers/remote/inference/ollama/ollama.py b/src/llama_stack/providers/remote/inference/ollama/ollama.py index 6f95a621c7e..45b41f7b3e0 100644 --- a/src/llama_stack/providers/remote/inference/ollama/ollama.py +++ b/src/llama_stack/providers/remote/inference/ollama/ollama.py @@ -13,8 +13,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.remote.inference.ollama.config import OllamaImplConfig from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin @@ -22,7 +20,9 @@ HealthResponse, HealthStatus, Model, + OpenAIChatCompletionChunkWithReasoning, OpenAIChatCompletionRequestWithExtraBody, + OpenAIChatCompletionWithReasoning, UnsupportedModelError, ) diff --git a/src/llama_stack/providers/remote/inference/vllm/vllm.py b/src/llama_stack/providers/remote/inference/vllm/vllm.py index 6c491fc8236..045dfb83b36 100644 --- a/src/llama_stack/providers/remote/inference/vllm/vllm.py +++ b/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, ) diff --git a/src/llama_stack_api/__init__.py b/src/llama_stack_api/__init__.py index 9054f19b527..90c209598fa 100644 --- a/src/llama_stack_api/__init__.py +++ b/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", diff --git a/src/llama_stack_api/inference/__init__.py b/src/llama_stack_api/inference/__init__.py index f14f62c7c8d..18431f6a52c 100644 --- a/src/llama_stack_api/inference/__init__.py +++ b/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", diff --git a/src/llama_stack_api/inference/api.py b/src/llama_stack_api/inference/api.py index 7c68c1fb19d..b15009f6a91 100644 --- a/src/llama_stack_api/inference/api.py +++ b/src/llama_stack_api/inference/api.py @@ -15,7 +15,9 @@ ListOpenAIChatCompletionResponse, OpenAIChatCompletion, OpenAIChatCompletionChunk, + OpenAIChatCompletionChunkWithReasoning, OpenAIChatCompletionRequestWithExtraBody, + OpenAIChatCompletionWithReasoning, OpenAICompletion, OpenAICompletionRequestWithExtraBody, OpenAICompletionWithInputMessages, @@ -67,7 +69,7 @@ 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 @@ -75,7 +77,7 @@ async def openai_chat_completions_with_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. diff --git a/src/llama_stack_api/inference/models.py b/src/llama_stack_api/inference/models.py index 616cf7c4dc4..1b5b8d9da95 100644 --- a/src/llama_stack_api/inference/models.py +++ b/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",