zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-18 14:28:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(agent): keep think tags in persisted agent session messages (#18222)

Browse Source
This commit is contained in:
euvre
2026-08-14 01:00:08 -07:00
committed by GitHub
parent 14c5a61350
commit e1330423ac
4 changed files with 135 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
internal/handler/agent.go
View File

@@ -1395,6 +1395,14 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) {
if c, ok := evData["content"].(string); ok {
fullContent += c
}
// Mirror Python agent_api.py: the reasoning segment stays
// wrapped in <think> tags in the aggregated answer so the
// chat UI can render the "thought" section.
if st, _ := evData["start_to_think"].(bool); st {
fullContent += "<think>"
} else if et, _ := evData["end_to_think"].(bool); et {
fullContent += "</think>"
}
}
if ref, _ := evData["reference"].(map[string]any); ref != nil {
for k, v := range ref {

40
internal/handler/agent_test.go
View File

@@ -866,6 +866,46 @@ func TestAgentChatCompletions_DefaultBranchNonStreaming(t *testing.T) {
}
}
// TestAgentChatCompletions_NonStreamingPreservesThinkMarkers covers the
// non-streaming aggregation: start_to_think/end_to_think message events must
// survive as <think> tags in the final content, mirroring Python
// agent_api.py, so clients can render the "thought" section.
func TestAgentChatCompletions_NonStreamingPreservesThinkMarkers(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/api/v1/agents/chat/completions",
strings.NewReader(`{"agent_id":"a1","query":"hello","stream":false}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set("user", &entity.User{ID: "u1"})
c.Set("user_id", "u1")
runner := &stubChatRunner{events: []canvas.RunEvent{
{Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"","start_to_think":true}`},
{Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"reasoning trace"}`},
{Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"","end_to_think":true}`},
{Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"final answer"}`},
{Type: "message_end", MessageID: "msg-1", SessionID: "sess-1", Data: `{}`},
{Type: "done", Data: ""},
}}
h := &AgentHandler{chatRunner: runner}
h.AgentChatCompletions(c)
body := w.Body.String()
if !strings.Contains(body, `"code":0`) {
t.Fatalf("body should contain success code, got %q", body)
}
var response map[string]any
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
t.Fatalf("decode response: %v", err)
}
data, _ := response["data"].(map[string]any)
ansData, _ := data["data"].(map[string]any)
if got, _ := ansData["content"].(string); got != "<think>reasoning trace</think>final answer" {
t.Errorf("aggregated content = %q, want the think section preserved as %q", got, "<think>reasoning trace</think>final answer")
}
}
type emptySessionCaptureRunner struct {
sessionID string
}

22
internal/service/agent.go
View File

@@ -253,6 +253,19 @@ func splitInlineThink(answer, thinking string) (string, string) {
return answer, thinking
}
// agentSessionMessageContent mirrors how Python's canvas_service.completion
// accumulates the persisted assistant message: reasoning streams first,
// bracketed by start_to_think/end_to_think markers, so the stored content
// keeps the thinking segment wrapped in <think> tags. The chat UI refetches
// the session message list after streaming and needs those tags for
// MarkdownContent to render the "thought" section (chat.thought).
func agentSessionMessageContent(answer, thinking string) string {
if thinking == "" {
return answer
}
return "<think>" + thinking + "</think>" + answer
}
func splitMessageContent(content string) []string {
if content == "" {
return nil
@@ -2110,7 +2123,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
if answer != "" {
appendAssistantHistory(state, partialAssistantOutput(answer, downloads))
}
if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, answer != ""); persistErr != nil {
if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, thinking, referencePayload, dsl, state, answer != ""); persistErr != nil {
return nil, fmt.Errorf("persist interrupted agent session: %w: %w", persistErr, ErrAgentStorageError)
}
if answer != "" && shouldEmitMessage {
@@ -2127,7 +2140,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
}
if shouldTreatAsCompletedLoopRun(err, answer) {
appendAssistantHistory(state, assistantOutput)
if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil {
if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, thinking, referencePayload, dsl, state, true); persistErr != nil {
s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error())
return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError)
}
@@ -2163,7 +2176,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
// Emit message + message_end (mirrors Python's ans dict).
appendAssistantHistory(state, assistantOutput)
if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil {
if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, thinking, referencePayload, dsl, state, true); persistErr != nil {
s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error())
return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError)
}
@@ -2291,6 +2304,7 @@ func (s *AgentService) persistAgentRunSession(
agentID, userID, sessionID, messageID string,
userInput any,
answer string,
thinking string,
reference map[string]interface{},
runDSL map[string]any,
state *canvas.CanvasState,
@@ -2313,7 +2327,7 @@ func (s *AgentService) persistAgentRunSession(
messages = append(messages, map[string]interface{}{"role": "user", "content": text, "id": utility.GenerateToken(), "created_at": now})
}
if appendAssistantMessage {
messages = append(messages, map[string]interface{}{"role": "assistant", "content": answer, "id": messageID, "created_at": now})
messages = append(messages, map[string]interface{}{"role": "assistant", "content": agentSessionMessageContent(answer, thinking), "id": messageID, "created_at": now})
}
if raw, err := json.Marshal(messages); err == nil {
session.Message = raw

69
internal/service/agent_persistence_test.go
View File

@@ -90,3 +90,72 @@ func TestAgentRunSessionUpdateFailurePreventsSuccessEvents(t *testing.T) {
}
}
}
func TestAgentSessionMessageContent(t *testing.T) {
tests := []struct {
name string
answer string
thinking string
want string
}{
{"no thinking keeps answer as-is", "final answer", "", "final answer"},
{"thinking wrapped before answer", "final answer", "reasoning trace", "<think>reasoning trace</think>final answer"},
{"empty answer keeps thinking section", "", "reasoning trace", "<think>reasoning trace</think>"},
{"inline think tags preserved when no separate thinking", "pre <think>inline</think> post", "", "pre <think>inline</think> post"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := agentSessionMessageContent(tt.answer, tt.thinking); got != tt.want {
t.Errorf("agentSessionMessageContent(%q, %q) = %q, want %q", tt.answer, tt.thinking, got, tt.want)
}
})
}
}
// TestPersistAgentRunSessionPreservesThinking guards the chat.thought
// ("思考完成") indicator: the assistant message stored on the agent session
// must keep the reasoning segment wrapped in <think> tags, mirroring Python
// canvas_service.completion. The chat UI refetches the session message list
// after streaming and renders the thought section from those tags.
func TestPersistAgentRunSessionPreservesThinking(t *testing.T) {
testDB := setupServiceTestDB(t)
if err := testDB.AutoMigrate(&entity.API4Conversation{}); err != nil {
t.Fatalf("migrate: %v", err)
}
originalDB := dao.DB
dao.DB = testDB
t.Cleanup(func() { dao.DB = originalDB })
if err := testDB.Create(&entity.API4Conversation{
ID: "session-think",
DialogID: "canvas-think",
UserID: "user-1",
Message: json.RawMessage(`[]`),
Reference: json.RawMessage(`[]`),
}).Error; err != nil {
t.Fatalf("create session: %v", err)
}
svc := NewAgentService()
if err := svc.persistAgentRunSession(context.Background(), "canvas-think", "user-1", "session-think", "msg-think-1", "question", "final answer", "reasoning trace", map[string]interface{}{}, nil, nil, true); err != nil {
t.Fatalf("persist: %v", err)
}
var conv entity.API4Conversation
if err := testDB.First(&conv, "id = ?", "session-think").Error; err != nil {
t.Fatalf("reload session: %v", err)
}
var messages []map[string]any
if err := json.Unmarshal(conv.Message, &messages); err != nil {
t.Fatalf("decode messages: %v", err)
}
if len(messages) != 2 {
t.Fatalf("messages = %d, want 2 (user + assistant)", len(messages))
}
if role, _ := messages[1]["role"].(string); role != "assistant" {
t.Fatalf("second message role = %q, want assistant", role)
}
if got, _ := messages[1]["content"].(string); got != "<think>reasoning trace</think>final answer" {
t.Fatalf("persisted assistant content = %q, want %q", got, "<think>reasoning trace</think>final answer")
}
}