Fix chat thinking & Figure issue in GO (#16558)

### Summary

Fix chat thinking & Figure issue

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"^.*</think>", "", ...) + ^.*```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 </think> (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)^.*</think>`)
+	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
+}

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 == "<think>" {
-									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 == "</think>" {
-									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 == "</think>" {
-									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 == "</think>" {
+					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 </think> 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 <think> 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 == "<think>" {
-								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 == "</think>" {
-								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 == "</think>" {
+					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 == "</think>" {
+					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 </think> 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]
+}

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"])
 		}
 	}
 }

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("<think>")
 				} else if result.EndToThink {
 					fullAnswer.WriteString("</think>")
 				} 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)

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 <think> or </think> prefix.
 	carry string
+	// inReasoning tracks whether the model is currently streaming reasoning chunks
+	// (via a dedicated reason field).  Kept separate from inThink, which tracks
+	// <think>/</think> 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.