feat(reasoning): add unified kimi_k2 reasoning parser - #1992
feat(reasoning): add unified kimi_k2 reasoning parser#1992ighutake-debug wants to merge 2 commits into
Conversation
One prefill-robust parser for the whole Kimi K2 family, matching vLLM/SGLang's kimi_k2 semantics: starts in reasoning (a leading <think> is consumed when present), and reasoning ends on </think> or <|tool_calls_section_begin|> -- Kimi can go straight from reasoning into a tool section without closing the think block. The tool-section marker is forwarded as content so the downstream tool parser can parse it. Replaces the per-SKU guessing of kimi_k25/kimi_thinking (static always_in_reasoning flag) that mis-split K2.5/K2.6 output when thinking was on -- the regression behind the nightly tau2-bench 0/100 in smg-project#1873. Streaming holds back trailing partial end markers so split tokens never leak into reasoning text. Registered as kimi_k2 (selectable via --reasoning-parser kimi_k2). Factory pattern-map changes, deprecated aliases, and the tool-parser rename land in follow-up PRs. Refs: smg-project#1873 Signed-off-by: ishan <ishanvgf@gmail.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds a unified ChangesKimi K2 parser
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ParserFactory
participant KimiK2Parser
participant ParserResult
ParserFactory->>KimiK2Parser: construct kimi_k2 parser
KimiK2Parser->>ParserResult: split reasoning and normal_text
Possibly related issues
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d03bbda99a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let normal = match kind { | ||
| EndKind::ThinkEnd => self.buffer[idx + THINK_END.len()..].to_string(), | ||
| EndKind::ToolSection => self.buffer[idx..].to_string(), |
There was a problem hiding this comment.
When a streaming chunk ends with something like </think><|tool_calls_se, this branch immediately emits the incomplete tool-section prefix as normal text and marks reasoning as ended, so the continuation arrives separately. In the inspected gRPC streaming pipeline, that first fragment is passed directly to KimiK2Parser::parse_incremental, whose no-marker path drains it as user-visible text; consequently the marker never reassembles and the tool call is not parsed. Hold a trailing prefix of TOOL_SECTION_START across the transition just as the reasoning path does.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 71f1caa — verified the claim against KimiK2Parser::parse_incremental first (its no-marker path mem::takes the buffer into normal_text, so a split marker is indeed lost). The think-end transition and the post-reasoning flush now hold back trailing partial tool-section prefixes, same hold-back as inside reasoning; new test kimi_k2_streaming_preserves_tool_marker_split_at_transition covers the exact chunk split you described. Worth noting for context: in production the section marker is a single special token so detokenization never actually splits it — this is defensive conformance for synthetic chunk boundaries. Crate suite 103 passed, clippy clean.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/reasoning_parser/src/factory.rs (1)
165-190: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winWire a K2-family model-ID pattern to
kimi_k2.
kimi_k2is registered and replaces the per-SKU parsers, butbuild_default_patternsstill mapskimi-k2-thinking→kimi_thinkingandkimi-k2.5→kimi_k25, and no pattern maps tokimi_k2. Automatic model selection therefore falls through to passthrough; add/update the K2-family patterns to route these models to the unified parser now.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/reasoning_parser/src/factory.rs` around lines 165 - 190, Update build_default_patterns so K2-family model IDs, including kimi-k2-thinking and kimi-k2.5, map to the registered kimi_k2 parser instead of kimi_thinking or kimi_k25. Ensure automatic model selection routes these variants to KimiK2Parser rather than passthrough.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@crates/reasoning_parser/src/factory.rs`:
- Around line 165-190: Update build_default_patterns so K2-family model IDs,
including kimi-k2-thinking and kimi-k2.5, map to the registered kimi_k2 parser
instead of kimi_thinking or kimi_k25. Ensure automatic model selection routes
these variants to KimiK2Parser rather than passthrough.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 1af90c45-8901-4519-be41-aa7db8929df8
📒 Files selected for processing (3)
crates/reasoning_parser/src/factory.rscrates/reasoning_parser/src/parsers/kimi_k2.rscrates/reasoning_parser/src/parsers/mod.rs
Review on lightseekorg#1992: when a streaming chunk ends right after </think> with a partial <|tool_calls_section_begin|>, the fragment was emitted as normal text. Downstream tool parsers drain marker-less deltas as user-visible text, so the marker never reassembled and the tool call was lost. Hold back trailing partial tool-section prefixes at the think-end transition and in the post-reasoning flush, matching the hold-back already used inside reasoning. Refs: lightseekorg#1873 Signed-off-by: ishan <ishanvgf@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 71f1caa3fc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let hold = Self::trailing_prefix_of(&self.buffer, &[THINK_END, TOOL_SECTION_START]); | ||
| let end = self.buffer.len() - hold; | ||
| let reasoning: String = self.buffer.drain(..end).collect(); |
There was a problem hiding this comment.
When a streamed response ends with a proper prefix of either delimiter—for example, valid reasoning whose final character is <, or output truncated at </thi—this code retains the suffix in self.buffer. The inspected gRPC streaming paths only invoke parse_reasoning_streaming_incremental for decoded chunks and provide no end-of-stream parser flush, so the retained text is silently lost and streaming disagrees with the one-shot truncated-reasoning behavior. Add an EOF/finalization mechanism that emits the held prefix while continuing to hold it between live chunks.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid edge, and it is pre-existing ecosystem-wide rather than introduced by this parser: BaseReasoningParser drops a stream-ending partial identically (base.rs is_partial_token holds </th forever), so every base-derived parser has the same leak — the trait simply has no EOF signal. A proper fix is a trait-level finalization hook plus a call from the streaming paths, which would touch all parsers + model_gateway streaming — well beyond this PR. Filed as #1998 with a proposed flush() design; keeping this PR scoped per one-concern-per-PR.
2 participants
Description
Problem
#1873: SMG ships three Kimi reasoning parsers (
kimi,kimi_k25,kimi_thinking) with a static per-SKUalways_in_reasoningflag, while the rest of the ecosystem (vLLM, SGLang, Moonshot's deploy guidance) uses onekimi_k2parser for the whole K2 family. K2.5+ chat templates prefill<think>at the generation prompt (thinking on by default), so model output begins inside reasoning — andkimi_k25(always_in_reasoning: false) waits for an opening<think>that never comes, mis-splitting reasoning vs content (the nightly τ²-bench 0/100 regression).This is PR 1 of the sequence proposed in #1873 (comment): the new parser itself. Factory pattern-map changes, deprecated aliases, and the tool-parser rename land in follow-ups; the
chat_template_kwargs.thinkingtoggle is the open design question there.Refs: #1873
Solution
New
KimiK2Parserincrates/reasoning_parser, matching current vLLM main (vllm/parser/kimi_k2.py) semantics:<think>is consumed when present (vLLM's self-correcting start-token behavior, both directions covered)</think>or<|tool_calls_section_begin|>— Kimi can go straight from reasoning into a tool section without closing the think block. The tool-section marker is forwarded as content so SMG's downstream tool parser can parse itBaseReasoningParserhas)Registered as
kimi_k2in the reasoning factory — selectable via--reasoning-parser kimi_k2. No behavior change to existing parsers.Changes
crates/reasoning_parser/src/parsers/kimi_k2.rs— new parser + 8 testscrates/reasoning_parser/src/parsers/mod.rs— module exportcrates/reasoning_parser/src/factory.rs— registerkimi_k2Test Plan
kimi_k2_golden_k26_output_split): frozen K2.6-style output (starts mid-reasoning,</think>, then a full tool section) — asserts the exact reasoning/content split. This is the regression class that previously only surfaced as a benchmark scorekimi_k2_ends_reasoning_at_tool_section_without_think_end— reasoning runs straight into<|tool_calls_section_begin|>; marker forwarded verbatimkimi_k2_consumes_leading_think_start_when_present+ streaming variant — non-prefilled templates work tookimi_k2_streaming_chunked_matches_non_streaming— golden output fed in chunks that split both end markers; streamed split identical to one-shotkimi_k2_truncated_reasoning_is_all_reasoning,kimi_k2_reset_restores_initial_state,kimi_k2_model_typeGate output (macOS, rustc 1.97.1 stable):
cargo test -p reasoning-parser --lib— 102 passed (94 existing + 8 new)cargo test -p smg --lib— 1323 passed, 0 failedcargo +nightly fmt --all— silent successcargo clippy -p reasoning-parser --all-targets -- -D warnings— zero warningsChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses (changed crate clean)Summary by CodeRabbit
Summary by CodeRabbit
New Features
Tests
Chores