feat(reasoning): add unified kimi_k2 reasoning parser - #1992
feat(reasoning): add unified kimi_k2 reasoning parser#1992ighutake-debug wants to merge 5 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. |
📝 WalkthroughSummary by CodeRabbitNew Features
Bug Fixes
Tests
WalkthroughAdds ChangesKimi K2 parser
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ParserFactory
participant KimiK2Parser
participant ReasoningParser
ParserFactory->>KimiK2Parser: construct kimi_k2 parser
ReasoningParser->>KimiK2Parser: provide one-shot or streaming content
KimiK2Parser-->>ReasoningParser: return reasoning and normal content
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 smg-project#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: smg-project#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.
|
Friendly ping — this has been green and CodeRabbit-approved for a few days; anything you would like changed? Happy to adjust. (This is PR 1 of the #1873 sequence; PRs 2 and 3 are ready to follow once this lands.) |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
crates/reasoning_parser/src/factory.rs (1)
187-190: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add a factory-level registration test for
"kimi_k2".Direct
KimiK2Parser::model_type()coverage exists, but the--reasoning-parser kimi_k2contract is covered only byParserFactory::new()/ParserRegistryflow tests for the other parsers. Add a test that callsParserFactory::new().create("<kimi-k2 model>")and assertsmodel_type() == "kimi_k2".🤖 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 187 - 190, Add a factory-level test near the existing ParserFactory/ParserRegistry flow tests that calls ParserFactory::new().create with a Kimi K2 model identifier and asserts the returned parser’s model_type() equals "kimi_k2".Source: Coding guidelines
crates/reasoning_parser/src/parsers/kimi_k2.rs (1)
91-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit Test both buffer-overflow paths.
The test suite does not execute the new one-shot or streaming
ParseError::BufferOverflowbranches. Add one oversized one-shot input test. Add one streaming test that retains a partial start marker, then exceeds the cumulative limit. AssertErr(ParseError::BufferOverflow(_))in both cases.As per coding guidelines, “Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality.”
Also applies to: 129-131
🤖 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/parsers/kimi_k2.rs` around lines 91 - 93, Add coverage for both ParseError::BufferOverflow paths in the Kimi K2 parser: test an oversized one-shot input and assert Err(ParseError::BufferOverflow(_)), then test streaming input that preserves a partial start marker before exceeding the cumulative max_buffer_size and assert the same error. Run the pr-test-analyzer agent to verify coverage of these changed branches.Source: Coding guidelines
🤖 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.
Inline comments:
In `@crates/reasoning_parser/src/parsers/kimi_k2.rs`:
- Around line 200-202: Update mark_reasoning_started() to reset reasoning_ended
and clear any stale buffered marker state whenever a new reasoning block begins,
while preserving the existing in_reasoning transition. Add a regression test
covering one completed reasoning block followed by mark_reasoning_started() and
parsing a second block, verifying the second block is classified as reasoning
rather than normal content.
---
Nitpick comments:
In `@crates/reasoning_parser/src/factory.rs`:
- Around line 187-190: Add a factory-level test near the existing
ParserFactory/ParserRegistry flow tests that calls ParserFactory::new().create
with a Kimi K2 model identifier and asserts the returned parser’s model_type()
equals "kimi_k2".
In `@crates/reasoning_parser/src/parsers/kimi_k2.rs`:
- Around line 91-93: Add coverage for both ParseError::BufferOverflow paths in
the Kimi K2 parser: test an oversized one-shot input and assert
Err(ParseError::BufferOverflow(_)), then test streaming input that preserves a
partial start marker before exceeding the cumulative max_buffer_size and assert
the same error. Run the pr-test-analyzer agent to verify coverage of these
changed branches.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ee21ac9f-1c22-49b5-8e76-f8141f02d00e
📒 Files selected for processing (3)
crates/reasoning_parser/src/factory.rscrates/reasoning_parser/src/parsers/kimi_k2.rscrates/reasoning_parser/src/parsers/mod.rs
CodeRabbit on smg-project#1992: after reasoning ended, mark_reasoning_started() set in_reasoning but left reasoning_ended latched, so a reused parser reported reasoning mode while classifying all text as normal content. mark_reasoning_started() now fully restarts turn state: clears the terminal latch and stale buffer, and re-resolves the leading <think> question for the new output. Refs: smg-project#1873 Signed-off-by: ishan <ishanvgf@gmail.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
crates/reasoning_parser/src/parsers/kimi_k2.rs (1)
91-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit Cover both
BufferOverflowbranches.No test reaches the one-shot guard at Lines 91-93 or the cumulative streaming guard at Lines 129-131. Add one oversized one-shot input test and one streaming test that buffers a partial
<think>prefix before exceedingmax_buffer_size. AssertErr(ParseError::BufferOverflow(_))in both cases.As per coding guidelines, “Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality.”
Also applies to: 129-131
🤖 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/parsers/kimi_k2.rs` around lines 91 - 93, Add tests covering both BufferOverflow paths: an oversized one-shot input exercising the guard in the parser method containing the text.len() check, and a streaming case that buffers a partial “<think>” prefix before exceeding max_buffer_size and reaches the cumulative guard. Assert Err(ParseError::BufferOverflow(_)) for both, then run the pr-test-analyzer agent to verify coverage.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@crates/reasoning_parser/src/parsers/kimi_k2.rs`:
- Around line 91-93: Add tests covering both BufferOverflow paths: an oversized
one-shot input exercising the guard in the parser method containing the
text.len() check, and a streaming case that buffers a partial “<think>” prefix
before exceeding max_buffer_size and reaches the cumulative guard. Assert
Err(ParseError::BufferOverflow(_)) for both, then run the pr-test-analyzer agent
to verify coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 55140b07-e2db-401a-88c8-e183740471b3
📒 Files selected for processing (1)
crates/reasoning_parser/src/parsers/kimi_k2.rs
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)