From 8f780799c88b85b0ff1ead69652891d248d6e381 Mon Sep 17 00:00:00 2001 From: Sbaaoui Idriss <112825897+6ba3i@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:15:30 +0800 Subject: [PATCH] fix(ingestor): initialize shared retrieval services (#22172) --- internal/server/api/server.go | 61 +------------ internal/server/bootstrap/retrieval.go | 94 +++++++++++++++++++++ internal/server/bootstrap/retrieval_test.go | 82 ++++++++++++++++++ internal/server/ingestor/server.go | 10 ++- 4 files changed, 185 insertions(+), 62 deletions(-) create mode 100644 internal/server/bootstrap/retrieval.go create mode 100644 internal/server/bootstrap/retrieval_test.go diff --git a/internal/server/api/server.go b/internal/server/api/server.go index c993848bde..5c7e669263 100644 --- a/internal/server/api/server.go +++ b/internal/server/api/server.go @@ -26,14 +26,8 @@ import ( "ragflow/internal/agent/audio" "ragflow/internal/agent/canvas" - "ragflow/internal/agent/retrievalbridge" - "ragflow/internal/agent/runtime" - agenttool "ragflow/internal/agent/tool" - agentic_rag "ragflow/internal/agentic_rag" "ragflow/internal/channels" - modelModule "ragflow/internal/entity/models" "ragflow/internal/handler" - "ragflow/internal/permission" "ragflow/internal/router" "ragflow/internal/server/local" "ragflow/internal/service" @@ -145,60 +139,7 @@ func startServer(ctx context.Context, serverName string, arguments *bootstrap.Se // Initialize doc engine for skill search docEngine := engine.Get() documentDAO := dao.NewDocumentDAO() - retrievalEnhancer := retrievalbridge.NewEnhancer(docEngine, metadataService) - // Keep the concrete adapter: it is both the canvas/agent-tool backend and the - // target of the agentic-RAG bridge wired below. The bridge must hold it - // directly (not look it up in the registry) because the registry is shared — - // agenttool.SetRetrievalService and runtime.SetRetrievalService write the same - // singleton, so a bridge that resolved its target per call would find itself. - retrievalAdapter := agenttool.NewNLPRetrievalAdapterFromDeps( - docEngine, - documentDAO, - agenttool.RetrievalModelConstructors{ - Embedding: func(ctx context.Context, tenantID, modelRef string) (*modelModule.EmbeddingModel, error) { - if strings.TrimSpace(modelRef) == "" { - return modelFactory.NewDefaultEmbeddingModel(permission.WithTenantID(ctx, tenantID)) - } - return modelFactory.NewEmbeddingModel(permission.WithTenantID(ctx, tenantID), modelRef) - }, - Chat: func(ctx context.Context, tenantID, modelRef string) (*modelModule.ChatModel, error) { - if strings.TrimSpace(modelRef) == "" { - return modelFactory.NewDefaultChatModel(permission.WithTenantID(ctx, tenantID)) - } - return modelFactory.NewChatModel(permission.WithTenantID(ctx, tenantID), modelRef) - }, - Rerank: func(ctx context.Context, tenantID, modelRef string) (*modelModule.RerankModel, error) { - return modelFactory.NewRerankModel(permission.WithTenantID(ctx, tenantID), modelRef) - }, - }, - retrievalEnhancer, - ) - agenttool.SetRetrievalService(retrievalAdapter) - agenttool.SetMemoryRetrievalService(retrievalbridge.NewMemoryAdapter(memoryService, retrievalEnhancer)) - common.Info("agent: retrieval service adapter installed") - - // The smart-reasoning agent's corpus tools (internal/agentic_rag): regex - // pushdown and lexical BM25 over the same doc engine the retrieval adapter - // wraps. The locate tools resolve these services per call, so they must be - // registered before the server starts serving. - runtime.SetGrepService(agentic_rag.NewGrepAdapter(docEngine)) - bm25Adapter := agentic_rag.NewBm25Adapter(docEngine) - bm25Adapter.SetQueryBuilder(nlp.GetQueryBuilder()) - runtime.SetBm25Service(bm25Adapter) - common.Info("agent: smart-reasoning corpus services installed (grep + bm25)") - - // The compiled-knowledge graph walk (graph_explore) resolves its service per - // call the same way. Without this registration GetKgService() returns the - // failing stub, so the tool could only ever report an infrastructure error. - // The seed encoder is left unset: the walk then seeds by keyword match on the - // entity rows, the documented fallback on a deployment without an embedder. - agentic_rag.SetKgService(agentic_rag.NewKgAdapter(docEngine)) - common.Info("agent: smart-reasoning compiled-knowledge graph service installed") - // The smart-reasoning agent's corpus tools resolve the shared retrieval - // service per call, so the production adapter must be registered before the - // server starts serving. - runtime.SetRetrievalService(retrievalbridge.NewRuntimeAdapter(retrievalAdapter)) - common.Info("agent: runtime retrieval service wired") + bootstrap.InitializeRetrievalServices(metadataService, memoryService, modelFactory) // Initialize handler layer authHandler := handler.NewAuthHandler() diff --git a/internal/server/bootstrap/retrieval.go b/internal/server/bootstrap/retrieval.go new file mode 100644 index 0000000000..7e3dc29985 --- /dev/null +++ b/internal/server/bootstrap/retrieval.go @@ -0,0 +1,94 @@ +// +// 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 bootstrap + +import ( + "context" + "strings" + + "ragflow/internal/agent/retrievalbridge" + "ragflow/internal/agent/runtime" + agenttool "ragflow/internal/agent/tool" + "ragflow/internal/agentic_rag" + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine" + modelModule "ragflow/internal/entity/models" + "ragflow/internal/permission" + "ragflow/internal/service" + "ragflow/internal/service/nlp" +) + +// InitializeRetrievalServices installs corpus tools for API and ingestor processes. +func InitializeRetrievalServices(metadataService *service.MetadataService, memoryService *service.MemoryService, modelFactory *service.ModelFactory) { + docEngine := engine.Get() + documentDAO := dao.NewDocumentDAO() + retrievalEnhancer := retrievalbridge.NewEnhancer(docEngine, metadataService) + // Keep the concrete adapter: it is both the canvas/agent-tool backend and the + // target of the agentic-RAG bridge wired below. The bridge must hold it + // directly (not look it up in the registry) because the registry is shared — + // agenttool.SetRetrievalService and runtime.SetRetrievalService write the same + // singleton, so a bridge that resolved its target per call would find itself. + retrievalAdapter := agenttool.NewNLPRetrievalAdapterFromDeps( + docEngine, + documentDAO, + agenttool.RetrievalModelConstructors{ + Embedding: func(ctx context.Context, tenantID, modelRef string) (*modelModule.EmbeddingModel, error) { + if strings.TrimSpace(modelRef) == "" { + return modelFactory.NewDefaultEmbeddingModel(permission.WithTenantID(ctx, tenantID)) + } + return modelFactory.NewEmbeddingModel(permission.WithTenantID(ctx, tenantID), modelRef) + }, + Chat: func(ctx context.Context, tenantID, modelRef string) (*modelModule.ChatModel, error) { + if strings.TrimSpace(modelRef) == "" { + return modelFactory.NewDefaultChatModel(permission.WithTenantID(ctx, tenantID)) + } + return modelFactory.NewChatModel(permission.WithTenantID(ctx, tenantID), modelRef) + }, + Rerank: func(ctx context.Context, tenantID, modelRef string) (*modelModule.RerankModel, error) { + return modelFactory.NewRerankModel(permission.WithTenantID(ctx, tenantID), modelRef) + }, + }, + retrievalEnhancer, + ) + agenttool.SetRetrievalService(retrievalAdapter) + agenttool.SetMemoryRetrievalService(retrievalbridge.NewMemoryAdapter(memoryService, retrievalEnhancer)) + common.Info("agent: retrieval service adapter installed") + + // The smart-reasoning agent's corpus tools (internal/agentic_rag): regex + // pushdown and lexical BM25 over the same doc engine the retrieval adapter + // wraps. The locate tools resolve these services per call, so they must be + // registered before the server starts serving. + runtime.SetGrepService(agentic_rag.NewGrepAdapter(docEngine)) + bm25Adapter := agentic_rag.NewBm25Adapter(docEngine) + bm25Adapter.SetQueryBuilder(nlp.GetQueryBuilder()) + runtime.SetBm25Service(bm25Adapter) + common.Info("agent: smart-reasoning corpus services installed (grep + bm25)") + + // The compiled-knowledge graph walk (graph_explore) resolves its service per + // call the same way. Without this registration GetKgService() returns the + // failing stub, so the tool could only ever report an infrastructure error. + // The seed encoder is left unset: the walk then seeds by keyword match on the + // entity rows, the documented fallback on a deployment without an embedder. + agentic_rag.SetKgService(agentic_rag.NewKgAdapter(docEngine)) + common.Info("agent: smart-reasoning compiled-knowledge graph service installed") + // The smart-reasoning agent's corpus tools resolve the shared retrieval + // service per call, so the production adapter must be registered before the + // server starts serving. + runtime.SetRetrievalService(retrievalbridge.NewRuntimeAdapter(retrievalAdapter)) + common.Info("agent: runtime retrieval service wired") +} diff --git a/internal/server/bootstrap/retrieval_test.go b/internal/server/bootstrap/retrieval_test.go new file mode 100644 index 0000000000..4718502241 --- /dev/null +++ b/internal/server/bootstrap/retrieval_test.go @@ -0,0 +1,82 @@ +// +// 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 bootstrap + +import ( + "errors" + "testing" + + "ragflow/internal/agent/retrievalbridge" + "ragflow/internal/agent/runtime" + agenttool "ragflow/internal/agent/tool" + "ragflow/internal/agentic_rag" + "ragflow/internal/service" +) + +func TestInitializeRetrievalServices(t *testing.T) { + retrieval := runtime.GetRetrievalService() + memory := runtime.GetMemoryRetrievalService() + grep := runtime.GetGrepService() + bm25 := runtime.GetBm25Service() + kg := agentic_rag.GetKgService() + t.Cleanup(func() { + runtime.SetRetrievalService(retrieval) + runtime.SetMemoryRetrievalService(memory) + runtime.SetGrepService(grep) + runtime.SetBm25Service(bm25) + agentic_rag.SetKgService(kg) + }) + + memoryService := service.NewMemoryService(nil, nil) + InitializeRetrievalServices(service.NewMetadataService(), memoryService, service.NewModelFactory()) + + if _, ok := runtime.GetRetrievalService().(*retrievalbridge.RuntimeAdapter); !ok { + t.Fatalf("retrieval service = %T, want runtime bridge", runtime.GetRetrievalService()) + } + if _, ok := runtime.GetMemoryRetrievalService().(*retrievalbridge.MemoryAdapter); !ok { + t.Errorf("memory service = %T, want memory adapter", runtime.GetMemoryRetrievalService()) + } + if _, ok := runtime.GetGrepService().(*agentic_rag.GrepAdapter); !ok { + t.Errorf("grep service = %T, want grep adapter", runtime.GetGrepService()) + } + if _, ok := runtime.GetBm25Service().(*agentic_rag.Bm25Adapter); !ok { + t.Errorf("BM25 service = %T, want BM25 adapter", runtime.GetBm25Service()) + } + if _, ok := agentic_rag.GetKgService().(*agentic_rag.KgAdapter); !ok { + t.Errorf("graph service = %T, want graph adapter", agentic_rag.GetKgService()) + } + + // Both registries share the runtime bridge, which must call the concrete NLP + // adapter rather than look itself up and recurse. + for _, svc := range []runtime.RetrievalService{runtime.GetRetrievalService(), agenttool.GetRetrievalService()} { + chunks, err := svc.Search(t.Context(), nil, runtime.RetrievalRequest{}) + if err != nil || len(chunks) != 0 { + t.Errorf("empty retrieval = %v, %v; want no chunks or error", chunks, err) + } + _, err = svc.Search(t.Context(), nil, runtime.RetrievalRequest{UseKG: true}) + if !errors.Is(err, agenttool.ErrGraphRAGNotSupported) { + t.Errorf("GraphRAG retrieval error = %v, want NLP adapter validation", err) + } + } + + // A configured memory service reaches request validation, not the missing + // service sentinel returned when the adapter has no memory backend. + _, err := runtime.GetMemoryRetrievalService().Search(t.Context(), nil, runtime.RetrievalRequest{}) + if err == nil || err.Error() != "memory retrieval: memory_ids is required" { + t.Errorf("memory retrieval error = %v, want missing memory IDs", err) + } +} diff --git a/internal/server/ingestor/server.go b/internal/server/ingestor/server.go index 09ad677b11..8698726790 100644 --- a/internal/server/ingestor/server.go +++ b/internal/server/ingestor/server.go @@ -71,11 +71,15 @@ func Run(ctx context.Context, cancel context.CancelFunc, serverName string, args // Initialize tokenizer (rag_analyzer) // tokenizer.Init handles DictPath fallback: env var → /usr/share/infinity/resource - err := tokenizer.Init(&tokenizer.PoolConfig{}) + tokenizerCfg := &tokenizer.PoolConfig{} + err := tokenizer.Init(tokenizerCfg) if err != nil { common.Fatal("Failed to initialize tokenizer", zap.Error(err)) } defer tokenizer.Close() + if err := nlp.InitQueryBuilderFromTokenizer(tokenizerCfg.DictPath); err != nil { + return fmt.Errorf("initialize query builder: %w", err) + } // Fail fast if the cl100k_base BPE table is missing. NumTokensFromString / // TrimContentToTokenLimit now panic rather than degrading silently, so this @@ -141,7 +145,9 @@ func Run(ctx context.Context, cancel context.CancelFunc, serverName string, args // Memory extraction runs on the Ingestor's shared NATS consumer + worker // pool (task_type="memory" dispatched by handleAndExecute -> executeMemoryTask), // so there is no longer a dedicated Redis memory consumer to start. - ingestor.SetMemoryMessageService(service.NewMemoryMessageService(service.NewMemoryService(dao.DB, engine.Get(), service.WithMemoryKvrocks(kvrocks.Get()), service.WithMemoryMQEngine(engine.GetMessageQueueEngine())))) + memoryService := service.NewMemoryService(dao.DB, engine.Get(), service.WithMemoryKvrocks(kvrocks.Get()), service.WithMemoryMQEngine(engine.GetMessageQueueEngine())) + ingestor.SetMemoryMessageService(service.NewMemoryMessageService(memoryService)) + bootstrap.InitializeRetrievalServices(service.NewMetadataService(), memoryService, service.NewModelFactory()) // Bind the relational handle and document engine into the production // embedder and knowledge-compiler resolvers. The task package is the // composition root for ingestion runs; it reads these from the explicit