Port agent PRs to GO - 5 (#16667)

### Summary Port https://github.com/infiniflow/ragflow/pull/15376 https://github.com/infiniflow/ragflow/pull/16401 https://github.com/infiniflow/ragflow/pull/15484 https://github.com/infiniflow/ragflow/pull/16685
2026-07-12 06:35:44 +08:00 · 2026-07-08 19:54:29 +08:00
parent 0f08dc070d
commit 8e3bbad4da
5 changed files with 220 additions and 84 deletions
--- a/internal/handler/agent.go
+++ b/internal/handler/agent.go
@@ -808,15 +808,17 @@ func (h *AgentHandler) DeleteAgentSession(c *gin.Context) {

 // AgentChatCompletions POST /api/v1/agents/chat/completions
 //
-// Runs the canvas against `agent_id` and streams the result as SSE.
+// Runs the canvas against `agent_id`. Mirrors the Python contract in
+// api/db/services/canvas_service.py:313 (`completion()`) and
+// api/apps/restful_apis/agent_api.py:1440-1676.
 //
-// Behaviour matches the Python reference at
-// api/db/services/canvas_service.py:313 (`completion()`):
-//
-//   - Non-openai path: always streams SSE — one `data: {...}\n\n` frame per
-//     canvas RunEvent, terminated by `data: [DONE]\n\n`. The `stream` field
-//     is ignored on this path because Python's `completion()` always yields
-//     SSE frames regardless of the flag.
+//   - When `stream` is true: streams SSE — one `data: {...}\n\n` frame per
+//     canvas RunEvent, terminated by `data: [DONE]\n\n`.
+//   - When `stream` is omitted or false (default, matches the Python
+//     agent_chat_completion contract where `req.get("stream", False)`
+//     defaults to non-streaming): collects all canvas events and returns a
+//     plain JSON response with `data.content` set to the concatenated
+//     message content (matching Python's final_ans["data"]["content"]).
 //   - Openai-compatible path: requires `messages` (a non-empty list with at
 //     least one user message is needed to derive the question). The full
 //     OpenAI wire framing (delta + reference + token counts — see
@@ -1004,62 +1006,164 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) {
 		return
 	}

-	c.Writer.Header().Set("Content-Type", "text/event-stream")
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	// SSE wire format is the flat Python agent-canvas envelope:
-	// {event,message_id,task_id,session_id,created_at,data}. One
-	// frame is emitted per canvas event through service.WriteChatbotRunEvent.
-	// The channel close is signalled by `data: [DONE]\n\n`. We do NOT
-	// emit an SSE `event:` line — the front-end's use-send-message.ts
-	// parser feeds each `data:` line directly into JSON.parse and
-	// expects the event type in the JSON object's top-level `event`
-	// field.
-	emitted := false
-	for ev := range events {
-		emitted = true
-		common.Debug("agent chat completions: streaming event",
+	if req.Stream {
+		// SSE streaming: one `data: {...}\n\n` frame per canvas RunEvent,
+		// terminated by `data: [DONE]\n\n`. We do NOT emit an SSE `event:`
+		// line — the front-end's use-send-message.ts parser feeds each
+		// `data:` line directly into JSON.parse and expects the event type
+		// in the JSON object's top-level `event` field.
+		c.Writer.Header().Set("Content-Type", "text/event-stream")
+		c.Writer.Header().Set("Cache-Control", "no-cache")
+		c.Writer.Header().Set("Connection", "keep-alive")
+		emitted := false
+		for ev := range events {
+			emitted = true
+			common.Debug("agent chat completions: streaming event",
+				zap.String("agent_id", req.AgentID),
+				zap.String("session_id", req.SessionID),
+				zap.String("event_type", ev.Type),
+				zap.String("message_id", ev.MessageID),
+				zap.String("task_id", ev.TaskID),
+			)
+			if err := service.WriteChatbotRunEvent(c.Writer, ev); err != nil {
+				common.Debug("agent chat completions: client disconnected",
+					zap.String("agent_id", req.AgentID),
+					zap.Error(err),
+				)
+				return
+			}
+		}
+		if !emitted {
+			// Canvas produced no events (e.g. empty query). Echo the
+			// session_id so the client can resume the conversation
+			// (fixes #15169). The [DONE] terminator must be emitted
+			// here explicitly because the canvas never sends a
+			// "done" event on this path.
+			common.Info("empty agent output - returning session_id",
+				zap.String("agent_id", req.AgentID),
+				zap.String("session_id", req.SessionID),
+			)
+			event := canvas.RunEvent{
+				Type:      "",
+				Data:      "{}",
+				CreatedAt: time.Now().Unix(),
+				SessionID: req.SessionID,
+			}
+			_ = service.WriteChatbotRunEvent(c.Writer, event)
+			if _, err := c.Writer.Write([]byte("data: [DONE]\n\n")); err != nil {
+				common.Debug("agent chat completions: failed to write [DONE]",
+					zap.Error(err),
+				)
+			}
+		}
+		common.Debug("agent chat completions: stream closed",
 			zap.String("agent_id", req.AgentID),
 			zap.String("session_id", req.SessionID),
-			zap.String("event_type", ev.Type),
-			zap.String("message_id", ev.MessageID),
-			zap.String("task_id", ev.TaskID),
 		)
-		if err := service.WriteChatbotRunEvent(c.Writer, ev); err != nil {
-			common.Debug("agent chat completions: client disconnected",
-				zap.String("agent_id", req.AgentID),
-				zap.Error(err),
-			)
-			return
+		return
+	}
+
+	// Non-streaming: collect all canvas events, accumulate message
+	// content and reference, then return a plain JSON response matching
+	// Python's get_result(data=final_ans) contract (agent_api.py:1616-1676).
+	var fullContent string
+	reference := map[string]any{}
+	structuredOutput := map[string]any{}
+	var runUsage any
+	var finalAns *canvas.RunEvent
+	hasEvents := false
+	for ev := range events {
+		hasEvents = true
+		var evData map[string]any
+		if err := json.Unmarshal([]byte(ev.Data), &evData); err == nil {
+			if ev.Type == "message" {
+				if c, ok := evData["content"].(string); ok {
+					fullContent += c
+				}
+			}
+			if ref, _ := evData["reference"].(map[string]any); ref != nil {
+				for k, v := range ref {
+					reference[k] = v
+				}
+			}
+			// #4: collect structured_output from node_finished events
+			// (Python: agent_api.py:1628-1633).
+			if ev.Type == "node_finished" {
+				if outputs, ok := evData["outputs"].(map[string]any); ok {
+					if structured, ok := outputs["structured"]; ok {
+						if componentID, ok := evData["component_id"].(string); ok {
+							structuredOutput[componentID] = structured
+						}
+					}
+				}
+			}
+			// #5: capture run_usage from workflow_finished events
+			// (Python: agent_api.py:1641-1644).
+			if ev.Type == "workflow_finished" {
+				if u := evData["usage"]; u != nil {
+					runUsage = u
+				}
+			}
+		}
+
+		// Python final_ans selection: prefer "message_end" event,
+		// skip "workflow_finished" (set by canvas as the last event
+		// but should not be the final answer), use "user_inputs" as
+		// fallback.
+		if ev.Type == "message_end" {
+			copy := ev
+			finalAns = &copy
+		}
+		if ev.Type == "workflow_finished" {
+			continue
+		}
+		if finalAns == nil {
+			copy := ev
+			finalAns = &copy
 		}
 	}
-	if !emitted {
+
+	if !hasEvents || finalAns == nil {
 		// Canvas produced no events (e.g. empty query). Echo the
 		// session_id so the client can resume the conversation
-		// (fixes #15169). The [DONE] terminator must be emitted
-		// here explicitly because the canvas never sends a
-		// "done" event on this path.
+		// (fixes #15169). Python returns get_result(data={"session_id": session_id}).
 		common.Info("empty agent output - returning session_id",
 			zap.String("agent_id", req.AgentID),
 			zap.String("session_id", req.SessionID),
 		)
-		event := canvas.RunEvent{
-			Type:      "",
-			Data:      "{}",
-			CreatedAt: time.Now().Unix(),
-			SessionID: req.SessionID,
-		}
-		_ = service.WriteChatbotRunEvent(c.Writer, event)
-		if _, err := c.Writer.Write([]byte("data: [DONE]\n\n")); err != nil {
-			common.Debug("agent chat completions: failed to write [DONE]",
-				zap.Error(err),
-			)
-		}
+		common.SuccessWithData(c, gin.H{"session_id": req.SessionID}, "success")
+		return
 	}
-	common.Debug("agent chat completions: stream closed",
-		zap.String("agent_id", req.AgentID),
-		zap.String("session_id", req.SessionID),
-	)
+
+	// Build the final answer payload:
+	//   final_ans["data"]["content"] = full_content
+	//   final_ans["data"]["reference"] = reference
+	//   final_ans["data"]["usage"] = run_usage (if present)
+	//   final_ans["data"]["structured"] = structured_output (if present)
+	var ansData map[string]any
+	if err := json.Unmarshal([]byte(finalAns.Data), &ansData); err != nil || ansData == nil {
+		ansData = map[string]any{}
+	}
+	ansData["content"] = fullContent
+	if len(reference) > 0 {
+		ansData["reference"] = reference
+	}
+	if runUsage != nil {
+		ansData["usage"] = runUsage
+	}
+	if len(structuredOutput) > 0 {
+		ansData["structured"] = structuredOutput
+	}
+
+	result := gin.H{
+		"event":      finalAns.Type,
+		"data":       ansData,
+		"message_id": finalAns.MessageID,
+		"created_at": finalAns.CreatedAt,
+		"task_id":    finalAns.TaskID,
+		"session_id": finalAns.SessionID,
+	}
+	common.SuccessWithData(c, result, "success")
 }

 // RerunAgent POST /api/v1/agents/rerun — requires id, dsl, and
--- a/internal/handler/agent_test.go
+++ b/internal/handler/agent_test.go
@@ -770,12 +770,12 @@ func TestAgentChatCompletions_StreamSetsContentType(t *testing.T) {
 	}
 }

-// TestAgentChatCompletions_DefaultBranchStreamsSSE covers the
-// scenario the user actually hit: `openai-compatible: false` with no
-// `stream` field on the body. The handler must still invoke the
-// canvas runner and stream the result as SSE — the SSE envelope is
-// the flat Python agent-canvas shape regardless of the stream flag.
-func TestAgentChatCompletions_DefaultBranchStreamsSSE(t *testing.T) {
+// TestAgentChatCompletions_DefaultBranchNonStreaming covers the
+// scenario where `stream` is omitted from the request body. When
+// `stream` is absent, the handler must return a plain JSON response
+// (non-streaming), matching the Python contract where
+// `req.get("stream", False)` defaults to non-streaming.
+func TestAgentChatCompletions_DefaultBranchNonStreaming(t *testing.T) {
 	gin.SetMode(gin.TestMode)
 	w := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(w)
@@ -792,17 +792,20 @@ func TestAgentChatCompletions_DefaultBranchStreamsSSE(t *testing.T) {
 	h := &AgentHandler{chatRunner: runner}
 	h.AgentChatCompletions(c)

-	if got := w.Header().Get("Content-Type"); !strings.Contains(got, "text/event-stream") {
-		t.Errorf("Content-Type = %q, want text/event-stream (default branch must stream)", got)
+	if got := w.Header().Get("Content-Type"); !strings.Contains(got, "application/json") {
+		t.Errorf("Content-Type = %q, want application/json (default branch must not stream)", got)
 	}
 	body := w.Body.String()
+	if !strings.Contains(body, `"code":0`) {
+		t.Errorf("body should contain success code, got %q", body)
+	}
 	if !strings.Contains(body, `"event":"message"`) ||
 		!strings.Contains(body, `"message_id":"msg-2"`) ||
-		!strings.Contains(body, `"content":"hello back"`) {
-		t.Errorf("body should contain flat agent event with content, got %q", body)
+		!strings.Contains(body, `"hello back"`) {
+		t.Errorf("body should contain agent event with content in data, got %q", body)
 	}
-	if !strings.HasSuffix(body, "data: [DONE]\n\n") {
-		t.Errorf("body should end with [DONE] terminator, got %q", body)
+	if strings.Contains(body, "data: [DONE]") {
+		t.Errorf("body should not contain [DONE] terminator in non-streaming mode, got %q", body)
 	}
 }

@@ -911,8 +914,12 @@ func TestAgentChatCompletions_OpenAICompat_NonStreamReturnsChoices(t *testing.T)

 	var resp map[string]interface{}
 	_ = json.Unmarshal(w.Body.Bytes(), &resp)
-	if _, ok := resp["choices"]; !ok {
-		t.Errorf("response should contain top-level 'choices', got keys: %v", resp)
+	data, _ := resp["data"].(map[string]interface{})
+	if data == nil {
+		t.Fatalf("response should contain 'data', got keys: %v", resp)
+	}
+	if _, ok := data["choices"]; !ok {
+		t.Errorf("response data should contain 'choices', got keys: %v", data)
 	}
 }

--- a/internal/handler/memory.go
+++ b/internal/handler/memory.go
@@ -789,6 +789,10 @@ func (h *MemoryHandler) SearchMessage(c *gin.Context) {
 	similarityThreshold, _ := strconv.ParseFloat(c.DefaultQuery("similarity_threshold", "0.2"), 64)
 	keywordsSimilarityWeight, _ := strconv.ParseFloat(c.DefaultQuery("keywords_similarity_weight", "0.7"), 64)
 	topN, _ := strconv.Atoi(c.DefaultQuery("top_n", "5"))
+	if topN <= 0 || topN > 100 {
+		common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "top_n must be between 1 and 100")
+		return
+	}

 	agentID := c.DefaultQuery("agent_id", "")
 	sessionID := c.DefaultQuery("session_id", "")
@@ -856,6 +860,10 @@ func (h *MemoryHandler) GetMessages(c *gin.Context) {
 	agentID := c.DefaultQuery("agent_id", "")
 	sessionID := c.DefaultQuery("session_id", "")
 	limit, _ := strconv.Atoi(c.DefaultQuery("limit", "10"))
+	if limit <= 0 || limit > 100 {
+		common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "limit must be between 1 and 100")
+		return
+	}
 	if len(memoryIDs) == 0 {
 		common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "memory_ids is required.")
 		return