feat(neuron): strip reasoning from chat completions by default
Some checks failed
CI / CUDA type-check (push) Failing after 18s
build-prerelease / Resolve version stamps (push) Successful in 32s
CI / Format (push) Successful in 32s
CI / Clippy (push) Successful in 2m36s
build-prerelease / Build cortex binary (push) Successful in 4m29s
CI / Test (push) Successful in 5m19s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 5m56s
build-prerelease / Package cortex RPM (push) Successful in 1m21s
build-prerelease / Build neuron-ampere (push) Successful in 7m45s
build-prerelease / Build neuron-ada (push) Successful in 5m24s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m53s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m0s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m43s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m2s

Closes #8.

Reasoning-capable models (Qwen3, DeepSeek-R1, gpt-oss, Mistral
Magistral, …) emit `<think>...</think>` blocks inline in their
content stream. The chat-completions wire format has no slot for
reasoning, so until this change every consumer either parsed the
markers themselves (helexa-acp) or wrote the raw scratchpad
content into their UI (Zed's commit-message generator — visible
as the leaked reasoning block on every generated commit message
against benjy's Qwen3-8B).

## Implementation, model-agnostic by design

The neuron side now does token-level routing without any
hardcoded model knowledge:

1. **At load time** (`detect_reasoning_token_pair` in
   `wire::event`), probe the tokenizer's vocabulary for a known
   reasoning-marker pair: `<think>` / `</think>` (Qwen3,
   DeepSeek-R1, gpt-oss), `[THINK]` / `[/THINK]` (Mistral
   Magistral), and a couple of derivatives. Each marker must
   resolve to a single token id; if both open and close resolve,
   stash on `LoadedModel.reasoning_tokens` (similarly
   `TpLoadedModel`). Non-reasoning models get `None` and pass
   through unchanged.

2. **At inference time**, the three streaming paths
   (`run_inference_streaming` CPU, `stream_inference_via_worker`
   CUDA single-GPU, `chat_completion_tp_stream` CUDA TP) now
   check each sampled token against the pair via the new
   `handle_reasoning_marker` helper before feeding it to the
   detokeniser. Open marker → set `in_reasoning = true`, drop
   the marker. Close marker → unset, drop. Other tokens go
   through `emit_delta(_blocking)` which now picks
   `ReasoningDelta` or `TextDelta` based on state. Markers
   never appear in the streamed output.

3. **In `wire::openai_chat`**, the projector splits into:
   - `project_chat_stream` (unchanged signature; default
     behaviour — drops `ReasoningDelta`)
   - `project_chat_stream_with(rx, …, ChatProjectionConfig)` —
     when `include_thinking: true` and `reasoning_markers:
     Some(_)`, re-wraps reasoning content with the literal
     open/close marker text and emits as content deltas.
     Preserves the on-the-wire shape that helexa-acp's
     `ThinkParser` expects.

4. **HTTP handler** reads `x-include-thinking: true` (case-
   insensitive `1`/`true`/`yes`) from the request headers and
   threads it into the projection config. cortex-gateway already
   forwards arbitrary headers verbatim, so the opt-in works
   end-to-end without gateway changes.

5. **helexa-acp's `openai_chat` provider** sets
   `x-include-thinking: true` on every request so its existing
   `ThinkParser` keeps receiving the marked content stream.
   `ThinkParser` itself is unchanged — needed for endpoints that
   aren't reasoning-aware (OpenRouter, OpenAI directly, etc.).

## Acceptance

- Zed's commit-message generator (vanilla chat-completions
  client, no `x-include-thinking`) gets clean commit messages
  with no `<think>` block.
- helexa-acp sessions continue to render thinking in Zed's
  thought UI via the opt-in path.
- Models without reasoning tokens declared in their tokenizer
  pass through unchanged.
- Implementation contains zero references to "qwen3" or any
  specific model — entirely driven by tokenizer metadata.

## Tests

