From a55a438b4262c75c42fff99cd93957fa0be3d5d6 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Wed, 5 Aug 2026 15:48:54 +0800 Subject: [PATCH] fix(go-agent): finish retrieval component (#17845) - Prefer canonical `dataset_ids` over legacy `kb_ids` in Canvas retrieval components. - Cover selected and explicitly cleared dataset IDs with regression tests. - Add cross-language retrieval with dataset-specific embedding and rerank models. - Support vector and keyword similarity controls, metadata filtering, TOC enhancement, and child-chunk expansion. - Route Canvas retrieval across datasets and memories with compatible embedding validation. --- cmd/ragflow_server.go | 10 +- .../component/production_chain_fixes_test.go | 56 +++ .../agent/component/retrieval_params_test.go | 149 +++++++ .../agent/component/retrieval_swap_test.go | 4 +- .../agent/component/universe_a_wrappers.go | 72 +++- internal/agent/retrievalbridge/enhancer.go | 143 +++++++ internal/agent/retrievalbridge/memory.go | 104 +++++ internal/agent/tool/agentic_search.go | 3 +- internal/agent/tool/registry.go | 36 +- internal/agent/tool/retrieval.go | 225 +++++++++- internal/agent/tool/retrieval_nlp.go | 392 ++++++++++++++++-- internal/agent/tool/retrieval_nlp_test.go | 222 ++++++++-- internal/agent/tool/retrieval_service.go | 76 +++- internal/agent/tool/retrieval_test.go | 106 ++++- internal/agent/tool/retrieval_wiring_test.go | 4 +- internal/service/memory.go | 20 + internal/service/memory_message_test.go | 18 + 17 files changed, 1500 insertions(+), 140 deletions(-) create mode 100644 internal/agent/component/retrieval_params_test.go create mode 100644 internal/agent/retrievalbridge/enhancer.go create mode 100644 internal/agent/retrievalbridge/memory.go diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 928c08d09b..57266f78b6 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -26,6 +26,7 @@ import ( "ragflow/internal/admin" "ragflow/internal/agent/audio" "ragflow/internal/agent/canvas" + "ragflow/internal/agent/retrievalbridge" agenttool "ragflow/internal/agent/tool" "ragflow/internal/channels" "ragflow/internal/handler" @@ -728,7 +729,14 @@ func startServer(ctx context.Context) { // Initialize doc engine for skill search docEngine := engine.Get() documentDAO := dao.NewDocumentDAO() - agenttool.SetRetrievalService(agenttool.NewNLPRetrievalAdapterFromDeps(docEngine, documentDAO)) + retrievalEnhancer := retrievalbridge.NewEnhancer(docEngine, metadataService) + agenttool.SetRetrievalService(agenttool.NewNLPRetrievalAdapterFromDeps( + docEngine, + documentDAO, + modelProviderService, + retrievalEnhancer, + )) + agenttool.SetMemoryRetrievalService(retrievalbridge.NewMemoryAdapter(memoryService)) common.Info("agent: retrieval service adapter installed") // Initialize handler layer diff --git a/internal/agent/component/production_chain_fixes_test.go b/internal/agent/component/production_chain_fixes_test.go index 8336a84253..615764e18b 100644 --- a/internal/agent/component/production_chain_fixes_test.go +++ b/internal/agent/component/production_chain_fixes_test.go @@ -159,6 +159,62 @@ func TestRetrieval_KbIDsTranslatedToDatasetIDs(t *testing.T) { if !ok || len(ds3) != 1 || ds3[0] != "kb-new" { t.Errorf("dataset_ids should keep call-time value %v, got %v", "kb-new", merged3["dataset_ids"]) } + + // Case 4: canonical node-level dataset_ids override stale kb_ids. + canonical, err := newRetrievalComponent(map[string]any{ + "dataset_ids": []any{"kb-current"}, + "kb_ids": []any{"kb-stale"}, + }) + if err != nil { + t.Fatalf("newRetrievalComponent with canonical dataset_ids: %v", err) + } + merged4 := canonical.(*retrievalComponent).applyDefaults(nil) + if ds, ok := merged4["dataset_ids"].([]any); !ok || len(ds) != 1 || ds[0] != "kb-current" { + t.Errorf("node dataset_ids should override stale kb_ids, got %v", merged4["dataset_ids"]) + } + + // Case 5: an explicitly empty canonical list clears stale kb_ids. + cleared, err := newRetrievalComponent(map[string]any{ + "dataset_ids": []any{}, + "kb_ids": []any{"kb-stale"}, + }) + if err != nil { + t.Fatalf("newRetrievalComponent with empty dataset_ids: %v", err) + } + merged5 := cleared.(*retrievalComponent).applyDefaults(nil) + if _, ok := merged5["dataset_ids"]; ok { + t.Errorf("empty dataset_ids should clear stale kb_ids, got %v", merged5["dataset_ids"]) + } +} + +func TestRetrieval_NodeQueryResolvedFromCanvasState(t *testing.T) { + previous := agenttool.GetRetrievalService() + agenttool.SetSimpleRetrievalService() + t.Cleanup(func() { agenttool.SetRetrievalService(previous) }) + + c, err := newRetrievalComponent(map[string]any{ + "query": "{sys.query}", + "kb_ids": []any{"kb-1"}, + }) + if err != nil { + t.Fatalf("newRetrievalComponent: %v", err) + } + + state := runtime.NewCanvasState("run-1", "session-1") + state.Sys["query"] = "AirPure X200 vs AirPure X300" + ctx := runtime.WithState(context.Background(), state) + out, err := c.Invoke(ctx, nil, map[string]any{ + "category": "Product Feature Comparison", + "category_name": "Product Feature Comparison", + "_next": []string{"Retrieval:EightyDaysHappen"}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + formalizedContent, _ := out["formalized_content"].(string) + if !strings.Contains(formalizedContent, "AirPure X200 vs AirPure X300") { + t.Fatalf("formalized_content = %q, want resolved canvas query", formalizedContent) + } } func TestRetrieval_LegacyQueryStringNormalized(t *testing.T) { diff --git a/internal/agent/component/retrieval_params_test.go b/internal/agent/component/retrieval_params_test.go new file mode 100644 index 0000000000..f5b036d935 --- /dev/null +++ b/internal/agent/component/retrieval_params_test.go @@ -0,0 +1,149 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package component + +import ( + "context" + "reflect" + "testing" + + agenttool "ragflow/internal/agent/tool" + + "gorm.io/gorm" +) + +type retrievalRequestRecorder struct { + request agenttool.RetrievalRequest +} + +func (r *retrievalRequestRecorder) Search(_ context.Context, _ *gorm.DB, request agenttool.RetrievalRequest) ([]agenttool.RetrievalChunk, error) { + r.request = request + return []agenttool.RetrievalChunk{}, nil +} + +func TestRetrievalComponentAppliesAdvancedNodeParams(t *testing.T) { + component, err := newRetrievalComponent(map[string]any{ + "dataset_ids": []any{"dataset-1"}, + "memory_ids": []any{"memory-1", "memory-2"}, + "cross_languages": []any{"English", "Chinese", "Spanish"}, + "toc_enhance": true, + "use_kg": true, + "meta_data_filter": map[string]any{"method": "manual"}, + "retrieval_from": "memory", + }) + if err != nil { + t.Fatalf("newRetrievalComponent: %v", err) + } + + merged := component.(*retrievalComponent).applyDefaults(nil) + if got := merged["memory_ids"]; !reflect.DeepEqual(got, []string{"memory-1", "memory-2"}) { + t.Errorf("memory_ids = %#v", got) + } + if got := merged["cross_languages"]; !reflect.DeepEqual(got, []string{"English", "Chinese", "Spanish"}) { + t.Errorf("cross_languages = %#v", got) + } + if got := merged["toc_enhance"]; got != true { + t.Errorf("toc_enhance = %#v", got) + } + if got := merged["use_kg"]; got != true { + t.Errorf("use_kg = %#v", got) + } + if got := merged["meta_data_filter"]; !reflect.DeepEqual(got, map[string]any{"method": "manual"}) { + t.Errorf("meta_data_filter = %#v", got) + } + if got := merged["retrieval_from"]; got != "memory" { + t.Errorf("retrieval_from = %#v", got) + } +} + +func TestRetrievalComponentForwardsAdvancedNodeParams(t *testing.T) { + previous := agenttool.GetRetrievalService() + recorder := &retrievalRequestRecorder{} + agenttool.SetRetrievalService(recorder) + t.Cleanup(func() { agenttool.SetRetrievalService(previous) }) + + component, err := newRetrievalComponent(map[string]any{ + "dataset_ids": []any{"dataset-1"}, + "memory_ids": []any{"memory-1"}, + "cross_languages": []any{"English", "Chinese", "Spanish"}, + "toc_enhance": true, + "use_kg": false, + "meta_data_filter": map[string]any{"method": "manual"}, + "retrieval_from": "dataset", + }) + if err != nil { + t.Fatalf("newRetrievalComponent: %v", err) + } + + if _, err := component.Invoke(context.Background(), nil, map[string]any{"query": "爱"}); err != nil { + t.Fatalf("Invoke: %v", err) + } + + request := recorder.request + if !reflect.DeepEqual(request.MemoryIDs, []string{"memory-1"}) { + t.Errorf("MemoryIDs = %#v", request.MemoryIDs) + } + if !reflect.DeepEqual(request.CrossLanguages, []string{"English", "Chinese", "Spanish"}) { + t.Errorf("CrossLanguages = %#v", request.CrossLanguages) + } + if !request.TOCEnhance { + t.Error("TOCEnhance = false") + } + if request.UseKG { + t.Error("UseKG = true") + } + if !reflect.DeepEqual(request.MetaDataFilter, map[string]any{"method": "manual"}) { + t.Errorf("MetaDataFilter = %#v", request.MetaDataFilter) + } + if request.RetrievalFrom != "dataset" { + t.Errorf("RetrievalFrom = %q", request.RetrievalFrom) + } +} + +func TestRetrievalComponentForwardsExplicitZeroSimilarityParams(t *testing.T) { + previous := agenttool.GetRetrievalService() + recorder := &retrievalRequestRecorder{} + agenttool.SetRetrievalService(recorder) + t.Cleanup(func() { agenttool.SetRetrievalService(previous) }) + + component, err := newRetrievalComponent(map[string]any{ + "dataset_ids": []any{"dataset-1"}, + "similarity_threshold": 0, + "keywords_similarity_weight": 0, + }) + if err != nil { + t.Fatalf("newRetrievalComponent: %v", err) + } + + merged := component.(*retrievalComponent).applyDefaults(nil) + if value, ok := merged["similarity_threshold"]; !ok || value != float64(0) { + t.Fatalf("similarity_threshold = %#v, present = %v; want explicit zero", value, ok) + } + if value, ok := merged["keywords_similarity_weight"]; !ok || value != float64(0) { + t.Fatalf("keywords_similarity_weight = %#v, present = %v; want explicit zero", value, ok) + } + + if _, err := component.Invoke(context.Background(), nil, map[string]any{"query": "zero"}); err != nil { + t.Fatalf("Invoke: %v", err) + } + if recorder.request.SimilarityThreshold == nil || *recorder.request.SimilarityThreshold != 0 { + t.Fatalf("SimilarityThreshold = %v; want explicit zero", recorder.request.SimilarityThreshold) + } + if recorder.request.KeywordsSimilarityWeight == nil || *recorder.request.KeywordsSimilarityWeight != 0 { + t.Fatalf("KeywordsSimilarityWeight = %v; want explicit zero", recorder.request.KeywordsSimilarityWeight) + } +} diff --git a/internal/agent/component/retrieval_swap_test.go b/internal/agent/component/retrieval_swap_test.go index a3691ee194..a39ffe4146 100644 --- a/internal/agent/component/retrieval_swap_test.go +++ b/internal/agent/component/retrieval_swap_test.go @@ -76,7 +76,9 @@ func TestSearchMyDataset_AliasDelegatesToRealWrapper(t *testing.T) { agenttool.SetSimpleRetrievalService() t.Cleanup(func() { agenttool.SetRetrievalService(prev) }) - c, err := New("SearchMyDataset", nil) + c, err := New("SearchMyDataset", map[string]any{ + "kb_ids": []any{"kb-1"}, + }) if err != nil { t.Fatalf("New(SearchMyDataset) errored: %v", err) } diff --git a/internal/agent/component/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index 0fae38a723..e5e77f76eb 100644 --- a/internal/agent/component/universe_a_wrappers.go +++ b/internal/agent/component/universe_a_wrappers.go @@ -65,13 +65,20 @@ func anySlice(v any) []any { // defaults to the per-invocation RetrievalRequest. The fields are // the same the Python agent/component/retrieval.py exposes. type retrievalParams struct { + Query string KbIDs []string + MemoryIDs []string TopN int TopK int - SimilarityThreshold float64 - KeywordsSimilarityWeight float64 + SimilarityThreshold *float64 + KeywordsSimilarityWeight *float64 RerankID string EmptyResponse string + CrossLanguages []string + TOCEnhance bool + UseKG bool + MetaDataFilter map[string]any + RetrievalFrom string } // parseRetrievalParams reads the v1 DSL node params for Retrieval. @@ -84,6 +91,12 @@ func parseRetrievalParams(params map[string]any) retrievalParams { if params == nil { return out } + if ids, ok := params["dataset_ids"]; ok { + params["kb_ids"] = ids + } + if v, ok := params["query"].(string); ok { + out.Query = v + } if v, ok := params["kb_ids"].([]any); ok { for _, x := range v { if s, ok := x.(string); ok { @@ -94,6 +107,7 @@ func parseRetrievalParams(params map[string]any) retrievalParams { if v, ok := params["kb_ids"].([]string); ok { out.KbIDs = append(out.KbIDs, v...) } + out.MemoryIDs = toStringSlice(params["memory_ids"]) if v, ok := params["top_n"]; ok { out.TopN = toIntParam(v) } @@ -101,10 +115,12 @@ func parseRetrievalParams(params map[string]any) retrievalParams { out.TopK = toIntParam(v) } if v, ok := params["similarity_threshold"]; ok { - out.SimilarityThreshold = toFloatParam(v) + value := toFloatParam(v) + out.SimilarityThreshold = &value } if v, ok := params["keywords_similarity_weight"]; ok { - out.KeywordsSimilarityWeight = toFloatParam(v) + value := toFloatParam(v) + out.KeywordsSimilarityWeight = &value } if v, ok := params["rerank_id"].(string); ok { out.RerankID = v @@ -112,14 +128,25 @@ func parseRetrievalParams(params map[string]any) retrievalParams { if v, ok := params["empty_response"].(string); ok { out.EmptyResponse = v } + out.CrossLanguages = toStringSlice(params["cross_languages"]) + if v, ok := params["toc_enhance"].(bool); ok { + out.TOCEnhance = v + } + if v, ok := params["use_kg"].(bool); ok { + out.UseKG = v + } + if v, ok := params["meta_data_filter"].(map[string]any); ok { + out.MetaDataFilter = cloneAnyMap(v) + } + if v, ok := params["retrieval_from"].(string); ok { + out.RetrievalFrom = v + } return out } // retrievalComponent delegates to internal/agent/tool/RetrievalTool. -// The wrapper captures the v1 DSL node params (kb_ids, top_n, -// top_k, similarity_threshold, keywords_similarity_weight, -// rerank_id, empty_response) at build time and applies them as -// defaults to each invocation. Per-call inputs override the +// The wrapper captures the Retrieval node's DSL params at build time and +// applies them as defaults to each invocation. Per-call inputs override the // defaults. type retrievalComponent struct { inner *agenttool.RetrievalTool @@ -213,6 +240,9 @@ func (c *retrievalComponent) applyDefaults(inputs map[string]any) map[string]any for k, v := range inputs { out[k] = v } + if _, ok := out["query"]; !ok && c.params.Query != "" { + out["query"] = c.params.Query + } if _, ok := out["kb_ids"]; !ok && len(c.params.KbIDs) > 0 { ids := make([]any, len(c.params.KbIDs)) for i, s := range c.params.KbIDs { @@ -226,11 +256,11 @@ func (c *retrievalComponent) applyDefaults(inputs map[string]any) map[string]any if _, ok := out["top_k"]; !ok && c.params.TopK > 0 { out["top_k"] = c.params.TopK } - if _, ok := out["similarity_threshold"]; !ok && c.params.SimilarityThreshold > 0 { - out["similarity_threshold"] = c.params.SimilarityThreshold + if _, ok := out["similarity_threshold"]; !ok && c.params.SimilarityThreshold != nil { + out["similarity_threshold"] = *c.params.SimilarityThreshold } - if _, ok := out["keywords_similarity_weight"]; !ok && c.params.KeywordsSimilarityWeight > 0 { - out["keywords_similarity_weight"] = c.params.KeywordsSimilarityWeight + if _, ok := out["keywords_similarity_weight"]; !ok && c.params.KeywordsSimilarityWeight != nil { + out["keywords_similarity_weight"] = *c.params.KeywordsSimilarityWeight } if _, ok := out["rerank_id"]; !ok && c.params.RerankID != "" { out["rerank_id"] = c.params.RerankID @@ -238,6 +268,24 @@ func (c *retrievalComponent) applyDefaults(inputs map[string]any) map[string]any if _, ok := out["empty_response"]; !ok && c.params.EmptyResponse != "" { out["empty_response"] = c.params.EmptyResponse } + if _, ok := out["memory_ids"]; !ok && len(c.params.MemoryIDs) > 0 { + out["memory_ids"] = append([]string(nil), c.params.MemoryIDs...) + } + if _, ok := out["cross_languages"]; !ok && len(c.params.CrossLanguages) > 0 { + out["cross_languages"] = append([]string(nil), c.params.CrossLanguages...) + } + if _, ok := out["toc_enhance"]; !ok && c.params.TOCEnhance { + out["toc_enhance"] = true + } + if _, ok := out["use_kg"]; !ok && c.params.UseKG { + out["use_kg"] = true + } + if _, ok := out["meta_data_filter"]; !ok && c.params.MetaDataFilter != nil { + out["meta_data_filter"] = cloneAnyMap(c.params.MetaDataFilter) + } + if _, ok := out["retrieval_from"]; !ok && c.params.RetrievalFrom != "" { + out["retrieval_from"] = c.params.RetrievalFrom + } // Translate v1 DSL name `kb_ids` to the tool's expected // name `dataset_ids`. dataset_ids already-set wins; kb_ids // is consumed and removed so the marshalled JSON carries a diff --git a/internal/agent/retrievalbridge/enhancer.go b/internal/agent/retrievalbridge/enhancer.go new file mode 100644 index 0000000000..cb5dc2e7ac --- /dev/null +++ b/internal/agent/retrievalbridge/enhancer.go @@ -0,0 +1,143 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +// Package retrievalbridge connects the agent retrieval interfaces to the +// parent service package without making agent/tool import internal/service. +package retrievalbridge + +import ( + "context" + "fmt" + + "ragflow/internal/engine" + "ragflow/internal/entity" + modelModule "ragflow/internal/entity/models" + "ragflow/internal/service" + "ragflow/internal/service/nlp" +) + +// Enhancer delegates retrieval query and result enhancement to the existing +// service-layer implementations used by chunk and chat retrieval. +type Enhancer struct { + docEngine engine.DocEngine + metadataSvc *service.MetadataService +} + +// NewEnhancer creates an Enhancer for the configured document engine. +func NewEnhancer(docEngine engine.DocEngine, metadataSvc *service.MetadataService) *Enhancer { + if metadataSvc == nil { + metadataSvc = service.NewMetadataService() + } + return &Enhancer{docEngine: docEngine, metadataSvc: metadataSvc} +} + +// CrossLanguages translates the query into the configured languages using the +// tenant's default chat model. +func (e *Enhancer) CrossLanguages( + ctx context.Context, + tenantID, query string, + languages []string, +) (string, error) { + return service.CrossLanguages(ctx, tenantID, "", query, languages) +} + +// FilterDocuments applies auto, semi-auto, or manual metadata filtering and +// constrains the result by any document scope supplied by upstream tools. +func (e *Enhancer) FilterDocuments( + ctx context.Context, + filter map[string]any, + query string, + chatModel *modelModule.ChatModel, + baseDocIDs []string, + kbIDs []string, +) ([]string, error) { + if e == nil || e.metadataSvc == nil { + return nil, fmt.Errorf("metadata service is not configured") + } + metadata, err := e.metadataSvc.GetFlattedMetaByKBs(ctx, kbIDs) + if err != nil { + return nil, err + } + docIDs, noMatches := service.ApplyMetaDataFilter( + ctx, + filter, + metadata, + query, + chatModel, + baseDocIDs, + kbIDs, + ) + if noMatches { + return []string{service.NoMatchDocIDSentinel}, nil + } + return docIDs, nil +} + +// LabelQuestion returns tag-based rank features for NLP reranking. +func (e *Enhancer) LabelQuestion( + ctx context.Context, + question string, + kbs []*entity.Knowledgebase, +) map[string]float64 { + if e == nil || e.metadataSvc == nil { + return nil + } + return e.metadataSvc.LabelQuestion(ctx, question, kbs) +} + +// EnhanceTOC adds or boosts chunks selected through the document table of +// contents. The service enhancer mutates the supplied kbinfos map. +func (e *Enhancer) EnhanceTOC( + ctx context.Context, + chatModel *modelModule.ChatModel, + tenantIDs, kbIDs []string, + question string, + topN int, + chunks []map[string]any, +) ([]map[string]any, error) { + if e == nil { + return nil, fmt.Errorf("retrieval enhancer is not configured") + } + kbinfos := map[string]any{"chunks": chunks} + enhancer := service.NewTOCEnhancer( + e.docEngine, + chatModel, + tenantIDs, + kbIDs, + question, + topN, + ) + if _, err := enhancer.Enhance(ctx, kbinfos); err != nil { + return nil, err + } + enhanced, ok := kbinfos["chunks"].([]map[string]any) + if !ok { + return nil, fmt.Errorf("TOC enhancer returned invalid chunks type %T", kbinfos["chunks"]) + } + return enhanced, nil +} + +// RetrieveByChildren aggregates child chunks under their parent chunks. +func (e *Enhancer) RetrieveByChildren( + ctx context.Context, + chunks []map[string]any, + tenantIDs []string, +) []map[string]any { + if e == nil || e.docEngine == nil { + return chunks + } + return nlp.RetrievalByChildren(chunks, tenantIDs, e.docEngine, ctx) +} diff --git a/internal/agent/retrievalbridge/memory.go b/internal/agent/retrievalbridge/memory.go new file mode 100644 index 0000000000..503805957d --- /dev/null +++ b/internal/agent/retrievalbridge/memory.go @@ -0,0 +1,104 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package retrievalbridge + +import ( + "context" + "fmt" + "strings" + + agenttool "ragflow/internal/agent/tool" + "ragflow/internal/service" + + "gorm.io/gorm" +) + +// MemoryAdapter exposes MemoryService.SearchMessage through the agent tool's +// retrieval interface. +type MemoryAdapter struct { + svc *service.MemoryService +} + +// NewMemoryAdapter creates a memory retrieval adapter. +func NewMemoryAdapter(svc *service.MemoryService) *MemoryAdapter { + return &MemoryAdapter{svc: svc} +} + +// Search performs hybrid memory-message retrieval and translates messages to +// the common RetrievalChunk result shape. +func (a *MemoryAdapter) Search( + ctx context.Context, + _ *gorm.DB, + req agenttool.RetrievalRequest, +) ([]agenttool.RetrievalChunk, error) { + if a == nil || a.svc == nil { + return nil, agenttool.ErrMemoryRetrievalServiceMissing + } + if strings.TrimSpace(req.TenantID) == "" { + return nil, fmt.Errorf("memory retrieval: tenant id is required") + } + memoryIDs := compactStrings(req.MemoryIDs) + if len(memoryIDs) == 0 { + return nil, fmt.Errorf("memory retrieval: memory_ids is required") + } + keywordWeight := 0.7 + if req.KeywordsSimilarityWeight != nil { + keywordWeight = *req.KeywordsSimilarityWeight + } + messages, _, err := a.svc.SearchMessage( + ctx, + req.TenantID, + map[string]any{"memory_id": memoryIDs}, + map[string]any{ + "query": req.Query, + "similarity_threshold": req.SimilarityThreshold, + "keywords_similarity_weight": keywordWeight, + "top_n": req.TopN, + }, + ) + if err != nil { + return nil, err + } + chunks := make([]agenttool.RetrievalChunk, 0, len(messages)) + for _, message := range messages { + memoryID := fmt.Sprint(message["memory_id"]) + chunks = append(chunks, agenttool.RetrievalChunk{ + ID: fmt.Sprint(message["message_id"]), + Content: fmt.Sprint(message["content"]), + DocumentID: memoryID, + DatasetID: memoryID, + }) + } + return chunks, nil +} + +func compactStrings(values []string) []string { + seen := make(map[string]struct{}, len(values)) + result := make([]string, 0, len(values)) + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" { + continue + } + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + result = append(result, value) + } + return result +} diff --git a/internal/agent/tool/agentic_search.go b/internal/agent/tool/agentic_search.go index 3a5de007f9..8510a6c431 100644 --- a/internal/agent/tool/agentic_search.go +++ b/internal/agent/tool/agentic_search.go @@ -118,12 +118,13 @@ func (a *AgenticSearchTool) InvokableRun(ctx context.Context, argumentsInJSON st } weight := a.weight + similarityThreshold := 0.2 req := RetrievalRequest{ Query: strings.TrimSpace(args.Query + " " + args.Keywords), DatasetIDs: datasetIDs, TopN: args.TopN, TopK: args.TopN * 4, - SimilarityThreshold: 0.2, + SimilarityThreshold: &similarityThreshold, KeywordsSimilarityWeight: &weight, DocScope: args.DocScope, } diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index 2487dce735..89fafeefd7 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -415,6 +415,11 @@ func buildRetrievalTool(params map[string]any) (einotool.BaseTool, error) { defaults.DatasetIDs = ids } } + if ids, ok, err := stringSliceParam(params, "memory_ids"); err != nil { + return nil, fmt.Errorf("agent tool: retrieval config: %w", err) + } else if ok { + defaults.MemoryIDs = ids + } if v, ok := intParam(params, "top_n"); ok { defaults.TopN = v } @@ -428,8 +433,11 @@ func buildRetrievalTool(params map[string]any) (einotool.BaseTool, error) { if v, ok := boolParam(params, "use_kg"); ok { defaults.UseKG = v } + if v, ok := boolParam(params, "toc_enhance"); ok { + defaults.TOCEnhance = v + } if v, ok := floatParam(params, "similarity_threshold"); ok { - defaults.SimilarityThreshold = v + defaults.SimilarityThreshold = &v } if v, ok := floatParam(params, "keywords_similarity_weight"); ok { if v < 0 || v > 1 { @@ -444,6 +452,32 @@ func buildRetrievalTool(params map[string]any) (einotool.BaseTool, error) { } defaults.EmptyResponse = emptyResponse } + if value, ok := params["rerank_id"]; ok { + rerankID, ok := value.(string) + if !ok { + return nil, fmt.Errorf("agent tool: retrieval tool requires string node-level param rerank_id") + } + defaults.RerankID = rerankID + } + if languages, ok, err := stringSliceParam(params, "cross_languages"); err != nil { + return nil, fmt.Errorf("agent tool: retrieval config: %w", err) + } else if ok { + defaults.CrossLanguages = languages + } + if value, ok := params["meta_data_filter"]; ok { + filter, ok := value.(map[string]any) + if !ok { + return nil, fmt.Errorf("agent tool: retrieval tool requires object node-level param meta_data_filter") + } + defaults.MetaDataFilter = cloneStringAnyMap(filter) + } + if value, ok := params["retrieval_from"]; ok { + retrievalFrom, ok := value.(string) + if !ok || (retrievalFrom != "dataset" && retrievalFrom != "memory") { + return nil, fmt.Errorf("agent tool: retrieval tool requires retrieval_from to be dataset or memory") + } + defaults.RetrievalFrom = retrievalFrom + } return NewRetrievalToolWithDefaults(defaults), nil } diff --git a/internal/agent/tool/retrieval.go b/internal/agent/tool/retrieval.go index a24fde6db5..7f9dc34993 100644 --- a/internal/agent/tool/retrieval.go +++ b/internal/agent/tool/retrieval.go @@ -56,15 +56,21 @@ 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"` - 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"` - EmptyResponse string `json:"empty_response,omitempty"` + Query string `json:"query"` + DatasetIDs []string `json:"dataset_ids,omitempty"` + KBIDs []string `json:"kb_ids,omitempty"` + MemoryIDs []string `json:"memory_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"` + RerankID string `json:"rerank_id,omitempty"` + CrossLanguages []string `json:"cross_languages,omitempty"` + TOCEnhance bool `json:"toc_enhance,omitempty"` + MetaDataFilter map[string]any `json:"meta_data_filter,omitempty"` + RetrievalFrom string `json:"retrieval_from,omitempty"` + EmptyResponse string `json:"empty_response,omitempty"` } // retrievalResult is the JSON shape returned to the model. The `_ERROR` @@ -139,6 +145,11 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string } } args = r.mergeDefaults(args) + resolvedQuery, err := resolveRetrievalQuery(ctx, args.Query) + if err != nil { + return "", err + } + args.Query = resolvedQuery common.Debug("agent retrieval tool: parsed arguments", zap.String("query", args.Query), zap.Strings("dataset_ids", args.DatasetIDs), @@ -150,32 +161,63 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string if args.Query == "" { return stubJSONWithErr(retrievalResult{FormalizedContent: args.EmptyResponse}) } - if args.UseKG { - // Plan + §9 Q3: GraphRAG is out of scope for the Go - // Canvas. Return the structured error so the model can react. return stubJSON(retrievalResult{ Stub: true, Error: ErrGraphRAGNotSupported.Error(), }), ErrGraphRAGNotSupported } + if args.RetrievalFrom == "" { + return stubJSONWithErr(retrievalResult{FormalizedContent: args.EmptyResponse}) + } + if args.RetrievalFrom != "dataset" && args.RetrievalFrom != "memory" { + return "", fmt.Errorf("retrieval: unsupported retrieval_from %q", args.RetrievalFrom) + } + if args.RetrievalFrom == "dataset" && len(args.DatasetIDs) == 0 { + return "", fmt.Errorf("retrieval: dataset_ids is required") + } + if args.RetrievalFrom == "memory" && len(args.MemoryIDs) == 0 { + return "", fmt.Errorf("retrieval: memory_ids is required") + } + resolvedDatasetIDs, err := resolveRetrievalDatasetIDs(ctx, args.DatasetIDs) + if err != nil { + return "", err + } + args.DatasetIDs = resolvedDatasetIDs + resolvedFilter, err := resolveRetrievalFilter(ctx, args.MetaDataFilter) + if err != nil { + return "", err + } + args.MetaDataFilter = resolvedFilter // Dispatch to the registered RetrievalService. When the // default stub is in place, the call surfaces // ErrRetrievalServiceMissing; once a real impl is installed // via SetRetrievalService (or SetSimpleRetrievalService for // dev), the chunks flow through normally. - svc := GetRetrievalService() - chunks, err := svc.Search(ctx, dao.DB, RetrievalRequest{ + searchReq := RetrievalRequest{ Query: args.Query, DatasetIDs: args.DatasetIDs, + MemoryIDs: args.MemoryIDs, TopN: args.TopN, TopK: args.TopK, KeywordsSimilarityWeight: args.KeywordsSimilarityWeight, UseKG: args.UseKG, SimilarityThreshold: args.SimilarityThreshold, + RerankID: args.RerankID, + CrossLanguages: append([]string(nil), args.CrossLanguages...), + TOCEnhance: args.TOCEnhance, + MetaDataFilter: cloneStringAnyMap(args.MetaDataFilter), + RetrievalFrom: args.RetrievalFrom, TenantID: retrievalTenantID(ctx), - }) + } + + var chunks []RetrievalChunk + if args.RetrievalFrom == "memory" { + chunks, err = GetMemoryRetrievalService().Search(ctx, dao.DB, searchReq) + } else { + chunks, err = GetRetrievalService().Search(ctx, dao.DB, searchReq) + } if err != nil { return stubJSON(retrievalResult{ Stub: true, @@ -198,6 +240,9 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string }) } formalizedContent := renderChunks(chunks, args.Query) + if args.RetrievalFrom == "memory" { + formalizedContent = renderMemoryChunks(chunks) + } if len(chunks) == 0 { formalizedContent = args.EmptyResponse } @@ -206,7 +251,7 @@ func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string // citation grounding call can read them. The recording is // 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 { + if state, _, sErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); sErr == nil && state != nil && len(chunks) > 0 && args.RetrievalFrom == "dataset" { state.SetRetrievalReferences(referenceChunksFromRetrieval(chunks), referenceDocAggsFromRetrieval(chunks)) } result, err := stubJSONWithErr(out) @@ -223,6 +268,9 @@ func (r *RetrievalTool) mergeDefaults(args retrievalArgs) retrievalArgs { if len(args.DatasetIDs) == 0 && len(r.defaults.DatasetIDs) != 0 { args.DatasetIDs = append([]string(nil), r.defaults.DatasetIDs...) } + if len(args.MemoryIDs) == 0 && len(r.defaults.MemoryIDs) != 0 { + args.MemoryIDs = append([]string(nil), r.defaults.MemoryIDs...) + } if args.TopN <= 0 { args.TopN = r.defaults.TopN } @@ -232,16 +280,148 @@ func (r *RetrievalTool) mergeDefaults(args retrievalArgs) retrievalArgs { if args.KeywordsSimilarityWeight == nil { args.KeywordsSimilarityWeight = r.defaults.KeywordsSimilarityWeight } - if args.SimilarityThreshold <= 0 { + if args.SimilarityThreshold == nil { args.SimilarityThreshold = r.defaults.SimilarityThreshold } if args.EmptyResponse == "" { args.EmptyResponse = r.defaults.EmptyResponse } + if args.RerankID == "" { + args.RerankID = r.defaults.RerankID + } + if len(args.CrossLanguages) == 0 && len(r.defaults.CrossLanguages) != 0 { + args.CrossLanguages = append([]string(nil), r.defaults.CrossLanguages...) + } + if args.MetaDataFilter == nil && r.defaults.MetaDataFilter != nil { + args.MetaDataFilter = cloneStringAnyMap(r.defaults.MetaDataFilter) + } + if args.RetrievalFrom == "" { + args.RetrievalFrom = r.defaults.RetrievalFrom + } + if args.RetrievalFrom == "" && len(args.DatasetIDs) > 0 { + args.RetrievalFrom = "dataset" + } + if args.RetrievalFrom == "" && len(args.MemoryIDs) > 0 { + args.RetrievalFrom = "memory" + } + args.TOCEnhance = args.TOCEnhance || r.defaults.TOCEnhance args.UseKG = args.UseKG || r.defaults.UseKG return args } +func cloneStringAnyMap(src map[string]any) map[string]any { + if src == nil { + return nil + } + dst := make(map[string]any, len(src)) + for key, value := range src { + dst[key] = value + } + return dst +} + +func resolveRetrievalQuery(ctx context.Context, query string) (string, error) { + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return query, nil + } + resolved, err := runtime.ResolveTemplateAuto(query, state) + if err != nil { + return "", fmt.Errorf("retrieval: resolve query variables: %w", err) + } + return resolved, nil +} + +func resolveRetrievalDatasetIDs(ctx context.Context, datasetIDs []string) ([]string, error) { + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return compactStrings(datasetIDs), nil + } + resolved := make([]string, 0, len(datasetIDs)) + for _, datasetID := range datasetIDs { + if !strings.Contains(datasetID, "@") { + resolved = append(resolved, datasetID) + continue + } + value, getErr := state.GetVar(datasetID) + if getErr != nil { + return nil, fmt.Errorf("retrieval: resolve dataset variable %q: %w", datasetID, getErr) + } + if value == nil { + return nil, fmt.Errorf("retrieval: dataset variable %q is empty", datasetID) + } + switch typed := value.(type) { + case string: + resolved = append(resolved, typed) + case []string: + resolved = append(resolved, typed...) + case []any: + for _, item := range typed { + text, ok := item.(string) + if !ok { + return nil, fmt.Errorf("retrieval: dataset variable %q contains non-string value", datasetID) + } + resolved = append(resolved, text) + } + default: + return nil, fmt.Errorf("retrieval: dataset variable %q must be a string or string list", datasetID) + } + } + return compactStrings(resolved), nil +} + +func resolveRetrievalFilter(ctx context.Context, filter map[string]any) (map[string]any, error) { + if filter == nil { + return nil, nil + } + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return cloneStringAnyMap(filter), nil + } + resolved, err := resolveRetrievalValue(filter, state) + if err != nil { + return nil, err + } + result, ok := resolved.(map[string]any) + if !ok { + return nil, fmt.Errorf("retrieval: metadata filter must be an object") + } + return result, nil +} + +func resolveRetrievalValue(value any, state *runtime.CanvasState) (any, error) { + switch typed := value.(type) { + case string: + resolved, err := runtime.ResolveTemplateAuto(typed, state) + if err != nil { + return nil, fmt.Errorf("retrieval: resolve metadata filter value: %w", err) + } + return resolved, nil + case map[string]any: + result := make(map[string]any, len(typed)) + for key, item := range typed { + resolved, err := resolveRetrievalValue(item, state) + if err != nil { + return nil, err + } + result[key] = resolved + } + return result, nil + case []any: + result := make([]any, len(typed)) + for index, item := range typed { + resolved, err := resolveRetrievalValue(item, state) + if err != nil { + return nil, err + } + result[index] = resolved + } + return result, nil + default: + return value, nil + } +} + // renderChunks concatenates the retrieved chunks into a human- // readable content string. Mirrors Python's // `kb_prompt(kbinfos, ...)` format: each chunk gets a header @@ -254,6 +434,17 @@ func renderChunks(chunks []RetrievalChunk, query string) string { return sb.String() } +func renderMemoryChunks(chunks []RetrievalChunk) string { + var builder strings.Builder + for index, chunk := range chunks { + if index > 0 { + builder.WriteByte('\n') + } + builder.WriteString(chunk.Content) + } + return builder.String() +} + func retrievalTenantID(ctx context.Context) string { state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) if err != nil || state == nil { diff --git a/internal/agent/tool/retrieval_nlp.go b/internal/agent/tool/retrieval_nlp.go index 2a15193b3d..7434c5fd6e 100644 --- a/internal/agent/tool/retrieval_nlp.go +++ b/internal/agent/tool/retrieval_nlp.go @@ -24,12 +24,6 @@ // the same service that powers chat / dataset search / chunk // retrieval across the rest of the codebase. // -// Wiring is one line at boot: -// -// tool.SetRetrievalService(tool.NewNLPRetrievalAdapter( -// nlp.NewRetrievalService(docEngine, documentDAO), -// )) -// // Translation rules: // // tool.RetrievalRequest.Query → nlp.RetrievalRequest.Question @@ -41,8 +35,10 @@ // tool.RetrievalRequest.KeywordsSimilarityWeight // → nlp.RetrievalRequest.VectorSimilarityWeight // as 1-keyword weight +// resolved knowledge-base model → nlp.RetrievalRequest.EmbeddingModel +// tool.RetrievalRequest.RerankID → nlp.RetrievalRequest.RerankModel // tool.RetrievalRequest.UseKG → ErrGraphRAGNotSupported (out of -// scope per plan + §9 Q3) +// scope for the Go Agent tool) // // Chunk shape translation: nlp's Chunks are []map[string]any with // keys chunk_id, doc_id, docnm_kwd, content_with_weight, @@ -68,50 +64,129 @@ package tool import ( "context" "fmt" + "regexp" "strings" + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/engine" "ragflow/internal/entity" + modelModule "ragflow/internal/entity/models" "ragflow/internal/service/nlp" + "go.uber.org/zap" "gorm.io/gorm" ) +var retrievalUserPrefixPattern = regexp.MustCompile(`(?i)^user[:：\s]*`) + // NLPRetrievalAdapter wraps *nlp.RetrievalService behind the // agent-tool RetrievalService interface. The adapter is safe to // share across goroutines — the wrapped service is stateless // beyond its docEngine + documentDAO handles, both of which the // nlp package treats as concurrent-safe. type NLPRetrievalAdapter struct { - svc *nlp.RetrievalService - kbDAO knowledgebaseLookup + svc *nlp.RetrievalService + kbDAO knowledgebaseLookup + modelResolver modelResolver + enhancer retrievalEnhancer } type knowledgebaseLookup interface { GetByIDs(ctx context.Context, sqlDB *gorm.DB, ids []string) ([]*entity.Knowledgebase, error) + GetByName(ctx context.Context, sqlDB *gorm.DB, name, tenantID string) (*entity.Knowledgebase, error) +} + +// modelResolver is the narrow model-provider surface needed by retrieval. +// Keeping this interface in the tool package avoids importing the parent +// internal/service package, which already depends on agent/tool. +type modelResolver interface { + GetModelConfigByID( + ctx context.Context, + tenantID string, + modelType entity.ModelType, + modelID string, + ) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) + ResolveModelConfig( + ctx context.Context, + tenantID string, + modelType entity.ModelType, + modelRef string, + ) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) + GetTenantDefaultModelByType( + ctx context.Context, + tenantID string, + modelType entity.ModelType, + ) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) +} + +// retrievalEnhancer exposes the service-layer query and result enhancements +// without making agent/tool import the parent internal/service package. +type retrievalEnhancer interface { + CrossLanguages( + ctx context.Context, + tenantID, query string, + languages []string, + ) (string, error) + FilterDocuments( + ctx context.Context, + filter map[string]any, + query string, + chatModel *modelModule.ChatModel, + baseDocIDs []string, + kbIDs []string, + ) ([]string, error) + LabelQuestion( + ctx context.Context, + question string, + kbs []*entity.Knowledgebase, + ) map[string]float64 + EnhanceTOC( + ctx context.Context, + chatModel *modelModule.ChatModel, + tenantIDs, kbIDs []string, + question string, + topN int, + chunks []map[string]any, + ) ([]map[string]any, error) + RetrieveByChildren( + ctx context.Context, + chunks []map[string]any, + tenantIDs []string, + ) []map[string]any } // NewNLPRetrievalAdapter wraps an already-constructed // *nlp.RetrievalService. -func NewNLPRetrievalAdapter(svc *nlp.RetrievalService) *NLPRetrievalAdapter { +func NewNLPRetrievalAdapter( + svc *nlp.RetrievalService, + resolver modelResolver, + enhancer retrievalEnhancer, +) *NLPRetrievalAdapter { return &NLPRetrievalAdapter{ - svc: svc, - kbDAO: dao.NewKnowledgebaseDAO(), + svc: svc, + kbDAO: dao.NewKnowledgebaseDAO(), + modelResolver: resolver, + enhancer: enhancer, } } // NewNLPRetrievalAdapterFromDeps is the convenience constructor // for the common boot path: // -// tool.SetRetrievalService(tool.NewNLPRetrievalAdapterFromDeps(docEngine, docDAO)) -// -// matches chat_session.go's newChatSessionServiceWithRetrieval -// call site. -func NewNLPRetrievalAdapterFromDeps(docEngine engine.DocEngine, documentDAO *dao.DocumentDAO) *NLPRetrievalAdapter { +// The boot path supplies the model provider and service enhancement bridge so +// the adapter can build a complete hybrid-retrieval request. +func NewNLPRetrievalAdapterFromDeps( + docEngine engine.DocEngine, + documentDAO *dao.DocumentDAO, + resolver modelResolver, + enhancer retrievalEnhancer, +) *NLPRetrievalAdapter { return &NLPRetrievalAdapter{ - svc: nlp.NewRetrievalService(docEngine, documentDAO), - kbDAO: dao.NewKnowledgebaseDAO(), + svc: nlp.NewRetrievalService(docEngine, documentDAO), + kbDAO: dao.NewKnowledgebaseDAO(), + modelResolver: resolver, + enhancer: enhancer, } } @@ -122,11 +197,7 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, db *gorm.DB, req Retri return nil, ErrRetrievalServiceMissing } if req.UseKG { - // Plan + §9 Q3: GraphRAG is out of scope for the - // Go Canvas. The tool layer also returns the error; we - // surface it here so any future direct caller of the - // adapter (bypassing the tool envelope) sees the same - // contract. + // Keep direct adapter callers consistent with RetrievalTool. return nil, ErrGraphRAGNotSupported } if req.Query == "" { @@ -137,11 +208,69 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, db *gorm.DB, req Retri topN = 8 } - tenantIDs, err := a.resolveTenantIDs(ctx, db, req) + datasets, err := a.resolveDatasets(ctx, db, req) if err != nil { return nil, err } - nlpReq := nlpRequestFromRetrieval(req, tenantIDs, topN) + if len(datasets.tenantIDs) != 1 { + return nil, fmt.Errorf("retrieval: datasets span multiple tenants") + } + if err := validateEmbeddingModels(datasets.kbs); err != nil { + return nil, err + } + embeddingModel, err := a.resolveEmbeddingModel(ctx, datasets.kbs[0]) + if err != nil { + return nil, err + } + + chatModel, err := a.resolveChatModel(ctx, req, datasets.kbs[0].TenantID) + if err != nil { + return nil, err + } + query := req.Query + docIDs := compactStrings(req.DocScope) + if len(req.MetaDataFilter) > 0 { + if a.enhancer == nil { + return nil, fmt.Errorf("retrieval: metadata filter service is not configured") + } + docIDs, err = a.enhancer.FilterDocuments( + ctx, req.MetaDataFilter, query, chatModel, docIDs, datasets.kbIDs, + ) + if err != nil { + return nil, fmt.Errorf("retrieval: filter documents: %w", err) + } + } + if len(req.CrossLanguages) > 0 { + if a.enhancer == nil { + return nil, fmt.Errorf("retrieval: cross-language service is not configured") + } + translated, translateErr := a.enhancer.CrossLanguages( + ctx, datasets.kbs[0].TenantID, query, req.CrossLanguages, + ) + if translateErr != nil { + common.Warn("agent retrieval: cross-language query failed; using original query", zap.Error(translateErr)) + } else if strings.TrimSpace(translated) != "" { + query = translated + } + } + query = retrievalUserPrefixPattern.ReplaceAllString(query, "") + var rankFeature map[string]float64 + if a.enhancer != nil { + rankFeature = a.enhancer.LabelQuestion(ctx, query, datasets.kbs) + } + rerankModel, err := a.resolveRerankModel(ctx, req, datasets.kbs[0].TenantID) + if err != nil { + return nil, err + } + preparedReq := req + preparedReq.Query = query + preparedReq.DocScope = docIDs + preparedReq.DatasetIDs = append([]string(nil), datasets.kbIDs...) + nlpReq := nlpRequestFromRetrieval(preparedReq, datasets.tenantIDs, topN, embeddingModel) + nlpReq.RerankModel = rerankModel + if rankFeature != nil { + nlpReq.RankFeature = &rankFeature + } res, err := a.svc.Retrieval(ctx, nlpReq) if err != nil { @@ -150,23 +279,51 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, db *gorm.DB, req Retri if res == nil || len(res.Chunks) == 0 { return []RetrievalChunk{}, nil } - out := make([]RetrievalChunk, 0, len(res.Chunks)) - for _, raw := range res.Chunks { + rawChunks := res.Chunks + if req.TOCEnhance { + if a.enhancer == nil { + return nil, fmt.Errorf("retrieval: TOC enhancement service is not configured") + } + rawChunks, err = a.enhancer.EnhanceTOC( + ctx, + chatModel, + datasets.tenantIDs, + datasets.kbIDs, + query, + topN, + rawChunks, + ) + if err != nil { + common.Warn("agent retrieval: TOC enhancement failed; using retrieval results", zap.Error(err)) + rawChunks = res.Chunks + } + } + if a.enhancer != nil { + rawChunks = a.enhancer.RetrieveByChildren(ctx, rawChunks, datasets.tenantIDs) + } + out := make([]RetrievalChunk, 0, len(rawChunks)) + for _, raw := range rawChunks { out = append(out, translateChunk(raw)) } return out, nil } -func nlpRequestFromRetrieval(req RetrievalRequest, tenantIDs []string, topN int) *nlp.RetrievalRequest { +func nlpRequestFromRetrieval( + req RetrievalRequest, + tenantIDs []string, + topN int, + embeddingModel *modelModule.EmbeddingModel, +) *nlp.RetrievalRequest { nlpReq := &nlp.RetrievalRequest{ - Question: req.Query, - TenantIDs: append([]string(nil), tenantIDs...), - KbIDs: append([]string(nil), req.DatasetIDs...), - DocIDs: append([]string(nil), compactStrings(req.DocScope)...), - Page: 1, - PageSize: topN, - Aggs: boolPtr(false), - Highlight: boolPtr(false), + Question: req.Query, + TenantIDs: append([]string(nil), tenantIDs...), + KbIDs: append([]string(nil), req.DatasetIDs...), + DocIDs: append([]string(nil), compactStrings(req.DocScope)...), + Page: 1, + PageSize: topN, + EmbeddingModel: embeddingModel, + Aggs: boolPtr(false), + Highlight: boolPtr(false), } if req.TopK > 0 { nlpReq.Top = &req.TopK @@ -174,8 +331,8 @@ func nlpRequestFromRetrieval(req RetrievalRequest, tenantIDs []string, topN int) rerankBudget := topN * 4 nlpReq.Top = &rerankBudget } - if req.SimilarityThreshold > 0 { - nlpReq.SimilarityThreshold = &req.SimilarityThreshold + if req.SimilarityThreshold != nil { + nlpReq.SimilarityThreshold = req.SimilarityThreshold } if req.KeywordsSimilarityWeight != nil { vectorSimilarityWeight := 1 - *req.KeywordsSimilarityWeight @@ -184,7 +341,17 @@ func nlpRequestFromRetrieval(req RetrievalRequest, tenantIDs []string, topN int) return nlpReq } -func (a *NLPRetrievalAdapter) resolveTenantIDs(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]string, error) { +type resolvedDatasets struct { + kbs []*entity.Knowledgebase + kbIDs []string + tenantIDs []string +} + +func (a *NLPRetrievalAdapter) resolveDatasets( + ctx context.Context, + db *gorm.DB, + req RetrievalRequest, +) (*resolvedDatasets, error) { seen := map[string]struct{}{} tenantIDs := make([]string, 0, 1) appendTenantID := func(tenantID string) { @@ -199,16 +366,40 @@ func (a *NLPRetrievalAdapter) resolveTenantIDs(ctx context.Context, db *gorm.DB, tenantIDs = append(tenantIDs, tenantID) } - appendTenantID(req.TenantID) datasetIDs := compactStrings(req.DatasetIDs) - if len(datasetIDs) == 0 || a == nil || a.kbDAO == nil { - return tenantIDs, nil + if len(datasetIDs) == 0 { + return nil, fmt.Errorf("retrieval: dataset_ids is required") + } + if a == nil || a.kbDAO == nil { + return nil, fmt.Errorf("retrieval: knowledge base lookup is not configured") } - kbs, err := a.kbDAO.GetByIDs(ctx, db, datasetIDs) + foundKBs, err := a.kbDAO.GetByIDs(ctx, db, datasetIDs) if err != nil { return nil, fmt.Errorf("retrieval: resolve dataset tenants: %w", err) } + kbsByID := make(map[string]*entity.Knowledgebase, len(foundKBs)) + for _, kb := range foundKBs { + if kb != nil { + kbsByID[kb.ID] = kb + } + } + kbs := make([]*entity.Knowledgebase, 0, len(datasetIDs)) + resolvedKBIDs := make([]string, 0, len(datasetIDs)) + for _, datasetID := range datasetIDs { + kb := kbsByID[datasetID] + if kb == nil && strings.TrimSpace(req.TenantID) != "" { + kb, err = a.kbDAO.GetByName(ctx, db, datasetID, req.TenantID) + if err != nil && err != gorm.ErrRecordNotFound { + return nil, fmt.Errorf("retrieval: resolve dataset %q by name: %w", datasetID, err) + } + } + if kb == nil { + return nil, fmt.Errorf("retrieval: dataset %q was not found", datasetID) + } + kbs = append(kbs, kb) + resolvedKBIDs = append(resolvedKBIDs, kb.ID) + } for _, kb := range kbs { if kb == nil { continue @@ -218,7 +409,118 @@ func (a *NLPRetrievalAdapter) resolveTenantIDs(ctx context.Context, db *gorm.DB, if len(tenantIDs) == 0 { return nil, fmt.Errorf("retrieval: no valid knowledge bases found for dataset_ids %v", datasetIDs) } - return tenantIDs, nil + return &resolvedDatasets{kbs: kbs, kbIDs: resolvedKBIDs, tenantIDs: tenantIDs}, nil +} + +func validateEmbeddingModels(kbs []*entity.Knowledgebase) error { + if len(kbs) == 0 { + return fmt.Errorf("retrieval: no datasets selected") + } + for _, kb := range kbs { + if kb == nil { + return fmt.Errorf("retrieval: dataset record is nil") + } + } + firstKey := knowledgebaseEmbeddingKey(kbs[0]) + for _, kb := range kbs[1:] { + if knowledgebaseEmbeddingKey(kb) != firstKey { + return fmt.Errorf("retrieval: datasets use different embedding models") + } + } + return nil +} + +func knowledgebaseEmbeddingKey(kb *entity.Knowledgebase) string { + if kb.TenantEmbdID != nil && strings.TrimSpace(*kb.TenantEmbdID) != "" { + return "tenant:" + strings.TrimSpace(*kb.TenantEmbdID) + } + if strings.TrimSpace(kb.EmbdID) != "" { + return "embedding:" + strings.TrimSpace(kb.EmbdID) + } + return "default:" + strings.TrimSpace(kb.TenantID) +} + +func (a *NLPRetrievalAdapter) resolveEmbeddingModel( + ctx context.Context, + kb *entity.Knowledgebase, +) (*modelModule.EmbeddingModel, error) { + if a == nil || a.modelResolver == nil { + return nil, fmt.Errorf("retrieval: embedding model resolver is not configured") + } + + var ( + driver modelModule.ModelDriver + modelName string + apiConfig *modelModule.APIConfig + maxTokens int + err error + ) + switch { + case kb.TenantEmbdID != nil && strings.TrimSpace(*kb.TenantEmbdID) != "": + driver, modelName, apiConfig, maxTokens, err = a.modelResolver.GetModelConfigByID( + ctx, kb.TenantID, entity.ModelTypeEmbedding, *kb.TenantEmbdID, + ) + case strings.TrimSpace(kb.EmbdID) != "": + driver, modelName, apiConfig, maxTokens, err = a.modelResolver.ResolveModelConfig( + ctx, kb.TenantID, entity.ModelTypeEmbedding, kb.EmbdID, + ) + default: + driver, modelName, apiConfig, maxTokens, err = a.modelResolver.GetTenantDefaultModelByType( + ctx, kb.TenantID, entity.ModelTypeEmbedding, + ) + } + if err != nil { + return nil, fmt.Errorf("retrieval: resolve embedding model for dataset %s: %w", kb.ID, err) + } + return modelModule.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens), nil +} + +func (a *NLPRetrievalAdapter) resolveChatModel( + ctx context.Context, + req RetrievalRequest, + tenantID string, +) (*modelModule.ChatModel, error) { + method, _ := req.MetaDataFilter["method"].(string) + needsChatModel := req.TOCEnhance || method == "auto" || method == "semi_auto" + if !needsChatModel { + return nil, nil + } + if a == nil || a.modelResolver == nil { + return nil, fmt.Errorf("retrieval: model resolver is not configured") + } + driver, modelName, apiConfig, _, err := a.modelResolver.GetTenantDefaultModelByType( + ctx, tenantID, entity.ModelTypeChat, + ) + if err != nil { + return nil, fmt.Errorf("retrieval: resolve default chat model: %w", err) + } + return modelModule.NewChatModel(driver, &modelName, apiConfig), nil +} + +func (a *NLPRetrievalAdapter) resolveRerankModel( + ctx context.Context, + req RetrievalRequest, + tenantID string, +) (*modelModule.RerankModel, error) { + if req.RerankID == "" { + return nil, nil + } + if a == nil || a.modelResolver == nil { + return nil, fmt.Errorf("retrieval: model resolver is not configured") + } + var ( + driver modelModule.ModelDriver + modelName string + apiConfig *modelModule.APIConfig + err error + ) + driver, modelName, apiConfig, _, err = a.modelResolver.ResolveModelConfig( + ctx, tenantID, entity.ModelTypeRerank, req.RerankID, + ) + if err != nil { + return nil, fmt.Errorf("retrieval: resolve rerank model: %w", err) + } + return modelModule.NewRerankModel(driver, &modelName, apiConfig), nil } // translateChunk converts one nlp chunk map into a RetrievalChunk. diff --git a/internal/agent/tool/retrieval_nlp_test.go b/internal/agent/tool/retrieval_nlp_test.go index 5e7c7cb983..ed15ee6cd8 100644 --- a/internal/agent/tool/retrieval_nlp_test.go +++ b/internal/agent/tool/retrieval_nlp_test.go @@ -33,6 +33,8 @@ import ( "testing" "ragflow/internal/entity" + modelModule "ragflow/internal/entity/models" + "ragflow/internal/service/nlp" "gorm.io/gorm" ) @@ -223,7 +225,7 @@ func TestTranslateChunk_MissingAllScores(t *testing.T) { // must produce an adapter whose Search returns the missing-service // error, not a panic. func TestNewNLPRetrievalAdapter_NilService(t *testing.T) { - a := NewNLPRetrievalAdapter(nil) + a := NewNLPRetrievalAdapter(nil, nil, nil) _, err := a.Search(context.TODO(), nil, RetrievalRequest{Query: "hi"}) if err == nil { t.Fatal("expected error from nil-service adapter") @@ -235,14 +237,16 @@ func TestNewNLPRetrievalAdapter_NilService(t *testing.T) { func TestNLPRequestFromRetrieval_ThreadsSearchControls(t *testing.T) { keywordWeight := 0.25 + similarityThreshold := 0.42 + embeddingModel := &modelModule.EmbeddingModel{} got := nlpRequestFromRetrieval(RetrievalRequest{ Query: "hi", DatasetIDs: []string{"kb-1"}, TopN: 3, TopK: 99, KeywordsSimilarityWeight: &keywordWeight, - SimilarityThreshold: 0.42, - }, []string{"tenant-a"}, 3) + SimilarityThreshold: &similarityThreshold, + }, []string{"tenant-a"}, 3, embeddingModel) if got.Question != "hi" { t.Fatalf("Question=%q want hi", got.Question) @@ -265,6 +269,9 @@ func TestNLPRequestFromRetrieval_ThreadsSearchControls(t *testing.T) { if got.VectorSimilarityWeight == nil || !floatEqual(*got.VectorSimilarityWeight, 0.75) { t.Fatalf("VectorSimilarityWeight=%v want 0.75", got.VectorSimilarityWeight) } + if got.EmbeddingModel != embeddingModel { + t.Fatal("EmbeddingModel was not passed to nlp retrieval request") + } } func TestNLPRequestFromRetrieval_FallsBackToTopNHeadroom(t *testing.T) { @@ -272,7 +279,7 @@ func TestNLPRequestFromRetrieval_FallsBackToTopNHeadroom(t *testing.T) { Query: "hi", DatasetIDs: []string{"kb-1"}, TopN: 3, - }, []string{"tenant-a"}, 3) + }, []string{"tenant-a"}, 3, &modelModule.EmbeddingModel{}) if got.Top == nil || *got.Top != 12 { t.Fatalf("Top=%v want 12", got.Top) @@ -282,25 +289,57 @@ func TestNLPRequestFromRetrieval_FallsBackToTopNHeadroom(t *testing.T) { } } -func TestNLPRetrievalAdapter_ResolveTenantIDsStaysWithinRequestTenant(t *testing.T) { - a := &NLPRetrievalAdapter{} - got, err := a.resolveTenantIDs(nil, nil, RetrievalRequest{ - TenantID: "tenant-a", - DatasetIDs: []string{"kb-1", "kb-2", "kb-missing"}, - }) - if err != nil { - t.Fatalf("resolveTenantIDs: %v", err) - } +func TestNLPRequestFromRetrieval_PreservesExplicitZeroSimilarityThreshold(t *testing.T) { + similarityThreshold := 0.0 + got := nlpRequestFromRetrieval(RetrievalRequest{ + Query: "hi", + DatasetIDs: []string{"kb-1"}, + SimilarityThreshold: &similarityThreshold, + }, []string{"tenant-a"}, 3, &modelModule.EmbeddingModel{}) - if len(got) != 1 { - t.Fatalf("tenantIDs len=%d want 1, got=%v", len(got), got) - } - if got[0] != "tenant-a" { - t.Fatalf("tenantIDs=%v want [tenant-a]", got) + if got.SimilarityThreshold == nil || *got.SimilarityThreshold != 0 { + t.Fatalf("SimilarityThreshold = %v; want explicit zero", got.SimilarityThreshold) } } -func TestNLPRetrievalAdapter_ResolveTenantIDsFromDatasetIDs(t *testing.T) { +func TestNLPRetrievalAdapter_ResolveDatasetsRequiresEveryDataset(t *testing.T) { + a := &NLPRetrievalAdapter{ + kbDAO: fakeKnowledgebaseLookup{ + kbs: []*entity.Knowledgebase{ + {ID: "kb-1", TenantID: "tenant-a"}, + {ID: "kb-2", TenantID: "tenant-a"}, + }, + }, + } + _, err := a.resolveDatasets(nil, nil, RetrievalRequest{ + TenantID: "tenant-a", + DatasetIDs: []string{"kb-1", "kb-2", "kb-missing"}, + }) + if err == nil { + t.Fatal("resolveDatasets: expected missing dataset error") + } +} + +func TestNLPRetrievalAdapter_SearchRejectsDatasetsFromMultipleTenants(t *testing.T) { + a := &NLPRetrievalAdapter{ + svc: &nlp.RetrievalService{}, + kbDAO: fakeKnowledgebaseLookup{ + kbs: []*entity.Knowledgebase{ + {ID: "kb-1", TenantID: "tenant-a", EmbdID: "embedding@provider"}, + {ID: "kb-2", TenantID: "tenant-b", EmbdID: "embedding@provider"}, + }, + }, + } + _, err := a.Search(t.Context(), nil, RetrievalRequest{ + Query: "hi", + DatasetIDs: []string{"kb-1", "kb-2"}, + }) + if err == nil || err.Error() != "retrieval: datasets span multiple tenants" { + t.Fatalf("Search error = %v, want multiple-tenant rejection", err) + } +} + +func TestNLPRetrievalAdapter_ResolveDatasetsFromDatasetIDs(t *testing.T) { a := &NLPRetrievalAdapter{ kbDAO: fakeKnowledgebaseLookup{ kbs: []*entity.Knowledgebase{ @@ -310,18 +349,18 @@ func TestNLPRetrievalAdapter_ResolveTenantIDsFromDatasetIDs(t *testing.T) { }, }, } - got, err := a.resolveTenantIDs(nil, nil, RetrievalRequest{ + got, err := a.resolveDatasets(nil, nil, RetrievalRequest{ DatasetIDs: []string{"kb-1", "kb-2", "kb-3", " "}, }) if err != nil { - t.Fatalf("resolveTenantIDs: %v", err) + t.Fatalf("resolveDatasets: %v", err) } - if len(got) != 2 || got[0] != "tenant-a" || got[1] != "tenant-b" { - t.Fatalf("tenantIDs=%v want [tenant-a tenant-b]", got) + if len(got.tenantIDs) != 2 || got.tenantIDs[0] != "tenant-a" || got.tenantIDs[1] != "tenant-b" { + t.Fatalf("tenantIDs=%v want [tenant-a tenant-b]", got.tenantIDs) } } -func TestNLPRetrievalAdapter_ResolveTenantIDsKeepsRequestTenantFirst(t *testing.T) { +func TestNLPRetrievalAdapter_ResolveDatasetsUsesDatasetTenants(t *testing.T) { a := &NLPRetrievalAdapter{ kbDAO: fakeKnowledgebaseLookup{ kbs: []*entity.Knowledgebase{ @@ -329,15 +368,81 @@ func TestNLPRetrievalAdapter_ResolveTenantIDsKeepsRequestTenantFirst(t *testing. }, }, } - got, err := a.resolveTenantIDs(nil, nil, RetrievalRequest{ + got, err := a.resolveDatasets(nil, nil, RetrievalRequest{ TenantID: "tenant-a", DatasetIDs: []string{"kb-1"}, }) if err != nil { - t.Fatalf("resolveTenantIDs: %v", err) + t.Fatalf("resolveDatasets: %v", err) } - if len(got) != 2 || got[0] != "tenant-a" || got[1] != "tenant-b" { - t.Fatalf("tenantIDs=%v want [tenant-a tenant-b]", got) + if len(got.tenantIDs) != 1 || got.tenantIDs[0] != "tenant-b" { + t.Fatalf("tenantIDs=%v want [tenant-b]", got.tenantIDs) + } +} + +func TestNLPRetrievalAdapter_ResolveEmbeddingModelPriority(t *testing.T) { + tenantEmbeddingID := "tenant-embedding-1" + tests := []struct { + name string + kb *entity.Knowledgebase + wantCall string + }{ + { + name: "tenant embedding id", + kb: &entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + EmbdID: "embedding@provider", + TenantEmbdID: &tenantEmbeddingID, + }, + wantCall: "id:tenant-embedding-1", + }, + { + name: "knowledge base embedding reference", + kb: &entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + EmbdID: "embedding@provider", + }, + wantCall: "resolve:embedding@provider", + }, + { + name: "tenant default only when knowledge base has no model", + kb: &entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + }, + wantCall: "default", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + resolver := &fakeModelResolver{modelName: "resolved-model"} + adapter := &NLPRetrievalAdapter{modelResolver: resolver} + model, err := adapter.resolveEmbeddingModel(t.Context(), test.kb) + if err != nil { + t.Fatalf("resolveEmbeddingModel: %v", err) + } + if resolver.call != test.wantCall { + t.Fatalf("resolver call = %q, want %q", resolver.call, test.wantCall) + } + if model == nil || model.ModelName == nil || *model.ModelName != "resolved-model" { + t.Fatalf("resolved model = %#v", model) + } + }) + } +} + +func TestValidateEmbeddingModelsRejectsDifferentTenantModels(t *testing.T) { + firstID := "tenant-embedding-1" + secondID := "tenant-embedding-2" + err := validateEmbeddingModels([]*entity.Knowledgebase{ + {ID: "kb-1", TenantID: "tenant-1", TenantEmbdID: &firstID}, + {ID: "kb-2", TenantID: "tenant-1", TenantEmbdID: &secondID}, + }) + if err == nil { + t.Fatal("expected different tenant embedding models to be rejected") } } @@ -346,6 +451,63 @@ type fakeKnowledgebaseLookup struct { err error } -func (f fakeKnowledgebaseLookup) GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Knowledgebase, error) { - return f.kbs, f.err +type fakeModelResolver struct { + call string + modelName string + err error +} + +func (f *fakeModelResolver) GetModelConfigByID( + _ context.Context, + _ string, + _ entity.ModelType, + modelID string, +) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + f.call = "id:" + modelID + return nil, f.modelName, &modelModule.APIConfig{}, 512, f.err +} + +func (f *fakeModelResolver) ResolveModelConfig( + _ context.Context, + _ string, + _ entity.ModelType, + modelRef string, +) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + f.call = "resolve:" + modelRef + return nil, f.modelName, &modelModule.APIConfig{}, 512, f.err +} + +func (f *fakeModelResolver) GetTenantDefaultModelByType( + _ context.Context, + _ string, + _ entity.ModelType, +) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + f.call = "default" + return nil, f.modelName, &modelModule.APIConfig{}, 512, f.err +} + +func (f fakeKnowledgebaseLookup) GetByIDs(ctx context.Context, db *gorm.DB, ids []string) ([]*entity.Knowledgebase, error) { + requested := make(map[string]struct{}, len(ids)) + for _, id := range ids { + requested[id] = struct{}{} + } + kbs := make([]*entity.Knowledgebase, 0, len(f.kbs)) + for _, kb := range f.kbs { + if kb == nil { + continue + } + if _, ok := requested[kb.ID]; ok { + kbs = append(kbs, kb) + } + } + return kbs, f.err +} + +func (f fakeKnowledgebaseLookup) GetByName(_ context.Context, _ *gorm.DB, name, tenantID string) (*entity.Knowledgebase, error) { + for _, kb := range f.kbs { + if kb != nil && kb.Name == name && kb.TenantID == tenantID { + return kb, nil + } + } + return nil, gorm.ErrRecordNotFound } diff --git a/internal/agent/tool/retrieval_service.go b/internal/agent/tool/retrieval_service.go index 9b44733acc..0ca0ae2c84 100644 --- a/internal/agent/tool/retrieval_service.go +++ b/internal/agent/tool/retrieval_service.go @@ -50,29 +50,40 @@ type RetrievalChunk struct { type RetrievalRequest struct { Query string DatasetIDs []string + MemoryIDs []string TopN int TopK int KeywordsSimilarityWeight *float64 UseKG bool - SimilarityThreshold float64 + SimilarityThreshold *float64 + RerankID string + CrossLanguages []string + TOCEnhance bool + MetaDataFilter map[string]any + RetrievalFrom string // DocScope restricts retrieval to a set of document ids (the doc_id list // routed by the dataset_navigation_by_tree tool). Empty = no doc filter. DocScope []string // 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 + // It is used for dataset-name resolution and memory access. Reads from // CanvasState.Sys["user_id"] when empty (set by the Begin component at // internal/agent/component/begin.go:82). TenantID string } -// RetrievalService is the interface the Retrieval tool uses. -// Today only the stub impl exists; production code can register -// a real impl via SetRetrievalService during boot. +// RetrievalService is the knowledge-base search interface used by the tool. +// The server installs NLPRetrievalAdapter during boot. type RetrievalService interface { Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) } +// MemoryRetrievalService is the memory-message retrieval surface used when +// retrieval_from=memory. It is separate from knowledge-base retrieval because +// memory messages live in different indices and have a different result shape. +type MemoryRetrievalService interface { + Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) +} + // KGRetrievalService is the GraphRAG retrieval surface. The // KBRetrieval service and the KGRetrieval service are kept // separate on purpose: the kg backend's signature requires @@ -85,16 +96,18 @@ type KGRetrievalService interface { Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) } -// ErrRetrievalServiceMissing is declared in retrieval.go (kept -// for backward compat with retrieval_test.go). Callers using -// the RetrievalService interface here can match against the -// same sentinel by referring to the same package-level var. - +// ErrRetrievalServiceMissing is declared in retrieval.go so callers and the +// default stub share the same sentinel. var ( retrievalServiceMu sync.RWMutex retrievalServiceImpl RetrievalService = stubRetrievalService{} ) +var ( + memoryRetrievalServiceMu sync.RWMutex + memoryRetrievalServiceImpl MemoryRetrievalService = stubMemoryRetrievalService{} +) + func SetRetrievalService(svc RetrievalService) { retrievalServiceMu.Lock() defer retrievalServiceMu.Unlock() @@ -111,17 +124,36 @@ func GetRetrievalService() RetrievalService { return retrievalServiceImpl } +func SetMemoryRetrievalService(svc MemoryRetrievalService) { + memoryRetrievalServiceMu.Lock() + defer memoryRetrievalServiceMu.Unlock() + if svc == nil { + memoryRetrievalServiceImpl = stubMemoryRetrievalService{} + return + } + memoryRetrievalServiceImpl = svc +} + +func GetMemoryRetrievalService() MemoryRetrievalService { + memoryRetrievalServiceMu.RLock() + defer memoryRetrievalServiceMu.RUnlock() + return memoryRetrievalServiceImpl +} + type stubRetrievalService struct{} func (stubRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { return nil, ErrRetrievalServiceMissing } -// simpleRetrievalService is a deterministic test/demo impl that -// returns synthetic chunks based on the query. Useful for -// development and integration tests; the production impl lands -// when the boot path wires internal/service.ChunkService into -// SetRetrievalService. +type stubMemoryRetrievalService struct{} + +func (stubMemoryRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { + return nil, ErrMemoryRetrievalServiceMissing +} + +// simpleRetrievalService is a deterministic test implementation that returns +// synthetic chunks based on the query. type simpleRetrievalService struct{} func (simpleRetrievalService) Search(_ context.Context, _ *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) { @@ -154,12 +186,8 @@ func (simpleRetrievalService) Search(_ context.Context, _ *gorm.DB, req Retrieva return chunks, nil } -// SetSimpleRetrievalService installs the simpleRetrievalService -// (deterministic synthetic chunks). Useful for development and -// integration tests. Production code should call -// SetRetrievalService with a real implementation backed by -// internal/service.ChunkService — see design doc §4.2 -// RetrievalService. +// SetSimpleRetrievalService installs deterministic synthetic retrieval for +// tests and local demos. func SetSimpleRetrievalService() { SetRetrievalService(simpleRetrievalService{}) } @@ -175,6 +203,10 @@ var ErrKGRetrievalServiceMissing = errors.New( "call tool.SetKGRetrievalService(tool.NewKGRetrievalAdapter(...)) at boot", ) +var ErrMemoryRetrievalServiceMissing = errors.New( + "memory retrieval service not registered", +) + var ( kgRetrievalServiceMu sync.RWMutex kgRetrievalServiceImpl KGRetrievalService = stubKGRetrievalService{} diff --git a/internal/agent/tool/retrieval_test.go b/internal/agent/tool/retrieval_test.go index baac86a8c0..924ab1a258 100644 --- a/internal/agent/tool/retrieval_test.go +++ b/internal/agent/tool/retrieval_test.go @@ -32,7 +32,7 @@ func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) { t.Parallel() rt := NewRetrievalTool() - out, err := rt.InvokableRun(context.Background(), `{"query":"hello"}`) + out, err := rt.InvokableRun(context.Background(), `{"query":"hello","dataset_ids":["kb-1"]}`) if err == nil { t.Fatal("expected stub error, got nil") } @@ -55,7 +55,6 @@ func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) { func TestRetrieval_RejectsUseKG(t *testing.T) { t.Parallel() - rt := NewRetrievalTool() out, err := rt.InvokableRun(context.Background(), `{"query":"x","use_kg":true}`) if !errors.Is(err, ErrGraphRAGNotSupported) { @@ -171,6 +170,10 @@ func TestRetrieval_UsesNodeParamsAsDefaults(t *testing.T) { "top_k": float64(99), "keywords_similarity_weight": 0.7, "similarity_threshold": 0.42, + "rerank_id": "rerank@provider", + "cross_languages": []any{"English", "Chinese"}, + "toc_enhance": true, + "meta_data_filter": map[string]any{"method": "manual"}, }) if err != nil { t.Fatalf("BuildByName(retrieval): %v", err) @@ -196,12 +199,46 @@ func TestRetrieval_UsesNodeParamsAsDefaults(t *testing.T) { 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 { + if svc.req.SimilarityThreshold == nil || *svc.req.SimilarityThreshold != 0.42 { t.Fatalf("SimilarityThreshold=%v want 0.42", svc.req.SimilarityThreshold) } + if svc.req.RerankID != "rerank@provider" { + t.Fatalf("RerankID=%q", svc.req.RerankID) + } + if len(svc.req.CrossLanguages) != 2 || svc.req.CrossLanguages[1] != "Chinese" { + t.Fatalf("CrossLanguages=%#v", svc.req.CrossLanguages) + } + if !svc.req.TOCEnhance || svc.req.MetaDataFilter["method"] != "manual" { + t.Fatalf("enhancement request=%#v", svc.req) + } } -func TestRetrieval_AcceptsEmptyResponseNodeParam(t *testing.T) { +func TestRetrieval_ExplicitZeroSimilarityArgsOverrideDefaults(t *testing.T) { + previous := GetRetrievalService() + service := &capturingRetrievalService{} + SetRetrievalService(service) + t.Cleanup(func() { SetRetrievalService(previous) }) + + similarityThreshold := 0.42 + keywordsSimilarityWeight := 0.7 + retrievalTool := NewRetrievalToolWithDefaults(retrievalArgs{ + DatasetIDs: []string{"kb-1"}, + SimilarityThreshold: &similarityThreshold, + KeywordsSimilarityWeight: &keywordsSimilarityWeight, + }) + _, err := retrievalTool.InvokableRun(context.Background(), `{"query":"hello","similarity_threshold":0,"keywords_similarity_weight":0}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if service.req.SimilarityThreshold == nil || *service.req.SimilarityThreshold != 0 { + t.Fatalf("SimilarityThreshold = %v; want explicit zero", service.req.SimilarityThreshold) + } + if service.req.KeywordsSimilarityWeight == nil || *service.req.KeywordsSimilarityWeight != 0 { + t.Fatalf("KeywordsSimilarityWeight = %v; want explicit zero", service.req.KeywordsSimilarityWeight) + } +} + +func TestRetrieval_ParsesEnhancementNodeParams(t *testing.T) { t.Parallel() built, err := BuildByName("retrieval", map[string]any{ @@ -209,7 +246,7 @@ func TestRetrieval_AcceptsEmptyResponseNodeParam(t *testing.T) { "toc_enhance": true, "meta_data_filter": map[string]any{"method": "manual"}, "empty_response": "empty", - "retrieval_from": "database", + "retrieval_from": "dataset", "memory_ids": []any{"memory-1"}, "kb_vars": map[string]any{"x": "y"}, "cross_languages": []any{"English"}, @@ -221,8 +258,14 @@ func TestRetrieval_AcceptsEmptyResponseNodeParam(t *testing.T) { 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("unimplemented params should not mutate retrieval defaults: %#v", rt.defaults) + if rt.defaults.RerankID != "rerank-1" || !rt.defaults.TOCEnhance { + t.Fatalf("enhancement defaults were not parsed: %#v", rt.defaults) + } + if len(rt.defaults.CrossLanguages) != 1 || rt.defaults.CrossLanguages[0] != "English" { + t.Fatalf("CrossLanguages = %#v", rt.defaults.CrossLanguages) + } + if rt.defaults.MetaDataFilter["method"] != "manual" { + t.Fatalf("MetaDataFilter = %#v", rt.defaults.MetaDataFilter) } if rt.defaults.EmptyResponse != "empty" { t.Fatalf("EmptyResponse = %q, want empty", rt.defaults.EmptyResponse) @@ -246,6 +289,53 @@ func TestRetrieval_UsesEmptyResponseForEmptyQuery(t *testing.T) { } } +func TestRetrieval_RoutesMemoryRequests(t *testing.T) { + previous := GetMemoryRetrievalService() + memoryService := &capturingRetrievalService{} + SetMemoryRetrievalService(memoryService) + t.Cleanup(func() { SetMemoryRetrievalService(previous) }) + + rt := NewRetrievalToolWithDefaults(retrievalArgs{ + MemoryIDs: []string{"memory-1"}, + }) + if _, err := rt.InvokableRun(t.Context(), `{"query":"remember this"}`); err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if memoryService.req.RetrievalFrom != "memory" { + t.Fatalf("RetrievalFrom = %q, want memory", memoryService.req.RetrievalFrom) + } + if len(memoryService.req.MemoryIDs) != 1 || memoryService.req.MemoryIDs[0] != "memory-1" { + t.Fatalf("MemoryIDs = %#v", memoryService.req.MemoryIDs) + } +} + +func TestRetrieval_ResolvesCanvasVariables(t *testing.T) { + state := runtime.NewCanvasState("run-1", "session-1") + state.SetVar("source", "ids", []any{"kb-1", "kb-2"}) + state.SetVar("source", "query", "semantic question") + state.SetVar("source", "value", "2026") + ctx := runtime.WithState(context.Background(), state) + + ids, err := resolveRetrievalDatasetIDs(ctx, []string{"source@ids"}) + if err != nil { + t.Fatalf("resolveRetrievalDatasetIDs: %v", err) + } + if len(ids) != 2 || ids[0] != "kb-1" || ids[1] != "kb-2" { + t.Fatalf("resolved dataset IDs = %#v", ids) + } + query, err := resolveRetrievalQuery(ctx, "{{source@query}}") + if err != nil || query != "semantic question" { + t.Fatalf("resolved query = %q, err = %v", query, err) + } + filter, err := resolveRetrievalFilter(ctx, map[string]any{ + "method": "manual", + "value": "{{source@value}}", + }) + if err != nil || filter["value"] != "2026" { + t.Fatalf("resolved filter = %#v, err = %v", filter, err) + } +} + func TestRetrieval_OmitsUnsetEmptyResponseFromArguments(t *testing.T) { arguments, err := json.Marshal(retrievalArgs{Query: "love"}) if err != nil { @@ -261,7 +351,7 @@ func TestRetrieval_UsesEmptyResponseWhenSearchHasNoChunks(t *testing.T) { SetRetrievalService(staticRetrievalService{}) t.Cleanup(func() { SetRetrievalService(prev) }) - rt := NewRetrievalToolWithDefaults(retrievalArgs{EmptyResponse: "No matching chunk."}) + rt := NewRetrievalToolWithDefaults(retrievalArgs{DatasetIDs: []string{"kb-1"}, EmptyResponse: "No matching chunk."}) out, err := rt.InvokableRun(context.Background(), `{"query":"love"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) diff --git a/internal/agent/tool/retrieval_wiring_test.go b/internal/agent/tool/retrieval_wiring_test.go index fdc86829e3..347b5661e4 100644 --- a/internal/agent/tool/retrieval_wiring_test.go +++ b/internal/agent/tool/retrieval_wiring_test.go @@ -31,7 +31,7 @@ func TestRetrievalTool_StubReturnsServiceMissing(t *testing.T) { SetRetrievalService(nil) tool := NewRetrievalTool() ctx := t.Context() - _, err := tool.InvokableRun(ctx, `{"query":"hi"}`) + _, err := tool.InvokableRun(ctx, `{"query":"hi","dataset_ids":["kb-1"]}`) if !errors.Is(err, ErrRetrievalServiceMissing) { t.Errorf("err=%v, want ErrRetrievalServiceMissing", err) } @@ -46,7 +46,7 @@ func TestRetrievalTool_SimpleServiceReturnsChunks(t *testing.T) { tool := NewRetrievalTool() ctx := t.Context() - out, err := tool.InvokableRun(ctx, `{"query":"hello world","top_n":2}`) + out, err := tool.InvokableRun(ctx, `{"query":"hello world","dataset_ids":["kb-1"],"top_n":2}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } diff --git a/internal/service/memory.go b/internal/service/memory.go index c6599b6dea..6a7b6a8103 100644 --- a/internal/service/memory.go +++ b/internal/service/memory.go @@ -1042,10 +1042,30 @@ func (s *MemoryService) SearchMessage(ctx context.Context, userID string, filter if len(memories) == 0 { return []map[string]interface{}{}, common.CodeSuccess, nil } + if err := validateMemorySearchModels(memories); err != nil { + return nil, common.CodeArgumentError, err + } return s.queryMessage(ctx, memories, filterDict, params) } +func validateMemorySearchModels(memories []*entity.Memory) error { + if len(memories) == 0 { + return nil + } + firstKey := memorySearchEmbeddingKey(memories[0]) + for _, memory := range memories[1:] { + if memorySearchEmbeddingKey(memory) != firstKey { + return fmt.Errorf("memories use different embedding models") + } + } + return nil +} + +func memorySearchEmbeddingKey(memory *entity.Memory) string { + return "embedding:" + strings.TrimSpace(memory.EmbdID) +} + func (s *MemoryService) queryMessage(ctx context.Context, memories []*entity.Memory, filterDict, params map[string]interface{}) ([]map[string]interface{}, common.ErrorCode, error) { if s.docEngine == nil { return nil, common.CodeServerError, errors.New("message store is not initialized") diff --git a/internal/service/memory_message_test.go b/internal/service/memory_message_test.go index a8d1887ac7..e0a21937a9 100644 --- a/internal/service/memory_message_test.go +++ b/internal/service/memory_message_test.go @@ -27,6 +27,24 @@ func TestIsMessageDocumentNotFound(t *testing.T) { } } +func TestValidateMemorySearchModels(t *testing.T) { + firstTenantEmbeddingID := "tenant-embedding-1" + secondTenantEmbeddingID := "tenant-embedding-2" + if err := validateMemorySearchModels([]*entity.Memory{ + {ID: "memory-1", EmbdID: "embedding@provider", TenantEmbdID: &firstTenantEmbeddingID}, + {ID: "memory-2", EmbdID: "embedding@provider", TenantEmbdID: &secondTenantEmbeddingID}, + }); err != nil { + t.Fatalf("same memory embedding model rejected: %v", err) + } + + if err := validateMemorySearchModels([]*entity.Memory{ + {ID: "memory-1", EmbdID: "embedding-a@provider"}, + {ID: "memory-2", EmbdID: "embedding-b@provider"}, + }); err == nil { + t.Fatal("different memory embedding models were accepted") + } +} + func TestRequireMemoryAccessReturnsCanceledContext(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel()