From 24118ac0d169d5f3da75d41c5f81af69086eb6c3 Mon Sep 17 00:00:00 2001
From: qinling0210 <88864212+qinling0210@users.noreply.github.com>
Date: Thu, 2 Jul 2026 20:19:50 +0800
Subject: [PATCH] Fix chat thinking & Figure issue in GO (#16558)
### Summary
Fix chat thinking & Figure issue
---
internal/agent/component/llm.go | 47 ++++-
internal/service/chat_pipeline.go | 227 ++++++++++++++++---------
internal/service/chat_pipeline_test.go | 6 +-
internal/service/chat_session.go | 56 +++---
internal/service/think_tag.go | 24 +++
5 files changed, 260 insertions(+), 100 deletions(-)
diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go
index b06f7d0c0..46ccfe238 100644
--- a/internal/agent/component/llm.go
+++ b/internal/agent/component/llm.go
@@ -251,6 +251,18 @@ func (e *einoChatInvoker) Invoke(ctx context.Context, req ChatInvokeRequest) (*C
TopP: req.TopP,
MaxTokens: req.MaxTokens,
}
+ // Propagate the agent-level Thinking setting to the driver so
+ // providers like DeepSeek can send thinking: {type: "disabled"}
+ // and prevent chain-of-thought from leaking into the answer.
+ // Mirrors the Python agent/component/llm.py behaviour.
+ switch req.Thinking {
+ case "enabled":
+ t := true
+ chatCfg.Thinking = &t
+ case "disabled":
+ f := false
+ chatCfg.Thinking = &f
+ }
wrapper := models.NewEinoChatModel(cm, chatCfg)
out, err := wrapper.Generate(ctx, toEinoMessages(req.Messages))
if err != nil {
@@ -517,8 +529,18 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s
return nil, fmt.Errorf("component: LLM.Invoke: %w", err)
}
+ // Strip think blocks + JSON fences from the response.
+ // Mirrors Python's clean_formated_answer() exactly
+ // (re.sub(r"^.*", "", ...) + ^.*```json + trailing ```).
+ // Python only cleans for structured output — keep raw content for
+ // regular responses (llm.py:483: self.set_output("content", ans)).
+ cleaned := resp.Content
+ if p.OutputStructure != nil || p.JSONOutput {
+ cleaned = cleanFormattedAnswer(resp.Content)
+ }
+
out := map[string]any{
- "content": resp.Content,
+ "content": cleaned,
"model": resp.Model,
"stopped": resp.Stopped,
"tokens": resp.Tokens,
@@ -564,7 +586,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s
out["structured"] = parsed
// Also update content to the validated response so
// downstream consumers reading "content" get the JSON text.
- out["content"] = resp.Content
+ out["content"] = cleanFormattedAnswer(resp.Content)
} else {
common.Warn("component: LLM: output_structure set but no parseable JSON after retry")
}
@@ -1353,3 +1375,24 @@ func init() {
return NewLLMComponent(p), nil
})
}
+
+// cleanFormattedAnswer mirrors Python's clean_formated_answer():
+//
+// 1. Strip everything up to and including (dotall).
+// 2. Strip everything up to and including ```json (dotall).
+// 3. Strip trailing ``` and optional newlines.
+//
+// This removes DeepSeek-R1-style thinking blocks and JSON-fence
+// prefixes/suffixes from the raw model response.
+var (
+ reThinkPrefix = regexp.MustCompile(`(?s)^.*`)
+ reJSONFencePrefix = regexp.MustCompile(`(?s)^.*` + "```json")
+ reJSONFenceSuffix = regexp.MustCompile("```\n*$")
+)
+
+func cleanFormattedAnswer(ans string) string {
+ ans = reThinkPrefix.ReplaceAllString(ans, "")
+ ans = reJSONFencePrefix.ReplaceAllString(ans, "")
+ ans = reJSONFenceSuffix.ReplaceAllString(ans, "")
+ return ans
+}
diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go
index 256d918eb..0be17e034 100644
--- a/internal/service/chat_pipeline.go
+++ b/internal/service/chat_pipeline.go
@@ -1102,27 +1102,19 @@ func (s *ChatPipelineService) AsyncChat(
*chatDriver.ModelName, chatMessages, chatDriver.APIConfig, chatCfg,
func(answer *string, reason *string) error {
if reason != nil && *reason != "" {
+ if thinkState.EnterReasoning() {
+ out <- AsyncChatResult{
+ Answer: "",
+ Reference: map[string]interface{}{},
+ AudioBinary: nil,
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ StartToThink: true,
+ }
+ }
deltas := NextThinkDelta(thinkState, *reason, 16)
for _, d := range deltas {
- if d.Kind == ThinkDeltaMarker && d.Value == "" {
- out <- AsyncChatResult{
- Answer: "",
- Reference: map[string]interface{}{},
- AudioBinary: nil,
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
- StartToThink: true,
- }
- } else if d.Kind == ThinkDeltaMarker && d.Value == "" {
- out <- AsyncChatResult{
- Answer: "",
- Reference: map[string]interface{}{},
- AudioBinary: nil,
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
- EndToThink: true,
- }
- } else if d.Kind == ThinkDeltaText && d.Value != "" {
+ if d.Kind == ThinkDeltaText && d.Value != "" {
fullAnswer += d.Value
out <- AsyncChatResult{
Answer: d.Value,
@@ -1135,19 +1127,32 @@ func (s *ChatPipelineService) AsyncChat(
}
}
if isContentDelta(answer) {
+ if thinkState.ExitReasoning() {
+ for _, d := range FlushRemaining(thinkState) {
+ if d.Kind == ThinkDeltaText && d.Value != "" {
+ fullAnswer += d.Value
+ out <- AsyncChatResult{
+ Answer: d.Value,
+ Reference: map[string]interface{}{},
+ AudioBinary: s.synthesizeTTS(ttsModel, d.Value),
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ }
+ }
+ }
+ out <- AsyncChatResult{
+ Answer: "",
+ Reference: map[string]interface{}{},
+ AudioBinary: nil,
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ EndToThink: true,
+ }
+ }
fullAnswer += *answer
deltas := BufferAnswerDelta(thinkState, *answer, 16)
for _, d := range deltas {
- if d.Kind == ThinkDeltaMarker && d.Value == "" {
- out <- AsyncChatResult{
- Answer: "",
- Reference: map[string]interface{}{},
- AudioBinary: nil,
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
- EndToThink: true,
- }
- } else if d.Kind == ThinkDeltaText && d.Value != "" {
+ if d.Kind == ThinkDeltaText && d.Value != "" {
out <- AsyncChatResult{
Answer: d.Value,
Reference: map[string]interface{}{},
@@ -1173,15 +1178,17 @@ func (s *ChatPipelineService) AsyncChat(
// Flush remaining state matching Python's final flush order
// (dialog_service.py:1601-1612): think_buffer → marker → answer_buffer → pending_after_close
// Python has no Reasoning field — all text is Answer.
+ hadThinkClose := false
for _, d := range FlushRemaining(thinkState) {
if d.Kind == ThinkDeltaMarker && d.Value == "" {
+ hadThinkClose = true
out <- AsyncChatResult{
- Answer: "",
- Reference: map[string]interface{}{},
- AudioBinary: nil,
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
- EndToThink: true,
+ Answer: "",
+ Reference: map[string]interface{}{},
+ AudioBinary: nil,
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ EndToThink: true,
}
} else if d.Kind == ThinkDeltaText && d.Value != "" {
out <- AsyncChatResult{
@@ -1193,6 +1200,19 @@ func (s *ChatPipelineService) AsyncChat(
}
}
}
+ // Close reasoning if the stream ended while still in reasoning mode
+ // (e.g. model returned only reasoning chunks with no content delta).
+ // Skip when FlushRemaining already emitted a marker.
+ if !hadThinkClose && thinkState.ExitReasoning() {
+ out <- AsyncChatResult{
+ Answer: "",
+ Reference: map[string]interface{}{},
+ AudioBinary: nil,
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ EndToThink: true,
+ }
+ }
// Decorate and yield the final answer.
// Python uses state.full_text (raw text with tags) as input
@@ -1399,27 +1419,19 @@ func (s *ChatPipelineService) AsyncChatSolo(
*chatModel.ModelName, chatMessages, chatModel.APIConfig, chatCfg,
func(answer *string, reason *string) error {
if reason != nil && *reason != "" {
+ if thinkState.EnterReasoning() {
+ out <- AsyncChatResult{
+ Answer: "",
+ Reference: map[string]interface{}{},
+ AudioBinary: nil,
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ StartToThink: true,
+ }
+ }
deltas := NextThinkDelta(thinkState, *reason, 16)
for _, d := range deltas {
- if d.Kind == ThinkDeltaMarker && d.Value == "" {
- out <- AsyncChatResult{
- Answer: "",
- Reference: map[string]interface{}{},
- AudioBinary: nil,
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
- StartToThink: true,
- }
- } else if d.Kind == ThinkDeltaMarker && d.Value == "" {
- out <- AsyncChatResult{
- Answer: "",
- Reference: map[string]interface{}{},
- AudioBinary: nil,
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
- EndToThink: true,
- }
- } else if d.Kind == ThinkDeltaText && d.Value != "" {
+ if d.Kind == ThinkDeltaText && d.Value != "" {
fullAnswer += d.Value
out <- AsyncChatResult{
Answer: d.Value,
@@ -1431,11 +1443,20 @@ func (s *ChatPipelineService) AsyncChatSolo(
}
}
}
- if isContentDelta(answer) {
- fullAnswer += *answer
- deltas := BufferAnswerDelta(thinkState, *answer, 16)
- for _, d := range deltas {
- if d.Kind == ThinkDeltaMarker && d.Value == "" {
+ if isContentDelta(answer) {
+ if thinkState.ExitReasoning() {
+ for _, d := range FlushRemaining(thinkState) {
+ if d.Kind == ThinkDeltaText && d.Value != "" {
+ fullAnswer += d.Value
+ out <- AsyncChatResult{
+ Answer: d.Value,
+ Reference: map[string]interface{}{},
+ AudioBinary: s.synthesizeTTS(ttsModel, d.Value),
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ }
+ }
+ }
out <- AsyncChatResult{
Answer: "",
Reference: map[string]interface{}{},
@@ -1444,17 +1465,21 @@ func (s *ChatPipelineService) AsyncChatSolo(
Final: false,
EndToThink: true,
}
- } else if d.Kind == ThinkDeltaText && d.Value != "" {
- out <- AsyncChatResult{
- Answer: d.Value,
- Reference: map[string]interface{}{},
- AudioBinary: s.synthesizeTTS(ttsModel, d.Value),
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
+ }
+ fullAnswer += *answer
+ deltas := BufferAnswerDelta(thinkState, *answer, 16)
+ for _, d := range deltas {
+ if d.Kind == ThinkDeltaText && d.Value != "" {
+ out <- AsyncChatResult{
+ Answer: d.Value,
+ Reference: map[string]interface{}{},
+ AudioBinary: s.synthesizeTTS(ttsModel, d.Value),
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ }
}
}
}
- }
return nil
},
)
@@ -1466,15 +1491,17 @@ func (s *ChatPipelineService) AsyncChatSolo(
return
}
timer.Exit(common.PhaseGenerateAnswer)
+ hadThinkClose := false
for _, d := range FlushRemaining(thinkState) {
if d.Kind == ThinkDeltaMarker && d.Value == "" {
+ hadThinkClose = true
out <- AsyncChatResult{
- Answer: "",
- Reference: map[string]interface{}{},
- AudioBinary: nil,
- CreatedAt: float64(time.Now().Unix()),
- Final: false,
- EndToThink: true,
+ Answer: "",
+ Reference: map[string]interface{}{},
+ AudioBinary: nil,
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ EndToThink: true,
}
} else if d.Kind == ThinkDeltaText && d.Value != "" {
out <- AsyncChatResult{
@@ -1486,6 +1513,19 @@ func (s *ChatPipelineService) AsyncChatSolo(
}
}
}
+ // Close reasoning if the stream ended while still in reasoning mode
+ // (e.g. model returned only reasoning chunks with no content delta).
+ // Skip when FlushRemaining already emitted a marker.
+ if !hadThinkClose && thinkState.ExitReasoning() {
+ out <- AsyncChatResult{
+ Answer: "",
+ Reference: map[string]interface{}{},
+ AudioBinary: nil,
+ CreatedAt: float64(time.Now().Unix()),
+ Final: false,
+ EndToThink: true,
+ }
+ }
finalAnswer := ExtractVisibleAnswer(thinkState.fullText)
if finalAnswer == "" {
finalAnswer = fullAnswer
@@ -2758,7 +2798,7 @@ func (s *ChatPipelineService) decorateAnswer(
}
newChunks = append(newChunks, newChunk)
}
- refs["chunks"] = newChunks
+ refs["chunks"] = chunksFormat(newChunks)
}
}
@@ -4050,7 +4090,7 @@ func (s *ChatPipelineService) buildSQLReference(
"count": d["count"],
})
}
- ref["chunks"] = chunks
+ ref["chunks"] = chunksFormat(chunks)
ref["doc_aggs"] = docAggs
return answer, ref
}
@@ -4059,7 +4099,7 @@ func (s *ChatPipelineService) buildSQLReference(
if isAggregateSQL(originalSQL) {
chunks, docAggs := s.fetchAggregateChunks(ctx, docEngine, tableName, originalSQL, expectedCol, kbIDs)
if len(chunks) > 0 {
- ref["chunks"] = chunks
+ ref["chunks"] = chunksFormat(chunks)
ref["doc_aggs"] = docAggs
}
return answer, ref
@@ -4210,3 +4250,38 @@ func kbIDStrings(kbs []*entity.Knowledgebase) []string {
}
return out
}
+
+// chunksFormat normalizes raw chunk maps to the frontend-expected field names.
+// Mirrors Python's chunks_format in rag/prompts/generator.py:41-65.
+func chunksFormat(chunksRaw []map[string]interface{}) []map[string]interface{} {
+ result := make([]map[string]interface{}, 0, len(chunksRaw))
+ for _, chunk := range chunksRaw {
+ formatted := map[string]interface{}{
+ "id": getChunkValue(chunk, "chunk_id", "id"),
+ "content": getChunkValue(chunk, "content", "content_with_weight"),
+ "document_id": getChunkValue(chunk, "doc_id", "document_id"),
+ "document_name": getChunkValue(chunk, "docnm_kwd", "document_name"),
+ "dataset_id": getChunkValue(chunk, "kb_id", "dataset_id"),
+ "image_id": getChunkValue(chunk, "image_id", "img_id"),
+ "positions": getChunkValue(chunk, "positions", "position_int"),
+ "url": chunk["url"],
+ "similarity": chunk["similarity"],
+ "vector_similarity": chunk["vector_similarity"],
+ "term_similarity": chunk["term_similarity"],
+ "row_id": chunk["row_id"],
+ "doc_type": getChunkValue(chunk, "doc_type_kwd", "doc_type"),
+ "document_metadata": chunk["document_metadata"],
+ }
+ result = append(result, formatted)
+ }
+ return result
+}
+
+// getChunkValue returns the first non-nil value from a chunk map, trying k1 first then k2.
+// Mirrors Python's get_value helper in rag/prompts/generator.py:37-38.
+func getChunkValue(chunk map[string]interface{}, k1, k2 string) interface{} {
+ if v, ok := chunk[k1]; ok && v != nil {
+ return v
+ }
+ return chunk[k2]
+}
diff --git a/internal/service/chat_pipeline_test.go b/internal/service/chat_pipeline_test.go
index f4f3d6752..0bb202517 100644
--- a/internal/service/chat_pipeline_test.go
+++ b/internal/service/chat_pipeline_test.go
@@ -1317,10 +1317,10 @@ func TestBuildSQLReference_NonAggregateWithSourceColumns(t *testing.T) {
if len(docAggs) != 2 {
t.Fatalf("docAggs len = %d, want 2", len(docAggs))
}
- // Each chunk must carry kb_id from the single-KB dialog.
+ // Each chunk must carry dataset_id (remapped from kb_id by chunksFormat) from the single-KB dialog.
for i, cm := range chunks {
- if cm["kb_id"] != "kb_a" {
- t.Errorf("chunks[%d].kb_id = %v, want kb_a", i, cm["kb_id"])
+ if cm["dataset_id"] != "kb_a" {
+ t.Errorf("chunks[%d].dataset_id = %v, want kb_a", i, cm["dataset_id"])
}
}
}
diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go
index 49ae9a735..786296fa5 100644
--- a/internal/service/chat_session.go
+++ b/internal/service/chat_session.go
@@ -1609,20 +1609,31 @@ func (s *ChatSessionService) ChatCompletions(
ans["chat_id"] = chatID
}
sendOrCancel(fmt.Sprintf("data:%s\n\n", sseMarshalChunk(sanitizeJSONFloats(ans).(map[string]interface{}), chatID)))
+ } else {
+ ans := s.structureAnswer(session, "", messageID, sessionID, reference)
+ ans["final"] = true
+ if chatID != "" {
+ ans["chat_id"] = chatID
+ }
+ sendOrCancel(fmt.Sprintf("data:%s\n\n", sseMarshalChunk(sanitizeJSONFloats(ans).(map[string]interface{}), chatID)))
}
continue
}
+ deltaAnswer := ""
if result.StartToThink {
fullAnswer.WriteString("")
} else if result.EndToThink {
fullAnswer.WriteString("")
} else if result.Answer != "" {
fullAnswer.WriteString(result.Answer)
+ deltaAnswer = result.Answer
}
if session != nil {
s.appendAssistantToSession(session, fullAnswer.String(), messageID)
}
- ans := s.structureAnswer(session, result.Answer, messageID, sessionID, reference)
+ ans := s.structureAnswer(session, deltaAnswer, messageID, sessionID, reference)
+ ans["start_to_think"] = result.StartToThink
+ ans["end_to_think"] = result.EndToThink
if chatID != "" {
ans["chat_id"] = chatID
}
@@ -1928,15 +1939,17 @@ func (s *ChatSessionService) checkTenantLLMAPIKey(tenantID, modelName string) (b
// sseAnswerChunk has deterministic JSON field order matching Python's structure_answer output.
type sseAnswerChunk struct {
- Answer string `json:"answer"`
- Reference map[string]interface{} `json:"reference"`
- AudioBinary interface{} `json:"audio_binary"`
- Prompt string `json:"prompt"`
- CreatedAt float64 `json:"created_at"`
- Final bool `json:"final"`
- ID string `json:"id"`
- SessionID string `json:"session_id"`
- ChatID string `json:"chat_id,omitempty"`
+ Answer string `json:"answer"`
+ Reference map[string]interface{} `json:"reference"`
+ AudioBinary interface{} `json:"audio_binary"`
+ Prompt string `json:"prompt"`
+ CreatedAt float64 `json:"created_at"`
+ Final bool `json:"final"`
+ ID string `json:"id"`
+ SessionID string `json:"session_id"`
+ ChatID string `json:"chat_id,omitempty"`
+ StartToThink bool `json:"start_to_think,omitempty"`
+ EndToThink bool `json:"end_to_think,omitempty"`
}
// sseWrapper wraps the SSE response with deterministic field order matching Python:
@@ -2036,16 +2049,21 @@ func sseMarshalChunk(ans map[string]interface{}, chatID string) string {
createdAt, _ := ans["created_at"].(float64)
final, _ := ans["final"].(bool)
+ startToThink, _ := ans["start_to_think"].(bool)
+ endToThink, _ := ans["end_to_think"].(bool)
+
chunk := sseAnswerChunk{
- Answer: answer,
- Reference: ref,
- AudioBinary: ans["audio_binary"],
- Prompt: prompt,
- CreatedAt: createdAt,
- Final: final,
- ID: id,
- SessionID: sessionID,
- ChatID: chatID,
+ Answer: answer,
+ Reference: ref,
+ AudioBinary: ans["audio_binary"],
+ Prompt: prompt,
+ CreatedAt: createdAt,
+ Final: final,
+ ID: id,
+ SessionID: sessionID,
+ ChatID: chatID,
+ StartToThink: startToThink,
+ EndToThink: endToThink,
}
wrapper := sseWrapper{Code: 0, Message: "", Data: chunk}
return marshalJSONWithSpaces(wrapper)
diff --git a/internal/service/think_tag.go b/internal/service/think_tag.go
index 582238f66..db8736a79 100644
--- a/internal/service/think_tag.go
+++ b/internal/service/think_tag.go
@@ -46,6 +46,30 @@ type ThinkStreamState struct {
answerBuffer string
// carry holds text at the end of a chunk that may be a partial or prefix.
carry string
+ // inReasoning tracks whether the model is currently streaming reasoning chunks
+ // (via a dedicated reason field). Kept separate from inThink, which tracks
+ // / tag boundaries parsed from the text stream itself.
+ inReasoning bool
+}
+
+// EnterReasoning marks the start of a reasoning block (model-level, not tag-based).
+// 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
+ return true
+}
+
+// ExitReasoning marks the end of a reasoning block (model-level, not tag-based).
+// Returns true when this is a new transition (reasoning was active).
+func (s *ThinkStreamState) ExitReasoning() bool {
+ if !s.inReasoning {
+ return false
+ }
+ s.inReasoning = false
+ return true
}
// ThinkDeltaKind describes the type of a think-tag delta event.