9 new tests in `wire::event` (token-pair detection across 4
marker conventions, edge cases) and `wire::openai_chat` (default
drop, opt-in re-wrap with multi-chunk reasoning, close-marker on
Finish, fallback when markers absent, off-switch with markers
present). All 213 workspace tests pass; fmt + clippy clean.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-05-31 17:55:04 +03:00
parent fdc0adb738
commit 7733eecba5
6 changed files with 645 additions and 67 deletions

View File

@@ -97,3 +97,134 @@ impl FinishReason {
}
}
}
/// Open/close token IDs for the reasoning marker a loaded model uses
/// (or `None` for non-reasoning models). The harness reads this once
/// at load time from the tokenizer's added-tokens table, then the
/// inference loop checks `next_token` against the pair to flip
/// between [`InferenceEvent::TextDelta`] and
/// [`InferenceEvent::ReasoningDelta`].
///
/// `open` and `close` text are kept alongside the IDs so wire
/// projectors that want to re-emit the literal markers (the
/// opt-in `include_thinking` path on chat completions) don't have
/// to reach back into the tokenizer for the strings.
#[derive(Debug, Clone)]
pub struct ReasoningTokenPair {
pub open_id: u32,
pub close_id: u32,
pub open_text: String,
pub close_text: String,
}
/// Known reasoning-marker conventions. Each is a `(open, close)`
/// pair of literal token strings. Each modern reasoning model
/// declares its markers in the tokenizer's `added_tokens` table;
/// at load time we probe for whichever pair the loaded tokenizer
/// has and stash both IDs.
///
/// Ordering matters only for tie-breaking when a model declares
/// multiple pairs (shouldn't happen in practice); the first hit
/// wins.
const KNOWN_REASONING_MARKERS: &[(&str, &str)] = &[
// Qwen3, DeepSeek-R1, gpt-oss, and most other open-weight
// reasoning models.
("<think>", "</think>"),
// Mistral Magistral.
("[THINK]", "[/THINK]"),
// Some older derivatives; harmless to probe.
("<thought>", "</thought>"),
("<reasoning>", "</reasoning>"),
];
/// Inspect a tokenizer for known reasoning-marker pairs and return
/// the first match. The tokenizer types this trait is defined over
/// just need to expose `token_to_id(&str) -> Option<u32>` so this
/// stays decoupled from the candle crate — the production caller
/// passes a `tokenizers::Tokenizer`, but tests can fake one.
///
/// Returns `None` when no known marker pair is fully declared
/// (both open AND close token ids must resolve). That's the
/// pass-through case — non-reasoning models, or reasoning models
/// whose tokenizer split the markers across multiple tokens (rare
/// in practice; modern reasoning tokenizers list them as
/// `added_tokens`).
pub fn detect_reasoning_token_pair<F>(token_to_id: F) -> Option<ReasoningTokenPair>
where
F: Fn(&str) -> Option<u32>,
{
for (open_text, close_text) in KNOWN_REASONING_MARKERS {
let open_id = token_to_id(open_text);
let close_id = token_to_id(close_text);
if let (Some(open_id), Some(close_id)) = (open_id, close_id) {
return Some(ReasoningTokenPair {
open_id,
close_id,
open_text: (*open_text).into(),
close_text: (*close_text).into(),
});
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
fn lookup<'a>(map: &'a HashMap<&'static str, u32>) -> impl Fn(&str) -> Option<u32> + 'a {
|s| map.get(s).copied()
}
#[test]
fn detects_qwen3_style_think_markers() {
let mut m = HashMap::new();
m.insert("<think>", 151648);
m.insert("</think>", 151649);
let pair = detect_reasoning_token_pair(lookup(&m)).expect("pair detected");
assert_eq!(pair.open_id, 151648);
assert_eq!(pair.close_id, 151649);
assert_eq!(pair.open_text, "<think>");
assert_eq!(pair.close_text, "</think>");
}
#[test]
fn detects_mistral_magistral_markers() {
let mut m = HashMap::new();
m.insert("[THINK]", 100);
m.insert("[/THINK]", 101);
let pair = detect_reasoning_token_pair(lookup(&m)).expect("pair detected");
assert_eq!(pair.open_text, "[THINK]");
}
#[test]
fn returns_none_when_only_open_marker_present() {
// A pathological tokenizer that has `<think>` but not
// `</think>` shouldn't half-detect. Pass-through.
let mut m = HashMap::new();
m.insert("<think>", 1);
assert!(detect_reasoning_token_pair(lookup(&m)).is_none());
}
#[test]
fn returns_none_for_non_reasoning_tokenizer() {
let m: HashMap<&'static str, u32> = HashMap::new();
assert!(detect_reasoning_token_pair(lookup(&m)).is_none());
}
#[test]
fn first_match_wins_when_multiple_pairs_declared() {
// Hypothetical tokenizer with both Qwen-style AND Mistral-style
// markers — the `<think>` pair is earlier in the convention
// table so it wins.
let mut m = HashMap::new();
m.insert("<think>", 1);
m.insert("</think>", 2);
m.insert("[THINK]", 3);
m.insert("[/THINK]", 4);
let pair = detect_reasoning_token_pair(lookup(&m)).unwrap();
assert_eq!(pair.open_id, 1);
assert_eq!(pair.close_id, 2);
}
}

View File

@@ -21,4 +21,4 @@ pub mod event;
pub mod openai_chat;
pub mod openai_responses;
pub use event::{FinishReason, InferenceEvent};
pub use event::{FinishReason, InferenceEvent, ReasoningTokenPair, detect_reasoning_token_pair};

View File

@@ -30,13 +30,42 @@ use cortex_core::openai::{ChatCompletionChunk, ChunkChoice};
use serde_json::json;
use tokio::sync::mpsc;
use super::event::{FinishReason, InferenceEvent};
use super::event::{FinishReason, InferenceEvent, ReasoningTokenPair};
/// Output channel buffer size. Mirrors the input side's bound; one
/// event maps to at most one chunk, so equal capacity keeps the
/// two ends in sync without surprising memory growth.
const CHUNK_CHANNEL_CAPACITY: usize = 32;
/// Per-stream config for the chat projector. Used by the
/// production handler to thread per-request choices (currently:
/// whether to surface reasoning content) into the projection
/// without bloating the function signature.
#[derive(Debug, Clone, Default)]
pub struct ChatProjectionConfig {
/// When `true`, reasoning content is re-wrapped with the
/// model's literal open/close markers and emitted as content
/// deltas — preserving the on-the-wire shape that
/// reasoning-aware clients like helexa-acp's `ThinkParser`
/// expect.
///
/// When `false` (the default), [`InferenceEvent::ReasoningDelta`]s
/// are dropped entirely so consumers that don't know about
/// reasoning (Zed's commit-message generator, any vanilla
/// OpenAI client) don't have model-internal scratchpad
/// material leaking into their UI. The chat-completions wire
/// format has no slot for reasoning, so the default chooses
/// the safer-for-naïve-clients behaviour.
pub include_thinking: bool,
/// Open/close marker strings to re-emit when `include_thinking`
/// is set. Sourced from the loaded model's
/// [`ReasoningTokenPair`]; `None` for non-reasoning models or
/// when the caller doesn't have the pair handy (in which case
/// `include_thinking` becomes equivalent to dropping reasoning
/// because there's nothing to wrap).
pub reasoning_markers: Option<ReasoningTokenPair>,
}
/// Project an [`InferenceEvent`] receiver into a
/// [`ChatCompletionChunk`] receiver. Spawns one tokio task that
/// owns the input receiver for the stream's lifetime and exits
@@ -46,15 +75,55 @@ const CHUNK_CHANNEL_CAPACITY: usize = 32;
/// chunk so the receiver can stay generic (decoupled from
/// per-request metadata).
pub fn project_chat_stream(
mut rx: mpsc::Receiver<InferenceEvent>,
rx: mpsc::Receiver<InferenceEvent>,
id: String,
created: u64,
model_id: String,
) -> mpsc::Receiver<ChatCompletionChunk> {
// Default config: include_thinking off, no marker rewrap.
project_chat_stream_with(rx, id, created, model_id, ChatProjectionConfig::default())
}
/// Same as [`project_chat_stream`] but with a per-stream config
/// (currently controlling reasoning surfacing). Production
/// callers that need the opt-in path call this directly; the
/// shorter wrapper above stays as the no-config convenience.
pub fn project_chat_stream_with(
mut rx: mpsc::Receiver<InferenceEvent>,
id: String,
created: u64,
model_id: String,
config: ChatProjectionConfig,
) -> mpsc::Receiver<ChatCompletionChunk> {
let (tx, out_rx) = mpsc::channel::<ChatCompletionChunk>(CHUNK_CHANNEL_CAPACITY);
tokio::spawn(async move {
// Track whether the previous event was inside a reasoning
// block — used to decide when to emit the literal close
// marker on the include_thinking re-wrap path. When this
// flips from true → false (a TextDelta or Finish lands
// after one or more ReasoningDeltas), we emit the close
// marker exactly once.
let mut was_in_reasoning = false;
while let Some(event) = rx.recv().await {
// Close-marker insertion: if we're leaving a reasoning
// chain, emit the literal close marker before the
// current event.
if was_in_reasoning && !matches!(event, InferenceEvent::ReasoningDelta(_)) {
if let Some(marker) = config
.include_thinking
.then_some(())
.and(config.reasoning_markers.as_ref())
{
let chunk = content_chunk(&id, created, &model_id, &marker.close_text);
if tx.send(chunk).await.is_err() {
return;
}
}
was_in_reasoning = false;
}
let chunks = match event {
InferenceEvent::Start => vec![role_chunk(&id, created, &model_id)],
InferenceEvent::TextDelta(text) => {
@@ -66,12 +135,42 @@ pub fn project_chat_stream(
}
vec![content_chunk(&id, created, &model_id, &text)]
}
InferenceEvent::ReasoningDelta(_) => {
// Reasoning isn't representable in OpenAI chat
// streaming today. The o-series uses a separate
// `summary` event but it's gated by the
// Responses API; chat-completions just drops it.
continue;
InferenceEvent::ReasoningDelta(text) => {
if !config.include_thinking {
// Default path — reasoning has no slot in
// chat completions, so it's dropped. Naïve
// clients (Zed commit-message generator,
// any vanilla OpenAI client) get clean
// output.
continue;
}
let Some(markers) = config.reasoning_markers.as_ref() else {
// Caller asked to include thinking but
// didn't supply markers — best we can do
// is emit the content as visible text.
// Skip the wrap entirely.
if text.is_empty() {
continue;
}
let chunk = content_chunk(&id, created, &model_id, &text);
if tx.send(chunk).await.is_err() {
return;
}
continue;
};
// First chunk of a reasoning block → open
// marker prelude. Subsequent reasoning deltas
// in the same block reuse `was_in_reasoning`
// to skip the prelude.
let mut chunks = Vec::new();
if !was_in_reasoning {
chunks.push(content_chunk(&id, created, &model_id, &markers.open_text));
}
if !text.is_empty() {
chunks.push(content_chunk(&id, created, &model_id, &text));
}
was_in_reasoning = true;
chunks
}
InferenceEvent::Finish { reason } => {
vec![final_chunk(&id, created, &model_id, reason)]
@@ -238,4 +337,165 @@ mod tests {
assert_eq!(out.len(), 1);
assert_eq!(out[0].choices[0].delta["content"], "real");
}
fn pair() -> ReasoningTokenPair {
ReasoningTokenPair {
open_id: 0,
close_id: 1,
open_text: "<think>".into(),
close_text: "</think>".into(),
}
}
#[tokio::test]
async fn include_thinking_rewraps_reasoning_with_literal_markers() {
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
let out_rx = project_chat_stream_with(
rx,
"id".into(),
1,
"m".into(),
ChatProjectionConfig {
include_thinking: true,
reasoning_markers: Some(pair()),
},
);
tx.send(InferenceEvent::ReasoningDelta("first ".into()))
.await
.unwrap();
tx.send(InferenceEvent::ReasoningDelta("second".into()))
.await
.unwrap();
tx.send(InferenceEvent::TextDelta("answer".into()))
.await
.unwrap();
tx.send(InferenceEvent::Finish {
reason: FinishReason::Stop,
})
.await
.unwrap();
drop(tx);
let out = collect(out_rx).await;
// Expected sequence: open marker → reasoning content (2 chunks)
// → close marker → visible answer → final chunk.
let contents: Vec<&str> = out
.iter()
.filter_map(|c| c.choices[0].delta["content"].as_str())
.collect();
assert_eq!(
contents,
vec!["<think>", "first ", "second", "</think>", "answer"]
);
assert_eq!(
out.last().unwrap().choices[0].finish_reason.as_deref(),
Some("stop")
);
}
#[tokio::test]
async fn include_thinking_closes_marker_at_finish_when_no_trailing_text() {
// Edge case: stream ends inside a reasoning block (model
// hit max_tokens mid-thought, no visible answer ever).
// The Finish event still triggers the close marker so the
// stream is balanced.
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
let out_rx = project_chat_stream_with(
rx,
"id".into(),
1,
"m".into(),
ChatProjectionConfig {
include_thinking: true,
reasoning_markers: Some(pair()),
},
);
tx.send(InferenceEvent::ReasoningDelta("thinking...".into()))
.await
.unwrap();
tx.send(InferenceEvent::Finish {
reason: FinishReason::Length,
})
.await
.unwrap();
drop(tx);
let out = collect(out_rx).await;
let contents: Vec<&str> = out
.iter()
.filter_map(|c| c.choices[0].delta["content"].as_str())
.collect();
assert_eq!(contents, vec!["<think>", "thinking...", "</think>"]);
assert_eq!(
out.last().unwrap().choices[0].finish_reason.as_deref(),
Some("length")
);
}
#[tokio::test]
async fn include_thinking_without_markers_emits_content_directly() {
// Defensive: if the caller asks for thinking but the
// model declared no markers, we still emit the content
// rather than dropping it. Better to leak than to lose.
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
let out_rx = project_chat_stream_with(
rx,
"id".into(),
1,
"m".into(),
ChatProjectionConfig {
include_thinking: true,
reasoning_markers: None,
},
);
tx.send(InferenceEvent::ReasoningDelta("raw".into()))
.await
.unwrap();
tx.send(InferenceEvent::Finish {
reason: FinishReason::Stop,
})
.await
.unwrap();
drop(tx);
let out = collect(out_rx).await;
let contents: Vec<&str> = out
.iter()
.filter_map(|c| c.choices[0].delta["content"].as_str())
.collect();
assert_eq!(contents, vec!["raw"]);
}
#[tokio::test]
async fn include_thinking_off_drops_reasoning_even_with_markers() {
// Default behaviour even when markers happen to be
// configured. The flag is the gate, not the marker
// presence.
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
let out_rx = project_chat_stream_with(
rx,
"id".into(),
1,
"m".into(),
ChatProjectionConfig {
include_thinking: false,
reasoning_markers: Some(pair()),
},
);
tx.send(InferenceEvent::ReasoningDelta("hidden".into()))
.await
.unwrap();
tx.send(InferenceEvent::TextDelta("visible".into()))
.await
.unwrap();
tx.send(InferenceEvent::Finish {
reason: FinishReason::Stop,
})
.await
.unwrap();
drop(tx);
let out = collect(out_rx).await;
let contents: Vec<&str> = out
.iter()
.filter_map(|c| c.choices[0].delta["content"].as_str())
.collect();
assert_eq!(contents, vec!["visible"]);
}
}