fix[go]: support retrieval kb in agent chat (#16944)

### Summary As title - [x] fix tool & component `Retrieval KB` - [x] fix agent cannot use `Retrieval KB` in agent chat #### Main cause Model provider do not set up: ```Go if chatModelConfig.Tools != nil { reqBody["tools"] = chatModelConfig.Tools } ``` #### Working Now <img width="3774" height="2128" alt="image" src="https://github.com/user-attachments/assets/400a349d-0211-43e5-a7ec-7a014acf77a6" /> ``` ____________________________________ < This PR takes me all day to do it. > ------------------------------------ \ \ (\__/) (•ㅅ•) / づ ```
2026-07-17 05:07:23 +08:00 · 2026-07-15 21:44:28 +08:00
parent a7da78d0d7
commit ff3d566a4f
15 changed files with 1559 additions and 95 deletions
--- a/internal/agent/component/agent.go
+++ b/internal/agent/component/agent.go
@@ -1017,11 +1017,12 @@ func mergeAgentParam(base AgentParam, inputs map[string]any) AgentParam {
 	if v, ok := stringFrom(inputs, "base_url"); ok {
 		p.BaseURL = v
 	}
-	if v, ok := sliceFrom(inputs, "tools"); ok {
-		p.Tools = v
+	if tools, params, ok := agentToolsFrom(inputs, "tools"); ok {
+		p.Tools = tools
+		p.ToolParams = mergeToolParams(p.ToolParams, params)
 	}
 	if v, ok := nestedMapFrom(inputs, "tool_params"); ok {
-		p.ToolParams = v
+		p.ToolParams = mergeToolParams(p.ToolParams, v)
 	}
 	if v, ok := boolFrom(inputs, "optimize_multi_turn"); ok {
 		p.OptimizeMultiTurn = v
@@ -1035,6 +1036,107 @@ func mergeAgentParam(base AgentParam, inputs map[string]any) AgentParam {
 	return p
 }

+// agentToolsFrom extracts the Agent tools list. The Go-native shape is
+// []string; the canvas DSL shape stores tool component objects with
+// component_name and params.
+func agentToolsFrom(inputs map[string]any, name string) ([]string, map[string]map[string]any, bool) {
+	v, ok := inputs[name]
+	if !ok {
+		return nil, nil, false
+	}
+	switch x := v.(type) {
+	case []string:
+		return x, nil, true
+	case []any:
+		out := make([]string, 0, len(x))
+		params := make(map[string]map[string]any)
+		for _, item := range x {
+			switch tool := item.(type) {
+			case string:
+				if strings.TrimSpace(tool) == "" {
+					continue
+				}
+				out = append(out, tool)
+			case map[string]any:
+				toolName, toolParams, ok := agentToolObject(tool)
+				if !ok {
+					continue
+				}
+				out = append(out, toolName)
+				if len(toolParams) != 0 {
+					params[strings.ToLower(strings.TrimSpace(toolName))] = toolParams
+				}
+			}
+		}
+		return out, params, true
+	}
+	return nil, nil, false
+}
+
+func agentToolObject(item map[string]any) (string, map[string]any, bool) {
+	toolName, ok := stringFrom(item, "component_name")
+	if !ok || strings.TrimSpace(toolName) == "" {
+		toolName, ok = stringFrom(item, "tool_name")
+	}
+	if !ok || strings.TrimSpace(toolName) == "" {
+		toolName, ok = stringFrom(item, "name")
+	}
+	if !ok || strings.TrimSpace(toolName) == "" {
+		return "", nil, false
+	}
+	toolName = strings.TrimSpace(toolName)
+
+	rawParams, _ := item["params"].(map[string]any)
+	toolParams := cloneMap(rawParams)
+	if fn, ok := stringFrom(item, "function_name"); ok && strings.TrimSpace(fn) != "" {
+		if toolParams == nil {
+			toolParams = make(map[string]any)
+		}
+		toolParams["function_name"] = strings.TrimSpace(fn)
+	}
+	return toolName, toolParams, true
+}
+
+func cloneMap(in map[string]any) map[string]any {
+	if len(in) == 0 {
+		return nil
+	}
+	out := make(map[string]any, len(in))
+	for k, v := range in {
+		out[k] = v
+	}
+	return out
+}
+
+func mergeToolParams(base, overrides map[string]map[string]any) map[string]map[string]any {
+	if len(base) == 0 && len(overrides) == 0 {
+		return nil
+	}
+	out := make(map[string]map[string]any, len(base)+len(overrides))
+	for name, params := range base {
+		out[name] = cloneMap(params)
+		if lower := strings.ToLower(strings.TrimSpace(name)); lower != "" && lower != name {
+			out[lower] = cloneMap(params)
+		}
+	}
+	for name, params := range overrides {
+		if len(params) == 0 {
+			continue
+		}
+		lower := strings.ToLower(strings.TrimSpace(name))
+		for k := range out {
+			if strings.ToLower(strings.TrimSpace(k)) == lower {
+				delete(out, k)
+			}
+		}
+		out[name] = cloneMap(params)
+		if lower != "" && lower != name {
+			out[lower] = cloneMap(params)
+		}
+	}
+	return out
+}
+
 // sliceFrom extracts []string from inputs[name].
 func sliceFrom(inputs map[string]any, name string) ([]string, bool) {
 	v, ok := inputs[name]
@@ -1102,11 +1204,12 @@ func init() {
 			f := v
 			p.TopP = &f
 		}
-		if v, ok := sliceFrom(params, "tools"); ok {
-			p.Tools = v
+		if tools, toolParams, ok := agentToolsFrom(params, "tools"); ok {
+			p.Tools = tools
+			p.ToolParams = mergeToolParams(p.ToolParams, toolParams)
 		}
 		if v, ok := nestedMapFrom(params, "tool_params"); ok {
-			p.ToolParams = v
+			p.ToolParams = mergeToolParams(p.ToolParams, v)
 		}
 		if v, ok := intFrom(params, "max_rounds"); ok {
 			p.MaxRounds = v
--- a/internal/agent/component/agent_test.go
+++ b/internal/agent/component/agent_test.go
@@ -306,6 +306,74 @@ func TestAgent_AllRegisteredToolsConfigPassesToRunner(t *testing.T) {
 	}
 }

+func TestAgent_AcceptsCanvasToolObjects(t *testing.T) {
+	var captured AgentParam
+	withAgentRunner(t, func(_ context.Context, p AgentParam) (*schema.Message, error) {
+		captured = p
+		return &schema.Message{Role: schema.Assistant, Content: "ok"}, nil
+	})
+
+	c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1})
+	_, err := c.Invoke(context.Background(), map[string]any{
+		"user_prompt": "x",
+		"tools": []any{
+			map[string]any{
+				"component_name": "Retrieval",
+				"name":           "Docs Retrieval",
+				"params": map[string]any{
+					"kb_ids": []any{"kb-1"},
+					"top_n":  float64(3),
+				},
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Invoke: %v", err)
+	}
+	if len(captured.Tools) != 1 || captured.Tools[0] != "Retrieval" {
+		t.Fatalf("captured.Tools = %#v, want [Retrieval]", captured.Tools)
+	}
+	params := captured.ToolParams["retrieval"]
+	if params == nil {
+		t.Fatalf("captured.ToolParams missing retrieval: %#v", captured.ToolParams)
+	}
+	ids, ok := params["kb_ids"].([]any)
+	if !ok || len(ids) != 1 || ids[0] != "kb-1" {
+		t.Fatalf("retrieval kb_ids = %#v, want [kb-1]", params["kb_ids"])
+	}
+}
+
+func TestAgent_NewAcceptsCanvasToolObjects(t *testing.T) {
+	cmp, err := New("Agent", map[string]any{
+		"model_id":    "stub",
+		"user_prompt": "x",
+		"tools": []any{
+			map[string]any{
+				"component_name": "Retrieval",
+				"params": map[string]any{
+					"dataset_ids": []any{"kb-1"},
+				},
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("New(Agent): %v", err)
+	}
+	agent, ok := cmp.(*AgentComponent)
+	if !ok {
+		t.Fatalf("New(Agent) returned %T, want *AgentComponent", cmp)
+	}
+	if len(agent.param.Tools) != 1 || agent.param.Tools[0] != "Retrieval" {
+		t.Fatalf("agent.param.Tools = %#v, want [Retrieval]", agent.param.Tools)
+	}
+	if agent.param.ToolParams["retrieval"] == nil {
+		t.Fatalf("agent.param.ToolParams missing retrieval: %#v", agent.param.ToolParams)
+	}
+	if _, err := buildAgentTools(agent.param); err != nil {
+		t.Fatalf("buildAgentTools: %v", err)
+	}
+}
+
 type fakeToolCallingChatModel struct {
 	tools []*schema.ToolInfo
 }
--- a/internal/agent/component/message.go
+++ b/internal/agent/component/message.go
@@ -35,6 +35,7 @@ import (
 	"fmt"
 	"maps"
 	"regexp"
+	"strings"

 	"ragflow/internal/agent/audio"
 	"ragflow/internal/agent/runtime"
@@ -179,6 +180,9 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m
 	if text == "" {
 		text = m.text
 	}
+	if text == "" {
+		text = fallbackMessageText(inputs)
+	}
 	// Message is a display node, not parameter binding. Use the
 	// tolerant resolver (nil refs render as empty string) instead
 	// of runtime.ResolveTemplate — matches the Python canvas.py
@@ -326,6 +330,46 @@ func extractMemoryIDs(inputs map[string]any) []string {
 	return nil
 }

+func fallbackMessageText(inputs map[string]any) string {
+	if inputs == nil {
+		return ""
+	}
+	if text, _ := inputs["formalized_content"].(string); strings.TrimSpace(text) != "" {
+		return text
+	}
+
+	var only string
+	count := 0
+	for key, value := range inputs {
+		if isMessageInfraInput(key) {
+			continue
+		}
+		text, ok := value.(string)
+		if !ok || strings.TrimSpace(text) == "" {
+			continue
+		}
+		only = text
+		count++
+		if count > 1 {
+			return ""
+		}
+	}
+	if count == 1 {
+		return only
+	}
+	return ""
+}
+
+func isMessageInfraInput(key string) bool {
+	switch key {
+	case "state", "__cpn_id__", "__legacy_noop__", "_created_time", "_elapsed_time",
+		"output_format", "voice", "lang", "auto_play", "memory_save", "stream":
+		return true
+	default:
+		return false
+	}
+}
+
 // extractMemoryIDsFromParams looks for a "_memory_ids" hint in
 // the component's stored text — used as a last-ditch fallback
 // when the orchestrator does not re-pass memory_ids. Returns nil
--- a/internal/agent/component/message_test.go
+++ b/internal/agent/component/message_test.go
@@ -112,6 +112,41 @@ func TestMessage_RuntimeContentInput(t *testing.T) {
 	}
 }

+func TestMessage_FormalizedContentFallback(t *testing.T) {
+	c, _ := NewMessageComponent(nil)
+	state := canvas.NewCanvasState("run-5", "task-5")
+	ctx := withStateForTest(context.Background(), state)
+
+	out, err := c.Invoke(ctx, map[string]any{
+		"formalized_content": "retrieved answer",
+		"_created_time":      "2026-07-15T00:00:00Z",
+		"_elapsed_time":      0.01,
+	})
+	if err != nil {
+		t.Fatalf("Invoke: %v", err)
+	}
+	if got, _ := out["content"].(string); got != "retrieved answer" {
+		t.Errorf("content: got %q, want %q", got, "retrieved answer")
+	}
+}
+
+func TestMessage_SingleStringFallback(t *testing.T) {
+	c, _ := NewMessageComponent(nil)
+	state := canvas.NewCanvasState("run-6", "task-6")
+	ctx := withStateForTest(context.Background(), state)
+
+	out, err := c.Invoke(ctx, map[string]any{
+		"value":         "single upstream text",
+		"_elapsed_time": 0.01,
+	})
+	if err != nil {
+		t.Fatalf("Invoke: %v", err)
+	}
+	if got, _ := out["content"].(string); got != "single upstream text" {
+		t.Errorf("content: got %q, want %q", got, "single upstream text")
+	}
+}
+
 // withStateForTest is a thin alias for canvas.WithState kept for
 // readability at the test call sites (also defined in begin_test.go).
 // Same package → same symbol; defining it twice is a compile error,
--- a/internal/agent/component/universe_a_wrappers.go
+++ b/internal/agent/component/universe_a_wrappers.go
@@ -148,6 +148,15 @@ func (c *retrievalComponent) Inputs() map[string]string {
 	}
 }

+func (c *retrievalComponent) GetInputForm() map[string]any {
+	return map[string]any{
+		"query": map[string]any{
+			"name": "Query",
+			"type": "line",
+		},
+	}
+}
+
 func (c *retrievalComponent) Outputs() map[string]string {
 	return map[string]string{
 		"formalized_content": "Rendered chunks for downstream LLM prompts.",