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-16 12:47:19 +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
+}
--- a/internal/entity/models/deepseek.go
+++ b/internal/entity/models/deepseek.go
@@ -24,6 +24,7 @@ import (
 	"io"
 	"net/http"
 	"ragflow/internal/common"
+	"sort"
 	"strconv"
 	"strings"
 )
@@ -70,10 +71,17 @@ func (d *DeepSeekModel) ChatWithMessages(modelName string, messages []Message, a
 	// Convert messages to the format expected by API
 	apiMessages := make([]map[string]interface{}, len(messages))
 	for i, msg := range messages {
-		apiMessages[i] = map[string]interface{}{
+		apiMsg := map[string]interface{}{
 			"role":    msg.Role,
 			"content": msg.Content,
 		}
+		if msg.ToolCallID != "" {
+			apiMsg["tool_call_id"] = msg.ToolCallID
+		}
+		if len(msg.ToolCalls) > 0 {
+			apiMsg["tool_calls"] = msg.ToolCalls
+		}
+		apiMessages[i] = apiMsg
 	}

 	// Build request body
@@ -105,6 +113,13 @@ func (d *DeepSeekModel) ChatWithMessages(modelName string, messages []Message, a
 			reqBody["stop"] = *chatModelConfig.Stop
 		}

+		if chatModelConfig.Tools != nil {
+			reqBody["tools"] = chatModelConfig.Tools
+		}
+		if chatModelConfig.ToolChoice != nil {
+			reqBody["tool_choice"] = *chatModelConfig.ToolChoice
+		}
+
 		if chatModelConfig.Thinking != nil {
 			if *chatModelConfig.Thinking {
 				var thinkingFlag string
@@ -198,26 +213,38 @@ func (d *DeepSeekModel) ChatWithMessages(modelName string, messages []Message, a
 		return nil, fmt.Errorf("invalid message format")
 	}

-	content, ok := messageMap["content"].(string)
-	if !ok {
-		return nil, fmt.Errorf("invalid content format")
+	var content string
+	if c, ok := messageMap["content"].(string); ok {
+		content = c
 	}

 	var reasonContent string
-	if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking {
-		reasonContent, ok = messageMap["reasoning_content"].(string)
-		if !ok {
-			return nil, fmt.Errorf("invalid content format")
-		}
+	if rc, ok := messageMap["reasoning_content"].(string); ok {
+		reasonContent = rc
 		// if first char of reasonContent is \n remove the '\n'
 		if reasonContent != "" && reasonContent[0] == '\n' {
 			reasonContent = reasonContent[1:]
 		}
 	}

+	var toolCalls []map[string]interface{}
+	if tcs, ok := messageMap["tool_calls"].([]interface{}); ok {
+		for _, tc := range tcs {
+			if tcMap, ok := tc.(map[string]interface{}); ok {
+				toolCalls = append(toolCalls, tcMap)
+			}
+		}
+	}
+
 	chatResponse := &ChatResponse{
 		Answer:        &content,
 		ReasonContent: &reasonContent,
+		ToolCalls:     toolCalls,
+	}
+	if pt, ct, tt := extractUsageFromMap(result); tt > 0 {
+		chatResponse.Usage = &ChatUsage{
+			PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt,
+		}
 	}

 	return chatResponse, nil
@@ -242,10 +269,17 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 	// Convert messages to API format
 	apiMessages := make([]map[string]interface{}, len(messages))
 	for i, msg := range messages {
-		apiMessages[i] = map[string]interface{}{
+		apiMsg := map[string]interface{}{
 			"role":    msg.Role,
 			"content": msg.Content,
 		}
+		if msg.ToolCallID != "" {
+			apiMsg["tool_call_id"] = msg.ToolCallID
+		}
+		if len(msg.ToolCalls) > 0 {
+			apiMsg["tool_calls"] = msg.ToolCalls
+		}
+		apiMessages[i] = apiMsg
 	}

 	// Build request body with streaming enabled
@@ -280,6 +314,12 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 		if chatModelConfig.Stop != nil {
 			reqBody["stop"] = *chatModelConfig.Stop
 		}
+		if chatModelConfig.Tools != nil {
+			reqBody["tools"] = chatModelConfig.Tools
+		}
+		if chatModelConfig.ToolChoice != nil {
+			reqBody["tool_choice"] = *chatModelConfig.ToolChoice
+		}

 		if chatModelConfig.Thinking != nil {
 			if *chatModelConfig.Thinking {
@@ -351,8 +391,8 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 		return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}

-	// SSE parsing: read line by line
 	sawTerminal := false
+	accumulatedToolCalls := make(map[int]map[string]interface{})
 	done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
 		common.Info(fmt.Sprintf("%v", event))

@@ -371,6 +411,37 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 			return nil
 		}

+		if tcs, ok := delta["tool_calls"].([]interface{}); ok {
+			for _, tc := range tcs {
+				tcMap, ok := tc.(map[string]interface{})
+				if !ok {
+					continue
+				}
+				idxF, ok := tcMap["index"].(float64)
+				if !ok {
+					continue
+				}
+				idx := int(idxF)
+				existing, hasExisting := accumulatedToolCalls[idx]
+				if hasExisting {
+					if fn, ok := tcMap["function"].(map[string]interface{}); ok {
+						if args, ok := fn["arguments"].(string); ok {
+							if ef, ok := existing["function"].(map[string]interface{}); ok {
+								if ea, ok := ef["arguments"].(string); ok {
+									ef["arguments"] = ea + args
+								} else {
+									ef["arguments"] = args
+								}
+							}
+						}
+					}
+				} else {
+					accumulatedToolCalls[idx] = cloneMap(tcMap)
+				}
+			}
+			return nil
+		}
+
 		content, ok := delta["content"].(string)
 		if ok && content != "" {
 			if err := sender(&content, nil); err != nil {
@@ -398,6 +469,19 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(modelName string, messages []Mess
 		return fmt.Errorf("deepseek: stream ended before [DONE] or finish_reason")
 	}

+	if len(accumulatedToolCalls) > 0 && chatModelConfig != nil {
+		indices := make([]int, 0, len(accumulatedToolCalls))
+		for idx := range accumulatedToolCalls {
+			indices = append(indices, idx)
+		}
+		sort.Ints(indices)
+		tcs := make([]map[string]interface{}, 0, len(accumulatedToolCalls))
+		for _, idx := range indices {
+			tcs = append(tcs, accumulatedToolCalls[idx])
+		}
+		chatModelConfig.ToolCallsResult = &tcs
+	}
+
 	// Send [DONE] marker for OpenAI compatibility
 	endOfStream := "[DONE]"
 	return sender(&endOfStream, nil)
--- a/internal/entity/models/deepseek_test.go
+++ b/internal/entity/models/deepseek_test.go
@@ -0,0 +1,161 @@
+package models
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func newDeepSeekForTest(baseURL string) *DeepSeekModel {
+	return NewDeepSeekModel(
+		map[string]string{"default": baseURL},
+		URLSuffix{
+			Chat:   "chat/completions",
+			Models: "models",
+		},
+	)
+}
+
+func TestDeepSeekChatWithMessagesSupportsToolCalls(t *testing.T) {
+	var requestBody map[string]interface{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/chat/completions" {
+			t.Errorf("path=%s, want /chat/completions", r.URL.Path)
+		}
+		if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
+			t.Errorf("Authorization=%q, want Bearer test-key", got)
+		}
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Errorf("decode request: %v", err)
+			http.Error(w, "bad request", http.StatusBadRequest)
+			return
+		}
+		_ = json.NewEncoder(w).Encode(map[string]interface{}{
+			"choices": []map[string]interface{}{
+				{
+					"message": map[string]interface{}{
+						"role":    "assistant",
+						"content": nil,
+						"tool_calls": []map[string]interface{}{
+							{
+								"id":   "call-1",
+								"type": "function",
+								"function": map[string]interface{}{
+									"name":      "search_my_dateset",
+									"arguments": `{"query":"marigold"}`,
+								},
+							},
+						},
+					},
+				},
+			},
+			"usage": map[string]interface{}{
+				"prompt_tokens":     1,
+				"completion_tokens": 2,
+				"total_tokens":      3,
+			},
+		})
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	toolChoice := "auto"
+	resp, err := newDeepSeekForTest(srv.URL).ChatWithMessages(
+		"deepseek-chat",
+		[]Message{
+			{Role: "user", Content: "what is marigold"},
+		},
+		&APIConfig{ApiKey: &apiKey},
+		&ChatConfig{
+			Tools: []map[string]interface{}{
+				{
+					"type": "function",
+					"function": map[string]interface{}{
+						"name":        "search_my_dateset",
+						"description": "Search datasets.",
+					},
+				},
+			},
+			ToolChoice: &toolChoice,
+		},
+	)
+	if err != nil {
+		t.Fatalf("ChatWithMessages: %v", err)
+	}
+	if requestBody["tool_choice"] != "auto" {
+		t.Fatalf("tool_choice=%#v, want auto", requestBody["tool_choice"])
+	}
+	if _, ok := requestBody["tools"].([]interface{}); !ok {
+		t.Fatalf("tools missing or wrong type: %#v", requestBody["tools"])
+	}
+	if resp.Answer == nil || *resp.Answer != "" {
+		t.Fatalf("Answer=%#v, want empty string for tool-call response", resp.Answer)
+	}
+	if len(resp.ToolCalls) != 1 {
+		t.Fatalf("ToolCalls len=%d, want 1", len(resp.ToolCalls))
+	}
+	fn, _ := resp.ToolCalls[0]["function"].(map[string]interface{})
+	if fn["name"] != "search_my_dateset" {
+		t.Fatalf("tool call function=%#v, want search_my_dateset", fn)
+	}
+	if resp.Usage == nil || resp.Usage.TotalTokens != 3 {
+		t.Fatalf("Usage=%#v, want total tokens 3", resp.Usage)
+	}
+}
+
+func TestDeepSeekChatWithMessagesForwardsToolHistory(t *testing.T) {
+	var requestBody map[string]interface{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Errorf("decode request: %v", err)
+			http.Error(w, "bad request", http.StatusBadRequest)
+			return
+		}
+		_ = json.NewEncoder(w).Encode(map[string]interface{}{
+			"choices": []map[string]interface{}{
+				{"message": map[string]interface{}{"role": "assistant", "content": "answer"}},
+			},
+		})
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	_, err := newDeepSeekForTest(srv.URL).ChatWithMessages(
+		"deepseek-chat",
+		[]Message{
+			{
+				Role:    "assistant",
+				Content: nil,
+				ToolCalls: []map[string]interface{}{
+					{
+						"id":   "call-1",
+						"type": "function",
+						"function": map[string]interface{}{
+							"name":      "search_my_dateset",
+							"arguments": `{"query":"marigold"}`,
+						},
+					},
+				},
+			},
+			{Role: "tool", ToolCallID: "call-1", Content: `{"formalized_content":"marigold"}`},
+		},
+		&APIConfig{ApiKey: &apiKey},
+		nil,
+	)
+	if err != nil {
+		t.Fatalf("ChatWithMessages: %v", err)
+	}
+	messages, ok := requestBody["messages"].([]interface{})
+	if !ok || len(messages) != 2 {
+		t.Fatalf("messages=%#v, want 2 messages", requestBody["messages"])
+	}
+	assistant, _ := messages[0].(map[string]interface{})
+	if _, ok := assistant["tool_calls"].([]interface{}); !ok {
+		t.Fatalf("assistant tool_calls missing: %#v", assistant)
+	}
+	toolMsg, _ := messages[1].(map[string]interface{})
+	if toolMsg["tool_call_id"] != "call-1" {
+		t.Fatalf("tool_call_id=%#v, want call-1", toolMsg["tool_call_id"])
+	}
+}
--- a/internal/entity/models/llm.go
+++ b/internal/entity/models/llm.go
@@ -32,6 +32,7 @@ package models

 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"sync"

@@ -89,7 +90,14 @@ func toInternalMessages(msgs []*schema.Message) []Message {
 		if role == "" {
 			role = "user"
 		}
-		out = append(out, Message{Role: role, Content: mm.Content})
+		msg := Message{Role: role, Content: mm.Content}
+		if len(mm.ToolCalls) > 0 {
+			msg.ToolCalls = toolCallsToInternal(mm.ToolCalls)
+		}
+		if mm.ToolCallID != "" {
+			msg.ToolCallID = mm.ToolCallID
+		}
+		out = append(out, msg)
 	}
 	return out
 }
@@ -105,7 +113,14 @@ func fromInternalResponse(resp *ChatResponse) *schema.Message {
 	if resp.Answer != nil {
 		content = *resp.Answer
 	}
-	return &schema.Message{Role: schema.Assistant, Content: content}
+	msg := &schema.Message{Role: schema.Assistant, Content: content}
+	if resp.ReasonContent != nil {
+		msg.ReasoningContent = *resp.ReasonContent
+	}
+	if len(resp.ToolCalls) > 0 {
+		msg.ToolCalls = toolCallsFromInternal(resp.ToolCalls)
+	}
+	return msg
 }

 // Generate blocks until the model returns a complete response. Mirrors
@@ -128,7 +143,11 @@ func (m *EinoChatModel) Generate(ctx context.Context, msgs []*schema.Message, op
 	// without a usage block doesn't leak the previous call's data.
 	// Mirrors Python's LLMBundle._reset_last_usage().
 	m.inner.LastUsage = nil
-	resp, err := m.inner.ModelDriver.ChatWithMessages(*m.inner.ModelName, internal, m.inner.APIConfig, m.chatCfg)
+	chatCfg, err := m.chatConfigForGenerate()
+	if err != nil {
+		return nil, err
+	}
+	resp, err := m.inner.ModelDriver.ChatWithMessages(*m.inner.ModelName, internal, m.inner.APIConfig, chatCfg)
 	if err != nil {
 		return nil, fmt.Errorf("models: EinoChatModel.Generate(%s): %w", *m.inner.ModelName, err)
 	}
@@ -144,6 +163,106 @@ func (m *EinoChatModel) Generate(ctx context.Context, msgs []*schema.Message, op
 	return fromInternalResponse(resp), nil
 }

+func (m *EinoChatModel) chatConfigForGenerate() (*ChatConfig, error) {
+	if len(m.tools) == 0 {
+		return m.chatCfg, nil
+	}
+	cfg := &ChatConfig{}
+	if m.chatCfg != nil {
+		cp := *m.chatCfg
+		cfg = &cp
+	}
+	tools, err := openAIToolsFromEino(m.tools)
+	if err != nil {
+		return nil, err
+	}
+	cfg.Tools = tools
+	choice := "auto"
+	cfg.ToolChoice = &choice
+	return cfg, nil
+}
+
+func openAIToolsFromEino(infos []*schema.ToolInfo) ([]map[string]any, error) {
+	tools := make([]map[string]any, 0, len(infos))
+	for _, info := range infos {
+		if info == nil {
+			continue
+		}
+		fn := map[string]any{
+			"name":        info.Name,
+			"description": info.Desc,
+		}
+		if info.ParamsOneOf != nil {
+			params, err := info.ParamsOneOf.ToJSONSchema()
+			if err != nil {
+				return nil, fmt.Errorf("models: convert tool %q schema: %w", info.Name, err)
+			}
+			fn["parameters"] = params
+		} else {
+			fn["parameters"] = map[string]any{"type": "object", "properties": map[string]any{}}
+		}
+		tools = append(tools, map[string]any{
+			"type":     "function",
+			"function": fn,
+		})
+	}
+	return tools, nil
+}
+
+func toolCallsToInternal(calls []schema.ToolCall) []map[string]interface{} {
+	out := make([]map[string]interface{}, 0, len(calls))
+	for _, call := range calls {
+		fn := map[string]interface{}{
+			"name":      call.Function.Name,
+			"arguments": call.Function.Arguments,
+		}
+		out = append(out, map[string]interface{}{
+			"id":       call.ID,
+			"type":     call.Type,
+			"function": fn,
+		})
+	}
+	return out
+}
+
+func toolCallsFromInternal(calls []map[string]interface{}) []schema.ToolCall {
+	out := make([]schema.ToolCall, 0, len(calls))
+	for i, call := range calls {
+		id, _ := call["id"].(string)
+		if id == "" {
+			id = fmt.Sprintf("call_%d", i)
+		}
+		callType, _ := call["type"].(string)
+		if callType == "" {
+			callType = "function"
+		}
+		var fnName, fnArgs string
+		if fn, ok := call["function"].(map[string]interface{}); ok {
+			fnName, _ = fn["name"].(string)
+			switch args := fn["arguments"].(type) {
+			case string:
+				fnArgs = args
+			case nil:
+				fnArgs = "{}"
+			default:
+				b, err := json.Marshal(args)
+				if err == nil {
+					fnArgs = string(b)
+				}
+			}
+		}
+		out = append(out, schema.ToolCall{
+			ID:   id,
+			Type: callType,
+			Function: schema.FunctionCall{
+				Name:      fnName,
+				Arguments: fnArgs,
+			},
+		})
+	}
+	return out
+}
+
 // Stream returns a schema.StreamReader that yields message chunks
 // incrementally. Uses the existing ChatStreamlyWithSender pathway; the
 // sender callback pushes the streamed delta into the StreamReader.
@@ -157,6 +276,13 @@ func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts
 	if m.inner.ModelName == nil {
 		return nil, fmt.Errorf("models: EinoChatModel: nil model name")
 	}
+	if len(m.tools) > 0 {
+		msg, err := m.Generate(ctx, msgs, opts...)
+		if err != nil {
+			return nil, err
+		}
+		return schema.StreamReaderFromArray([]*schema.Message{msg}), nil
+	}
 	internal := toInternalMessages(msgs)

 	sr, sw := schema.Pipe[*schema.Message](1)
--- a/internal/entity/models/llm_test.go
+++ b/internal/entity/models/llm_test.go
@@ -0,0 +1,196 @@
+package models
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cloudwego/eino/schema"
+)
+
+func TestEinoChatModelGenerateSendsBoundTools(t *testing.T) {
+	apiKey := "key"
+	modelName := "chat"
+	driver := &captureToolDriver{
+		resp: &ChatResponse{
+			ToolCalls: []map[string]interface{}{
+				{
+					"id":   "call-1",
+					"type": "function",
+					"function": map[string]interface{}{
+						"name":      "search_my_dateset",
+						"arguments": `{"query":"hello"}`,
+					},
+				},
+			},
+		},
+	}
+	base := NewChatModel(driver, &modelName, &APIConfig{ApiKey: &apiKey})
+	model := NewEinoChatModel(base, nil)
+	bound, err := model.WithTools([]*schema.ToolInfo{
+		{
+			Name: "search_my_dateset",
+			Desc: "Search datasets.",
+			ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
+				"query": &schema.ParameterInfo{Type: schema.String, Required: true},
+			}),
+		},
+	})
+	if err != nil {
+		t.Fatalf("WithTools: %v", err)
+	}
+
+	msg, err := bound.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")})
+	if err != nil {
+		t.Fatalf("Generate: %v", err)
+	}
+	if driver.lastConfig == nil || driver.lastConfig.Tools == nil {
+		t.Fatal("Generate did not send tools to driver")
+	}
+	tools, ok := driver.lastConfig.Tools.([]map[string]any)
+	if !ok || len(tools) != 1 {
+		t.Fatalf("driver tools = %#v, want one OpenAI-style tool", driver.lastConfig.Tools)
+	}
+	fn, _ := tools[0]["function"].(map[string]any)
+	if fn["name"] != "search_my_dateset" {
+		t.Fatalf("tool function name = %#v, want search_my_dateset", fn["name"])
+	}
+	if driver.lastConfig.ToolChoice == nil || *driver.lastConfig.ToolChoice != "auto" {
+		t.Fatalf("ToolChoice = %#v, want auto", driver.lastConfig.ToolChoice)
+	}
+	if len(msg.ToolCalls) != 1 {
+		t.Fatalf("msg.ToolCalls len = %d, want 1", len(msg.ToolCalls))
+	}
+	if msg.ToolCalls[0].Function.Name != "search_my_dateset" || msg.ToolCalls[0].Function.Arguments != `{"query":"hello"}` {
+		t.Fatalf("tool call = %#v, want search_my_dateset query call", msg.ToolCalls[0])
+	}
+}
+
+func TestEinoChatModelStreamWithToolsYieldsToolCalls(t *testing.T) {
+	apiKey := "key"
+	modelName := "chat"
+	driver := &captureToolDriver{
+		resp: &ChatResponse{
+			ToolCalls: []map[string]interface{}{
+				{
+					"id":   "call-1",
+					"type": "function",
+					"function": map[string]interface{}{
+						"name":      "search_my_dateset",
+						"arguments": `{"query":"hello"}`,
+					},
+				},
+			},
+		},
+	}
+	base := NewChatModel(driver, &modelName, &APIConfig{ApiKey: &apiKey})
+	model := NewEinoChatModel(base, nil)
+	bound, err := model.WithTools([]*schema.ToolInfo{
+		{
+			Name: "search_my_dateset",
+			Desc: "Search datasets.",
+			ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
+				"query": &schema.ParameterInfo{Type: schema.String, Required: true},
+			}),
+		},
+	})
+	if err != nil {
+		t.Fatalf("WithTools: %v", err)
+	}
+
+	stream, err := bound.Stream(context.Background(), []*schema.Message{schema.UserMessage("hello")})
+	if err != nil {
+		t.Fatalf("Stream: %v", err)
+	}
+	msg, err := stream.Recv()
+	if err != nil {
+		t.Fatalf("stream.Recv: %v", err)
+	}
+	if msg == nil {
+		t.Fatal("stream ended before yielding message")
+	}
+	if len(msg.ToolCalls) != 1 || msg.ToolCalls[0].Function.Name != "search_my_dateset" {
+		t.Fatalf("stream message tool calls = %#v, want search_my_dateset", msg.ToolCalls)
+	}
+}
+
+func TestToInternalMessagesPreservesToolMessages(t *testing.T) {
+	internal := toInternalMessages([]*schema.Message{
+		{
+			Role: schema.Assistant,
+			ToolCalls: []schema.ToolCall{{
+				ID:   "call-1",
+				Type: "function",
+				Function: schema.FunctionCall{
+					Name:      "search_my_dateset",
+					Arguments: `{"query":"hello"}`,
+				},
+			}},
+		},
+		{
+			Role:       schema.Tool,
+			Content:    `{"formalized_content":"answer"}`,
+			ToolCallID: "call-1",
+		},
+	})
+	if len(internal) != 2 {
+		t.Fatalf("len(internal) = %d, want 2", len(internal))
+	}
+	if len(internal[0].ToolCalls) != 1 {
+		t.Fatalf("assistant ToolCalls = %#v, want one tool call", internal[0].ToolCalls)
+	}
+	if internal[1].ToolCallID != "call-1" || internal[1].Role != "tool" {
+		t.Fatalf("tool message = %#v, want tool role with call id", internal[1])
+	}
+}
+
+type captureToolDriver struct {
+	resp       *ChatResponse
+	lastConfig *ChatConfig
+}
+
+func (d *captureToolDriver) NewInstance(baseURL map[string]string) ModelDriver { return d }
+func (d *captureToolDriver) Name() string                                      { return "capture" }
+func (d *captureToolDriver) ChatWithMessages(_ string, _ []Message, _ *APIConfig, cfg *ChatConfig) (*ChatResponse, error) {
+	d.lastConfig = cfg
+	return d.resp, nil
+}
+func (d *captureToolDriver) ChatStreamlyWithSender(_ string, _ []Message, _ *APIConfig, _ *ChatConfig, _ func(*string, *string) error) error {
+	return nil
+}
+func (d *captureToolDriver) Embed(_ *string, _ []string, _ *APIConfig, _ *EmbeddingConfig) ([]EmbeddingData, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) Rerank(_ *string, _ string, _ []string, _ *APIConfig, _ *RerankConfig) (*RerankResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) TranscribeAudio(_ *string, _ *string, _ *APIConfig, _ *ASRConfig) (*ASRResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) TranscribeAudioWithSender(_ *string, _ *string, _ *APIConfig, _ *ASRConfig, _ func(*string, *string) error) error {
+	return nil
+}
+func (d *captureToolDriver) AudioSpeech(_ *string, _ *string, _ *APIConfig, _ *TTSConfig) (*TTSResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) AudioSpeechWithSender(_ *string, _ *string, _ *APIConfig, _ *TTSConfig, _ func(*string, *string) error) error {
+	return nil
+}
+func (d *captureToolDriver) OCRFile(_ *string, _ []byte, _ *string, _ *APIConfig, _ *OCRConfig) (*OCRFileResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) ParseFile(_ *string, _ []byte, _ *string, _ *APIConfig, _ *ParseFileConfig) (*ParseFileResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) ListModels(_ *APIConfig) ([]ListModelResponse, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) Balance(_ *APIConfig) (map[string]interface{}, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) CheckConnection(_ *APIConfig) error { return nil }
+func (d *captureToolDriver) ListTasks(_ *APIConfig) ([]ListTaskStatus, error) {
+	return nil, nil
+}
+func (d *captureToolDriver) ShowTask(_ string, _ *APIConfig) (*TaskResponse, error) {
+	return nil, nil
+}