diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index db565e538b..216c8a75ff 100644 --- 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 diff --git a/internal/agent/component/agent_test.go b/internal/agent/component/agent_test.go index f801d89960..63d9dd5a19 100644 --- 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 } diff --git a/internal/agent/component/message.go b/internal/agent/component/message.go index 5a834b8b84..b4544e864e 100644 --- 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 diff --git a/internal/agent/component/message_test.go b/internal/agent/component/message_test.go index 5983bb5fa1..19a659d24a 100644 --- 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, diff --git a/internal/agent/component/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index 2d3a272342..a2d016e033 100644 --- 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.", diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index 76cc6edc0b..e4751c1434 100644 --- 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: diff --git a/internal/agent/tool/retrieval.go b/internal/agent/tool/retrieval.go index 81d39d74af..03c097e125 100644 --- 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) { diff --git a/internal/agent/tool/retrieval_nlp.go b/internal/agent/tool/retrieval_nlp.go index 7b5285d0e8..8b1abe2ae2 100644 --- 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) { diff --git a/internal/agent/tool/retrieval_nlp_test.go b/internal/agent/tool/retrieval_nlp_test.go index 666dc23dc3..7771073781 100644 --- 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 +} diff --git a/internal/agent/tool/retrieval_service.go b/internal/agent/tool/retrieval_service.go index 32e3e7840a..86d40e26f5 100644 --- 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 diff --git a/internal/agent/tool/retrieval_test.go b/internal/agent/tool/retrieval_test.go index a127dd83ca..97aa4b1926 100644 --- 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 +} diff --git a/internal/entity/models/deepseek.go b/internal/entity/models/deepseek.go index 47be9d9e41..5e8b7b81f6 100644 --- 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) diff --git a/internal/entity/models/deepseek_test.go b/internal/entity/models/deepseek_test.go new file mode 100644 index 0000000000..36340033e3 --- /dev/null +++ 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"]) + } +} diff --git a/internal/entity/models/llm.go b/internal/entity/models/llm.go index c78ad02862..1e5fdd0bf4 100644 --- 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) diff --git a/internal/entity/models/llm_test.go b/internal/entity/models/llm_test.go new file mode 100644 index 0000000000..ecaedf2ae1 --- /dev/null +++ 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 +}