Route Harmony reasoning separately from final output (#21378)

Differential Revision: D112991739

Pull Request resolved: #21378

Author: Gasoonjia

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
 
 

examples/llm_server/python/serving_chat.py

@@ -67,13 +67,17 @@ 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
         self._model_id = model_id
         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 <tool_call> 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,