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-20 06:31:02 +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.",
--- a/internal/agent/tool/registry.go
+++ b/internal/agent/tool/registry.go
@@ -49,9 +49,9 @@ var registry = map[string]Factory{
 	"keenable":              buildKeenableTool,
 	"pubmed":                buildPubMedTool,
 	"qweather":              noConfig("qweather", func() einotool.BaseTool { return NewQWeatherTool() }),
-	"retrieval":             noConfig("retrieval", func() einotool.BaseTool { return NewRetrievalTool() }),
-	"search_my_dataset":     noConfig("search_my_dataset", func() einotool.BaseTool { return NewRetrievalTool() }),
-	"search_my_dateset":     noConfig("search_my_dateset", func() einotool.BaseTool { return NewRetrievalTool() }),
+	"retrieval":             buildRetrievalTool,
+	"search_my_dataset":     buildRetrievalTool,
+	"search_my_dateset":     buildRetrievalTool,
 	"searxng":               buildSearXNGTool,
 	"tavily":                buildTavilyTool,
 	"tavily_extract":        buildTavilyExtractTool,
@@ -350,6 +350,58 @@ func buildPubMedTool(params map[string]any) (einotool.BaseTool, error) {
 	return NewPubMedToolWithDefaults(nil, defaults), nil
 }

+func buildRetrievalTool(params map[string]any) (einotool.BaseTool, error) {
+	defaults := retrievalArgs{}
+	for key := range params {
+		switch key {
+		case "dataset_ids", "kb_ids", "top_n", "top_k", "similarity_threshold",
+			"keywords_similarity_weight", "use_kg", "rerank_id", "empty_response",
+			"toc_enhance", "meta_data_filter", "retrieval_from", "memory_ids",
+			"kb_vars", "cross_languages", "function_name", "description", "meta",
+			"inputs", "outputs":
+		default:
+			return nil, fmt.Errorf("agent tool: retrieval tool does not accept node-level param %s", key)
+		}
+	}
+
+	if ids, ok, err := stringSliceParam(params, "dataset_ids"); err != nil {
+		return nil, fmt.Errorf("agent tool: retrieval config: %w", err)
+	} else if ok {
+		defaults.DatasetIDs = ids
+	}
+	if ids, ok, err := stringSliceParam(params, "kb_ids"); err != nil {
+		return nil, fmt.Errorf("agent tool: retrieval config: %w", err)
+	} else if ok {
+		defaults.KBIDs = ids
+		if len(defaults.DatasetIDs) == 0 {
+			defaults.DatasetIDs = ids
+		}
+	}
+	if v, ok := intParam(params, "top_n"); ok {
+		defaults.TopN = v
+	}
+	if raw, exists := params["top_k"]; exists {
+		value, ok := strictInt(raw)
+		if !ok || value <= 0 {
+			return nil, fmt.Errorf("agent tool: retrieval tool requires positive integer node-level param top_k")
+		}
+		defaults.TopK = value
+	}
+	if v, ok := boolParam(params, "use_kg"); ok {
+		defaults.UseKG = v
+	}
+	if v, ok := floatParam(params, "similarity_threshold"); ok {
+		defaults.SimilarityThreshold = v
+	}
+	if v, ok := floatParam(params, "keywords_similarity_weight"); ok {
+		if v < 0 || v > 1 {
+			return nil, fmt.Errorf("agent tool: retrieval tool requires node-level param keywords_similarity_weight in [0,1]")
+		}
+		defaults.KeywordsSimilarityWeight = &v
+	}
+	return NewRetrievalToolWithDefaults(defaults), nil
+}
+
 func buildSearXNGTool(params map[string]any) (einotool.BaseTool, error) {
 	defaults := defaultSearXNGParams()
 	if value, ok := params["top_n"]; ok {
@@ -585,6 +637,52 @@ func intParam(params map[string]any, key string) (int, bool) {
 	}
 }

+func floatParam(params map[string]any, key string) (float64, bool) {
+	v, ok := params[key]
+	if !ok {
+		return 0, false
+	}
+	switch x := v.(type) {
+	case float64:
+		return x, true
+	case float32:
+		return float64(x), true
+	case int:
+		return float64(x), true
+	case int32:
+		return float64(x), true
+	case int64:
+		return float64(x), true
+	default:
+		return 0, false
+	}
+}
+
+func stringSliceParam(params map[string]any, key string) ([]string, bool, error) {
+	v, ok := params[key]
+	if !ok {
+		return nil, false, nil
+	}
+	switch x := v.(type) {
+	case []string:
+		return append([]string(nil), x...), true, nil
+	case []any:
+		out := make([]string, 0, len(x))
+		for _, item := range x {
+			s, ok := item.(string)
+			if !ok {
+				return nil, true, fmt.Errorf("%s must be a string list", key)
+			}
+			if strings.TrimSpace(s) != "" {
+				out = append(out, strings.TrimSpace(s))
+			}
+		}
+		return out, true, nil
+	default:
+		return nil, true, fmt.Errorf("%s must be a string list", key)
+	}
+}
+
 func strictInt(value any) (int, bool) {
 	switch x := value.(type) {
 	case int:
--- a/internal/agent/tool/retrieval.go
+++ b/internal/agent/tool/retrieval.go
@@ -55,10 +55,14 @@ const retrievalToolDescription = "This tool can be utilized for relevant content
 // accept both `query` (canonical) and `dataset_ids` / `use_kg` etc. to
 // match the Python ToolMeta field set.
 type retrievalArgs struct {
-	Query      string   `json:"query"`
-	DatasetIDs []string `json:"dataset_ids,omitempty"`
-	TopN       int      `json:"top_n,omitempty"`
-	UseKG      bool     `json:"use_kg,omitempty"`
+	Query                    string   `json:"query"`
+	DatasetIDs               []string `json:"dataset_ids,omitempty"`
+	KBIDs                    []string `json:"kb_ids,omitempty"`
+	TopN                     int      `json:"top_n,omitempty"`
+	TopK                     int      `json:"top_k,omitempty"`
+	KeywordsSimilarityWeight *float64 `json:"keywords_similarity_weight,omitempty"`
+	UseKG                    bool     `json:"use_kg,omitempty"`
+	SimilarityThreshold      float64  `json:"similarity_threshold,omitempty"`
 }

 // retrievalResult is the JSON shape returned to the model. The `_ERROR`
@@ -86,12 +90,23 @@ type chunkPayload struct {
 // dispatches to the registered RetrievalService via
 // SetRetrievalService. When no service is registered, the call
 // surfaces ErrRetrievalServiceMissing.
-type RetrievalTool struct{}
+type RetrievalTool struct {
+	defaults retrievalArgs
+}

 // NewRetrievalTool returns a RetrievalTool implementing eino's
 // tool.InvokableTool interface.
 func NewRetrievalTool() *RetrievalTool {
-	return &RetrievalTool{}
+	return NewRetrievalToolWithDefaults(retrievalArgs{})
+}
+
+// NewRetrievalToolWithDefaults returns a RetrievalTool with node-level
+// defaults from the Agent tool configuration.
+func NewRetrievalToolWithDefaults(defaults retrievalArgs) *RetrievalTool {
+	if len(defaults.DatasetIDs) == 0 && len(defaults.KBIDs) != 0 {
+		defaults.DatasetIDs = append([]string(nil), defaults.KBIDs...)
+	}
+	return &RetrievalTool{defaults: defaults}
 }

 // Info returns the tool's metadata for the chat model. The schema mirrors
@@ -111,16 +126,36 @@ func (r *RetrievalTool) Info(_ context.Context) (*schema.ToolInfo, error) {
 				Desc:     "Optional list of dataset IDs to restrict the search to.",
 				Required: false,
 			},
+			"kb_ids": {
+				Type:     schema.Array,
+				Desc:     "Optional list of knowledge base IDs to restrict the search to.",
+				Required: false,
+			},
 			"top_n": {
 				Type:     schema.Integer,
 				Desc:     "Number of top chunks to return. Defaults to 8 if omitted.",
 				Required: false,
 			},
+			"top_k": {
+				Type:     schema.Integer,
+				Desc:     "Maximum candidate chunks retrieved before final top_n trimming.",
+				Required: false,
+			},
+			"keywords_similarity_weight": {
+				Type:     schema.Number,
+				Desc:     "Keyword similarity weight in [0,1]; vector similarity weight is 1 - this value.",
+				Required: false,
+			},
 			"use_kg": {
 				Type:     schema.Boolean,
 				Desc:     "GraphRAG toggle. Not supported in Go Canvas (plan ); must be false.",
 				Required: false,
 			},
+			"similarity_threshold": {
+				Type:     schema.Number,
+				Desc:     "Minimum similarity threshold for dataset retrieval.",
+				Required: false,
+			},
 		}),
 	}, nil
 }
@@ -136,10 +171,13 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string
 			return "", fmt.Errorf("retrieval: parse arguments: %w", err)
 		}
 	}
