From a8cc9afb9191b1e9a7cd8fb9e1053c6f86f2c672 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Sun, 26 Jul 2026 19:55:12 -0700 Subject: [PATCH] Route Harmony reasoning separately from final output (#21378) Summary: Teach the shared OpenAI chat adapter to expose optional reasoning_content. Reviewed By: metascroy Differential Revision: D112991739 --- examples/llm_server/python/protocol.py | 2 + examples/llm_server/python/serving_chat.py | 131 ++++++++++++++++----- 2 files changed, 101 insertions(+), 32 deletions(-) diff --git a/examples/llm_server/python/protocol.py b/examples/llm_server/python/protocol.py index 1a6e4771b9f..d8a8edc9cbb 100644 --- a/examples/llm_server/python/protocol.py +++ b/examples/llm_server/python/protocol.py @@ -102,6 +102,7 @@ class Usage(BaseModel): class ResponseMessage(BaseModel): role: str = "assistant" content: Optional[str] = None + reasoning_content: Optional[str] = None tool_calls: Optional[list[ToolCall]] = None @@ -123,6 +124,7 @@ class ChatCompletionResponse(BaseModel): class DeltaMessage(BaseModel): role: Optional[str] = None content: Optional[str] = None + reasoning_content: Optional[str] = None tool_calls: Optional[list[ToolCall]] = None diff --git a/examples/llm_server/python/serving_chat.py b/examples/llm_server/python/serving_chat.py index 28551f6bff3..95e3272d6fd 100644 --- a/examples/llm_server/python/serving_chat.py +++ b/examples/llm_server/python/serving_chat.py @@ -67,6 +67,9 @@ def __init__( prompt_token_offset: int = 0, content_filter: Optional[Callable[[str], str]] = None, content_filter_specials: Optional[set[str]] = None, + reasoning_extractor: Optional[ + Callable[[str], tuple[Optional[str], str]] + ] = None, ): self._runtime = runtime self._template = template @@ -74,6 +77,7 @@ def __init__( self._max_context = max_context self._prompt_token_offset = prompt_token_offset self._content_filter = content_filter + self._reasoning_extractor = reasoning_extractor # Detector CLASS; a fresh instance is created per request so streaming # state is never shared across concurrent requests. self._tool_detector_cls = tool_detector_cls @@ -122,6 +126,17 @@ def _visible_content(self, text: str) -> str: text = self._content_filter(text) return self._strip_specials(text) + def _split_reasoning(self, text: str) -> tuple[Optional[str], str]: + if self._reasoning_extractor is None: + return None, text + return self._reasoning_extractor(text) + + @staticmethod + def _return_reasoning(req: ChatCompletionRequest) -> bool: + kwargs = req.chat_template_kwargs or {} + value = kwargs.get("return_reasoning", False) + return value if isinstance(value, bool) else False + @staticmethod def _to_openai_tool_call(item: ToolCallItem) -> ToolCall: return ToolCall( @@ -172,17 +187,20 @@ async def _collect_until_stop(self, stream: AsyncIterator[str], stops: list[str] break return text, stopped - def _extract_tools(self, req: ChatCompletionRequest, text: str): - """Returns (tool_calls | None, content_text). Falls back to plain text.""" + def _extract_response(self, req: ChatCompletionRequest, text: str): + """Return tool calls, optional reasoning, and user-visible content.""" + tool_calls = None if self._tools_active(req): parsed = self._tool_detector_cls().detect_and_parse( text, self._tool_schemas(req) ) if parsed.calls: - content = self._visible_content(parsed.normal_text) or None - return [self._to_openai_tool_call(c) for c in parsed.calls], content + tool_calls = [self._to_openai_tool_call(c) for c in parsed.calls] text = parsed.normal_text - return None, self._visible_content(text) + reasoning, text = self._split_reasoning(text) + if not self._return_reasoning(req): + reasoning = None + return tool_calls, reasoning, self._visible_content(text) or None @staticmethod def _log_generation_stats( @@ -408,8 +426,10 @@ async def create(self, req: ChatCompletionRequest): # the tool schemas, the model can emit a that we'd surface as # plain text (parsing is disabled), instead of a normal answer. template_tools = None if req.tool_choice == "none" else req.tools + template_kwargs = dict(req.chat_template_kwargs or {}) + template_kwargs.pop("return_reasoning", None) prompt = self._template.render( - req.messages, tools=template_tools, template_kwargs=req.chat_template_kwargs + req.messages, tools=template_tools, template_kwargs=template_kwargs ) # Token-ID segments splice prior assistant turns' exact ids so warm resume # survives the template's lossy tool-call re-render; plain text when @@ -420,7 +440,7 @@ async def create(self, req: ChatCompletionRequest): messages=req.messages, rendered_prompt=prompt, tools=template_tools, - template_kwargs=req.chat_template_kwargs, + template_kwargs=template_kwargs, ) # Pre-flight context check against the tokens the worker will actually # assemble: for segments that is sum(len(ids)) + tokenized text, not the @@ -450,7 +470,7 @@ async def create(self, req: ChatCompletionRequest): # reproduces the exact resident scaffold even if the mode changes between # requests. preamble = self._template.generation_preamble( - req.chat_template_kwargs, tools=template_tools + template_kwargs, tools=template_tools ) # Admit the session up front (before the stream's first chunk) so a # capacity refusal is an HTTP status, not a mid-stream error event. @@ -488,7 +508,9 @@ async def _complete( raise GenerationError(str(e)) # Bound the raw output at the first stop/special token BEFORE tool # parsing, so a call after the stop boundary is not parsed/emitted. - tool_calls, content = self._extract_tools(req, self._truncate_raw(text, req)) + tool_calls, reasoning, content = self._extract_response( + req, self._truncate_raw(text, req) + ) # Record after the response is finalized: the fingerprint is of exactly # what we return (content + tool_calls), so the next turn can confirm the # client echoed this turn before splicing its ids. @@ -508,7 +530,11 @@ async def _complete( model=self._model_id, choices=[ Choice( - message=ResponseMessage(content=content, tool_calls=tool_calls), + message=ResponseMessage( + content=content, + reasoning_content=reasoning, + tool_calls=tool_calls, + ), finish_reason=finish, ) ], @@ -549,6 +575,42 @@ def on_stop(): ): yield token + async def _stream_final_chunks( + self, + req: ChatCompletionRequest, + stats: GenStats, + use_tools: bool, + tool_calls, + reasoning: Optional[str], + content: Optional[str], + stopped: bool, + chunk, + ) -> AsyncIterator[str]: + if use_tools: + if reasoning: + yield chunk(DeltaMessage(reasoning_content=reasoning)) + if content: + yield chunk(DeltaMessage(content=content)) + for tool_call in tool_calls or []: + yield chunk(DeltaMessage(tool_calls=[tool_call])) + finish = self._finish_reason( + req, + stats.completion_tokens, + tool_calls, + stopped, + stats.finish_reason, + ) + else: + finish = self._finish_reason( + req, + stats.completion_tokens, + stopped=stopped, + worker_finish=stats.finish_reason, + ) + + self._log_generation_stats(req.session_id, stats, finish) + yield chunk(DeltaMessage(), finish=finish) + async def _stream( self, req: ChatCompletionRequest, @@ -571,6 +633,7 @@ def chunk(delta: DeltaMessage, finish=None) -> str: error: Optional[Exception] = None use_tools = self._tools_active(req) tool_calls = None + reasoning = None content = None stats = GenStats() @@ -594,9 +657,23 @@ def chunk(delta: DeltaMessage, finish=None) -> str: ), stops, ) - tool_calls, content = self._extract_tools( + tool_calls, reasoning, content = self._extract_response( req, self._truncate_raw(raw, req) ) + elif self._reasoning_extractor is not None: + raw, stop_hit[0] = await self._collect_until_stop( + self._runtime.generate_stream( + req.session_id, prompt, options, stats + ), + stops, + ) + tool_calls, reasoning, content = self._extract_response( + req, self._apply_stop(raw, stops) + ) + if reasoning: + yield chunk(DeltaMessage(reasoning_content=reasoning)) + if content: + yield chunk(DeltaMessage(content=content)) else: streamed: list[str] = [] async for token in self._stream_plain_content( @@ -628,27 +705,17 @@ def chunk(delta: DeltaMessage, finish=None) -> str: preamble=preamble, ) - if use_tools: - if content: - yield chunk(DeltaMessage(content=content)) - for tc in tool_calls or []: - yield chunk(DeltaMessage(tool_calls=[tc])) - finish = self._finish_reason( - req, - stats.completion_tokens, - tool_calls, - stop_hit[0], - stats.finish_reason, - ) - else: - finish = self._finish_reason( - req, - stats.completion_tokens, - stopped=stop_hit[0], - worker_finish=stats.finish_reason, - ) - self._log_generation_stats(req.session_id, stats, finish) - yield chunk(DeltaMessage(), finish=finish) + async for final_chunk in self._stream_final_chunks( + req, + stats, + use_tools, + tool_calls, + reasoning, + content, + stop_hit[0], + chunk, + ): + yield final_chunk if req.stream_options and req.stream_options.include_usage: usage_chunk = ChatCompletionChunk( id=cid,