Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 12 additions & 7 deletions internal/service/bot_completion.go
Original file line number Diff line number Diff line change
Expand Up @@ -429,11 +429,14 @@ func (s *BotService) ChatbotCompletion(
}

// 5. Translate pipeline results into chatbot frames. The
// pipeline streams deltas; the python iframe contract sends the
// full accumulated answer in every frame, so we accumulate here
// (the front-end accepts both shapes, but full-text frames are
// byte-parity with python). The final pipeline result carries
// the decorated full answer plus the retrieval reference.
// pipeline streams deltas; the python iframe contract also yields
// deltas (async_chat yields the value from _stream_with_think_delta).
// We accumulate the full answer locally for persistence and for the
// final frame, but forward each result's own delta in the SSE frame
// so the front-end can build the <think>-wrapped answer correctly.
// Sending accumulated full text on every frame interacts badly with
// the front-end's start_to_think/end_to_think marker append, causing
// reasoning content to leak into the visible answer.
out := make(chan ChatbotSSEFrame, 16)
go func() {
defer close(out)
Expand Down Expand Up @@ -468,8 +471,10 @@ func (s *BotService) ChatbotCompletion(
continue
}
if res.StartToThink || res.EndToThink {
// Marker frames carry no text; the front-end appends
// <think> / </think> to the accumulated answer.
out <- ChatbotSSEFrame{
Data: fullAnswer,
Data: "",
Reference: map[string]any{},
SessionID: session.ID,
StartToThink: res.StartToThink,
Expand All @@ -489,7 +494,7 @@ func (s *BotService) ChatbotCompletion(
finalRef = res.Reference
}
out <- ChatbotSSEFrame{
Data: fullAnswer,
Data: res.Answer,
Reference: map[string]any{},
SessionID: session.ID,
}
Expand Down
7 changes: 7 additions & 0 deletions internal/service/think_tag.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,22 +53,29 @@ type ThinkStreamState struct {
}

// EnterReasoning marks the start of a reasoning block (model-level, not tag-based).
// It also writes a <think> marker into fullText so that ExtractVisibleAnswer can
// strip the reasoning segment from the final answer when the model exposes
// reasoning through a separate channel (e.g. delta.reasoning_content).
// Returns true when this is a new transition (reasoning was not already active).
func (s *ThinkStreamState) EnterReasoning() bool {
if s.inReasoning {
return false
}
s.inReasoning = true
s.fullText += "<think>"
return true
}

// ExitReasoning marks the end of a reasoning block (model-level, not tag-based).
// It writes a matching </think> marker into fullText so the reasoning segment is
// closed before the visible answer text begins.
// Returns true when this is a new transition (reasoning was active).
func (s *ThinkStreamState) ExitReasoning() bool {
if !s.inReasoning {
return false
}
s.inReasoning = false
s.fullText += "</think>"
return true
}

Expand Down
36 changes: 36 additions & 0 deletions internal/service/think_tag_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -301,3 +301,39 @@ func TestStreamThinkTagDelta_DeferredClose(t *testing.T) {
t.Errorf("second marker = %q", markers[1])
}
}

// TestReasoningChannelWrapsThinkTags verifies that model-level reasoning
// (arriving through a separate reason channel rather than <think> tags in
// the text) is bracketed by <think></think> in fullText. This prevents the
// reasoning content from leaking into the final visible answer.
func TestReasoningChannelWrapsThinkTags(t *testing.T) {
state := &ThinkStreamState{}

if !state.EnterReasoning() {
t.Fatal("EnterReasoning should return true on first call")
}
// Simulate the chat pipeline feeding reasoning chunks through NextThinkDelta.
deltas := NextThinkDelta(state, "reasoning step 1", 0)
if len(deltas) != 1 || deltas[0].Value != "reasoning step 1" {
t.Fatalf("unexpected reasoning deltas: %+v", deltas)
}
deltas = NextThinkDelta(state, " reasoning step 2", 0)
if len(deltas) != 1 || deltas[0].Value != " reasoning step 2" {
t.Fatalf("unexpected reasoning deltas: %+v", deltas)
}

if state.ExitReasoning() != true {
t.Fatal("ExitReasoning should return true when reasoning was active")
}

// Visible answer arrives after reasoning ends.
deltas = BufferAnswerDelta(state, "visible answer", 0)
if len(deltas) != 1 || deltas[0].Value != "visible answer" {
t.Fatalf("unexpected answer deltas: %+v", deltas)
}

wantVisible := "visible answer"
if got := ExtractVisibleAnswer(state.fullText); got != wantVisible {
t.Errorf("ExtractVisibleAnswer(%q) = %q, want %q", state.fullText, got, wantVisible)
}
}
Comment on lines +304 to +339

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Assert that both reasoning markers are persisted.

The current test can pass even if EnterReasoning() stops appending <think>, because ExtractVisibleAnswer only needs </think> to discard the preceding text. Assert the exact state.fullText value so the new opening-marker behavior is covered.

Proposed test assertion
  wantVisible := "visible answer"
  if got := ExtractVisibleAnswer(state.fullText); got != wantVisible {
    t.Errorf("ExtractVisibleAnswer(%q) = %q, want %q", state.fullText, got, wantVisible)
  }
+	wantFullText := "<think>reasoning step 1 reasoning step 2</think>visible answer"
+	if state.fullText != wantFullText {
+		t.Errorf("fullText = %q, want %q", state.fullText, wantFullText)
+	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// TestReasoningChannelWrapsThinkTags verifies that model-level reasoning
// (arriving through a separate reason channel rather than <think> tags in
// the text) is bracketed by <think></think> in fullText. This prevents the
// reasoning content from leaking into the final visible answer.
func TestReasoningChannelWrapsThinkTags(t *testing.T) {
state := &ThinkStreamState{}
if !state.EnterReasoning() {
t.Fatal("EnterReasoning should return true on first call")
}
// Simulate the chat pipeline feeding reasoning chunks through NextThinkDelta.
deltas := NextThinkDelta(state, "reasoning step 1", 0)
if len(deltas) != 1 || deltas[0].Value != "reasoning step 1" {
t.Fatalf("unexpected reasoning deltas: %+v", deltas)
}
deltas = NextThinkDelta(state, " reasoning step 2", 0)
if len(deltas) != 1 || deltas[0].Value != " reasoning step 2" {
t.Fatalf("unexpected reasoning deltas: %+v", deltas)
}
if state.ExitReasoning() != true {
t.Fatal("ExitReasoning should return true when reasoning was active")
}
// Visible answer arrives after reasoning ends.
deltas = BufferAnswerDelta(state, "visible answer", 0)
if len(deltas) != 1 || deltas[0].Value != "visible answer" {
t.Fatalf("unexpected answer deltas: %+v", deltas)
}
wantVisible := "visible answer"
if got := ExtractVisibleAnswer(state.fullText); got != wantVisible {
t.Errorf("ExtractVisibleAnswer(%q) = %q, want %q", state.fullText, got, wantVisible)
}
}
// TestReasoningChannelWrapsThinkTags verifies that model-level reasoning
// (arriving through a separate reason channel rather than <think> tags in
// the text) is bracketed by <think></think> in fullText. This prevents the
// reasoning content from leaking into the final visible answer.
func TestReasoningChannelWrapsThinkTags(t *testing.T) {
state := &ThinkStreamState{}
if !state.EnterReasoning() {
t.Fatal("EnterReasoning should return true on first call")
}
// Simulate the chat pipeline feeding reasoning chunks through NextThinkDelta.
deltas := NextThinkDelta(state, "reasoning step 1", 0)
if len(deltas) != 1 || deltas[0].Value != "reasoning step 1" {
t.Fatalf("unexpected reasoning deltas: %+v", deltas)
}
deltas = NextThinkDelta(state, " reasoning step 2", 0)
if len(deltas) != 1 || deltas[0].Value != " reasoning step 2" {
t.Fatalf("unexpected reasoning deltas: %+v", deltas)
}
if state.ExitReasoning() != true {
t.Fatal("ExitReasoning should return true when reasoning was active")
}
// Visible answer arrives after reasoning ends.
deltas = BufferAnswerDelta(state, "visible answer", 0)
if len(deltas) != 1 || deltas[0].Value != "visible answer" {
t.Fatalf("unexpected answer deltas: %+v", deltas)
}
wantVisible := "visible answer"
if got := ExtractVisibleAnswer(state.fullText); got != wantVisible {
t.Errorf("ExtractVisibleAnswer(%q) = %q, want %q", state.fullText, got, wantVisible)
}
wantFullText := "<think>reasoning step 1 reasoning step 2</think>visible answer"
if state.fullText != wantFullText {
t.Errorf("fullText = %q, want %q", state.fullText, wantFullText)
}
}
🤖 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 `@internal/service/think_tag_test.go` around lines 304 - 339, Strengthen
TestReasoningChannelWrapsThinkTags by asserting the exact state.fullText after
reasoning and visible-answer chunks are processed, including both the opening
<think> and closing </think> markers around the reasoning text followed by the
visible answer. Keep the existing ExtractVisibleAnswer assertion unchanged.

Loading