+	args = r.mergeDefaults(args)
 	common.Debug("agent retrieval tool: parsed arguments",
 		zap.String("query", args.Query),
 		zap.Strings("dataset_ids", args.DatasetIDs),
 		zap.Int("top_n", args.TopN),
+		zap.Int("top_k", args.TopK),
+		zap.Float64p("keywords_similarity_weight", args.KeywordsSimilarityWeight),
 		zap.Bool("use_kg", args.UseKG),
 	)

@@ -159,14 +197,14 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string
 	// dev), the chunks flow through normally.
 	svc := GetRetrievalService()
 	chunks, err := svc.Search(ctx, RetrievalRequest{
-		Query:          args.Query,
-		DatasetIDs:     args.DatasetIDs,
-		TopN:           args.TopN,
-		UseKG:          args.UseKG,
-		UseRerank:      false, // future enhancement
-		RerankID:       "",
-		TOCEnhance:     false, // future enhancement
-		MetadataFilter: nil,   // future enhancement
+		Query:                    args.Query,
+		DatasetIDs:               args.DatasetIDs,
+		TopN:                     args.TopN,
+		TopK:                     args.TopK,
+		KeywordsSimilarityWeight: args.KeywordsSimilarityWeight,
+		UseKG:                    args.UseKG,
+		SimilarityThreshold:      args.SimilarityThreshold,
+		TenantID:                 retrievalTenantID(ctx),
 	})
 	if err != nil {
 		return stubJSON(retrievalResult{
@@ -198,16 +236,7 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string
 	// best-effort — when the canvas state is not
 	// attached (e.g. unit tests), we skip silently.
 	if state, _, sErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); sErr == nil && state != nil && len(chunks) > 0 {
-		asMap := make([]map[string]any, 0, len(chunks))
-		for _, c := range chunks {
-			asMap = append(asMap, map[string]any{
-				"id":          c.ID,
-				"content":     c.Content,
-				"document_id": c.DocumentID,
-				"score":       c.Score,
-			})
-		}
-		state.SetRetrievalChunks(asMap)
+		state.SetRetrievalReferences(referenceChunksFromRetrieval(chunks), referenceDocAggsFromRetrieval(chunks))
 	}
 	result, err := stubJSONWithErr(out)
 	if err != nil {
@@ -216,6 +245,29 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string
 	return result, nil
 }

+func (r *RetrievalTool) mergeDefaults(args retrievalArgs) retrievalArgs {
+	if len(args.DatasetIDs) == 0 && len(args.KBIDs) != 0 {
+		args.DatasetIDs = append([]string(nil), args.KBIDs...)
+	}
+	if len(args.DatasetIDs) == 0 && len(r.defaults.DatasetIDs) != 0 {
+		args.DatasetIDs = append([]string(nil), r.defaults.DatasetIDs...)
+	}
+	if args.TopN <= 0 {
+		args.TopN = r.defaults.TopN
+	}
+	if args.TopK <= 0 {
+		args.TopK = r.defaults.TopK
+	}
+	if args.KeywordsSimilarityWeight == nil {
+		args.KeywordsSimilarityWeight = r.defaults.KeywordsSimilarityWeight
+	}
+	if args.SimilarityThreshold <= 0 {
+		args.SimilarityThreshold = r.defaults.SimilarityThreshold
+	}
+	args.UseKG = args.UseKG || r.defaults.UseKG
+	return args
+}
+
 // renderChunks concatenates the retrieved chunks into a human-
 // readable content string. Mirrors Python's
 // `kb_prompt(kbinfos, ...)` format: each chunk gets a header
@@ -228,6 +280,89 @@ func renderChunks(chunks []RetrievalChunk, query string) string {
 	return sb.String()
 }

+func retrievalTenantID(ctx context.Context) string {
+	state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx)
+	if err != nil || state == nil {
+		return ""
+	}
+	if tenantID, _ := state.Sys["tenant_id"].(string); tenantID != "" {
+		return tenantID
+	}
+	userID, _ := state.Sys["user_id"].(string)
+	return userID
+}
+
+func referenceChunksFromRetrieval(chunks []RetrievalChunk) []map[string]any {
+	out := make([]map[string]any, 0, len(chunks))
+	for idx, c := range chunks {
+		id := c.ID
+		if id == "" {
+			id = fmt.Sprint(idx)
+		}
+		chunk := map[string]any{
+			"id":                  id,
+			"chunk_id":            c.ID,
+			"content":             c.Content,
+			"content_with_weight": c.Content,
+			"document_id":         c.DocumentID,
+			"doc_id":              c.DocumentID,
+			"document_name":       c.DocumentName,
+			"docnm_kwd":           c.DocumentName,
+			"dataset_id":          c.DatasetID,
+			"kb_id":               c.DatasetID,
+			"image_id":            c.ImageID,
+			"img_id":              c.ImageID,
+			"similarity":          c.Score,
+			"term_similarity":     c.TermSimilarity,
+			"vector_similarity":   c.VectorSimilarity,
+		}
+		if c.URL != "" {
+			chunk["url"] = c.URL
+			chunk["document_url"] = c.URL
+		}
+		if c.Positions != nil {
+			chunk["positions"] = c.Positions
+			chunk["position_int"] = c.Positions
+		}
+		out = append(out, chunk)
+	}
+	return out
+}
+
+func referenceDocAggsFromRetrieval(chunks []RetrievalChunk) []map[string]any {
+	byDocID := make(map[string]map[string]any)
+	order := make([]string, 0, len(chunks))
+	for _, c := range chunks {
+		if c.DocumentID == "" && c.DocumentName == "" {
+			continue
+		}
+		key := c.DocumentID
+		if key == "" {
+			key = c.DocumentName
+		}
+		agg, exists := byDocID[key]
+		if !exists {
+			agg = map[string]any{
+				"count":    0,
+				"doc_id":   c.DocumentID,
+				"doc_name": c.DocumentName,
+			}
+			if c.URL != "" {
+				agg["url"] = c.URL
+			}
+			byDocID[key] = agg
+			order = append(order, key)
+		}
+		agg["count"] = agg["count"].(int) + 1
+	}
+
+	out := make([]map[string]any, 0, len(order))
+	for _, key := range order {
+		out = append(out, byDocID[key])
+	}
+	return out
+}
+
 // stubJSONWithErr is the (string, error) variant for call sites
 // that need to propagate marshal failures.
 func stubJSONWithErr(r retrievalResult) (string, error) {
--- a/internal/agent/tool/retrieval_nlp.go
+++ b/internal/agent/tool/retrieval_nlp.go
@@ -35,19 +35,27 @@
 //   tool.RetrievalRequest.Query      → nlp.RetrievalRequest.Question
 //   tool.RetrievalRequest.DatasetIDs → nlp.RetrievalRequest.KbIDs
 //   tool.RetrievalRequest.TopN       → nlp.RetrievalRequest.PageSize
-//                                       (Page=1, Top=TopN*4 so rerank
+//   tool.RetrievalRequest.TopK       → nlp.RetrievalRequest.Top
+//                                       (fallback Top=TopN*4 so rerank
 //                                        has headroom)
+//   tool.RetrievalRequest.KeywordsSimilarityWeight
+//                                    → nlp.RetrievalRequest.VectorSimilarityWeight
+//                                       as 1-keyword weight
 //   tool.RetrievalRequest.UseKG      → ErrGraphRAGNotSupported (out of
 //                                       scope per plan  + §9 Q3)
 //
 // Chunk shape translation: nlp's Chunks are []map[string]any with
 // keys chunk_id, doc_id, docnm_kwd, content_with_weight,
 // content_ltks, similarity, term_similarity, vector_similarity. The
-// tool side wants a flat RetrievalChunk{ID, Content, DocumentID,
-// Score}. We pick the most user-facing fields:
+// tool side wants a flat RetrievalChunk with the fields needed for both
+// display and the frontend reference strip. We pick the most user-facing fields:
 //   - ID         ← chunk_id
 //   - Content    ← content_with_weight (fallback to content_ltks)
 //   - DocumentID ← doc_id
+//   - DocumentName ← docnm_kwd
+//   - DatasetID  ← kb_id
+//   - ImageID    ← image_id/img_id
+//   - Positions  ← positions/position_int
 //   - Score      ← similarity (fallback to avg of term+vector)
 //
 // Defensive defaults: missing or wrong-typed chunk fields become
@@ -59,6 +67,8 @@ package tool

 import (
 	"context"
+	"fmt"
+	"strings"

 	"ragflow/internal/dao"
 	"ragflow/internal/engine"
@@ -125,27 +135,11 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, req RetrievalRequest)
 		topN = 8
 	}

-	// nlp.Retrieval applies its own defaults for SimilarityThreshold
-	// (0.2), VectorSimilarityWeight (0.3), RankFeature, etc. We
-	// surface only the fields the agent tool actually controls:
-	// Page=1, PageSize=TopN, KbIDs=DatasetIDs, Top=TopN*4 (rerank
-	// headroom — matches the chat_session.go call pattern).
-	nlpReq := &nlp.RetrievalRequest{
-		Question:  req.Query,
-		TenantIDs: a.resolveTenantIDs(req),
-		KbIDs:     append([]string(nil), req.DatasetIDs...),
-		Page:      1,
-		PageSize:  topN,
-		Aggs:      boolPtr(false),
-		Highlight: boolPtr(false),
-	}
-	if topN > 0 {
-		rerankBudget := topN * 4
-		nlpReq.Top = &rerankBudget
-	}
-	if req.SimilarityThreshold > 0 {
-		nlpReq.SimilarityThreshold = &req.SimilarityThreshold
+	tenantIDs, err := a.resolveTenantIDs(req)
+	if err != nil {
+		return nil, err
 	}
+	nlpReq := nlpRequestFromRetrieval(req, tenantIDs, topN)

 	res, err := a.svc.Retrieval(ctx, nlpReq)
 	if err != nil {
@@ -161,10 +155,37 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, req RetrievalRequest)
 	return out, nil
 }

-func (a *NLPRetrievalAdapter) resolveTenantIDs(req RetrievalRequest) []string {
+func nlpRequestFromRetrieval(req RetrievalRequest, tenantIDs []string, topN int) *nlp.RetrievalRequest {
+	nlpReq := &nlp.RetrievalRequest{
+		Question:  req.Query,
+		TenantIDs: append([]string(nil), tenantIDs...),
+		KbIDs:     append([]string(nil), req.DatasetIDs...),
+		Page:      1,
+		PageSize:  topN,
+		Aggs:      boolPtr(false),
+		Highlight: boolPtr(false),
+	}
+	if req.TopK > 0 {
+		nlpReq.Top = &req.TopK
+	} else if topN > 0 {
+		rerankBudget := topN * 4
+		nlpReq.Top = &rerankBudget
+	}
+	if req.SimilarityThreshold > 0 {
+		nlpReq.SimilarityThreshold = &req.SimilarityThreshold
+	}
+	if req.KeywordsSimilarityWeight != nil {
+		vectorSimilarityWeight := 1 - *req.KeywordsSimilarityWeight
+		nlpReq.VectorSimilarityWeight = &vectorSimilarityWeight
+	}
+	return nlpReq
+}
+
+func (a *NLPRetrievalAdapter) resolveTenantIDs(req RetrievalRequest) ([]string, error) {
 	seen := map[string]struct{}{}
 	tenantIDs := make([]string, 0, 1)
 	appendTenantID := func(tenantID string) {
+		tenantID = strings.TrimSpace(tenantID)
 		if tenantID == "" {
 			return
 		}
@@ -176,7 +197,25 @@ func (a *NLPRetrievalAdapter) resolveTenantIDs(req RetrievalRequest) []string {
 	}

 	appendTenantID(req.TenantID)
-	return tenantIDs
+	datasetIDs := compactStrings(req.DatasetIDs)
+	if len(datasetIDs) == 0 || a == nil || a.kbDAO == nil {
+		return tenantIDs, nil
+	}
+
+	kbs, err := a.kbDAO.GetByIDs(datasetIDs)
+	if err != nil {
+		return nil, fmt.Errorf("retrieval: resolve dataset tenants: %w", err)
+	}
+	for _, kb := range kbs {
+		if kb == nil {
+			continue
+		}
+		appendTenantID(kb.TenantID)
+	}
+	if len(tenantIDs) == 0 {
+		return nil, fmt.Errorf("retrieval: no valid knowledge bases found for dataset_ids %v", datasetIDs)
+	}
+	return tenantIDs, nil
 }

 // translateChunk converts one nlp chunk map into a RetrievalChunk.
@@ -185,10 +224,17 @@ func (a *NLPRetrievalAdapter) resolveTenantIDs(req RetrievalRequest) []string {
 // can't break the whole result list.
 func translateChunk(raw map[string]any) RetrievalChunk {
 	return RetrievalChunk{
-		ID:         stringFromMap(raw, "chunk_id"),
-		Content:    contentFromMap(raw),
-		DocumentID: stringFromMap(raw, "doc_id"),
-		Score:      scoreFromMap(raw),
+		ID:               stringFromMap(raw, "chunk_id"),
+		Content:          contentFromMap(raw),
+		DocumentID:       stringFromMap(raw, "doc_id"),
+		DocumentName:     stringFromMap(raw, "docnm_kwd"),
+		DatasetID:        stringFromMap(raw, "kb_id"),
+		ImageID:          firstStringFromMap(raw, "image_id", "img_id"),
+		URL:              firstStringFromMap(raw, "url", "document_url", "doc_url"),
+		Positions:        firstValueFromMap(raw, "positions", "position_int"),
+		Score:            scoreFromMap(raw),
+		TermSimilarity:   scoreValueFromMap(raw, "term_similarity"),
+		VectorSimilarity: scoreValueFromMap(raw, "vector_similarity"),
 	}
 }

@@ -203,6 +249,24 @@ func stringFromMap(raw map[string]any, key string) string {
 	return ""
 }

+func firstStringFromMap(raw map[string]any, keys ...string) string {
+	for _, key := range keys {
+		if value := stringFromMap(raw, key); value != "" {
+			return value
+		}
+	}
+	return ""
+}
+
+func firstValueFromMap(raw map[string]any, keys ...string) any {
+	for _, key := range keys {
+		if value, ok := raw[key]; ok && value != nil {
+			return value
+		}
+	}
+	return nil
+}
+
 // contentFromMap picks the most user-facing content field. nlp
 // chunks carry content_with_weight (the highlightable string) and
 // content_ltks (the tokenised form). content_with_weight is what
@@ -238,6 +302,11 @@ func scoreFromMap(raw map[string]any) float64 {
 	return 0
 }

+func scoreValueFromMap(raw map[string]any, key string) float64 {
+	value, _ := numberFromMap(raw, key)
+	return value
+}
+
 // numberFromMap returns raw[key].(float64) with a tolerant path
 // for ints. JSON unmarshaling can produce either.
 func numberFromMap(raw map[string]any, key string) (float64, bool) {
--- a/internal/agent/tool/retrieval_nlp_test.go
+++ b/internal/agent/tool/retrieval_nlp_test.go
@@ -30,6 +30,8 @@ import (
 	"context"
 	"math"
 	"testing"
+
+	"ragflow/internal/entity"
 )

 // floatEqual compares two floats with a small epsilon so
@@ -44,6 +46,8 @@ func TestTranslateChunk_FullFields(t *testing.T) {
 		"doc_id":              "doc-7",
 		"docnm_kwd":           "report.pdf",
 		"kb_id":               "kb-1",
+		"image_id":            "img-1",
+		"position_int":        [][]float64{{1, 2, 3, 4}},
 		"content_with_weight": "the answer is 42",
 		"content_ltks":        "answer 42",
 		"similarity":          0.87,
@@ -60,9 +64,27 @@ func TestTranslateChunk_FullFields(t *testing.T) {
 	if got.DocumentID != "doc-7" {
 		t.Errorf("DocumentID = %q, want \"doc-7\"", got.DocumentID)
 	}
+	if got.DocumentName != "report.pdf" {
+		t.Errorf("DocumentName = %q, want \"report.pdf\"", got.DocumentName)
+	}
+	if got.DatasetID != "kb-1" {
+		t.Errorf("DatasetID = %q, want \"kb-1\"", got.DatasetID)
+	}
+	if got.ImageID != "img-1" {
+		t.Errorf("ImageID = %q, want \"img-1\"", got.ImageID)
+	}
+	if got.Positions == nil {
+		t.Errorf("Positions is nil, want position_int payload")
+	}
 	if got.Score != 0.87 {
 		t.Errorf("Score = %v, want 0.87 (similarity preferred)", got.Score)
 	}
+	if got.TermSimilarity != 0.5 {
+		t.Errorf("TermSimilarity = %v, want 0.5", got.TermSimilarity)
+	}
+	if got.VectorSimilarity != 0.9 {
+		t.Errorf("VectorSimilarity = %v, want 0.9", got.VectorSimilarity)
+	}
 }

 // TestTranslateChunk_ContentFallback: when content_with_weight is
@@ -208,12 +230,64 @@ func TestNewNLPRetrievalAdapter_NilService(t *testing.T) {
 	}
 }

+func TestNLPRequestFromRetrieval_ThreadsSearchControls(t *testing.T) {
+	keywordWeight := 0.25
+	got := nlpRequestFromRetrieval(RetrievalRequest{
+		Query:                    "hi",
+		DatasetIDs:               []string{"kb-1"},
+		TopN:                     3,
+		TopK:                     99,
+		KeywordsSimilarityWeight: &keywordWeight,
+		SimilarityThreshold:      0.42,
+	}, []string{"tenant-a"}, 3)
+
+	if got.Question != "hi" {
+		t.Fatalf("Question=%q want hi", got.Question)
+	}
+	if len(got.TenantIDs) != 1 || got.TenantIDs[0] != "tenant-a" {
+		t.Fatalf("TenantIDs=%v want [tenant-a]", got.TenantIDs)
+	}
+	if len(got.KbIDs) != 1 || got.KbIDs[0] != "kb-1" {
+		t.Fatalf("KbIDs=%v want [kb-1]", got.KbIDs)
+	}
+	if got.Page != 1 || got.PageSize != 3 {
+		t.Fatalf("Page/PageSize=%d/%d want 1/3", got.Page, got.PageSize)
+	}
+	if got.Top == nil || *got.Top != 99 {
+		t.Fatalf("Top=%v want 99", got.Top)
+	}
+	if got.SimilarityThreshold == nil || *got.SimilarityThreshold != 0.42 {
+		t.Fatalf("SimilarityThreshold=%v want 0.42", got.SimilarityThreshold)
+	}
+	if got.VectorSimilarityWeight == nil || !floatEqual(*got.VectorSimilarityWeight, 0.75) {
+		t.Fatalf("VectorSimilarityWeight=%v want 0.75", got.VectorSimilarityWeight)
+	}
+}
+
+func TestNLPRequestFromRetrieval_FallsBackToTopNHeadroom(t *testing.T) {
+	got := nlpRequestFromRetrieval(RetrievalRequest{
+		Query:      "hi",
+		DatasetIDs: []string{"kb-1"},
+		TopN:       3,
+	}, []string{"tenant-a"}, 3)
+
+	if got.Top == nil || *got.Top != 12 {
+		t.Fatalf("Top=%v want 12", got.Top)
+	}
+	if got.VectorSimilarityWeight != nil {
+		t.Fatalf("VectorSimilarityWeight=%v want nil", got.VectorSimilarityWeight)
+	}
+}
+
 func TestNLPRetrievalAdapter_ResolveTenantIDsStaysWithinRequestTenant(t *testing.T) {
 	a := &NLPRetrievalAdapter{}
-	got := a.resolveTenantIDs(RetrievalRequest{
+	got, err := a.resolveTenantIDs(RetrievalRequest{
 		TenantID:   "tenant-a",
 		DatasetIDs: []string{"kb-1", "kb-2", "kb-missing"},
 	})
+	if err != nil {
+		t.Fatalf("resolveTenantIDs: %v", err)
+	}

 	if len(got) != 1 {
 		t.Fatalf("tenantIDs len=%d want 1, got=%v", len(got), got)
@@ -222,3 +296,53 @@ func TestNLPRetrievalAdapter_ResolveTenantIDsStaysWithinRequestTenant(t *testing
 		t.Fatalf("tenantIDs=%v want [tenant-a]", got)
 	}
 }
+
+func TestNLPRetrievalAdapter_ResolveTenantIDsFromDatasetIDs(t *testing.T) {
+	a := &NLPRetrievalAdapter{
+		kbDAO: fakeKnowledgebaseLookup{
+			kbs: []*entity.Knowledgebase{
+				{ID: "kb-1", TenantID: "tenant-a"},
+				{ID: "kb-2", TenantID: "tenant-b"},
+				{ID: "kb-3", TenantID: "tenant-a"},
+			},
+		},
+	}
+	got, err := a.resolveTenantIDs(RetrievalRequest{
+		DatasetIDs: []string{"kb-1", "kb-2", "kb-3", " "},
+	})
+	if err != nil {
+		t.Fatalf("resolveTenantIDs: %v", err)
+	}
+	if len(got) != 2 || got[0] != "tenant-a" || got[1] != "tenant-b" {
+		t.Fatalf("tenantIDs=%v want [tenant-a tenant-b]", got)
+	}
+}
+
+func TestNLPRetrievalAdapter_ResolveTenantIDsKeepsRequestTenantFirst(t *testing.T) {
+	a := &NLPRetrievalAdapter{
+		kbDAO: fakeKnowledgebaseLookup{
+			kbs: []*entity.Knowledgebase{
+				{ID: "kb-1", TenantID: "tenant-b"},
+			},
+		},
+	}
+	got, err := a.resolveTenantIDs(RetrievalRequest{
+		TenantID:   "tenant-a",
+		DatasetIDs: []string{"kb-1"},
+	})
+	if err != nil {
+		t.Fatalf("resolveTenantIDs: %v", err)
+	}
+	if len(got) != 2 || got[0] != "tenant-a" || got[1] != "tenant-b" {
+		t.Fatalf("tenantIDs=%v want [tenant-a tenant-b]", got)
+	}
+}
+
+type fakeKnowledgebaseLookup struct {
+	kbs []*entity.Knowledgebase
+	err error
+}
+
+func (f fakeKnowledgebaseLookup) GetByIDs(_ []string) ([]*entity.Knowledgebase, error) {
+	return f.kbs, f.err
+}
--- a/internal/agent/tool/retrieval_service.go
+++ b/internal/agent/tool/retrieval_service.go
@@ -15,11 +15,9 @@
 //

 // RetrievalService is the abstract interface for retrieval.
-// The full Python `Dealer.search()` surface (KB + memory +
-// rerank + cross-language + toc_enhance + metadata filter +
-// GraphRAG) lands incrementally as the corresponding
-// enhancements fill in. For now the interface is minimal —
-// just the entry point the RetrievalTool calls.
+// The Go path exposes only the parameters currently threaded into
+// internal/service/nlp. Extra Python-only options should be added
+// here only when the adapter can pass them through.
 package tool

 import (
@@ -33,23 +31,28 @@ import (
 // full Chunk type (with document_id, docnm_kwd, position, etc.)
 // lives in internal/entity and is wired in by a follow-up phase.
 type RetrievalChunk struct {
-	ID         string
-	Content    string
-	DocumentID string
-	Score      float64
+	ID               string
+	Content          string
+	DocumentID       string
+	DocumentName     string
+	DatasetID        string
+	ImageID          string
+	URL              string
+	Positions        any
+	Score            float64
+	TermSimilarity   float64
+	VectorSimilarity float64
 }

 // RetrievalRequest is the input to RetrievalService.Search.
 type RetrievalRequest struct {
-	Query               string
-	DatasetIDs          []string
-	TopN                int
-	UseKG               bool
-	UseRerank           bool
-	RerankID            string
-	TOCEnhance          bool
-	MetadataFilter      map[string]string
-	SimilarityThreshold float64
+	Query                    string
+	DatasetIDs               []string
+	TopN                     int
+	TopK                     int
+	KeywordsSimilarityWeight *float64
+	UseKG                    bool
+	SimilarityThreshold      float64
 	// TenantID is the calling tenant (== user_id in RAGFlow's data model).
 	// Optional for the nlp adapter; the KG adapter uses it to resolve the
 	// tenant's default chat + embedding models. Reads from
--- a/internal/agent/tool/retrieval_test.go
+++ b/internal/agent/tool/retrieval_test.go
@@ -22,6 +22,8 @@ import (
 	"errors"
 	"strings"
 	"testing"
+
+	"ragflow/internal/agent/runtime"
 )

 func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) {
@@ -105,3 +107,210 @@ func TestRetrieval_EmptyArgsIsHandled(t *testing.T) {
 		t.Fatalf("err = %v, want ErrRetrievalServiceMissing", err)
 	}
 }
+
+func TestRetrieval_PassesTenantIDFromCanvasState(t *testing.T) {
+	prev := GetRetrievalService()
+	svc := &capturingRetrievalService{}
+	SetRetrievalService(svc)
+	t.Cleanup(func() { SetRetrievalService(prev) })
+
+	state := runtime.NewCanvasState("run-1", "task-1")
+	state.Sys["tenant_id"] = "tenant-1"
+	ctx := runtime.WithState(context.Background(), state)
+
+	rt := NewRetrievalTool()
+	_, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`)
+	if err != nil {
+		t.Fatalf("InvokableRun: %v", err)
+	}
+	if svc.req.TenantID != "tenant-1" {
+		t.Fatalf("TenantID=%q want tenant-1", svc.req.TenantID)
+	}
+}
+
+func TestRetrieval_PassesUserIDWhenTenantIDMissing(t *testing.T) {
+	prev := GetRetrievalService()
+	svc := &capturingRetrievalService{}
+	SetRetrievalService(svc)
+	t.Cleanup(func() { SetRetrievalService(prev) })
+
+	state := runtime.NewCanvasState("run-1", "task-1")
+	state.Sys["user_id"] = "user-1"
+	ctx := runtime.WithState(context.Background(), state)
+
+	rt := NewRetrievalTool()
+	_, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`)
+	if err != nil {
+		t.Fatalf("InvokableRun: %v", err)
+	}
+	if svc.req.TenantID != "user-1" {
+		t.Fatalf("TenantID=%q want user-1", svc.req.TenantID)
+	}
+}
+
+func TestRetrieval_UsesNodeParamsAsDefaults(t *testing.T) {
+	prev := GetRetrievalService()
+	svc := &capturingRetrievalService{}
+	SetRetrievalService(svc)
+	t.Cleanup(func() { SetRetrievalService(prev) })
+
+	built, err := BuildByName("retrieval", map[string]any{
+		"kb_ids":                     []any{"kb-1"},
+		"top_n":                      float64(3),
+		"top_k":                      float64(99),
+		"keywords_similarity_weight": 0.7,
+		"similarity_threshold":       0.42,
+	})
+	if err != nil {
+		t.Fatalf("BuildByName(retrieval): %v", err)
+	}
+	rt, ok := built.(*RetrievalTool)
+	if !ok {
+		t.Fatalf("BuildByName(retrieval) returned %T, want *RetrievalTool", built)
+	}
+
+	_, err = rt.InvokableRun(context.Background(), `{"query":"hello"}`)
+	if err != nil {
+		t.Fatalf("InvokableRun: %v", err)
+	}
+	if len(svc.req.DatasetIDs) != 1 || svc.req.DatasetIDs[0] != "kb-1" {
+		t.Fatalf("DatasetIDs=%#v want [kb-1]", svc.req.DatasetIDs)
+	}
+	if svc.req.TopN != 3 {
+		t.Fatalf("TopN=%d want 3", svc.req.TopN)
+	}
+	if svc.req.TopK != 99 {
+		t.Fatalf("TopK=%d want 99", svc.req.TopK)
+	}
+	if svc.req.KeywordsSimilarityWeight == nil || *svc.req.KeywordsSimilarityWeight != 0.7 {
+		t.Fatalf("KeywordsSimilarityWeight=%v want 0.7", svc.req.KeywordsSimilarityWeight)
+	}
+	if svc.req.SimilarityThreshold != 0.42 {
+		t.Fatalf("SimilarityThreshold=%v want 0.42", svc.req.SimilarityThreshold)
+	}
+}
+
+func TestRetrieval_IgnoresPythonOnlyNodeParams(t *testing.T) {
+	t.Parallel()
+
+	built, err := BuildByName("retrieval", map[string]any{
+		"rerank_id":        "rerank-1",
+		"toc_enhance":      true,
+		"meta_data_filter": map[string]any{"method": "manual"},
+		"empty_response":   "empty",
+		"retrieval_from":   "database",
+		"memory_ids":       []any{"memory-1"},
+		"kb_vars":          map[string]any{"x": "y"},
+		"cross_languages":  []any{"English"},
+	})
+	if err != nil {
+		t.Fatalf("BuildByName(retrieval): %v", err)
+	}
+	rt, ok := built.(*RetrievalTool)
+	if !ok {
+		t.Fatalf("BuildByName(retrieval) returned %T, want *RetrievalTool", built)
+	}
+	if rt.defaults.TopN != 0 || rt.defaults.TopK != 0 || rt.defaults.KeywordsSimilarityWeight != nil {
+		t.Fatalf("python-only params should not mutate retrieval defaults: %#v", rt.defaults)
+	}
+}
+
+func TestRetrieval_IgnoresCanvasMetadataNodeParams(t *testing.T) {
+	built, err := BuildByName("retrieval", map[string]any{
+		"kb_ids":  []any{"kb-1"},
+		"inputs":  map[string]any{"query": "upstream"},
+		"outputs": map[string]any{"formalized_content": "downstream"},
+	})
+	if err != nil {
+		t.Fatalf("BuildByName(retrieval): %v", err)
+	}
+	rt, ok := built.(*RetrievalTool)
+	if !ok {
+		t.Fatalf("BuildByName(retrieval) returned %T, want *RetrievalTool", built)
+	}
+	if len(rt.defaults.DatasetIDs) != 1 || rt.defaults.DatasetIDs[0] != "kb-1" {
+		t.Fatalf("defaults.DatasetIDs=%#v want [kb-1]", rt.defaults.DatasetIDs)
+	}
+}
+
+func TestRetrieval_ModelArgsOverrideNodeDatasetIDs(t *testing.T) {
+	prev := GetRetrievalService()
+	svc := &capturingRetrievalService{}
+	SetRetrievalService(svc)
+	t.Cleanup(func() { SetRetrievalService(prev) })
+
+	rt := NewRetrievalToolWithDefaults(retrievalArgs{DatasetIDs: []string{"kb-default"}, TopN: 3})
+	_, err := rt.InvokableRun(context.Background(), `{"query":"hello","dataset_ids":["kb-call"],"top_n":5}`)
+	if err != nil {
+		t.Fatalf("InvokableRun: %v", err)
+	}
+	if len(svc.req.DatasetIDs) != 1 || svc.req.DatasetIDs[0] != "kb-call" {
+		t.Fatalf("DatasetIDs=%#v want [kb-call]", svc.req.DatasetIDs)
+	}
+	if svc.req.TopN != 5 {
+		t.Fatalf("TopN=%d want 5", svc.req.TopN)
+	}
+}
+
+func TestRetrieval_RecordsFrontendReferencePayload(t *testing.T) {
+	prev := GetRetrievalService()
+	SetRetrievalService(staticRetrievalService{chunks: []RetrievalChunk{
+		{
+			ID:               "ck-1",
+			Content:          "answer",
+			DocumentID:       "doc-1",
+			DocumentName:     "paper.pdf",
+			DatasetID:        "kb-1",
+			ImageID:          "img-1",
+			Positions:        [][]float64{{1, 2, 3, 4}},
+			Score:            0.9,
+			TermSimilarity:   0.7,
+			VectorSimilarity: 0.8,
+		},
+	}})
+	t.Cleanup(func() { SetRetrievalService(prev) })
+
+	state := runtime.NewCanvasState("run-1", "task-1")
+	ctx := runtime.WithState(context.Background(), state)
+
+	rt := NewRetrievalTool()
+	_, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`)
+	if err != nil {
+		t.Fatalf("InvokableRun: %v", err)
+	}
+
+	reference := state.GetRetrievalReference()
+	chunks, _ := reference["chunks"].([]any)
+	if len(chunks) != 1 {
+		t.Fatalf("reference chunks length = %d, want 1", len(chunks))
+	}
+	chunk, _ := chunks[0].(map[string]any)
+	if chunk["document_name"] != "paper.pdf" || chunk["image_id"] != "img-1" {
+		t.Fatalf("reference chunk = %#v, want document_name/image_id", chunk)
+	}
+	docAggs, _ := reference["doc_aggs"].([]any)
+	if len(docAggs) != 1 {
+		t.Fatalf("reference doc_aggs length = %d, want 1", len(docAggs))
+	}
+	docAgg, _ := docAggs[0].(map[string]any)
+	if docAgg["doc_id"] != "doc-1" || docAgg["doc_name"] != "paper.pdf" || docAgg["count"] != 1 {
+		t.Fatalf("reference doc_agg = %#v, want doc metadata", docAgg)
+	}
+}
+
+type capturingRetrievalService struct {
+	req RetrievalRequest
+}
+
+func (s *capturingRetrievalService) Search(_ context.Context, req RetrievalRequest) ([]RetrievalChunk, error) {
+	s.req = req
+	return []RetrievalChunk{{ID: "ck-1", Content: "answer"}}, nil
+}
+
+type staticRetrievalService struct {
+	chunks []RetrievalChunk
+}
+
+func (s staticRetrievalService) Search(_ context.Context, _ RetrievalRequest) ([]RetrievalChunk, error) {
+	return s.chunks, nil
+}