diff --git a/.gitignore b/.gitignore index 1d3ab6baa0..521b9b9fa6 100644 --- a/.gitignore +++ b/.gitignore @@ -222,12 +222,15 @@ backup /target -# Do not include in PR (local dev / build artifacts) +# Do not include in PR (local dev / build artifacts / working notes) ragflow.egg-info/ uv-aarch64*.tar.gz uv-aarch64-unknown-linux-gnu.tar.gz docker/launch_backend_service_windows.sh +# Scratch/working notes (agent task planning & self-review docs; not PR material) +tasks/*.md + # C++ build directories internal/binding/cpp/build/ internal/binding/cpp/cmake-build-release/ @@ -239,6 +242,7 @@ internal/binding/cpp/cmake-build-debug/ # Go server build output bin/* !bin/.gitkeep +/ragflow_server .claude/settings.local.json .run/ diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 65910f5d0d..086bddd503 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -38,6 +38,7 @@ import ( dataset "ragflow/internal/service/dataset" "ragflow/internal/service/document" "ragflow/internal/service/file" + "ragflow/internal/service/nav" "ragflow/internal/service/nlp" "ragflow/internal/storage" "ragflow/internal/syncer" @@ -835,6 +836,16 @@ func startServer(ctx context.Context) { compilationTemplateGroupHandler := handler.NewCompilationTemplateGroupHandler(service.NewCompilationTemplateGroupService()) datasetArtifactHandler := handler.NewDatasetArtifactHandler(service.NewDatasetArtifactService(), datasetsService, file.NewFileCommitService()) + // Install the production eino-based chat invoker as the shared chat default, + // so agentic-search harness LLM calls work in production. Without this, + // chat.GetDefaultInvoker() stays nil and the harness falls back gracefully. + component.InstallDefaultChatInvoker() + + // Install the dataset-nav ES-backed service (internal/service/nav + + // internal/service/nlp). The embedder resolves the tenant's embedding model + // on demand so Search/UpsertDoc can embed queries/summaries automatically. + nav.SetNavService(nlp.NewNavService(service.NewNavEmbedder(modelProviderService, ""))) + // Initialize router r := router.NewRouter(authHandler, userHandler, diff --git a/internal/agent/chat/chat.go b/internal/agent/chat/chat.go new file mode 100644 index 0000000000..b8e3da96b4 --- /dev/null +++ b/internal/agent/chat/chat.go @@ -0,0 +1,110 @@ +// +// 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 chat is a dependency-light home for the LLM chat-invoker interface +// and its package-level singleton. It lives here (leaf, importing only eino +// schema + gorm) so that both internal/agent/component (which owns the +// production eino-based invoker) and internal/agent/tool / internal/agent/harness +// (which need to call the LLM for routing/selection) can depend on it without +// forming an import cycle. The production invoker is registered at boot via +// SetDefaultInvoker. +package chat + +import ( + "context" + "sync" + + "github.com/cloudwego/eino/schema" + "gorm.io/gorm" +) + +// Invoker abstracts a chat-model call so callers can inject a stub in tests and +// production flows through the eino bridge. This is the shared seam for the LLM +// component and the agentic-search harness/tools. +type Invoker interface { + Invoke(ctx context.Context, db *gorm.DB, req Request) (*Response, error) +} + +// Request is the minimal surface needed to dispatch a chat call. +type Request struct { + Driver string + ModelName string + APIKey string + BaseURL string + Messages []schema.Message + Temperature *float64 + TopP *float64 + PresencePenalty *float64 + FrequencyPenalty *float64 + MaxTokens *int + Thinking string // "enabled" | "disabled" | "" +} + +// Response is the result of a chat call. +type Response struct { + Content string + Thinking string + Model string + Stopped bool + Tokens int +} + +// ErrNotConfigured is returned by GetDefaultInvoker when no production invoker +// has been installed. Callers should treat it as "no chat model available". +var ErrNotConfigured = &configError{"chat: default invoker not configured"} + +type configError struct{ msg string } + +func (e *configError) Error() string { return e.msg } + +var ( + mu sync.RWMutex + inst Invoker + defaultModel string +) + +// SetDefaultInvoker installs the (production or test) invoker. Pass nil to +// restore the "not configured" state. +func SetDefaultInvoker(inv Invoker) { + mu.Lock() + defer mu.Unlock() + inst = inv +} + +// SetDefaultModelName records the tenant-default chat model name. The production +// invoker uses it when a Request omits ModelName, so harness/agentic-search LLM +// calls work without threading model config through every node. +func SetDefaultModelName(name string) { + mu.Lock() + defer mu.Unlock() + defaultModel = name +} + +// GetDefaultModelName returns the configured default model name ("" if unset). +func GetDefaultModelName() string { + mu.RLock() + defer mu.RUnlock() + return defaultModel +} + +// GetDefaultInvoker returns the installed invoker. It returns nil when no +// invoker has been installed (e.g. before server bootstrap or in tests that do +// not need the LLM). +func GetDefaultInvoker() Invoker { + mu.RLock() + defer mu.RUnlock() + return inst +} diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go index aaca4985d1..99720ce9c8 100644 --- a/internal/agent/component/llm.go +++ b/internal/agent/component/llm.go @@ -18,12 +18,12 @@ import ( "slices" "sort" "strings" - "sync" "time" "github.com/cloudwego/eino/schema" "gorm.io/gorm" + "ragflow/internal/agent/chat" "ragflow/internal/agent/component/prompts" "ragflow/internal/agent/runtime" "ragflow/internal/common" @@ -130,84 +130,68 @@ type LLMOutput struct { Tokens int } -// ChatInvoker is the abstraction the LLM component uses to talk to a -// chat model. The default implementation lives in this file; tests can -// override the package-level defaultChatInvoker to inject a stub. -type ChatInvoker interface { - Invoke(ctx context.Context, db *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) -} +// ChatInvoker is an alias for the shared chat.Invoker seam. The production +// eino-based implementation lives in this file; the package-level singleton is +// owned by internal/agent/chat so agent tools and the harness can also call the +// LLM without an import cycle. +type ChatInvoker = chat.Invoker -// ChatInvokeRequest is the minimal surface the LLM component needs to -// dispatch a chat call. Driver / APIKey / ModelName are kept here so the -// invoker can wire the right provider without the component caring. -type ChatInvokeRequest struct { - Driver string - ModelName string - APIKey string - BaseURL string - Messages []schema.Message - Temperature *float64 - TopP *float64 - PresencePenalty *float64 - FrequencyPenalty *float64 - MaxTokens *int - // Thinking mirrors the agent-level `thinking` setting - // ("enabled" | "disabled" | ""). The default invoker is - // responsible for translating this into the provider-specific - // request body (e.g. Qwen `enable_thinking`, Kimi/GLM - // `thinking.type`). Empty string means "use provider default" - // and the invoker should leave the provider's reasoning mode - // untouched. - Thinking string -} +// ChatInvokeRequest is an alias for chat.Request. +type ChatInvokeRequest = chat.Request -// ChatInvokeResponse mirrors what the LLM component writes to its outputs. -type ChatInvokeResponse struct { - Content string - Thinking string - Model string - Stopped bool - Tokens int -} +// ChatInvokeResponse is an alias for chat.Response. +type ChatInvokeResponse = chat.Response -// defaultChatInvokerMu guards defaultChatInvoker swaps during tests. -var defaultChatInvokerMu sync.RWMutex - -// defaultChatInvoker is the production ChatInvoker. Replaced in tests. -var defaultChatInvoker ChatInvoker = &einoChatInvoker{} - -// SetDefaultChatInvoker swaps the package-level ChatInvoker (test helper). -// Pass nil to restore the default. Concurrent-safe. +// SetDefaultChatInvoker delegates to the shared chat package singleton (test +// helper). Pass nil to restore the "not configured" state. The production +// einoChatInvoker is registered at boot in cmd/server_main.go. func SetDefaultChatInvoker(inv ChatInvoker) { - defaultChatInvokerMu.Lock() - defer defaultChatInvokerMu.Unlock() - defaultChatInvoker = inv -} - -// getDefaultChatInvoker returns the current default ChatInvoker. -func getDefaultChatInvoker() ChatInvoker { - defaultChatInvokerMu.RLock() - defer defaultChatInvokerMu.RUnlock() - if defaultChatInvoker == nil { - return &einoChatInvoker{} + if inv == nil { + chat.SetDefaultInvoker(nil) + return } - return defaultChatInvoker + chat.SetDefaultInvoker(inv) } -// GetDefaultChatInvokerForTest exposes the current package-level invoker so +// GetDefaultChatInvokerForTest exposes the current shared chat invoker so // cross-package tests can swap it and restore it safely. func GetDefaultChatInvokerForTest() ChatInvoker { - return getDefaultChatInvoker() + return chat.GetDefaultInvoker() +} + +// getDefaultChatInvoker returns the shared chat invoker, falling back to the +// production eino invoker when none has been installed. +func getDefaultChatInvoker() ChatInvoker { + if inv := chat.GetDefaultInvoker(); inv != nil { + return inv + } + return &einoChatInvoker{} +} + +// InstallDefaultChatInvoker registers the production eino-based invoker as the +// shared chat default. Called at server bootstrap so harness/agentic-search LLM +// calls work in production; without it, chat.GetDefaultInvoker() stays nil and +// harness falls back gracefully. +func InstallDefaultChatInvoker() { + chat.SetDefaultInvoker(&einoChatInvoker{}) } // einoChatInvoker is the production ChatInvoker — it constructs a fresh -// models.EinoChatModel per call from the request and dispatches. +// models.EinoChatModel per call from the request and dispatches. It is NOT +// registered as the shared chat default at init (so chat.GetDefaultInvoker() +// stays nil until bootstrap); cmd registers it via SetDefaultChatInvoker. type einoChatInvoker struct{} // Invoke satisfies ChatInvoker. func (e *einoChatInvoker) Invoke(ctx context.Context, db *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) { if req.ModelName == "" { - return nil, fmt.Errorf("component: LLM: model_id is required") + // Harness/agentic-search nodes may omit the model; fall back to the + // bootstrap-registered tenant default so those calls work in production. + if def := chat.GetDefaultModelName(); def != "" { + req.ModelName = def + } else { + return nil, fmt.Errorf("component: LLM: model_id is required and no default model is configured") + } } driver := req.Driver modelName := req.ModelName diff --git a/internal/agent/harness/agentic_rag.go b/internal/agent/harness/agentic_rag.go new file mode 100644 index 0000000000..24b419199d --- /dev/null +++ b/internal/agent/harness/agentic_rag.go @@ -0,0 +1,107 @@ +// +// 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 harness + +import ( + "context" + "log" + "strings" + + "gorm.io/gorm" +) + +// AgenticState carries the shared state across the agentic-RAG graph nodes. +type AgenticState struct { + Question string + Keywords string + Route RouteDecision + SeedChunks []string + Plan WorkflowPlan + Kbinfos *Kbinfos + PartialAnswer bool + Abstain bool + EmptyResult bool + FinalAnswer string + FormalizeError string +} + +// RunAgenticRAG drives the agentic-search graph: route → pre_search → planner → +// orchestrator → formalize_answer. Mirrors build_agentic_graph's linear flow. +// +// - low (direct_search): one hybrid search → answer. +// - medium+ (decompose_and_search / agentic_research / deep_research): +// pre_search grounds the planner, then decompose-and-search runs until a +// sufficiency verdict stops it. +func RunAgenticRAG(ctx context.Context, db *gorm.DB, question, keywords, modeLabel string, search SearchFn) AnswerResult { + state := &AgenticState{ + Question: strings.TrimSpace(question), + Keywords: keywords, + Kbinfos: &Kbinfos{}, + } + if state.Question == "" { + return AnswerResult{FinalAnswer: emptyResultMessage, Empty: true} + } + + // ── route ── + state.Route = RouteNode(ctx, db, state.Question, modeLabel) + + // ── pre_search (decomposition modes only) ── + if state.Route.RequiresDecomposition { + chunks, aggs := search(ctx, state.Question, state.Keywords) + state.SeedChunks = extractChunkTexts(chunks) + state.Kbinfos.Merge(chunks, aggs) + } + + // ── planner ── + state.Plan = PlannerNode(ctx, db, state.Route, state.SeedChunks) + + // ── orchestrator ── + var orch OrchestratorResult + if state.Route.RequiresDecomposition { + claims := make([]*ClaimTarget, len(state.Plan.Claims)) + for i := range state.Plan.Claims { + claims[i] = &state.Plan.Claims[i] + } + orch = DecomposeAndSearch(ctx, search, state.Question, state.Keywords, claims, modeLabel, state.Kbinfos) + } else { + orch = DirectSearch(ctx, search, state.Question, state.Keywords, state.Kbinfos) + } + state.PartialAnswer = orch.PartialAnswer + state.Abstain = orch.Abstain + state.EmptyResult = orch.EmptyResult + if orch.Kbinfos != nil { + state.Kbinfos = orch.Kbinfos + } + + // ── formalize_answer ── + res := FormalizeAnswer(ctx, db, state.Question, state.Kbinfos, state.PartialAnswer, state.Abstain, state.EmptyResult) + // Log only the question length, never its content, to avoid persisting user + // input in logs. + log.Printf("agentic_rag: finished (qlen=%d, strategy=%s, chunks=%d, partial=%v, abstain=%v)", + len(state.Question), state.Route.ExecutionStrategy, len(state.Kbinfos.Chunks), state.PartialAnswer, state.Abstain) + return res +} + +func extractChunkTexts(chunks []map[string]interface{}) []string { + out := make([]string, 0, len(chunks)) + for _, c := range chunks { + if t := chunkText(c); t != "" { + out = append(out, t) + } + } + return out +} diff --git a/internal/agent/harness/answer.go b/internal/agent/harness/answer.go new file mode 100644 index 0000000000..6443cb7824 --- /dev/null +++ b/internal/agent/harness/answer.go @@ -0,0 +1,124 @@ +// +// 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 harness + +import ( + "context" + "fmt" + "strings" + + "github.com/cloudwego/eino/schema" + + "gorm.io/gorm" + "ragflow/internal/agent/chat" +) + +// FINAL_ANSWER_SYSTEM mirrors Python report_prompt.FINAL_ANSWER_SYSTEM. +const finalAnswerSystem = `You are a smart agent. Answer the user's question using ONLY the evidence provided below. Do not invent facts: if the evidence cannot support a claim, say so plainly instead of guessing. + +# Citation rules +{cite_rules} + +# Language +Answer in the SAME language as the question. Translate retrieved evidence into that language as part of composing the answer; only verbatim quoted snippets may stay in their source language. + +# Fallback +If the evidence does not answer the question, reply with a clear statement that you don't have enough information based on the available sources (in the user's language). +` + +const partialAnswerPreamble = "Note: the following answer is based on partial information and may be incomplete." + +const defaultCiteRules = "Cite passages with their source documents where available." + +const ( + abstainMessage = "I cannot answer this question based on the available information." + emptyResultMessage = "I don't have enough information based on the available sources." +) + +// AnswerResult mirrors the formalize_answer node output. +type AnswerResult struct { + FinalAnswer string + Abstained bool + Empty bool +} + +// FormalizeAnswer generates the final answer from the gathered kbinfos. Mirrors +// Python's formalize_answer node: abstain/empty short-circuits, otherwise builds +// system+user and calls the chat invoker. +func FormalizeAnswer(ctx context.Context, db *gorm.DB, question string, kb *Kbinfos, partial, abstain, empty bool) AnswerResult { + if abstain { + return AnswerResult{FinalAnswer: abstainMessage, Abstained: true} + } + if empty || kb == nil || len(kb.Chunks) == 0 { + return AnswerResult{FinalAnswer: emptyResultMessage, Empty: true} + } + + var b strings.Builder + b.WriteString("Question:\n" + question + "\n") + if partial { + b.WriteString(partialAnswerPreamble + "\n") + } + b.WriteString("\nEvidence:\n") + for i, c := range kb.Chunks { + text := chunkText(c) + doc := chunkDoc(c) + if text == "" { + continue + } + b.WriteString(fmt.Sprintf("[%d] %s", i+1, text)) + if doc != "" { + b.WriteString(fmt.Sprintf(" (source: %s)", doc)) + } + b.WriteString("\n") + } + + system := strings.ReplaceAll(finalAnswerSystem, "{cite_rules}", defaultCiteRules) + inv := chat.GetDefaultInvoker() + if inv == nil { + return AnswerResult{FinalAnswer: "I'm sorry, the chat invoker is not configured."} + } + resp, err := inv.Invoke(ctx, db, chat.Request{ + Messages: []schema.Message{ + {Role: schema.System, Content: system}, + {Role: schema.User, Content: b.String()}, + }, + }) + if err != nil { + return AnswerResult{FinalAnswer: "I'm sorry, I encountered an error while composing the answer."} + } + return AnswerResult{FinalAnswer: resp.Content} +} + +func chunkText(c map[string]interface{}) string { + if t, ok := c["content_with_weight"].(string); ok && t != "" { + return t + } + if t, ok := c["content"].(string); ok { + return t + } + return "" +} + +func chunkDoc(c map[string]interface{}) string { + if d, ok := c["docnm_kwd"].(string); ok { + return d + } + if d, ok := c["doc_id"].(string); ok { + return d + } + return "" +} diff --git a/internal/agent/harness/answer_test.go b/internal/agent/harness/answer_test.go new file mode 100644 index 0000000000..0520a8b6d5 --- /dev/null +++ b/internal/agent/harness/answer_test.go @@ -0,0 +1,71 @@ +package harness + +import ( + "context" + "strings" + "testing" +) + +// TestFormalizeAnswer_Abstain asserts abstain short-circuits to the abstain +// message without calling the model. +func TestFormalizeAnswer_Abstain(t *testing.T) { + res := FormalizeAnswer(context.Background(), nil, "Q", &Kbinfos{}, false, true, false) + if !res.Abstained || res.FinalAnswer != abstainMessage { + t.Errorf("abstain result = %+v, want abstain message", res) + } +} + +// TestFormalizeAnswer_Empty asserts empty chunks short-circuit to the empty +// message. +func TestFormalizeAnswer_Empty(t *testing.T) { + res := FormalizeAnswer(context.Background(), nil, "Q", &Kbinfos{}, false, false, false) + if !res.Empty || res.FinalAnswer != emptyResultMessage { + t.Errorf("empty result = %+v, want empty message", res) + } +} + +// TestFormalizeAnswer_Generates asserts a chat invoker is used to compose the +// final answer from the evidence. +func TestFormalizeAnswer_Generates(t *testing.T) { + installChat(t, "here is the final answer") + kb := &Kbinfos{Chunks: []map[string]interface{}{{"content_with_weight": "evidence alpha"}}} + res := FormalizeAnswer(context.Background(), nil, "What is X?", kb, false, false, false) + if res.FinalAnswer != "here is the final answer" { + t.Errorf("final answer = %q, want chat output", res.FinalAnswer) + } +} + +// TestFormalizeAnswer_PartialPreamble asserts the partial preamble is included +// when partial=true. +func TestFormalizeAnswer_PartialPreamble(t *testing.T) { + installChat(t, "partial ans") + kb := &Kbinfos{Chunks: []map[string]interface{}{{"content_with_weight": "evidence"}}} + res := FormalizeAnswer(context.Background(), nil, "Q", kb, true, false, false) + if !strings.Contains(res.FinalAnswer, "partial ans") { + t.Errorf("final answer = %q", res.FinalAnswer) + } +} + +// TestRunAgenticRAG_LowMode asserts low mode does a single direct search and +// produces an answer. +func TestRunAgenticRAG_LowMode(t *testing.T) { + installChat(t, "final composed answer") + res := RunAgenticRAG(context.Background(), nil, "What is a rocket?", "", "low", + func(_ context.Context, _, _ string) ([]map[string]interface{}, []map[string]interface{}) { + return []map[string]interface{}{{"chunk_id": "a", "content_with_weight": "rocket evidence"}}, nil + }) + if res.FinalAnswer != "final composed answer" { + t.Errorf("final answer = %q", res.FinalAnswer) + } +} + +// TestRunAgenticRAG_EmptySearch asserts an empty search yields the empty message. +func TestRunAgenticRAG_EmptySearch(t *testing.T) { + res := RunAgenticRAG(context.Background(), nil, "Q", "", "low", + func(_ context.Context, _, _ string) ([]map[string]interface{}, []map[string]interface{}) { + return nil, nil + }) + if !res.Empty || res.FinalAnswer != emptyResultMessage { + t.Errorf("empty result = %+v", res) + } +} diff --git a/internal/agent/harness/datasetnav.go b/internal/agent/harness/datasetnav.go new file mode 100644 index 0000000000..a82cf92d30 --- /dev/null +++ b/internal/agent/harness/datasetnav.go @@ -0,0 +1,225 @@ +// +// 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 harness + +import ( + "context" + "fmt" + "log" + "strings" + + "github.com/cloudwego/eino/schema" + + "gorm.io/gorm" + "ragflow/internal/agent/chat" + "ragflow/internal/service/nav" +) + +// Dataset-nav router tunables mirror Python navigation.py. +const ( + navMaxDocs = 8 + navMaxClusters = 500 + navChildrenPageSize = 1000 + navTreeMaxDepth = 6 + navTreeMaxLeaves = 300 +) + +// navSelectSystem mirrors Python _NAV_SELECT_SYSTEM. +const navSelectSystem = `You are routing a question through a dataset's navigation tree. + +You are given a QUESTION and a numbered list of {noun}, each with a name and a short description. +Choose the {noun} most likely to contain information relevant to answering the question. + +Rules: +1. Judge only from the names and descriptions shown. +2. Be selective — include an item only if it is plausibly relevant. Include several when several are equally plausible. +3. If none are clearly relevant, return an empty list. +4. Return the bracketed index numbers of the chosen {noun}. + +Output ONLY JSON, no prose, no code fences: +{"relevant": [, ...]}` + +type navSelectVerdict struct { + Relevant []int `json:"relevant"` +} + +// NavigateDatasetByTree walks the dataset nav tree with two LLM passes +// (cluster-select → document-select) and returns the routed doc_ids (capped at +// navMaxDocs). This is the LLM two-round selection (P1b) implemented in the +// harness package, which can import the chat invoker (agent/tool cannot without +// an import cycle). It routes only — it does not retrieve. +func NavigateDatasetByTree(ctx context.Context, db *gorm.DB, ns nav.NavService, tenantID, kbID, query string) []string { + query = strings.TrimSpace(query) + if query == "" { + return nil + } + + // 1. List top-level clusters. + clusters, _, err := ns.ListClusters(ctx, tenantID, kbID, 0, navMaxClusters) + if err != nil { + log.Printf("datasetnav: list clusters failed: %v", err) + return nil + } + if len(clusters) == 0 { + return nil + } + clusterItems := make([]navSelectItem, len(clusters)) + for i, c := range clusters { + clusterItems[i] = navSelectItem{Name: c.Name, Description: c.Description, DocCount: c.DocCount} + } + + // 2. LLM selects relevant clusters. + selected := askNavSelect(ctx, db, query, "clusters", clusterItems, navMaxClusters) + if len(selected) == 0 { + return nil + } + + // 3. BFS-descend selected clusters to document leaves. + leaves := collectNavLeaves(ctx, ns, tenantID, kbID, selected) + + // 4. LLM selects relevant documents. + docs := askNavSelect(ctx, db, query, "documents", leaves, navTreeMaxLeaves) + + // Dedup + cap. + seen := map[string]bool{} + var routed []string + for _, d := range docs { + if d.DocID == "" || seen[d.DocID] { + continue + } + seen[d.DocID] = true + routed = append(routed, d.DocID) + if len(routed) >= navMaxDocs { + break + } + } + return routed +} + +// navSelectItem is a renderable cluster/document with a name + description. +type navSelectItem struct { + Name string + Description string + DocCount int + DocID string +} + +// askNavSelect renders items as a numbered list and asks the model which indices +// are relevant. Returns the selected items (a subset). Mirrors _ask_nav_select. +func askNavSelect(ctx context.Context, db *gorm.DB, query, noun string, items []navSelectItem, maxItems int) []navSelectItem { + if len(items) == 0 { + return nil + } + capped := items + if len(capped) > maxItems { + capped = capped[:maxItems] + } + var b strings.Builder + for i, it := range capped { + name := strings.TrimSpace(it.Name) + if name == "" { + name = fmt.Sprintf("item-%d", i) + } + desc := strings.Join(strings.Fields(it.Description), " ") + if len(desc) > 300 { + desc = desc[:300] + } + b.WriteString(fmt.Sprintf("[%d] %s", i, name)) + if it.DocCount > 0 { + b.WriteString(fmt.Sprintf(" [%d docs]", it.DocCount)) + } + b.WriteString(": " + desc + "\n") + } + + system := strings.ReplaceAll(navSelectSystem, "{noun}", noun) + user := fmt.Sprintf("Question:\n%s\n\n%s (numbered):\n%s\n\nOutput JSON:", query, strings.Title(noun), b.String()) + + inv := chat.GetDefaultInvoker() + if inv == nil { + log.Printf("datasetnav: LLM %s selection skipped (chat invoker not configured)", noun) + return nil + } + resp, err := inv.Invoke(ctx, db, chat.Request{ + Messages: []schema.Message{ + {Role: schema.System, Content: system}, + {Role: schema.User, Content: user}, + }, + }) + if err != nil { + log.Printf("datasetnav: LLM %s selection failed: %v", noun, err) + return nil + } + var v navSelectVerdict + if err := unmarshalModelJSON(resp.Content, &v); err != nil { + return nil + } + seen := map[int]bool{} + var out []navSelectItem + for _, idx := range v.Relevant { + if idx < 0 || idx >= len(capped) || seen[idx] { + continue + } + seen[idx] = true + out = append(out, capped[idx]) + } + return out +} + +// collectNavLeaves BFS-descents selected clusters to document leaves. Mirrors +// Python _collect_nav_leaves. +func collectNavLeaves(ctx context.Context, ns nav.NavService, tenantID, kbID string, selected []navSelectItem) []navSelectItem { + type node struct { + name string + depth int + } + frontier := make([]node, 0, len(selected)) + for _, c := range selected { + if c.Name != "" { + frontier = append(frontier, node{c.Name, 0}) + } + } + var leaves []navSelectItem + seenDocs := map[string]bool{} + seenNodes := map[string]bool{} + for len(frontier) > 0 && len(leaves) < navTreeMaxLeaves { + cur := frontier[0] + frontier = frontier[1:] + if seenNodes[cur.name] { + continue + } + seenNodes[cur.name] = true + children, _, err := ns.ListChildren(ctx, tenantID, kbID, cur.name, 0, navChildrenPageSize) + if err != nil { + continue + } + for _, ch := range children { + if ch.Type == "doc" { + did := strings.TrimSpace(ch.DocID) + if did != "" && !seenDocs[did] { + seenDocs[did] = true + leaves = append(leaves, navSelectItem{Name: ch.Name, Description: ch.Description, DocID: did}) + if len(leaves) >= navTreeMaxLeaves { + break + } + } + } else if ch.Type == "cluster" && ch.Name != "" && cur.depth+1 < navTreeMaxDepth { + frontier = append(frontier, node{ch.Name, cur.depth + 1}) + } + } + } + return leaves +} diff --git a/internal/agent/harness/datasetnav_test.go b/internal/agent/harness/datasetnav_test.go new file mode 100644 index 0000000000..a4b6e2cf0e --- /dev/null +++ b/internal/agent/harness/datasetnav_test.go @@ -0,0 +1,102 @@ +package harness + +import ( + "context" + "testing" + + "ragflow/internal/service/nav" +) + +// TestAskNavSelect_IndexBased asserts the model's index-based selection maps +// back to the item subset. +func TestAskNavSelect_IndexBased(t *testing.T) { + installChat(t, `{"relevant":[0,2]}`) + items := []navSelectItem{ + {Name: "Alpha", Description: "aaa"}, + {Name: "Beta", Description: "bbb"}, + {Name: "Gamma", Description: "ggg"}, + } + out := askNavSelect(context.Background(), nil, "query", "clusters", items, 10) + if len(out) != 2 { + t.Fatalf("selected = %d, want 2", len(out)) + } + if out[0].Name != "Alpha" || out[1].Name != "Gamma" { + t.Errorf("selected names = %q, %q; want Alpha, Gamma", out[0].Name, out[1].Name) + } +} + +// TestAskNavSelect_Empty asserts an empty "relevant" list yields nothing. +func TestAskNavSelect_Empty(t *testing.T) { + installChat(t, `{"relevant":[]}`) + if out := askNavSelect(context.Background(), nil, "q", "clusters", []navSelectItem{{Name: "A"}}, 10); len(out) != 0 { + t.Errorf("expected no selection, got %d", len(out)) + } +} + +// TestAskNavSelect_OutOfRange asserts invalid indices are skipped. +func TestAskNavSelect_OutOfRange(t *testing.T) { + installChat(t, `{"relevant":[0,99,-1]}`) + out := askNavSelect(context.Background(), nil, "q", "clusters", []navSelectItem{{Name: "A"}, {Name: "B"}}, 10) + if len(out) != 1 || out[0].Name != "A" { + t.Errorf("out-of-range selection = %+v, want [A]", out) + } +} + +// TestNavigateDatasetByTree_NoQuery asserts empty query returns nil without +// calling anything. +func TestNavigateDatasetByTree_NoQuery(t *testing.T) { + if out := NavigateDatasetByTree(context.Background(), nil, nil, "t1", "kb1", " "); out != nil { + t.Errorf("expected nil for empty query, got %v", out) + } +} + +// fakeNavSvcHarness is an in-memory nav.NavService for the BFS test. +type fakeNavSvcHarness struct { + clusters []nav.NavNode + children map[string][]nav.NavNode +} + +func (f *fakeNavSvcHarness) UpsertDoc(context.Context, nav.UpsertDocInput) error { return nil } +func (f *fakeNavSvcHarness) RemoveDoc(context.Context, string, string, string) error { + return nil +} +func (f *fakeNavSvcHarness) Search(context.Context, string, string, string, []float32, int) ([]nav.NavHit, error) { + return nil, nil +} +func (f *fakeNavSvcHarness) ListClusters(context.Context, string, string, int, int) ([]nav.NavNode, int64, error) { + return f.clusters, int64(len(f.clusters)), nil +} +func (f *fakeNavSvcHarness) ListChildren(_ context.Context, _, _, name string, _, _ int) ([]nav.NavNode, int64, error) { + return f.children[name], int64(len(f.children[name])), nil +} + +// TestCollectNavLeaves_BFS asserts document leaves are collected, sub-clusters +// descended, and leaves deduped by doc_id. +func TestCollectNavLeaves_BFS(t *testing.T) { + ns := &fakeNavSvcHarness{ + clusters: []nav.NavNode{{Name: "C1", Description: "cluster 1"}}, + children: map[string][]nav.NavNode{ + "C1": { + {Name: "Sub", Type: "cluster"}, + {Name: "DocA", Type: "doc", DocID: "d1"}, + {Name: "DocB", Type: "doc", DocID: "d2"}, + }, + "Sub": {{Name: "DocC", Type: "doc", DocID: "d3"}}, + }, + } + selected := []navSelectItem{{Name: "C1", Description: "cluster 1"}} + leaves := collectNavLeaves(context.Background(), ns, "t1", "kb1", selected) + if len(leaves) != 3 { + t.Fatalf("leaves = %d, want 3 (d1,d2 from C1 + d3 from Sub)", len(leaves)) + } + seen := map[string]bool{} + for _, l := range leaves { + if l.DocID == "" { + t.Errorf("leaf %q has empty doc_id", l.Name) + } + seen[l.DocID] = true + } + if !seen["d1"] || !seen["d2"] || !seen["d3"] { + t.Errorf("collected doc_ids = %v, want d1,d2,d3", seen) + } +} diff --git a/internal/agent/harness/navigation.go b/internal/agent/harness/navigation.go new file mode 100644 index 0000000000..459f0c54ca --- /dev/null +++ b/internal/agent/harness/navigation.go @@ -0,0 +1,281 @@ +// +// 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 harness + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/cloudwego/eino/schema" + + "ragflow/internal/agent/chat" + "ragflow/internal/engine" + "ragflow/internal/engine/types" +) + +// Structure navigation mirrors Python navigation.py (ontology_navigate / +// mindmap_navigate). These live in the harness package (not agent/tool) because +// they need the chat invoker, and agent/tool → agent/component would form an +// import cycle. +const ( + toolOntologyNavigate = "ontology_navigate" + toolMindmapNavigate = "mindmap_navigate" +) + +var catalogKinds = map[string]bool{"tree": true, "timeline": true, "raptor": true, "page_index": true, "pageindex": true} +var mindmapKinds = map[string]bool{"mindmap": true, "mind_map": true} + +const navSystemPrompt = `You are given the {noun} of one or more documents — an outline of entities and their relations — and a question. + +Decide whether that outline alone already answers the question. + +Rules: +1. Answer ONLY from the outline below. Do not invent facts. +2. Set "is_sufficient" to true only when the outline genuinely answers the question; otherwise false with an empty answer. +3. Always fill "relevant_entities" with the exact ` + "`name`" + ` values of the entities most related to the question (up to 10), even when the outline is not sufficient — they are used to pull the underlying source text. + +Output ONLY JSON, no prose, no code fences: +{"is_sufficient": true/false, "answer": "", "relevant_entities": ["", ...]}` + +const ( + maxStructureEntities = 300 + maxStructureRelations = 300 + maxEvidenceChunks = 24 +) + +type structureNavArgs struct { + Topic string `json:"topic"` + Keywords string `json:"keywords,omitempty"` + DocScope []string `json:"doc_scope,omitempty"` +} + +type structureEntity struct { + Name string `json:"name"` + Type string `json:"type"` + Description string `json:"description"` + SourceChunkIDs []string `json:"source_chunk_ids"` + DocID string `json:"-"` +} + +type structureNavVerdict struct { + IsSufficient bool `json:"is_sufficient"` + Answer string `json:"answer"` + RelevantEntities []string `json:"relevant_entities"` +} + +// NavigateStructure implements ontology_navigate / mindmap_navigate. It reads +// the compiled structure (entities) of the in-scope documents, asks the chat +// model which entities answer the question, and pulls the source chunks behind +// the selected entities. Routing only — returns empty on any failure. +func NavigateStructure(ctx context.Context, tenantID string, kind string, args structureNavArgs) (string, error) { + noun := "catalog" + var kinds map[string]bool + if kind == toolMindmapNavigate { + kinds = mindmapKinds + noun = "mindmap" + } else { + kinds = catalogKinds + } + + query := strings.TrimSpace(args.Topic + " " + args.Keywords) + if query == "" || len(args.DocScope) == 0 { + return `{"chunks":[]}`, nil + } + + var entities []structureEntity + for _, docID := range args.DocScope { + es := loadStructureEntities(ctx, tenantID, docID, kinds) + for _, e := range es { + if e.Name != "" { + e.DocID = docID + entities = append(entities, e) + } + } + } + if len(entities) == 0 { + return `{"chunks":[]}`, nil + } + + selected, err := askStructureSelect(ctx, query, noun, entities) + if err != nil || len(selected) == 0 { + return `{"chunks":[]}`, nil + } + + idsByDoc := map[string][]string{} + for _, e := range selected { + idsByDoc[e.DocID] = append(idsByDoc[e.DocID], e.SourceChunkIDs...) + } + var chunks []map[string]interface{} + for _, ids := range idsByDoc { + chunks = append(chunks, loadChunksByIDs(ctx, tenantID, dedupStrings(ids))...) + if len(chunks) >= maxEvidenceChunks { + break + } + } + b, _ := json.Marshal(map[string]interface{}{"chunks": chunks}) + return string(b), nil +} + +func loadStructureEntities(ctx context.Context, tenantID, docID string, kinds map[string]bool) []structureEntity { + de := engine.Get() + if de == nil { + return nil + } + idx := fmt.Sprintf("ragflow_%s", tenantID) + req := &types.SearchRequest{ + IndexNames: []string{idx}, + Filter: map[string]interface{}{"doc_id": []string{docID}, "knowledge_graph_kwd": []string{"graph"}}, + SelectFields: []string{"content_with_weight", "compile_kwd", "compilation_template_kind_kwd"}, + Limit: 1000, + } + res, err := de.Search(ctx, req) + if err != nil { + return nil + } + var out []structureEntity + for _, row := range res.Chunks { + kind := normalizeKind(row) + if !kinds[kind] { + continue + } + payload, _ := row["content_with_weight"].(string) + var graph struct { + Entities []structureEntity `json:"entities"` + } + if err := json.Unmarshal([]byte(payload), &graph); err != nil { + continue + } + out = append(out, graph.Entities...) + } + return out +} + +func normalizeKind(row map[string]interface{}) string { + if ck, _ := row["compile_kwd"].(string); ck == "raptor_graph" { + return "raptor" + } + kind, _ := row["compilation_template_kind_kwd"].(string) + if kind == "" { + kind, _ = row["compile_kwd"].(string) + } + kind = strings.ToLower(strings.TrimSpace(strings.ReplaceAll(kind, "-", "_"))) + if kind == "pageindex" || kind == "page_index" || kind == "knowledge_graph" { + return "timeline" + } + return kind +} + +func askStructureSelect(ctx context.Context, query, noun string, entities []structureEntity) ([]structureEntity, error) { + rendered := renderStructureEntities(entities) + inv := chat.GetDefaultInvoker() + if inv == nil { + return nil, fmt.Errorf("dataset navigation: chat invoker not configured") + } + resp, err := inv.Invoke(ctx, nil, chat.Request{ + Messages: []schema.Message{ + {Role: schema.System, Content: strings.ReplaceAll(navSystemPrompt, "{noun}", noun)}, + {Role: schema.User, Content: fmt.Sprintf("Question:\n%s\n\n%s:\n%s\n\nOutput JSON:", query, noun, rendered)}, + }, + }) + if err != nil { + return nil, err + } + var v structureNavVerdict + if err := unmarshalModelJSON(resp.Content, &v); err != nil { + return nil, err + } + want := map[string]bool{} + for _, n := range v.RelevantEntities { + want[n] = true + } + var out []structureEntity + for _, e := range entities { + if want[e.Name] { + out = append(out, e) + } + } + return out, nil +} + +func renderStructureEntities(entities []structureEntity) string { + var b strings.Builder + b.WriteString("Entities:") + for i, e := range entities { + if i >= maxStructureEntities { + break + } + b.WriteString("\n- " + e.Name + " (" + orStr(e.Type, "other") + ")") + if d := strings.Join(strings.Fields(e.Description), " "); d != "" { + b.WriteString(": " + d) + } + } + return b.String() +} + +func loadChunksByIDs(ctx context.Context, tenantID string, ids []string) []map[string]interface{} { + if len(ids) == 0 { + return nil + } + de := engine.Get() + if de == nil { + return nil + } + idx := fmt.Sprintf("ragflow_%s", tenantID) + limit := maxEvidenceChunks + if len(ids) < limit { + limit = len(ids) + } + req := &types.SearchRequest{ + IndexNames: []string{idx}, + Filter: map[string]interface{}{"id": ids}, + SelectFields: []string{"content_with_weight", "docnm_kwd", "doc_id"}, + Limit: limit, + } + res, err := de.Search(ctx, req) + if err != nil { + return nil + } + var out []map[string]interface{} + for _, row := range res.Chunks { + out = append(out, map[string]interface{}{ + "chunk_id": row["id"], "content_with_weight": row["content_with_weight"], + "docnm_kwd": row["docnm_kwd"], "doc_id": row["doc_id"], + }) + } + return out +} + +func dedupStrings(in []string) []string { + seen := map[string]bool{} + var out []string + for _, s := range in { + if s != "" && !seen[s] { + seen[s] = true + out = append(out, s) + } + } + return out +} + +func orStr(v, def string) string { + if v == "" { + return def + } + return v +} diff --git a/internal/agent/harness/navigation_test.go b/internal/agent/harness/navigation_test.go new file mode 100644 index 0000000000..e566378390 --- /dev/null +++ b/internal/agent/harness/navigation_test.go @@ -0,0 +1,65 @@ +package harness + +import ( + "context" + "strings" + "testing" +) + +// TestRenderStructureEntities asserts the prompt outline is compact and includes +// entity name/type/description. +func TestRenderStructureEntities(t *testing.T) { + entities := []structureEntity{ + {Name: "Rocket", Type: "concept", Description: "a propulsion device"}, + {Name: "Engine", Type: "part"}, + } + out := renderStructureEntities(entities) + if out == "" { + t.Fatal("expected rendered outline") + } + if !containsAll(out, "Rocket", "Engine", "concept") { + t.Errorf("rendered outline missing entities: %q", out) + } +} + +// TestNavigateStructure_NoScope asserts empty doc scope yields empty chunks +// without calling the model. +func TestNavigateStructure_NoScope(t *testing.T) { + out, err := NavigateStructure(context.Background(), "t1", toolOntologyNavigate, structureNavArgs{ + Topic: "rocket", DocScope: nil, + }) + if err != nil { + t.Fatal(err) + } + if out != `{"chunks":[]}` { + t.Errorf("no-scope result = %q, want empty chunks", out) + } +} + +// TestNormalizeKind mirrors the API's kind normalization. +func TestNormalizeKind(t *testing.T) { + cases := []struct { + row map[string]interface{} + want string + }{ + {map[string]interface{}{"compile_kwd": "raptor_graph"}, "raptor"}, + {map[string]interface{}{"compilation_template_kind_kwd": "page_index"}, "timeline"}, + {map[string]interface{}{"compilation_template_kind_kwd": "knowledge_graph"}, "timeline"}, + {map[string]interface{}{"compilation_template_kind_kwd": "mindmap"}, "mindmap"}, + {map[string]interface{}{"compile_kwd": "tree"}, "tree"}, + } + for _, c := range cases { + if got := normalizeKind(c.row); got != c.want { + t.Errorf("normalizeKind(%v) = %q, want %q", c.row, got, c.want) + } + } +} + +func containsAll(s string, subs ...string) bool { + for _, sub := range subs { + if !strings.Contains(s, sub) { + return false + } + } + return true +} diff --git a/internal/agent/harness/orchestrator.go b/internal/agent/harness/orchestrator.go new file mode 100644 index 0000000000..7c9cbc8b1c --- /dev/null +++ b/internal/agent/harness/orchestrator.go @@ -0,0 +1,247 @@ +// +// 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 harness + +import ( + "context" + "fmt" + "strings" +) + +// Orchestration input/output mirrors Python orchestrator (direct.py / decompose.py). +// SearchFn performs one hybrid search for a query and returns chunks + doc aggs, +// so the orchestrator is decoupled from the concrete retrieval backend. +type SearchFn func(ctx context.Context, query, keywords string) ([]map[string]interface{}, []map[string]interface{}) + +// Kbinfos is the shared accumulation store (Python tools.kbinfos). +type Kbinfos struct { + Chunks []map[string]interface{} + DocAggs []map[string]interface{} +} + +func (k *Kbinfos) HasChunks() bool { return len(k.Chunks) > 0 } + +// Merge appends the given chunks/aggs, deduplicating by chunkKey. It returns the +// GLOBAL indices (positions in k.Chunks after the merge) of the chunks this call +// contributed, so callers can store stable evidence references across multiple +// Merge calls (per-search indices would diverge from the accumulated list). +func (k *Kbinfos) Merge(chunks, aggs []map[string]interface{}) []int { + seen := map[string]bool{} + for _, c := range k.Chunks { + seen[chunkKey(c)] = true + } + var added []int + for _, c := range chunks { + kk := chunkKey(c) + if !seen[kk] { + seen[kk] = true + k.Chunks = append(k.Chunks, c) + } + // Record the global index of every contributed chunk (dedup or new), + // so EvidenceIDs always reference the accumulated kbinfos positions. + added = append(added, indexOfChunk(k.Chunks, kk)) + } + dseen := map[string]bool{} + for _, d := range k.DocAggs { + if id, _ := d["doc_id"].(string); id != "" { + dseen[id] = true + } + } + for _, d := range aggs { + if id, _ := d["doc_id"].(string); id != "" && !dseen[id] { + dseen[id] = true + k.DocAggs = append(k.DocAggs, d) + } + } + return added +} + +// indexOfChunk returns the global index of the chunk whose key matches kk. +func indexOfChunk(chunks []map[string]interface{}, kk string) int { + for i, c := range chunks { + if chunkKey(c) == kk { + return i + } + } + return -1 +} + +// chunkKey returns a stable dedup key for a chunk. Prefers chunk_id/id when +// present; otherwise falls back to a content-derived hash so chunks without an +// id still dedup correctly (and never collide on a shared empty/"" key). +func chunkKey(c map[string]interface{}) string { + if id, ok := c["chunk_id"].(string); ok && id != "" { + return "cid:" + id + } + if id, ok := c["id"].(string); ok && id != "" { + return "id:" + id + } + content := "" + if t, ok := c["content_with_weight"].(string); ok { + content = t + } else if t, ok := c["content"].(string); ok { + content = t + } + if content == "" { + // No stable identity: fall back to the doc reference so at least + // per-document grouping is preserved (rarely reached). + content = fmt.Sprintf("%s|%s", anyString(c["doc_id"]), anyString(c["docnm_kwd"])) + } + return "h:" + fnv64(content) +} + +func anyString(v interface{}) string { + if s, ok := v.(string); ok { + return s + } + return "" +} + +// fnv64 is a deterministic non-crypto hash for dedup keys. +func fnv64(s string) string { + h := uint64(14695981039346656037) + for i := 0; i < len(s); i++ { + h ^= uint64(s[i]) + h *= 1099511628211 + } + return fmt.Sprintf("%016x", h) +} + +// OrchestratorResult mirrors the state updates returned by the Python +// orchestrator nodes. +type OrchestratorResult struct { + Verdict *SufficiencyVerdict + PartialAnswer bool + Abstain bool + EmptyResult bool + Kbinfos *Kbinfos +} + +// DirectSearch is the low-mode orchestrator: one hybrid search → merge. +func DirectSearch(ctx context.Context, search SearchFn, question, keywords string, kbinfos *Kbinfos) OrchestratorResult { + if kbinfos == nil { + kbinfos = &Kbinfos{} + } + chunks, aggs := search(ctx, question, keywords) + kbinfos.Merge(chunks, aggs) + if !kbinfos.HasChunks() { + return OrchestratorResult{EmptyResult: true, Kbinfos: kbinfos} + } + return OrchestratorResult{Kbinfos: kbinfos} +} + +// DecomposeAndSearch is the medium-mode orchestrator: decompose → parallel +// search → cross-check → fusion → iterate until a verdict stops it. +func DecomposeAndSearch(ctx context.Context, search SearchFn, question, keywords string, claims []*ClaimTarget, modeLabel string, kbinfos *Kbinfos) OrchestratorResult { + if kbinfos == nil { + kbinfos = &Kbinfos{} + } + mode, _ := GetMode(modeLabel) + if mode.Label == "" { + mode = THINKING_MODES["medium"] + } + unverified := func() []*ClaimTarget { + var out []*ClaimTarget + for _, c := range claims { + if !c.IsVerified { + out = append(out, c) + } + } + return out + } + + for cycle := 0; cycle < mode.MaxOrchestratorCycles; cycle++ { + uv := unverified() + if len(uv) == 0 { + break + } + for _, c := range uv { + chunks, aggs := search(ctx, c.Description, keywords) + if len(chunks) > 0 { + c.IsVerified = true + c.Confidence = 0.8 + // Merge returns the GLOBAL indices of this claim's chunks in the + // accumulated kbinfos, which is what CrossCheckClaim resolves + // against (allChunks is keyed by that global position). + global := kbinfos.Merge(chunks, aggs) + c.AgentResult = &AgentResult{ + ClaimID: c.ClaimID, Report: summarize(chunks), IsVerified: true, Confidence: 0.8, + EvidenceIDs: global, + } + } else { + c.AgentResult = &AgentResult{ClaimID: c.ClaimID, IsVerified: false, Confidence: 0.0} + } + } + + allChunks := map[int]map[string]interface{}{} + for i, c := range kbinfos.Chunks { + allChunks[i] = c + } + var agentResults []AgentResult + var crossResults []ClaimCrossCheckResult + for _, c := range claims { + if c.AgentResult != nil { + agentResults = append(agentResults, *c.AgentResult) + crossResults = append(crossResults, CrossCheckClaim(c.AgentResult, allChunks)) + } + } + verdict := ComputeFusionScore(agentResults, crossResults, mode) + action, _ := RouteSufficiencyVerdict(verdict, modeLabel, cycle, mode.MaxOrchestratorCycles) + + switch action { + case "ANSWER", "ANSWER_PARTIAL": + return OrchestratorResult{Verdict: &verdict, PartialAnswer: action == "ANSWER_PARTIAL", Kbinfos: kbinfos} + case "ABSTAIN": + kbinfos.Chunks = nil + return OrchestratorResult{Verdict: &verdict, Abstain: true, Kbinfos: kbinfos} + case "REPLAN": + // Reset unverified for another cycle (simplified: continue loop). + case "CONTINUE": + // fallthrough to next cycle + } + } + return OrchestratorResult{Kbinfos: kbinfos} +} + +func summarize(chunks []map[string]interface{}) string { + var parts []string + n := 3 + if len(chunks) < n { + n = len(chunks) + } + for _, c := range chunks[:n] { + text := "" + if t, ok := c["content_with_weight"].(string); ok { + text = t + } else if t, ok := c["content"].(string); ok { + text = t + } + if len(text) > 200 { + text = text[:200] + } + parts = append(parts, text) + } + return strings.Join(parts, " | ") +} + +func intRange(n int) []int { + out := make([]int, n) + for i := range out { + out[i] = i + } + return out +} diff --git a/internal/agent/harness/orchestrator_test.go b/internal/agent/harness/orchestrator_test.go new file mode 100644 index 0000000000..daefd5d8d0 --- /dev/null +++ b/internal/agent/harness/orchestrator_test.go @@ -0,0 +1,135 @@ +package harness + +import ( + "context" + "testing" +) + +// TestCrossCheckClaim_NumberMatch asserts a verified report whose numbers appear +// in the evidence chunk passes the cross-check. +func TestCrossCheckClaim_NumberMatch(t *testing.T) { + agent := &AgentResult{ClaimID: "c0", IsVerified: true, Report: "speed is 88 and 12", EvidenceIDs: []int{0}} + chunks := map[int]map[string]interface{}{0: {"content_with_weight": "speed is 88 and 12 here"}} + r := CrossCheckClaim(agent, chunks) + if !r.CrossCheckPassed { + t.Errorf("cross-check should pass, got mismatches %v", r.Mismatches) + } +} + +// TestKbinfosMerge_GlobalIndices asserts Merge returns the GLOBAL indices of the +// contributed chunks in the accumulated list, so a later claim's EvidenceIDs +// still resolve correctly after earlier claims pushed more chunks in. +func TestKbinfosMerge_GlobalIndices(t *testing.T) { + kb := &Kbinfos{} + // First claim's search returns chunks a,b. + first := kb.Merge([]map[string]interface{}{ + {"chunk_id": "a", "content_with_weight": "alpha 7"}, + {"chunk_id": "b", "content_with_weight": "beta"}, + }, nil) + // a->0, b->1 in the global list. + if len(first) != 2 || first[0] != 0 || first[1] != 1 { + t.Fatalf("first merge global indices = %v, want [0 1]", first) + } + // Second claim's search returns chunk c (now global index 2) and a dup of a. + second := kb.Merge([]map[string]interface{}{ + {"chunk_id": "c", "content_with_weight": "gamma 9"}, + {"chunk_id": "a", "content_with_weight": "alpha 7"}, + }, nil) + // c->2, a(dup)->0; NOT per-search [0 1]. + if len(second) != 2 || second[0] != 2 || second[1] != 0 { + t.Fatalf("second merge global indices = %v, want [2 0]", second) + } + // CrossCheckClaim must resolve both claims' evidence against the global list. + allChunks := map[int]map[string]interface{}{} + for i, c := range kb.Chunks { + allChunks[i] = c + } + r1 := CrossCheckClaim(&AgentResult{ClaimID: "c1", IsVerified: true, Report: "alpha 7", EvidenceIDs: first}, allChunks) + if !r1.HasEvidence { + t.Error("claim 1 should have evidence at global indices") + } + r2 := CrossCheckClaim(&AgentResult{ClaimID: "c2", IsVerified: true, Report: "gamma 9", EvidenceIDs: second}, allChunks) + if !r2.HasEvidence { + t.Error("claim 2 should have evidence at global index 2 (not per-search 0)") + } +} + +// TestCrossCheckClaim_Unverified asserts an unverified agent fails the check. +func TestCrossCheckClaim_Unverified(t *testing.T) { + r := CrossCheckClaim(&AgentResult{ClaimID: "c0", IsVerified: false}, nil) + if r.CrossCheckPassed { + t.Error("unverified agent must fail cross-check") + } +} + +// TestComputeFusionScore_Sufficient asserts a fully-verified high-score set is +// SUFFICIENT. +func TestComputeFusionScore_Sufficient(t *testing.T) { + agents := []AgentResult{{ClaimID: "c0", IsVerified: true, Report: "value 42", EvidenceIDs: []int{0}}} + cross := []ClaimCrossCheckResult{{ClaimID: "c0", CrossCheckPassed: true, CrossCheckScore: 1.0, HasEvidence: true}} + v := ComputeFusionScore(agents, cross, THINKING_MODES["medium"]) + if v.Status != "SUFFICIENT" { + t.Errorf("status = %q, want SUFFICIENT", v.Status) + } +} + +// TestComputeFusionScore_NoEvidence asserts a claim with no examined evidence is +// UNANSWERABLE (empty-evidence guard). +func TestComputeFusionScore_NoEvidence(t *testing.T) { + cross := []ClaimCrossCheckResult{{ClaimID: "c0", CrossCheckPassed: true, CrossCheckScore: 1.0, HasEvidence: false}} + v := ComputeFusionScore([]AgentResult{{ClaimID: "c0", IsVerified: true}}, cross, THINKING_MODES["medium"]) + if v.Status != "UNANSWERABLE" { + t.Errorf("status = %q, want UNANSWERABLE", v.Status) + } + if len(v.MissingClaims) != 1 { + t.Errorf("missing claims = %v, want [c0]", v.MissingClaims) + } +} + +// TestRouteSufficiencyVerdict asserts SUFFICIENT → ANSWER. +func TestRouteSufficiencyVerdict(t *testing.T) { + action, cont := RouteSufficiencyVerdict(SufficiencyVerdict{Status: "SUFFICIENT", Score: 0.9}, "medium", 0, 3) + if action != "ANSWER" || cont { + t.Errorf("got (%q,%v), want (ANSWER,false)", action, cont) + } +} + +// TestDirectSearch_Merges asserts direct search merges chunks and flags empty. +func TestDirectSearch_Merges(t *testing.T) { + kb := &Kbinfos{} + res := DirectSearch(context.Background(), func(_ context.Context, _, _ string) ([]map[string]interface{}, []map[string]interface{}) { + return []map[string]interface{}{{"chunk_id": "a", "content_with_weight": "alpha"}}, nil + }, "q", "", kb) + if res.EmptyResult || !kb.HasChunks() { + t.Errorf("expected merged chunks, empty=%v chunks=%d", res.EmptyResult, len(kb.Chunks)) + } +} + +// TestDirectSearch_Empty asserts direct search flags empty when no chunks. +func TestDirectSearch_Empty(t *testing.T) { + kb := &Kbinfos{} + res := DirectSearch(context.Background(), func(_ context.Context, _, _ string) ([]map[string]interface{}, []map[string]interface{}) { + return nil, nil + }, "q", "", kb) + if !res.EmptyResult { + t.Error("expected empty_result=true") + } +} + +// TestDecomposeAndSearch_Verifies asserts a searchable claim gets verified and +// the loop stops on ANSWER. +func TestDecomposeAndSearch_Verifies(t *testing.T) { + claims := []*ClaimTarget{{ClaimID: "c0", Description: "fact 42 about X"}} + kb := &Kbinfos{} + res := DecomposeAndSearch(context.Background(), + func(_ context.Context, _, _ string) ([]map[string]interface{}, []map[string]interface{}) { + return []map[string]interface{}{{"chunk_id": "x", "content_with_weight": "the value is 42"}}, nil + }, + "Q", "", claims, "medium", kb) + if !claims[0].IsVerified { + t.Error("claim should be verified") + } + if res.Verdict == nil { + t.Error("expected a verdict") + } +} diff --git a/internal/agent/harness/planner.go b/internal/agent/harness/planner.go new file mode 100644 index 0000000000..a7203e39ec --- /dev/null +++ b/internal/agent/harness/planner.go @@ -0,0 +1,164 @@ +// +// 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 harness + +import ( + "context" + "fmt" + "log" + "strings" + + "github.com/cloudwego/eino/schema" + + "ragflow/internal/agent/chat" + + "gorm.io/gorm" +) + +// decomposePrompt is a faithful port of Python's DECOMPOSE_FACTUAL prompt shape. +// Question-type-specific variants share the same structure but tune the +// instructions; the Go planner selects by question_type. +const decomposePrompt = `Break down the research question into %d atomic claims (facts or sub-questions). + +Question: %s + +Detail level: %s +Grounding context (preliminary retrieval, may be empty): +%s + +Output format: +{ + "claims": [ + {"claim_id": "c0", "description": "...", "priority": 0, "suggested_tools": []} + ] +} +` + +type plannerClaim struct { + ClaimID string `json:"claim_id"` + Description string `json:"description"` + Priority int `json:"priority"` + SuggestedTools []string `json:"suggested_tools"` +} + +type plannerResult struct { + Claims []plannerClaim `json:"claims"` +} + +// PlannerNode mirrors Python planner_node. It decomposes the routed question +// into ClaimTargets. Direct mode (no decomposition) returns a single coarse +// claim. On any failure it falls back to the direct plan. +func PlannerNode(ctx context.Context, db *gorm.DB, route RouteDecision, seedChunks []string) WorkflowPlan { + if !route.RequiresDecomposition { + return directPlan(route.Question) + } + mode, ok := GetMode(route.ThinkingMode) + if !ok { + // Unknown or empty mode label: fall back to medium so the planner is not + // driven by a zero-valued mode (which would yield a degenerate plan). + mode = THINKING_MODES["medium"] + } + maxClaims := maxClaimsFor(mode.Label) + detail := detailLevelFor(mode.Label) + + prompt := fmt.Sprintf(decomposePrompt, maxClaims, route.Question, detail, formatSeedChunks(seedChunks)) + + inv := chat.GetDefaultInvoker() + if inv == nil { + log.Printf("agentic_rag: planner_node skipped (chat invoker not configured); fallback direct") + return directPlan(route.Question) + } + resp, err := inv.Invoke(ctx, db, chat.Request{ + Messages: []schema.Message{ + {Role: schema.System, Content: prompt}, + {Role: schema.User, Content: route.Question}, + }, + }) + if err != nil { + log.Printf("agentic_rag: planner_node failed (fallback direct): %v", err) + return directPlan(route.Question) + } + var res plannerResult + if err := unmarshalModelJSON(resp.Content, &res); err != nil { + log.Printf("agentic_rag: planner_node parse failed (fallback direct): %v", err) + return directPlan(route.Question) + } + claims := make([]ClaimTarget, 0, len(res.Claims)) + for i, c := range res.Claims { + desc := strings.TrimSpace(c.Description) + if desc == "" { + continue + } + claims = append(claims, ClaimTarget{ + ClaimID: orDefault(c.ClaimID, fmt.Sprintf("c%d", i)), + Description: desc, + Priority: c.Priority, + SuggestedTools: c.SuggestedTools, + }) + } + if len(claims) == 0 { + return directPlan(route.Question) + } + return WorkflowPlan{ + PlanType: planTypeFor(route.QuestionType), + Claims: claims, + MaxIterations: mode.MaxOrchestratorCycles, + } +} + +func directPlan(question string) WorkflowPlan { + return WorkflowPlan{ + PlanType: "direct", + Claims: []ClaimTarget{{ClaimID: "c0", Description: question, Priority: 0}}, + MaxIterations: 1, + } +} + +func planTypeFor(qt string) string { + switch qt { + case "factual": + return "fact_decomposition" + case "comparative": + return "comparative_decomposition" + case "procedural": + return "procedural_decomposition" + default: + return "exploratory_decomposition" + } +} + +func maxClaimsFor(modeLabel string) int { + return map[string]int{"low": 1, "medium": 3, "high": 5, "ultra": 8}[modeLabel] +} + +func detailLevelFor(modeLabel string) string { + return map[string]string{"low": "coarse", "medium": "normal", "high": "fine", "ultra": "extra_fine"}[modeLabel] +} + +func formatSeedChunks(seedChunks []string) string { + if len(seedChunks) == 0 { + return "(no preliminary results)" + } + return strings.Join(seedChunks, "\n") +} + +func orDefault(v, def string) string { + if v == "" { + return def + } + return v +} diff --git a/internal/agent/harness/production.go b/internal/agent/harness/production.go new file mode 100644 index 0000000000..2a6f3a9139 --- /dev/null +++ b/internal/agent/harness/production.go @@ -0,0 +1,147 @@ +// +// 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 harness + +import ( + "context" + "encoding/json" + "fmt" + "log" + "strings" + + einotool "github.com/cloudwego/eino/components/tool" + + "gorm.io/gorm" + "ragflow/internal/agent/tool" + "ragflow/internal/service/nav" +) + +// ProductionRunner wires the real agentic-search tools (hybrid_search, +// dataset_navigation_by_tree) into the RunAgenticRAG flow, so the tools are +// actually invoked rather than merely registered. This is the production +// counterpart to the unit-testable SearchFn seam. +type ProductionRunner struct { + db *gorm.DB + tenantID string + datasetIDs []string + searchTool einotool.InvokableTool + navSvc nav.NavService // defaults to nav.GetNavService() when nil +} + +// NewProductionRunner builds a ProductionRunner backed by the real tools. The +// dataset-nav router (harness.NavigateDatasetByTree) resolves its NavService +// lazily via nav.GetNavService(). +func NewProductionRunner(db *gorm.DB, tenantID string, datasetIDs []string) (*ProductionRunner, error) { + searchBase, err := tool.BuildByName("hybrid_search", nil) + if err != nil { + return nil, err + } + search, ok := searchBase.(einotool.InvokableTool) + if !ok { + return nil, fmt.Errorf("hybrid_search is not invokable") + } + return &ProductionRunner{db: db, tenantID: tenantID, datasetIDs: datasetIDs, searchTool: search}, nil +} + +// newProductionRunnerWithTools builds a ProductionRunner with an injected +// search tool and nav service, for unit/E2E tests that want to fake the +// invocation surface without real services. +func newProductionRunnerWithTools(db *gorm.DB, tenantID string, datasetIDs []string, searchTool einotool.InvokableTool, navSvc nav.NavService) *ProductionRunner { + return &ProductionRunner{db: db, tenantID: tenantID, datasetIDs: datasetIDs, searchTool: searchTool, navSvc: navSvc} +} + +// Run executes the agentic-search graph with the real tools. It returns the +// final answer. +func (r *ProductionRunner) Run(ctx context.Context, question, keywords, modeLabel string) AnswerResult { + if r.searchTool == nil { + log.Printf("agentic_rag: production runner not fully wired (search tool missing)") + return AnswerResult{FinalAnswer: emptyResultMessage, Empty: true} + } + // The router tool returns a doc list; feed it as the search DocScope. + searchFn := func(ctx context.Context, query, kws string) ([]map[string]interface{}, []map[string]interface{}) { + return r.search(ctx, query, kws, nil) + } + + // For decomposition modes, route the doc scope first via the nav tool. + if mode, _ := GetMode(modeLabel); mode.RequiresDecomposition { + docs := r.routeDocs(ctx, question, keywords) + if len(docs) > 0 { + searchFn = func(ctx context.Context, query, kws string) ([]map[string]interface{}, []map[string]interface{}) { + return r.search(ctx, query, kws, docs) + } + } + } + return RunAgenticRAG(ctx, r.db, question, keywords, modeLabel, searchFn) +} + +// search invokes the hybrid_search tool and normalizes its chunk output. +func (r *ProductionRunner) search(ctx context.Context, query, keywords string, docScope []string) ([]map[string]interface{}, []map[string]interface{}) { + args := map[string]interface{}{"query": query, "keywords": keywords, "kb_ids": r.datasetIDs, "top_n": 12} + if len(docScope) > 0 { + args["doc_scope"] = docScope + } + raw, err := r.searchTool.InvokableRun(ctx, mustJSON(args)) + if err != nil { + log.Printf("agentic_rag: hybrid_search failed: %v", err) + return nil, nil + } + var res struct { + Chunks []map[string]interface{} `json:"chunks"` + } + if err := json.Unmarshal([]byte(raw), &res); err != nil { + return nil, nil + } + return res.Chunks, nil +} + +// routeDocs derives the doc scope via the canonical dataset-nav router +// (harness.NavigateDatasetByTree — the full LLM two-round selection). It routes +// across ALL bound datasets and merges the doc ids, so every KB contributes its +// own relevant docs to the shared scope (a multi-KB session must not collapse to +// the first KB only). +func (r *ProductionRunner) routeDocs(ctx context.Context, topic, keywords string) []string { + ns := r.navSvc + if ns == nil { + ns = nav.GetNavService() + } + if ns == nil { + log.Printf("agentic_rag: dataset nav service not initialized; skipping doc routing") + return nil + } + // Combine topic + keywords into the routing query so the nav router actually + // uses the full user signal (keywords must not be dropped). + query := strings.TrimSpace(topic + " " + keywords) + seen := map[string]bool{} + var docs []string + for _, kbID := range r.datasetIDs { + for _, id := range NavigateDatasetByTree(ctx, r.db, ns, r.tenantID, kbID, query) { + if id != "" && !seen[id] { + seen[id] = true + docs = append(docs, id) + } + } + } + return docs +} + +func mustJSON(v interface{}) string { + b, err := json.Marshal(v) + if err != nil { + return "{}" + } + return string(b) +} diff --git a/internal/agent/harness/production_test.go b/internal/agent/harness/production_test.go new file mode 100644 index 0000000000..cd53998d99 --- /dev/null +++ b/internal/agent/harness/production_test.go @@ -0,0 +1,144 @@ +package harness + +import ( + "context" + "strings" + "sync" + "testing" + + einotool "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" + "gorm.io/gorm" + + "ragflow/internal/agent/component" + "ragflow/internal/service/nav" +) + +// routeChat returns a model-style response depending on the stage: nav-selection +// calls return {"relevant":[...]}, all other calls return a plain final answer. +type routeChat struct{} + +func (routeChat) Invoke(_ context.Context, _ *gorm.DB, req component.ChatInvokeRequest) (*component.ChatInvokeResponse, error) { + content := "" + msg := "" + if len(req.Messages) > 0 { + msg = req.Messages[len(req.Messages)-1].Content + } + switch { + case strings.Contains(msg, "(numbered)") || strings.Contains(msg, "Entities"): + // nav-select pass: keep every item (all relevant). + content = `{"relevant":[0,1,2,3,4,5,6,7,8,9]}` + default: + content = "final scoped answer" + } + return &component.ChatInvokeResponse{Content: content}, nil +} + +// installRouteChat installs the stage-aware chat invoker for E2E tests. +func installRouteChat(t *testing.T) { + t.Helper() + component.SetDefaultChatInvoker(routeChat{}) + t.Cleanup(func() { component.SetDefaultChatInvoker(nil) }) +} + +// fakeInvokableTool is an einotool.InvokableTool double for E2E testing. +type fakeInvokableTool struct { + name string + fn func(ctx context.Context, argsJSON string) string + mu sync.Mutex + lastArgs string +} + +func (f *fakeInvokableTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: f.name}, nil +} + +func (f *fakeInvokableTool) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { + f.mu.Lock() + f.lastArgs = argsJSON + f.mu.Unlock() + return f.fn(context.Background(), argsJSON), nil +} + +func (f *fakeInvokableTool) args() string { + f.mu.Lock() + defer f.mu.Unlock() + return f.lastArgs +} + +// TestProductionRunner_E2E_RouteToScopedSearch proves the chain +// "user question → dataset nav routes doc scope → hybrid_search retrieval is +// scoped to those docs". It injects a fake nav service (returns docs d1,d2) and +// asserts the search tool receives doc_scope=[d1,d2]. +func TestProductionRunner_E2E_RouteToScopedSearch(t *testing.T) { + installRouteChat(t) + + // Fake nav service whose two-round router returns docs d1,d2. + navSvc := &fakeNavSvcHarness{ + clusters: []nav.NavNode{{Name: "C1", Description: "cluster"}}, + children: map[string][]nav.NavNode{ + "C1": { + {Name: "DocA", Type: "doc", DocID: "d1"}, + {Name: "DocB", Type: "doc", DocID: "d2"}, + }, + }, + } + // AskNavSelect returns all items (both docs) for the doc-select pass. + searchTool := &fakeInvokableTool{name: "hybrid_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[{"chunk_id":"c1","content_with_weight":"scoped evidence"}]}` + }} + + // Multi-KB session: routing must cover every bound dataset, not just the + // first KB. + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1", "kb2"}, searchTool, navSvc) + res := runner.Run(context.Background(), "Compare A and B", "", "medium") + + if res.FinalAnswer != "final scoped answer" { + t.Errorf("final answer = %q, want chat output", res.FinalAnswer) + } + // The search tool must receive doc_scope=[d1,d2] — the docs the nav router + // selected — proving retrieval is scoped to those docs across KBs. + if !strings.Contains(searchTool.args(), `"doc_scope":["d1","d2"]`) { + t.Errorf("hybrid_search not scoped to routed docs; args=%s", searchTool.args()) + } + // The search tool must receive all bound KBs (not collapsed to one). + if !strings.Contains(searchTool.args(), `"kb_ids":["kb1","kb2"]`) { + t.Errorf("hybrid_search kb_ids must include all bound KBs; args=%s", searchTool.args()) + } +} + +// TestProductionRunner_E2E_NoRoute_SearchUnscoped asserts low mode (no +// decomposition) skips nav routing and searches without a doc_scope. +func TestProductionRunner_E2E_NoRoute_SearchUnscoped(t *testing.T) { + installRouteChat(t) + searchTool := &fakeInvokableTool{name: "hybrid_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[{"chunk_id":"c1","content_with_weight":"evidence"}]}` + }} + // A nav service is provided but must NOT be consulted in low mode. + navSvc := &fakeNavSvcHarness{clusters: []nav.NavNode{{Name: "C1"}}} + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, searchTool, navSvc) + res := runner.Run(context.Background(), "What is X?", "", "low") + + if res.FinalAnswer == "" { + t.Error("expected a final answer in low mode") + } + if strings.Contains(searchTool.args(), "doc_scope") { + t.Errorf("low mode search should have no doc_scope; args=%s", searchTool.args()) + } +} + +// TestProductionRunner_E2E_EmptyRoute_SearchUnscoped asserts an empty doc route +// falls back to an unscoped search (no hard failure). +func TestProductionRunner_E2E_EmptyRoute_SearchUnscoped(t *testing.T) { + installRouteChat(t) + searchTool := &fakeInvokableTool{name: "hybrid_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[{"chunk_id":"c1","content_with_weight":"evidence"}]}` + }} + // Nav service with no clusters -> empty route. + navSvc := &fakeNavSvcHarness{clusters: nil} + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, searchTool, navSvc) + res := runner.Run(context.Background(), "Compare A and B", "", "medium") + if res.FinalAnswer == "" { + t.Error("expected an answer even when nav routing returns no docs") + } +} diff --git a/internal/agent/harness/route.go b/internal/agent/harness/route.go new file mode 100644 index 0000000000..5459cfd0d5 --- /dev/null +++ b/internal/agent/harness/route.go @@ -0,0 +1,133 @@ +// +// 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 harness + +import ( + "context" + "encoding/json" + "log" + "regexp" + "strings" + + "github.com/cloudwego/eino/schema" + + "ragflow/internal/agent/chat" + + "gorm.io/gorm" +) + +// routePrompt mirrors Python harness/prompts/route_prompt.py. +const routePrompt = `Analyze the following question and output a structured query analysis. + +Question: %s + +Analyze it across these dimensions: +1. Question type: factual / comparative / analytical / procedural / exploratory / verification / summarization. +2. Whether it needs decomposition into atomic facts, meaning whether multiple independent pieces of information must be retrieved separately before answering: true/false. +3. Suggested knowledge compilation tool: null (none) / toc (document table of contents) / graph (knowledge graph) / wiki (compiled domain knowledge). + +Output format (JSON): +{ + "question_type": "comparative", + "requires_decomposition": true, + "suggests_compilation": null, + "reasoning": "This is a comparative question, so it needs to be decomposed into two independent facts and one comparison relation." +} +` + +type routeResult struct { + QuestionType string `json:"question_type"` + RequiresDecomp *bool `json:"requires_decomposition"` + Reasoning string `json:"reasoning"` +} + +// RouteNode mirrors Python route_node. It classifies the question into a +// RouteDecision using the default chat invoker. Pure classification, no KB +// dependency. Never fails — falls back to a factual/direct decision. +func RouteNode(ctx context.Context, db *gorm.DB, question, modeLabel string) RouteDecision { + if strings.TrimSpace(question) == "" { + return fallbackRoute(question, modeLabel, "fallback: empty question") + } + inv := chat.GetDefaultInvoker() + if inv == nil { + return fallbackRoute(question, modeLabel, "fallback: chat invoker not configured") + } + resp, err := inv.Invoke(ctx, db, chat.Request{ + Messages: []schema.Message{ + {Role: schema.System, Content: strings.ReplaceAll(routePrompt, "%s", question)}, + {Role: schema.User, Content: question}, + }, + }) + if err != nil { + log.Printf("agentic_rag: route_node failed (fallback): %v", err) + return fallbackRoute(question, modeLabel, "fallback: LLM error") + } + var res routeResult + if err := unmarshalModelJSON(resp.Content, &res); err != nil { + log.Printf("agentic_rag: route_node parse failed (fallback): %v", err) + return fallbackRoute(question, modeLabel, "fallback: parse error") + } + return decide(question, modeLabel, res) +} + +// decide applies the mode's execution strategy to the LLM route output. +func decide(question, modeLabel string, res routeResult) RouteDecision { + mode, ok := GetMode(modeLabel) + if !ok { + mode = THINKING_MODES["medium"] + } + qType := res.QuestionType + if qType == "" { + qType = "factual" + } + needDecomp := true + if res.RequiresDecomp != nil { + needDecomp = *res.RequiresDecomp + } + return RouteDecision{ + Question: question, + ThinkingMode: modeLabel, + QuestionType: qType, + RequiresDecomposition: mode.RequiresDecomposition && needDecomp, + ExecutionStrategy: mode.Strategy, + Reasoning: res.Reasoning, + } +} + +func fallbackRoute(question, modeLabel, reason string) RouteDecision { + return RouteDecision{ + Question: question, ThinkingMode: modeLabel, QuestionType: "factual", + RequiresDecomposition: false, ExecutionStrategy: "direct_search", Reasoning: reason, + } +} + +var ( + reThinkTag = regexp.MustCompile(`(?s)^.*`) + reFence = regexp.MustCompile("```(?:json)?\\s*|\\s*```") +) + +// unmarshalModelJSON mirrors Python's _extract_json: strip thinking preamble and +// markdown fences, then parse JSON. +func unmarshalModelJSON(text string, out interface{}) error { + text = reThinkTag.ReplaceAllString(text, "") + text = reFence.ReplaceAllString(text, "") + text = strings.TrimSpace(text) + if text == "" { + return json.Unmarshal([]byte("{}"), out) + } + return json.Unmarshal([]byte(text), out) +} diff --git a/internal/agent/harness/route_test.go b/internal/agent/harness/route_test.go new file mode 100644 index 0000000000..03d9170079 --- /dev/null +++ b/internal/agent/harness/route_test.go @@ -0,0 +1,140 @@ +package harness + +import ( + "context" + "testing" + + "gorm.io/gorm" + + "ragflow/internal/agent/component" +) + +// fakeChatInvoker returns a fixed content for the chat call, so tests can drive +// the route/planner LLM output deterministically. +type fakeChatInvoker struct{ content string } + +func (f *fakeChatInvoker) Invoke(_ context.Context, _ *gorm.DB, _ component.ChatInvokeRequest) (*component.ChatInvokeResponse, error) { + return &component.ChatInvokeResponse{Content: f.content}, nil +} + +func installChat(t *testing.T, content string) { + t.Helper() + component.SetDefaultChatInvoker(&fakeChatInvoker{content: content}) + t.Cleanup(func() { component.SetDefaultChatInvoker(nil) }) +} + +// TestRouteNode_Classifies asserts route classification drives the execution +// strategy from the mode and the LLM's question_type. +func TestRouteNode_Classifies(t *testing.T) { + installChat(t, `{"question_type":"comparative","requires_decomposition":true,"reasoning":"cmp"}`) + r := RouteNode(context.Background(), nil, "Compare A and B", "medium") + if r.QuestionType != "comparative" { + t.Errorf("question_type = %q, want comparative", r.QuestionType) + } + // medium mode requires decomposition AND LLM says true -> true. + if !r.RequiresDecomposition { + t.Errorf("requires_decomposition = false, want true") + } + if r.ExecutionStrategy != "decompose_and_search" { + t.Errorf("execution_strategy = %q, want decompose_and_search", r.ExecutionStrategy) + } +} + +// TestRouteNode_LowModeDisablesDecomposition asserts low mode never decomposes +// even if the LLM requests it. +func TestRouteNode_LowModeDisablesDecomposition(t *testing.T) { + installChat(t, `{"question_type":"analytical","requires_decomposition":true}`) + r := RouteNode(context.Background(), nil, "Analyze X", "low") + if r.RequiresDecomposition { + t.Errorf("low mode must disable decomposition, got true") + } + if r.ExecutionStrategy != "direct_search" { + t.Errorf("low mode execution_strategy = %q, want direct_search", r.ExecutionStrategy) + } +} + +// TestRouteNode_FencedJSON asserts think-tag/fence stripping works. +func TestRouteNode_FencedJSON(t *testing.T) { + installChat(t, "Sure!\n```json\n{\"question_type\":\"factual\",\"requires_decomposition\":false}\n```") + r := RouteNode(context.Background(), nil, "What is X?", "medium") + if r.QuestionType != "factual" { + t.Errorf("question_type = %q, want factual", r.QuestionType) + } + if r.RequiresDecomposition { + t.Errorf("requires_decomposition = true, want false") + } +} + +// TestRouteNode_EmptyQuestionFallsBack asserts an empty question yields a +// direct factual decision without calling the LLM. +func TestRouteNode_EmptyQuestionFallsBack(t *testing.T) { + r := RouteNode(context.Background(), nil, "", "medium") + if r.QuestionType != "factual" || r.RequiresDecomposition { + t.Errorf("empty question must fall back to direct factual, got %+v", r) + } +} + +// TestPlannerNode_DirectMode asserts a non-decomposed route yields one coarse +// claim without calling the LLM. +func TestPlannerNode_DirectMode(t *testing.T) { + plan := PlannerNode(context.Background(), nil, RouteDecision{ + Question: "What is X?", RequiresDecomposition: false, + }, nil) + if plan.PlanType != "direct" || len(plan.Claims) != 1 { + t.Fatalf("direct plan = %+v, want 1 direct claim", plan) + } +} + +// TestPlannerNode_Decomposes asserts the planner builds claims from the LLM +// output and applies the mode's max iterations. +func TestPlannerNode_Decomposes(t *testing.T) { + installChat(t, `{"claims":[ + {"claim_id":"c0","description":"fact one","priority":0}, + {"claim_id":"c1","description":"fact two","priority":1} + ]}`) + plan := PlannerNode(context.Background(), nil, RouteDecision{ + Question: "Compare A and B", QuestionType: "comparative", RequiresDecomposition: true, ThinkingMode: "medium", + }, nil) + if plan.PlanType != "comparative_decomposition" { + t.Errorf("plan_type = %q, want comparative_decomposition", plan.PlanType) + } + if len(plan.Claims) != 2 { + t.Fatalf("claims = %d, want 2", len(plan.Claims)) + } + if plan.Claims[1].Priority != 1 { + t.Errorf("claim priority = %d, want 1", plan.Claims[1].Priority) + } + // medium maxOrchestratorCycles = 3 + if plan.MaxIterations != 3 { + t.Errorf("max_iterations = %d, want 3", plan.MaxIterations) + } +} + +// TestPlannerNode_UnknownModeFallsBack asserts an unknown (non-empty) mode label +// falls back to medium so the planner is not driven by a zero-valued mode +// (which would produce a degenerate plan with max_claims=0). +func TestPlannerNode_UnknownModeFallsBack(t *testing.T) { + installChat(t, `{"claims":[{"claim_id":"c0","description":"fact one","priority":0}]}`) + plan := PlannerNode(context.Background(), nil, RouteDecision{ + Question: "Q", RequiresDecomposition: true, ThinkingMode: "turbo-unknown", + }, nil) + // medium maxOrchestratorCycles = 3, and claims must still be built. + if plan.MaxIterations != 3 { + t.Errorf("max_iterations = %d, want 3 (medium fallback)", plan.MaxIterations) + } + if len(plan.Claims) != 1 { + t.Errorf("claims = %d, want 1", len(plan.Claims)) + } +} + +// TestPlannerNode_BadJSONFallsBack asserts unparseable planner output falls back +// to the direct plan. +func TestPlannerNode_BadJSONFallsBack(t *testing.T) { + installChat(t, "not json at all") + plan := PlannerNode(context.Background(), nil, RouteDecision{ + Question: "Q", RequiresDecomposition: true, ThinkingMode: "medium", + }, nil) + if plan.PlanType != "direct" || len(plan.Claims) != 1 { + t.Fatalf("fallback plan = %+v, want direct", plan) + } +} diff --git a/internal/agent/harness/sufficiency.go b/internal/agent/harness/sufficiency.go new file mode 100644 index 0000000000..beb4290133 --- /dev/null +++ b/internal/agent/harness/sufficiency.go @@ -0,0 +1,250 @@ +// +// 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 harness + +import ( + "fmt" + "regexp" + "strings" +) + +// Sufficiency scoring (code-only, mirrors Python sufficiency.py): cross-check an +// agent result against the evidence chunks, fuse agent confidence + cross-check +// pass rate, then route to a 5-way verdict. + +var reNumber = regexp.MustCompile(`\d+\.?\d*`) +var reEntities = regexp.MustCompile(`\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b`) + +// extractNumbers returns numeric values found in text. +func extractNumbers(text string) []string { + return reNumber.FindAllString(text, -1) +} + +// extractNamedEntities returns capitalized multi-word sequences. +func extractNamedEntities(text string) []string { + seen := map[string]bool{} + var out []string + for _, e := range reEntities.FindAllString(text, -1) { + if !seen[e] { + seen[e] = true + out = append(out, e) + } + } + return out +} + +// CrossCheckClaim performs a code-level cross-check of an agent result against +// the accumulated evidence chunks (number matching + entity presence). +func CrossCheckClaim(agent *AgentResult, allChunks map[int]map[string]interface{}) ClaimCrossCheckResult { + if agent == nil { + return ClaimCrossCheckResult{ClaimID: "", CrossCheckPassed: false, Mismatches: []string{"nil agent result"}} + } + if !agent.IsVerified { + return ClaimCrossCheckResult{ClaimID: agent.ClaimID, CrossCheckPassed: false, Mismatches: []string{"agent self-reported as unverified"}} + } + numbers := extractNumbers(agent.Report) + entities := extractNamedEntities(agent.Report) + + var matches, mismatches []string + for _, eid := range agent.EvidenceIDs { + chunk, ok := allChunks[eid] + if !ok { + mismatches = append(mismatches, fmt.Sprintf("evidence_id=%d: chunk not found", eid)) + continue + } + text := "" + if c, ok := chunk["content_with_weight"].(string); ok { + text = strings.ToLower(c) + } else if c, ok := chunk["content"].(string); ok { + text = strings.ToLower(c) + } + for _, num := range numbers { + if strings.Contains(text, num) { + matches = append(matches, fmt.Sprintf("number %s found in chunk %d", num, eid)) + } else { + mismatches = append(mismatches, fmt.Sprintf("number %s not found in chunk %d", num, eid)) + } + } + for _, ent := range entities { + if strings.Contains(text, strings.ToLower(ent)) { + matches = append(matches, fmt.Sprintf("entity '%s' found in chunk %d", ent, eid)) + } else { + mismatches = append(mismatches, fmt.Sprintf("entity '%s' not found in chunk %d", ent, eid)) + } + } + } + // HasEvidence is true when at least one evidence id resolved to a chunk with + // content. A claim with no resolvable evidence is never considered verified. + hasEvidence := len(matches)+len(mismatches) > 0 + total := len(matches) + len(mismatches) + crossScore := 0.0 + if total > 0 { + crossScore = float64(len(matches)) / float64(total) + } + // Entity presence now contributes to matches too, so it can raise the score; + // use a float comparison to avoid integer-division truncation bias. + crossPassed := hasEvidence && float64(len(mismatches)) < float64(len(matches))/2.0 + return ClaimCrossCheckResult{ + ClaimID: agent.ClaimID, CrossCheckPassed: crossPassed, CrossCheckScore: crossScore, + EvidenceMatches: matches, Mismatches: mismatches, HasEvidence: hasEvidence, + } +} + +// ComputeFusionScore fuses agent confidence + cross-check pass rate into a +// SufficiencyVerdict for the given mode. +func ComputeFusionScore(agentResults []AgentResult, crossResults []ClaimCrossCheckResult, mode ExecutionStrategy) SufficiencyVerdict { + verified := 0 + for _, r := range agentResults { + if r.IsVerified { + verified++ + } + } + agentScore := 0.0 + if len(agentResults) > 0 { + agentScore = float64(verified) / float64(len(agentResults)) + } + passed := 0 + for _, r := range crossResults { + if r.CrossCheckPassed { + passed++ + } + } + crossScore := 0.0 + if len(crossResults) > 0 { + crossScore = float64(passed) / float64(len(crossResults)) + } + + fusionScore := agentScore + if crossScore > fusionScore { + fusionScore = crossScore // low/medium default: max + } + switch mode.Label { + case "ultra": + fusionScore = min(agentScore, crossScore) + case "high": + fusionScore = (agentScore + crossScore) / 2 + } + + hasConflicts := false + for _, r := range crossResults { + if len(r.Mismatches) > 0 { + hasConflicts = true + break + } + } + + // Empty-evidence guard: if no claim examined any evidence chunk, the answer + // cannot be grounded at all — this is UNANSWERABLE, not merely incomplete. + anyEvidence := false + for _, r := range crossResults { + if r.HasEvidence { + anyEvidence = true + break + } + } + + status := "INSUFFICIENT" + switch { + case !anyEvidence: + status = "UNANSWERABLE" + case hasConflicts && fusionScore < mode.PartialThreshold: + status = "CONFLICTING" + case fusionScore >= mode.SufficiencyThreshold: + status = "SUFFICIENT" + case fusionScore >= mode.PartialThreshold: + status = "USEFUL_BUT_INCOMPLETE" + case func() bool { + for _, r := range crossResults { + if r.CrossCheckPassed { + return true + } + } + return false + }(): + status = "INSUFFICIENT" + default: + status = "UNANSWERABLE" + } + + var missing []string + for _, r := range crossResults { + if !r.CrossCheckPassed || !r.HasEvidence { + missing = append(missing, r.ClaimID) + } + } + + assessments := make([]map[string]interface{}, 0, len(crossResults)) + for _, r := range crossResults { + assessments = append(assessments, map[string]interface{}{ + "claim_id": r.ClaimID, "is_verified": r.CrossCheckPassed && r.HasEvidence, "score": r.CrossCheckScore, + "mismatches": r.Mismatches, "has_evidence": r.HasEvidence, + }) + } + + return SufficiencyVerdict{ + Status: status, Score: fusionScore, AgentScore: agentScore, CrossScore: crossScore, + ClaimAssessments: assessments, HasConflicts: hasConflicts, MissingClaims: missing, + Feedback: buildFeedback(missing, crossResults), OverallReason: fmt.Sprintf("%s score=%.2f missing=%v", status, fusionScore, missing), + } +} + +func buildFeedback(missing []string, results []ClaimCrossCheckResult) string { + if len(missing) == 0 { + return "all claims verified" + } + var hints []string + for _, r := range results { + if !r.CrossCheckPassed { + hints = append(hints, fmt.Sprintf("claim %s: %d mismatch(es)", r.ClaimID, len(r.Mismatches))) + } + } + return "missing: " + strings.Join(hints, "; ") +} + +// RouteSufficiencyVerdict returns (action, shouldContinue) from the verdict. +func RouteSufficiencyVerdict(v SufficiencyVerdict, modeLabel string, cycle, maxCycles int) (string, bool) { + mode, _ := GetMode(modeLabel) + if mode.Label == "" { + mode = THINKING_MODES["medium"] + } + switch v.Status { + case "SUFFICIENT": + return "ANSWER", false + case "USEFUL_BUT_INCOMPLETE": + if mode.RequiresSelectiveGen { + return "ANSWER_PARTIAL", false + } + return "CONTINUE", false + case "INSUFFICIENT": + if cycle >= int(float64(maxCycles)*0.8) { + return "ANSWER_PARTIAL", false + } + return "CONTINUE", true + case "CONFLICTING": + if mode.AllowsReplan && cycle < int(float64(maxCycles)*0.5) { + return "REPLAN", true + } + return "ANSWER_PARTIAL", false + case "UNANSWERABLE": + if mode.FallbackToDirectLLM { + return "FALLBACK_LLM", false + } + return "ABSTAIN", false + default: + return "CONTINUE", true + } +} diff --git a/internal/agent/harness/types.go b/internal/agent/harness/types.go new file mode 100644 index 0000000000..3b9c56d357 --- /dev/null +++ b/internal/agent/harness/types.go @@ -0,0 +1,151 @@ +// +// 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 harness ports Python's rag/advanced_rag/harness (agentic-RAG): +// route → planner → orchestrator → formalize_answer. This is a phased port; see +// tasks/agentic_search_port_plan.md §P4. +package harness + +// RouteDecision mirrors Python RouteDecision. +type RouteDecision struct { + Question string + ThinkingMode string + QuestionType string // factual | comparative | analytical | procedural | exploratory | verification | summarization + RequiresDecomposition bool + ExecutionStrategy string // direct_search | decompose_and_search | agentic_research | deep_research + Reasoning string +} + +// ClaimTarget mirrors Python ClaimTarget. +type ClaimTarget struct { + ClaimID string + Description string + Priority int + SuggestedTools []string + IsVerified bool + Confidence float64 + AgentResult *AgentResult +} + +// WorkflowPlan mirrors Python WorkflowPlan. +type WorkflowPlan struct { + PlanType string // direct | fact_decomposition | ... + Claims []ClaimTarget + MaxIterations int +} + +// ExecutionStrategy mirrors Python ExecutionStrategy and is instantiated by the +// four thinking modes. +type ExecutionStrategy struct { + Label string + Strategy string + RequiresDecomposition bool + MaxOrchestratorCycles int + MaxAgentCycles int + MaxParallelAgents int + AvailableTools []string + SufficiencyThreshold float64 + PartialThreshold float64 + FallbackToDirectLLM bool + RequiresSelectiveGen bool + AllowsReplan bool +} + +// AgentResult mirrors Python AgentResult. +type AgentResult struct { + ClaimID string + Report string + IsVerified bool + Confidence float64 + EvidenceIDs []int +} + +// ClaimCrossCheckResult mirrors Python ClaimCrossCheckResult. +type ClaimCrossCheckResult struct { + ClaimID string + CrossCheckPassed bool + CrossCheckScore float64 + EvidenceMatches []string + Mismatches []string + // HasEvidence reports whether any evidence chunk was actually examined. When + // false the claim had no resolvable evidence, so it cannot be considered + // verified regardless of the other fields. + HasEvidence bool +} + +// SufficiencyVerdict mirrors Python SufficiencyVerdict. +type SufficiencyVerdict struct { + Status string // SUFFICIENT | USEFUL_BUT_INCOMPLETE | INSUFFICIENT | CONFLICTING | UNANSWERABLE + Score float64 + AgentScore float64 + CrossScore float64 + ClaimAssessments []map[string]interface{} + HasConflicts bool + MissingClaims []string + Feedback string + OverallReason string +} + +// OrchestratorContext mirrors Python OrchestratorContext. +type OrchestratorContext struct { + Question string + Claims []*ClaimTarget + Mode string + Iteration int +} + +// THINKING_MODES mirrors Python config.THINKING_MODES. +var THINKING_MODES = map[string]ExecutionStrategy{ + "low": { + Label: "low", Strategy: "direct_search", RequiresDecomposition: false, + MaxOrchestratorCycles: 1, MaxAgentCycles: 0, MaxParallelAgents: 1, + AvailableTools: []string{"hybrid_search"}, SufficiencyThreshold: 0.85, PartialThreshold: 0.50, + }, + "medium": { + Label: "medium", Strategy: "decompose_and_search", RequiresDecomposition: true, + MaxOrchestratorCycles: 3, MaxAgentCycles: 0, MaxParallelAgents: 1, + AvailableTools: []string{"hybrid_search"}, SufficiencyThreshold: 0.75, PartialThreshold: 0.40, + RequiresSelectiveGen: true, + }, + "high": { + Label: "high", Strategy: "agentic_research", RequiresDecomposition: true, + MaxOrchestratorCycles: 3, MaxAgentCycles: 2, MaxParallelAgents: 2, + AvailableTools: []string{ + "hybrid_search", "web_search", "ontology_navigate", "dataset_navigation_by_tree", + "graph_explore", "inspector_open_context", "inspector_compare", + }, + SufficiencyThreshold: 0.65, PartialThreshold: 0.30, + RequiresSelectiveGen: true, AllowsReplan: true, + }, + "ultra": { + Label: "ultra", Strategy: "deep_research", RequiresDecomposition: true, + MaxOrchestratorCycles: 4, MaxAgentCycles: 2, MaxParallelAgents: 3, + AvailableTools: []string{ + "hybrid_search", "bm25_search", "web_search", "structured_query", + "ontology_navigate", "dataset_navigation_by_tree", "mindmap_navigate", + "graph_explore", "wiki_query", "inspector_open_context", "inspector_compare", + "inspector_grep_within", "inspector_request_adjacent", + }, + SufficiencyThreshold: 0.55, PartialThreshold: 0.20, FallbackToDirectLLM: true, + RequiresSelectiveGen: true, AllowsReplan: true, + }, +} + +// GetMode mirrors Python get_mode(label). +func GetMode(label string) (ExecutionStrategy, bool) { + m, ok := THINKING_MODES[label] + return m, ok +} diff --git a/internal/agent/tool/agentic_search.go b/internal/agent/tool/agentic_search.go new file mode 100644 index 0000000000..3a5de007f9 --- /dev/null +++ b/internal/agent/tool/agentic_search.go @@ -0,0 +1,345 @@ +// +// 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 tool + +import ( + "context" + "encoding/json" + "fmt" + "regexp" + "strings" + + einotool "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// Agentic search tool names mirror Python rag/advanced_rag/harness/tools/search.py. +const ( + toolHybridSearch = "hybrid_search" + toolVectorSearch = "vector_search" + toolBM25Search = "bm25_search" + toolWebSearch = "web_search" + toolStructuredQuery = "structured_query" +) + +// hybridSearchArgs is the shared JSON schema for the three retrieval tools. +type hybridSearchArgs struct { + Query string `json:"query"` + KbIDs []string `json:"kb_ids,omitempty"` + TopN int `json:"top_n,omitempty"` + DocScope []string `json:"doc_scope,omitempty"` + Keywords string `json:"keywords,omitempty"` + UseCompiled bool `json:"use_compiled,omitempty"` +} + +type agenticSearchResult struct { + Chunks []map[string]interface{} `json:"chunks"` + DocAggs []map[string]interface{} `json:"doc_aggs,omitempty"` +} + +// AgenticSearchTool is the hybrid/vector/bm25 retrieval tool. The search mode +// selects the vector-similarity weight used by the underlying retrieval service: +// +// hybrid: 0.3 (hybrid of keyword + vector) +// vector: 1.0 (vector-only) +// bm25: 0.0 (keyword-only) +// +// It backs onto GetRetrievalService() (the same singleton the agent Retrieval +// tool uses), so DocScope and KB scoping carry through automatically. +type AgenticSearchTool struct { + mode string // hybrid_search | vector_search | bm25_search + weight float64 + defaults hybridSearchArgs +} + +// NewAgenticSearchTool returns the retrieval tool for the given mode. +func NewAgenticSearchTool(mode string) *AgenticSearchTool { + weight := 0.3 + switch mode { + case toolVectorSearch: + weight = 1.0 + case toolBM25Search: + weight = 0.0 + } + return &AgenticSearchTool{mode: mode, weight: weight, defaults: hybridSearchArgs{TopN: 12}} +} + +func (a *AgenticSearchTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: a.mode, + Desc: fmt.Sprintf("Search the bound knowledge base(s) for the query (mode=%s). Returns relevant passages.", a.mode), + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": { + Type: schema.String, Required: true, Desc: "The search query.", + }, + "kb_ids": {Type: schema.Array, Desc: "Optional dataset ids to restrict to."}, + "top_n": {Type: schema.Number, Desc: "Number of passages to return (default 12)."}, + "doc_scope": {Type: schema.Array, Desc: "Optional doc ids to restrict to."}, + "keywords": {Type: schema.String, Desc: "Comma-separated keywords to narrow results."}, + "use_compiled": {Type: schema.Boolean, Desc: "Whether to enrich with compiled products."}, + }), + }, nil +} + +// InvokableRun executes the retrieval. It returns JSON with "chunks" (array of +// chunk maps). Never returns a hard error for retrieval failures — it returns an +// empty result so the agent can fall back. +func (a *AgenticSearchTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...einotool.Option) (string, error) { + var args hybridSearchArgs + if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil { + return "", fmt.Errorf("%s: parse arguments: %w", a.mode, err) + } + if args.TopN <= 0 { + args.TopN = a.defaults.TopN + } + + svc := GetRetrievalService() + tenantID := canvasTenantID(ctx) + datasetIDs := args.KbIDs + if len(datasetIDs) == 0 { + datasetIDs = canvasDatasetIDs(ctx, nil) + } + if svc == nil || tenantID == "" || len(datasetIDs) == 0 { + return jsonChunksEmpty(), nil + } + + weight := a.weight + req := RetrievalRequest{ + Query: strings.TrimSpace(args.Query + " " + args.Keywords), + DatasetIDs: datasetIDs, + TopN: args.TopN, + TopK: args.TopN * 4, + SimilarityThreshold: 0.2, + KeywordsSimilarityWeight: &weight, + DocScope: args.DocScope, + } + chunks, err := svc.Search(ctx, nil, req) + if err != nil { + return jsonChunksEmpty(), nil // agent falls back on failure + } + + // Keyword narrowing (mirrors Python _narrow_by_keywords). + if args.Keywords != "" { + chunks = narrowByKeywords(chunks, args.Keywords) + } + return marshalSearchResult(chunks), nil +} + +// narrowByKeywords narrows each chunk to keyword-bearing sentences (+/-1 +// neighbour) and drops keyword-less chunks. A simplified port of Python's +// _narrow_by_keywords. +func narrowByKeywords(chunks []RetrievalChunk, keywords string) []RetrievalChunk { + kwds := splitKeywords(keywords) + if len(kwds) == 0 { + return chunks + } + seen := map[string]struct{}{} + out := make([]RetrievalChunk, 0, len(chunks)) + for _, c := range chunks { + nc, ok := narrowContent(c.Content, kwds) + if !ok { + continue + } + hash := md5Hex(nc) + if _, dup := seen[hash]; dup { + continue + } + seen[hash] = struct{}{} + c.Content = nc + out = append(out, c) + } + return out +} + +// splitKeywords normalizes a keyword string into a list of terms. When fewer +// than 3 comma terms exist, falls back to space-split bigrams (mirrors Python). +func splitKeywords(keywords string) []string { + if strings.TrimSpace(keywords) == "" { + return nil + } + kwds := make([]string, 0, 8) + for _, k := range strings.Split(keywords, ",") { + if k = strings.TrimSpace(k); k != "" { + kwds = append(kwds, strings.ToLower(k)) + } + } + if len(kwds) < 3 { + words := make([]string, 0, 8) + for _, w := range strings.Split(keywords, " ") { + if w = strings.TrimSpace(w); w != "" { + words = append(words, strings.ToLower(w)) + } + } + bigrams := make([]string, 0, len(words)) + for i := 0; i+1 < len(words); i++ { + bigrams = append(bigrams, words[i]+" "+words[i+1]) + } + if len(bigrams) > 0 { + return bigrams + } + } + return kwds +} + +var sentEnd = regexp.MustCompile(`[。！？；!?;]+|\.`) + +// splitSentences splits text into sentences at sentence terminators, keeping a +// digit-guarded period intact ("3.14" / "v1.2" are not split). Implemented with +// a simple splitter because RE2 (Go) does not support lookbehind/lookahead. +func splitSentences(content string) []string { + raw := sentEnd.Split(content, -1) + sents := make([]string, 0, len(raw)) + for _, s := range raw { + if strings.TrimSpace(s) == "" { + continue + } + // Re-join a trailing digit-period-digit that the splitter cut apart: + // if s ends with a digit and content had "." following, reattach. + sents = append(sents, s) + } + return rejoinDigitPeriods(sents) +} + +// rejoinDigitPeriods merges "…1" + "2…" back into "…1.2…" when a decimal point +// separated two digit groups. +func rejoinDigitPeriods(sents []string) []string { + out := make([]string, 0, len(sents)) + for i := 0; i < len(sents); i++ { + cur := sents[i] + // If current ends with a digit and next begins with a digit, the split + // point was a decimal point — merge them. + for i+1 < len(sents) && hasTrailingDigit(cur) && hasLeadingDigit(sents[i+1]) { + cur = strings.TrimRight(cur, " \t") + "." + sents[i+1] + i++ + } + out = append(out, cur) + } + return out +} + +func hasTrailingDigit(s string) bool { + s = strings.TrimRight(s, " \t") + return len(s) > 0 && s[len(s)-1] >= '0' && s[len(s)-1] <= '9' +} + +func hasLeadingDigit(s string) bool { + s = strings.TrimLeft(s, " \t") + return len(s) > 0 && s[0] >= '0' && s[0] <= '9' +} + +// narrowContent returns the keyword-bearing sentences (+/-1 neighbour) with the +// keyword highlighted, or (_, false) if no keyword occurs. +func narrowContent(content string, kwds []string) (string, bool) { + if strings.TrimSpace(content) == "" { + return "", false + } + sents := splitSentences(content) + if len(sents) == 0 { + return "", false + } + keep := map[int]bool{} + matched := false + for i, s := range sents { + low := strings.ToLower(s) + for _, kw := range kwds { + if kw != "" && strings.Contains(low, kw) { + matched = true + if i > 0 { + keep[i-1] = true + } + keep[i] = true + if i+1 < len(sents) { + keep[i+1] = true + } + break + } + } + } + if !matched { + return "", false + } + var b strings.Builder + for i := 0; i < len(sents); i++ { + if keep[i] { + b.WriteString(sents[i]) + } + } + return "..." + highlightKeywords(b.String(), kwds) + "...", true +} + +// highlightKeywords wraps keyword occurrences in . +func highlightKeywords(text string, kwds []string) string { + if len(kwds) == 0 { + return text + } + // Sort by length desc so longer terms match first. + terms := make([]string, len(kwds)) + copy(terms, kwds) + for i := 1; i < len(terms); i++ { + for j := i; j > 0 && len(terms[j]) > len(terms[j-1]); j-- { + terms[j], terms[j-1] = terms[j-1], terms[j] + } + } + pattern := "(" + for i, t := range terms { + if t == "" { + continue + } + if i > 0 { + pattern += "|" + } + pattern += regexp.QuoteMeta(t) + } + pattern += ")" + re := regexp.MustCompile(`(?i)` + pattern) + return re.ReplaceAllString(text, "${1}") +} + +func md5Hex(s string) string { + h := uint32(2166136261) + for i := 0; i < len(s); i++ { + h ^= uint32(s[i]) + h *= 16777619 + } + return fmt.Sprintf("%08x", h) +} + +func jsonChunksEmpty() string { + return `{"chunks":[]}` +} + +func marshalSearchResult(chunks []RetrievalChunk) string { + type outChunk struct { + ID string `json:"id"` + Content string `json:"content"` + DocumentID string `json:"doc_id"` + DocName string `json:"docnm_kwd"` + Score float64 `json:"similarity"` + } + out := make([]outChunk, 0, len(chunks)) + for _, c := range chunks { + out = append(out, outChunk{ + ID: c.ID, Content: c.Content, DocumentID: c.DocumentID, + DocName: c.DocumentName, Score: c.Score, + }) + } + b, err := json.Marshal(map[string]interface{}{"chunks": out}) + if err != nil { + return jsonChunksEmpty() + } + return string(b) +} diff --git a/internal/agent/tool/agentic_search_test.go b/internal/agent/tool/agentic_search_test.go new file mode 100644 index 0000000000..c0c49baf91 --- /dev/null +++ b/internal/agent/tool/agentic_search_test.go @@ -0,0 +1,78 @@ +package tool + +import ( + "strings" + "testing" +) + +// TestSplitKeywords_BigramFallback asserts comma terms are used when >=3, and +// space-split bigrams otherwise. +func TestSplitKeywords_BigramFallback(t *testing.T) { + got := splitKeywords("alpha, beta, gamma") + if len(got) != 3 { + t.Fatalf("comma kwds = %d, want 3", len(got)) + } + got = splitKeywords("a b c d") + // 4 words -> 3 bigrams + if len(got) != 3 { + t.Fatalf("bigram kwds = %d, want 3", len(got)) + } + if got[0] != "a b" { + t.Errorf("bigram[0] = %q, want \"a b\"", got[0]) + } +} + +// TestNarrowContent_KeepsKeywordSentence asserts narrowing keeps the +// keyword-bearing sentence and its neighbours, and highlights the keyword. +func TestNarrowContent_KeepsKeywordSentence(t *testing.T) { + content := "The introduction is boring. The key insight about rocket engines is here. The conclusion is short." + nc, ok := narrowContent(content, []string{"rocket"}) + if !ok { + t.Fatal("expected keyword match") + } + if !strings.Contains(nc, "key insight") { + t.Errorf("narrowed content missing keyword sentence: %q", nc) + } + if !strings.Contains(nc, "rocket") { + t.Errorf("narrowed content missing highlight: %q", nc) + } +} + +// TestNarrowContent_NoKeyword asserts narrowing returns false when no keyword +// occurs. +func TestNarrowContent_NoKeyword(t *testing.T) { + if _, ok := narrowContent("no match here at all", []string{"zzz"}); ok { + t.Fatal("expected no match") + } +} + +// TestSplitSentences_DigitPeriod asserts "3.14" is not split into two sentences. +func TestSplitSentences_DigitPeriod(t *testing.T) { + sents := splitSentences("pi is 3.14 and e is 2.71. end.") + found := false + for _, s := range sents { + if strings.Contains(s, "3.14") { + found = true + } + } + if !found { + t.Fatalf("digit-period sentence lost: %v", sents) + } +} + +// TestNarrowByKeywords_DropsKeywordless asserts chunks without keywords are +// dropped and deduped by narrowed content hash. +func TestNarrowByKeywords_DropsKeywordless(t *testing.T) { + chunks := []RetrievalChunk{ + {ID: "c1", Content: "alpha talks about beta in detail here."}, + {ID: "c2", Content: "unrelated content entirely."}, + {ID: "c3", Content: "alpha talks about beta in detail here."}, // dup of c1 + } + out := narrowByKeywords(chunks, "beta") + if len(out) != 1 { + t.Fatalf("narrowed chunks = %d, want 1 (drops keywordless + dedups)", len(out)) + } + if out[0].ID != "c1" { + t.Errorf("kept chunk = %q, want c1", out[0].ID) + } +} diff --git a/internal/agent/tool/canvas_ctx.go b/internal/agent/tool/canvas_ctx.go new file mode 100644 index 0000000000..6fb4d78687 --- /dev/null +++ b/internal/agent/tool/canvas_ctx.go @@ -0,0 +1,59 @@ +// +// 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 tool + +import ( + "context" + + "ragflow/internal/agent/runtime" +) + +// canvasTenantID derives the tenant id from canvas state, falling back to +// user_id. Shared by agentic search and dataset-navigation tools. +func canvasTenantID(ctx context.Context) string { + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return "" + } + if tenantID, _ := state.Sys["tenant_id"].(string); tenantID != "" { + return tenantID + } + userID, _ := state.Sys["user_id"].(string) + return userID +} + +// canvasDatasetIDs returns the explicit dataset ids (all of them, preserving +// multi-KB sessions), else the canvas sys dataset_id as a single-element list. +func canvasDatasetIDs(ctx context.Context, explicit []string) []string { + if len(explicit) > 0 { + out := make([]string, 0, len(explicit)) + for _, id := range explicit { + if id != "" { + out = append(out, id) + } + } + return out + } + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return nil + } + if id, _ := state.Sys["dataset_id"].(string); id != "" { + return []string{id} + } + return nil +} diff --git a/internal/agent/tool/dataset_navigation.go b/internal/agent/tool/dataset_navigation.go new file mode 100644 index 0000000000..0f33b6749e --- /dev/null +++ b/internal/agent/tool/dataset_navigation.go @@ -0,0 +1,220 @@ +// +// 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 tool + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + einotool "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" + + "ragflow/internal/service/nav" +) + +// datasetNavigationToolName mirrors Python's dataset_navigation_by_tree router +// tool. It navigates the dataset nav tree and returns the doc_ids to read. +const datasetNavigationToolName = "dataset_navigation_by_tree" + +const datasetNavigationToolDescription = "Navigate a dataset's navigation tree by topic and return the document ids that are likely relevant." + +// datasetNavigationArgs is the JSON schema the model sends into InvokableRun. +type datasetNavigationArgs struct { + Topic string `json:"topic"` + Keywords string `json:"keywords,omitempty"` + DatasetIDs []string `json:"dataset_ids,omitempty"` + DocScope string `json:"doc_scope,omitempty"` + MaxDocs int `json:"max_docs,omitempty"` +} + +// datasetNavigationResult is the JSON shape returned to the model. +type datasetNavigationResult struct { + Docs []string `json:"docs,omitempty"` + Error string `json:"_ERROR,omitempty"` + NotFound bool `json:"not_found,omitempty"` +} + +// datasetNavigationDefaultMaxDocs caps the number of doc_ids returned. +const datasetNavigationDefaultMaxDocs = 8 + +// DatasetNavigationByTree is the dataset-navigation router tool. Minimal closed +// loop: one-level drill-down from the root clusters and deduplicated doc ids +// (max MaxDocs). LLM-guided multi-level selection is deferred. +type DatasetNavigationByTree struct { + defaults datasetNavigationArgs +} + +// NewDatasetNavigationByTree returns a DatasetNavigationByTree implementing +// eino's tool.InvokableTool interface. +func NewDatasetNavigationByTree() *DatasetNavigationByTree { + return NewDatasetNavigationByTreeWithDefaults(datasetNavigationArgs{}) +} + +// NewDatasetNavigationByTreeWithDefaults returns a DatasetNavigationByTree with +// node-level defaults. +func NewDatasetNavigationByTreeWithDefaults(defaults datasetNavigationArgs) *DatasetNavigationByTree { + if defaults.MaxDocs <= 0 { + defaults.MaxDocs = datasetNavigationDefaultMaxDocs + } + return &DatasetNavigationByTree{defaults: defaults} +} + +// Info returns the tool's metadata for the chat model. +func (d *DatasetNavigationByTree) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: datasetNavigationToolName, + Desc: datasetNavigationToolDescription, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "topic": { + Type: schema.String, + Desc: "The topic to navigate to. Use the core subject from the original request.", + Required: true, + }, + "keywords": { + Type: schema.String, + Desc: "Optional additional keywords to disambiguate the topic.", + }, + }), + }, nil +} + +// InvokableRun executes the tool. It navigates the nav tree via the registered +// NavService (internal/service datasetnav) and returns a deduplicated doc_id +// list (max MaxDocs). +func (d *DatasetNavigationByTree) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...einotool.Option) (string, error) { + var args datasetNavigationArgs + if argumentsInJSON != "" { + if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil { + return "", fmt.Errorf("dataset_navigation: parse arguments: %w", err) + } + } + args = d.mergeDefaults(args) + if args.Topic == "" { + return "", fmt.Errorf("dataset_navigation: topic is required") + } + // Per-request max_docs overrides the node default; default to a sane cap. + maxDocs := args.MaxDocs + if maxDocs <= 0 { + maxDocs = datasetNavigationDefaultMaxDocs + } + + ns := nav.GetNavService() + if ns == nil { + return datasetNavigationJSON(datasetNavigationResult{ + Error: "dataset navigation service not initialized (SetNavService must be called at bootstrap)", + }), nil + } + + tenantID := canvasTenantID(ctx) + datasetIDs := canvasDatasetIDs(ctx, args.DatasetIDs) + if tenantID == "" || len(datasetIDs) == 0 { + return datasetNavigationJSON(datasetNavigationResult{ + NotFound: true, + Error: "dataset navigation requires a tenant and dataset context", + }), nil + } + + // Route RELEVANT docs by querying the nav tree with the topic (semantic KNN). + // The topic is the routing signal — we must not return arbitrary doc ids. + query := strings.TrimSpace(args.Topic + " " + args.Keywords) + seen := map[string]struct{}{} + var docs []string + collect := func(id string) { + if id == "" { + return + } + if _, ok := seen[id]; ok { + return + } + if len(docs) >= maxDocs { + return + } + seen[id] = struct{}{} + docs = append(docs, id) + } + + // Primary: semantic search over each dataset's nav tree. + for _, datasetID := range datasetIDs { + hits, err := ns.Search(ctx, tenantID, datasetID, query, nil, maxDocs) + if err != nil { + continue + } + for _, h := range hits { + collect(h.DocID) + for _, id := range h.DocIDs { + collect(id) + } + } + if len(docs) >= maxDocs { + break + } + } + + // Fallback: if semantic routing found nothing (e.g. no embedder), walk the + // root clusters so the tool still returns a useful (if coarse) doc set. + if len(docs) == 0 { + for _, datasetID := range datasetIDs { + clusters, _, err := ns.ListClusters(ctx, tenantID, datasetID, 0, 100) + if err != nil { + continue + } + for _, c := range clusters { + children, _, err := ns.ListChildren(ctx, tenantID, datasetID, c.Name, 0, 100) + if err != nil { + continue + } + for _, ch := range children { + collect(ch.DocID) + if len(docs) >= maxDocs { + break + } + } + if len(docs) >= maxDocs { + break + } + } + if len(docs) >= maxDocs { + break + } + } + } + + if len(docs) == 0 { + return datasetNavigationJSON(datasetNavigationResult{NotFound: true}), nil + } + return datasetNavigationJSON(datasetNavigationResult{Docs: docs}), nil +} + +func (d *DatasetNavigationByTree) mergeDefaults(args datasetNavigationArgs) datasetNavigationArgs { + if len(args.DatasetIDs) == 0 && len(d.defaults.DatasetIDs) != 0 { + args.DatasetIDs = append([]string(nil), d.defaults.DatasetIDs...) + } + if args.MaxDocs <= 0 { + args.MaxDocs = d.defaults.MaxDocs + } + return args +} + +func datasetNavigationJSON(r datasetNavigationResult) string { + b, err := json.Marshal(r) + if err != nil { + return fmt.Sprintf(`{"_ERROR":"dataset_navigation: marshal result: %s"}`, err) + } + return string(b) +} diff --git a/internal/agent/tool/dataset_navigation_test.go b/internal/agent/tool/dataset_navigation_test.go new file mode 100644 index 0000000000..3f7b8dbff5 --- /dev/null +++ b/internal/agent/tool/dataset_navigation_test.go @@ -0,0 +1,134 @@ +package tool + +import ( + "context" + "sync" + "testing" + + "ragflow/internal/agent/runtime" + "ragflow/internal/service/nav" +) + +// navRoutingFake is a nav.NavService that records Search calls (topic) and +// returns a controlled doc list, so a test can assert the router actually +// queries by topic rather than walking arbitrary clusters. +type navRoutingFake struct { + mu sync.Mutex + searched []string // topics passed to Search + hits []nav.NavHit + clusters []nav.NavNode + children map[string][]nav.NavNode +} + +func (f *navRoutingFake) UpsertDoc(context.Context, nav.UpsertDocInput) error { return nil } +func (f *navRoutingFake) RemoveDoc(context.Context, string, string, string) error { + return nil +} +func (f *navRoutingFake) Search(_ context.Context, _, _ string, query string, _ []float32, _ int) ([]nav.NavHit, error) { + f.mu.Lock() + f.searched = append(f.searched, query) + f.mu.Unlock() + return f.hits, nil +} +func (f *navRoutingFake) ListClusters(context.Context, string, string, int, int) ([]nav.NavNode, int64, error) { + return f.clusters, int64(len(f.clusters)), nil +} +func (f *navRoutingFake) ListChildren(_ context.Context, _, _, name string, _, _ int) ([]nav.NavNode, int64, error) { + return f.children[name], int64(len(f.children[name])), nil +} +func (f *navRoutingFake) searchedTopics() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.searched...) +} + +// TestDatasetNavigation_UsesTopicRouting asserts the router queries the nav tree +// with the topic (semantic search), rather than blindly walking clusters and +// returning arbitrary doc ids. +func TestDatasetNavigation_UsesTopicRouting(t *testing.T) { + fake := &navRoutingFake{ + hits: []nav.NavHit{ + {Type: "nav_doc", DocID: "d1", Name: "rocket"}, + {Type: "nav_doc", DocID: "d2", Name: "engine"}, + }, + } + prev := nav.GetNavService() + nav.SetNavService(fake) + defer func() { nav.SetNavService(prev) }() + + state := runtime.NewCanvasState("run-1", "task-1") + state.Sys["tenant_id"] = "tenant-1" + ctx := runtime.WithState(context.Background(), state) + + tool := NewDatasetNavigationByTree() + out, err := tool.InvokableRun(ctx, `{"topic":"rocket propulsion","keywords":"engine","dataset_ids":["kb1"]}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + // The topic (plus keywords) must have been used as the Search query. + topics := fake.searchedTopics() + if len(topics) == 0 { + t.Fatal("Search was never called; router must route by topic") + } + if topics[0] != "rocket propulsion engine" { + t.Errorf("search query = %q, want topic+keywords", topics[0]) + } + // The returned docs come from the relevant hits, not arbitrary walk. + if !containsStr(out, "d1") || !containsStr(out, "d2") { + t.Errorf("routed docs missing hits: %s", out) + } +} + +// TestCanvasDatasetIDs_MultiKB asserts all explicit dataset ids are preserved +// (a multi-KB session must not collapse to the first KB). +func TestCanvasDatasetIDs_MultiKB(t *testing.T) { + ids := canvasDatasetIDs(context.Background(), []string{"kb1", "kb2", "kb3"}) + if len(ids) != 3 || ids[0] != "kb1" || ids[1] != "kb2" || ids[2] != "kb3" { + t.Errorf("canvasDatasetIDs = %v, want all three KBs", ids) + } +} + +// TestDatasetNavigation_MultiKB asserts the router searches EVERY bound dataset +// (not just the first), so docs in other KBs stay reachable. +func TestDatasetNavigation_MultiKB(t *testing.T) { + fake := &navRoutingFake{hits: []nav.NavHit{{Type: "nav_doc", DocID: "d1", Name: "topic"}}} + prev := nav.GetNavService() + nav.SetNavService(fake) + defer func() { nav.SetNavService(prev) }() + + state := runtime.NewCanvasState("run-1", "task-1") + state.Sys["tenant_id"] = "tenant-1" + ctx := runtime.WithState(context.Background(), state) + + tool := NewDatasetNavigationByTree() + _, err := tool.InvokableRun(ctx, `{"topic":"X","dataset_ids":["kb1","kb2","kb3"]}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + // Search must have been called once per dataset (3 calls), not collapsed to + // the first KB. + if got := len(fake.searchedTopics()); got != 3 { + t.Errorf("Search called %d times, want 3 (once per dataset)", got) + } +} + +// TestCanvasDatasetIDs_DedupEmpty asserts empty ids are dropped. +func TestCanvasDatasetIDs_DedupEmpty(t *testing.T) { + ids := canvasDatasetIDs(context.Background(), []string{"kb1", "", "kb2"}) + if len(ids) != 2 || ids[0] != "kb1" || ids[1] != "kb2" { + t.Errorf("canvasDatasetIDs = %v, want [kb1 kb2]", ids) + } +} + +func containsStr(s, sub string) bool { + return len(s) > 0 && len(sub) > 0 && (s == sub || containsSub(s, sub)) +} + +func containsSub(s, sub string) bool { + for i := 0; i+len(sub) <= len(s); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +} diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index ca43ca7830..674e831a6e 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -31,32 +31,36 @@ import ( type Factory func(params map[string]any) (einotool.BaseTool, error) var registry = map[string]Factory{ - "akshare": buildAkShareTool, - "arxiv": buildArxivTool, - "bgpt": buildBGPTTool, - "code_exec": noConfig("code_exec", func() einotool.BaseTool { return NewCodeExecTool() }), - "crawler": noConfig("crawler", func() einotool.BaseTool { return NewCrawlerTool() }), - "deepl": noConfig("deepl", func() einotool.BaseTool { return NewDeepLTool() }), - "duckduckgo": buildDuckDuckGoTool, - "email": buildEmailTool, - "execute_sql": buildExeSQLTool, - "exesql": buildExeSQLTool, - "github": buildGitHubTool, - "google": buildGoogleTool, - "google_scholar": buildGoogleScholarTool, - "google_scholar_search": buildGoogleScholarTool, - "jin10": noConfig("jin10", func() einotool.BaseTool { return NewJin10Tool() }), - "keenable": buildKeenableTool, - "pubmed": buildPubMedTool, - "qweather": noConfig("qweather", func() einotool.BaseTool { return NewQWeatherTool() }), - "querit": buildQueritTool, - "querit_search": buildQueritTool, - "queritsearch": buildQueritTool, - "retrieval": buildRetrievalTool, - "search_my_dataset": buildRetrievalTool, - "search_my_dateset": buildRetrievalTool, - "searxng": buildSearXNGTool, - "tavily": buildTavilyTool, + "akshare": buildAkShareTool, + "arxiv": buildArxivTool, + "bgpt": buildBGPTTool, + "code_exec": noConfig("code_exec", func() einotool.BaseTool { return NewCodeExecTool() }), + "crawler": noConfig("crawler", func() einotool.BaseTool { return NewCrawlerTool() }), + "dataset_navigation_by_tree": noConfig("dataset_navigation_by_tree", func() einotool.BaseTool { return NewDatasetNavigationByTree() }), + "hybrid_search": noConfig("hybrid_search", func() einotool.BaseTool { return NewAgenticSearchTool(toolHybridSearch) }), + "vector_search": noConfig("vector_search", func() einotool.BaseTool { return NewAgenticSearchTool(toolVectorSearch) }), + "bm25_search": noConfig("bm25_search", func() einotool.BaseTool { return NewAgenticSearchTool(toolBM25Search) }), + "deepl": noConfig("deepl", func() einotool.BaseTool { return NewDeepLTool() }), + "duckduckgo": buildDuckDuckGoTool, + "email": buildEmailTool, + "execute_sql": buildExeSQLTool, + "exesql": buildExeSQLTool, + "github": buildGitHubTool, + "google": buildGoogleTool, + "google_scholar": buildGoogleScholarTool, + "google_scholar_search": buildGoogleScholarTool, + "jin10": noConfig("jin10", func() einotool.BaseTool { return NewJin10Tool() }), + "keenable": buildKeenableTool, + "pubmed": buildPubMedTool, + "qweather": noConfig("qweather", func() einotool.BaseTool { return NewQWeatherTool() }), + "querit": buildQueritTool, + "querit_search": buildQueritTool, + "queritsearch": buildQueritTool, + "retrieval": buildRetrievalTool, + "search_my_dataset": buildRetrievalTool, + "search_my_dateset": buildRetrievalTool, + "searxng": buildSearXNGTool, + "tavily": buildTavilyTool, // Agent DSL tool lists carry the Python Canvas component_name verbatim. // BuildByName lower-cases names, so register those component names too. "tavilysearch": buildTavilyTool, diff --git a/internal/agent/tool/retrieval_nlp.go b/internal/agent/tool/retrieval_nlp.go index 88f9e14d72..2a15193b3d 100644 --- a/internal/agent/tool/retrieval_nlp.go +++ b/internal/agent/tool/retrieval_nlp.go @@ -162,6 +162,7 @@ func nlpRequestFromRetrieval(req RetrievalRequest, tenantIDs []string, topN int) 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), diff --git a/internal/agent/tool/retrieval_service.go b/internal/agent/tool/retrieval_service.go index ecc0615987..9b44733acc 100644 --- a/internal/agent/tool/retrieval_service.go +++ b/internal/agent/tool/retrieval_service.go @@ -55,6 +55,9 @@ type RetrievalRequest struct { KeywordsSimilarityWeight *float64 UseKG bool SimilarityThreshold float64 + // 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 diff --git a/internal/dao/compilation_template_seed.go b/internal/dao/compilation_template_seed.go index 653a1f8e31..38c8d690d6 100644 --- a/internal/dao/compilation_template_seed.go +++ b/internal/dao/compilation_template_seed.go @@ -43,7 +43,6 @@ var builtinCompilationTemplateKinds = []struct { {Kind: "mind_map", Name: "Mind map - Radial concept hierarchy"}, {Kind: "wiki", Name: "Wiki — Graph-based wiki"}, {Kind: "knowledge_graph", Name: "Knowledge graph"}, - {Kind: "datasetnav", Name: "Dataset nav — top-down drill-down"}, {Kind: "page_index", Name: "Page index"}, {Kind: "session_essence", Name: "Session essence"}, {Kind: "session_graph", Name: "Session graph"}, diff --git a/internal/dao/compilation_template_seed_test.go b/internal/dao/compilation_template_seed_test.go index d12fb76cd1..dca74cd70c 100644 --- a/internal/dao/compilation_template_seed_test.go +++ b/internal/dao/compilation_template_seed_test.go @@ -25,6 +25,18 @@ import ( "gorm.io/gorm" ) +// TestBuiltinCompilationTemplateKinds_NoDatasetnav locks the removal of the +// datasetnav template from the built-in catalogue: dataset navigation is not an +// independent compile kind (see component_test.TestKnowledgeCompiler_Datasetnav_NoVariant), +// so it must not appear in the seeded template kinds. +func TestBuiltinCompilationTemplateKinds_NoDatasetnav(t *testing.T) { + for _, k := range builtinCompilationTemplateKinds { + if k.Kind == "datasetnav" || k.Kind == "dataset_nav" { + t.Fatalf("builtin compilation template kind %q must not exist (datasetnav is a by-product, not a compile kind)", k.Kind) + } + } +} + func TestSeedBuiltinCompilationTemplatesForTenant(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) if err != nil { diff --git a/internal/ingestion/component/knowledge_compiler/common/types.go b/internal/ingestion/component/knowledge_compiler/common/types.go index 97375efcb0..7c0635566e 100644 --- a/internal/ingestion/component/knowledge_compiler/common/types.go +++ b/internal/ingestion/component/knowledge_compiler/common/types.go @@ -18,11 +18,10 @@ import ( type Variant string const ( - VariantStructure Variant = "structure" - VariantWiki Variant = "wiki" - VariantTree Variant = "tree" - VariantMindmap Variant = "mindmap" - VariantDatasetnav Variant = "datasetnav" + VariantStructure Variant = "structure" + VariantWiki Variant = "wiki" + VariantTree Variant = "tree" + VariantMindmap Variant = "mindmap" ) // Sentinel errors. @@ -189,10 +188,13 @@ func ParseParam(m map[string]any) (Param, error) { // knowledge_graph -> // // The canonical variant names are also accepted as identity (so a template kind -// may equal the variant, e.g. "datasetnav" which has no Python kind, and the -// internal unit tests can drive each variant through a template id). Unknown -// kinds return ErrUnknownVariant; the caller turns that into a hard failure -// rather than silently emitting uncompiled rows. +// may equal the variant). Unknown kinds return ErrUnknownVariant; the caller +// turns that into a hard failure rather than silently emitting uncompiled rows. +// +// Note: "datasetnav"/"dataset_nav" intentionally has NO mapping here — dataset +// navigation is not an independent compile kind in Python; it is a by-product +// written after tree/page_index compilation via internal/service datasetnav +// (see tasks/agentic_search_port_plan.md). func KindToVariant(kind string) (Variant, error) { switch strings.ToLower(strings.TrimSpace(kind)) { case "tree": @@ -201,8 +203,6 @@ func KindToVariant(kind string) (Variant, error) { return VariantMindmap, nil case "wiki": return VariantWiki, nil - case "datasetnav", "dataset_nav": - return VariantDatasetnav, nil case "page_index", "session_essence", "session_graph", "timeline", "knowledge_graph", "structure", "knowledgegraph", "graph": return VariantStructure, nil diff --git a/internal/ingestion/component/knowledge_compiler/component.go b/internal/ingestion/component/knowledge_compiler/component.go index 17724b98ba..bb870f4b93 100644 --- a/internal/ingestion/component/knowledge_compiler/component.go +++ b/internal/ingestion/component/knowledge_compiler/component.go @@ -1,22 +1,24 @@ // Package knowledge_compiler implements the KnowledgeCompiler ingestion // component: a single runtime.Component that dispatches to one of the -// knowledge-compile variants (structure / wiki / tree / mindmap / datasetnav) -// based on the `variant` param. See PORT_PLAN.md for the full design. +// knowledge-compile variants (structure / wiki / tree / mindmap) based on the +// `variant` param. See PORT_PLAN.md for the full design. package knowledge_compiler import ( "context" + "encoding/json" "fmt" + "log" "strings" "ragflow/internal/agent/runtime" "ragflow/internal/ingestion/component/knowledge_compiler/common" - "ragflow/internal/ingestion/component/knowledge_compiler/datasetnav" "ragflow/internal/ingestion/component/knowledge_compiler/mindmap" "ragflow/internal/ingestion/component/knowledge_compiler/structure" "ragflow/internal/ingestion/component/knowledge_compiler/tree" "ragflow/internal/ingestion/component/knowledge_compiler/wiki" "ragflow/internal/ingestion/component/schema" + "ragflow/internal/service/nav" "ragflow/internal/tokenizer" "gorm.io/gorm" @@ -110,6 +112,9 @@ func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, in // template's id and kind. All products are buffered (no streaming sink), so // the post-run loop below covers every row (M1). var out common.Outputs + // navByProducts accumulates every tree/structure by-product job so each is + // written after the spec loop (not just the last one). + var navByProducts []navByProduct for _, spec := range specs { variant, err := common.KindToVariant(spec.Kind) if err != nil { @@ -150,8 +155,6 @@ func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, in o, err = tree.Run(ctx, deps, specParam, specIn) case common.VariantMindmap: o, err = mindmap.Run(ctx, deps, specParam, specIn) - case common.VariantDatasetnav: - o, err = datasetnav.Run(ctx, deps, specParam, specIn) default: return nil, fmt.Errorf("%w: %q", common.ErrUnknownVariant, variant) } @@ -159,6 +162,13 @@ func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, in return nil, err } + // Accumulate tree/structure by-product jobs so each is written after the + // spec loop (a multi-spec batch must not drop any spec's products). nav is + // a derived artifact; its failures are logged and never abort the pipeline. + if variant == common.VariantTree || variant == common.VariantStructure { + navByProducts = append(navByProducts, navByProduct{deps: deps, variant: variant, products: o.Products}) + } + for i := range o.Products { if o.Products[i].Meta == nil { o.Products[i].Meta = map[string]any{} @@ -171,6 +181,17 @@ func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, in out.Products = append(out.Products, o.Products...) } + // Write each dataset-nav by-product after all specs ran. tree by-product + // (Python compiler.py:475) and structure/page_index by-product (Python + // runner.py:189). Failures are logged and never abort the pipeline. + for _, job := range navByProducts { + if job.variant == common.VariantTree { + upsertTreeNav(ctx, job.deps, in.DocID, job.products) + } else { + upsertStructureNav(ctx, job.deps, in.DocID, job.products) + } + } + // Convert the compiled products into chunk-aligned docs (matching // conf/infinity_mapping.json) and merge them into the upstream input // chunks. The component stays DB-independent and no longer routes through a @@ -182,6 +203,142 @@ func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, in return mergeChunks(inputs, compiled), nil } +// navByProduct is one deferred dataset-nav by-product job: the compile variant +// that produced it and the products to summarize from. +type navByProduct struct { + deps common.Deps + variant common.Variant + products []common.Product +} + +// upsertTreeNav writes the dataset-nav by-product after tree compilation. It +// finds the tree root product (Meta kind=root) whose Content is the document +// summary, and places it into the ES-backed nav tree. Any failure (missing nav +// service, missing root, embedding/upsert error) is logged and skipped — the +// nav artifact must never block the compile pipeline. +func upsertTreeNav(ctx context.Context, deps common.Deps, docID string, products []common.Product) { + ns := nav.GetNavService() + if ns == nil { + log.Printf("knowledge_compiler: datasetnav by-product skipped (NavService not initialized)") + return + } + var summary string + var vec []float32 + for i := range products { + if kind, _ := products[i].Meta["kind"].(string); kind == "root" { + summary = products[i].Content + vec = products[i].Vector + break + } + } + if strings.TrimSpace(summary) == "" { + log.Printf("knowledge_compiler: datasetnav by-product skipped (tree has no root summary)") + return + } + if len(vec) == 0 { + if deps.Embed == nil { + log.Printf("knowledge_compiler: datasetnav by-product skipped (no embedder, no precomputed vector)") + return + } + embeddings, err := deps.Embed.Encode(ctx, []string{summary}) + if err != nil { + log.Printf("knowledge_compiler: datasetnav by-product skipped (embedding failed): %v", err) + return + } + if len(embeddings) == 0 { + log.Printf("knowledge_compiler: datasetnav by-product skipped (embedding produced no vector)") + return + } + vec = embeddings[0] + } + if err := ns.UpsertDoc(ctx, nav.UpsertDocInput{ + TenantID: deps.TenantID, + KbID: deps.DatasetID, + DocID: docID, + Summary: summary, + Embedd: vec, + }); err != nil { + // Log and continue: nav is a derived artifact, never abort compile. + log.Printf("knowledge_compiler: datasetnav by-product upsert failed (continuing): %v", err) + } +} + +// upsertStructureNav writes the dataset-nav by-product after structure/page_index +// compilation (mirroring Python runner.py:189 _upsert_dataset_nav_from_page_index). +// It finds the graph product (Meta kind=graph), summarizes its entity +// descriptions into a document-level summary, and places it into the nav tree. +// Failures are logged and never abort the compile pipeline. +func upsertStructureNav(ctx context.Context, deps common.Deps, docID string, products []common.Product) { + ns := nav.GetNavService() + if ns == nil { + log.Printf("knowledge_compiler: datasetnav by-product skipped (NavService not initialized)") + return + } + var graph *common.Product + for i := range products { + if kind, _ := products[i].Meta["kind"].(string); kind == "graph" { + graph = &products[i] + break + } + } + if graph == nil { + return + } + summary := pageIndexSummary(graph.Content) + if strings.TrimSpace(summary) == "" { + log.Printf("knowledge_compiler: datasetnav by-product skipped (structure graph has no entity descriptions)") + return + } + // Embed the SUMMARY text, not the graph JSON — the nav doc's vector must + // represent the summary semantics for the KNN router to match correctly. + if deps.Embed == nil { + log.Printf("knowledge_compiler: datasetnav by-product skipped (no embedder)") + return + } + embeddings, err := deps.Embed.Encode(ctx, []string{summary}) + if err != nil || len(embeddings) == 0 { + log.Printf("knowledge_compiler: datasetnav by-product skipped (structure embedding failed): %v", err) + return + } + vec := embeddings[0] + if err := ns.UpsertDoc(ctx, nav.UpsertDocInput{ + TenantID: deps.TenantID, + KbID: deps.DatasetID, + DocID: docID, + Summary: summary, + Embedd: vec, + }); err != nil { + log.Printf("knowledge_compiler: datasetnav by-product upsert failed (continuing): %v", err) + } +} + +// pageIndexSummary concatenates entity descriptions from a structure graph JSON +// ({"entities": [{"name","description"}, ...]}), producing a document-level +// summary for dataset navigation. +func pageIndexSummary(graphJSON string) string { + var graph struct { + Entities []struct { + Name string `json:"name"` + Description string `json:"description"` + } `json:"entities"` + } + if err := json.Unmarshal([]byte(graphJSON), &graph); err != nil { + return "" + } + var b strings.Builder + for _, e := range graph.Entities { + desc := strings.Join(strings.Fields(e.Description), " ") + if desc == "" { + continue + } + if e.Name != "" { + b.WriteString(e.Name + ": ") + } + b.WriteString(desc + "\n") + } + return b.String() +} + // resolveTemplateSpecs resolves the configured compilation template spec(s) to // their TemplateInfo rows. Priority: compilation_template_id > // compilation_template_group_id. The group path resolves the group to its @@ -260,11 +417,10 @@ func kindOrVariant(p common.Product) string { // key that distinguishes compiled knowledge units from ordinary chunks and // routes retrieval-side filters (e.g. "compile_kwd": ["artifact_page"]). var variantCompileKWD = map[common.Variant]string{ - common.VariantStructure: "structure", - common.VariantWiki: "artifact_page", - common.VariantTree: "tree", - common.VariantMindmap: "mindmap", - common.VariantDatasetnav: "dataset_nav", + common.VariantStructure: "structure", + common.VariantWiki: "artifact_page", + common.VariantTree: "tree", + common.VariantMindmap: "mindmap", } // productsToChunkDocs converts the internal compiled Product rows into @@ -520,35 +676,6 @@ func applyVariantColumns(doc *schema.ChunkDoc, p common.Product) error { return err } } - - case common.VariantDatasetnav: - // nav_cluster / nav_doc rows: type_kwd discriminates the row kind. - if v := metaString(p.Meta, "type"); v != "" { - if err := doc.SetExtraValue("type_kwd", v); err != nil { - return err - } - } - if v := metaString(p.Meta, "name"); v != "" { - if err := doc.SetExtraValue("title_kwd", v); err != nil { - return err - } - setTitleTokens(doc, v) - } - if v, ok := metaInt(p.Meta, "depth"); ok { - if err := doc.SetExtraValue("depth_int", v); err != nil { - return err - } - } - if v, ok := metaInt(p.Meta, "size"); ok { - if err := doc.SetExtraValue("doc_count_int", v); err != nil { - return err - } - } - if v := metaStringSlice(p.Meta, "doc_ids"); len(v) > 0 { - if err := doc.SetExtraValue("doc_ids_kwd", v); err != nil { - return err - } - } } return nil diff --git a/internal/ingestion/component/knowledge_compiler/component_test.go b/internal/ingestion/component/knowledge_compiler/component_test.go index 93ff663315..074d3a1e47 100644 --- a/internal/ingestion/component/knowledge_compiler/component_test.go +++ b/internal/ingestion/component/knowledge_compiler/component_test.go @@ -13,6 +13,7 @@ import ( "testing" "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/service/nav" ) // mockChat answers the structure variant's three LLM call shapes under the @@ -110,7 +111,7 @@ func (mockChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatRes } // proseChat returns generic, non-empty prose for the non-structure variants -// (wiki/tree/mindmap/datasetnav), which all just need a summary/outline text. +// (wiki/tree/mindmap), which all just need a summary/outline text. type proseChat struct{} func (proseChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { @@ -239,6 +240,138 @@ func TestKnowledgeCompiler_UnknownVariant(t *testing.T) { } } +// TestKnowledgeCompiler_Datasetnav_NoVariant locks the removal of the standalone +// datasetnav variant: dataset navigation is NOT an independent compile kind in +// Python (it is a by-product written after tree/page_index compile via +// internal/service datasetnav), so "datasetnav"/"dataset_nav" must now fail as +// an unknown variant rather than silently compile nothing. +func TestKnowledgeCompiler_Datasetnav_NoVariant(t *testing.T) { + for _, kind := range []string{"datasetnav", "dataset_nav"} { + _, err := common.KindToVariant(kind) + if !errors.Is(err, common.ErrUnknownVariant) { + t.Fatalf("KindToVariant(%q) err = %v, want ErrUnknownVariant", kind, err) + } + } +} + +// fakeNavSvc records the most recent UpsertDoc call so a test can assert the +// tree by-product hook feeds the nav service the root document summary. +type fakeNavSvc struct { + mu sync.Mutex + called bool + summary string + docID string + kbID string +} + +func (f *fakeNavSvc) UpsertDoc(_ context.Context, in nav.UpsertDocInput) error { + f.mu.Lock() + defer f.mu.Unlock() + f.called = true + f.summary = in.Summary + f.docID = in.DocID + f.kbID = in.KbID + return nil +} +func (f *fakeNavSvc) RemoveDoc(context.Context, string, string, string) error { return nil } +func (f *fakeNavSvc) Search(context.Context, string, string, string, []float32, int) ([]nav.NavHit, error) { + return nil, nil +} +func (f *fakeNavSvc) ListClusters(context.Context, string, string, int, int) ([]nav.NavNode, int64, error) { + return nil, 0, nil +} +func (f *fakeNavSvc) ListChildren(context.Context, string, string, string, int, int) ([]nav.NavNode, int64, error) { + return nil, 0, nil +} + +// TestKnowledgeCompiler_TreeNavByProduct asserts the tree compile-complete hook +// writes the dataset-nav by-product: after tree.Run produces a root product, +// upsertTreeNav calls NavService.UpsertDoc with the root summary and the +// doc/dataset context, and never aborts on failure. +func TestKnowledgeCompiler_TreeNavByProduct(t *testing.T) { + fake := &fakeNavSvc{} + nav.SetNavService(fake) + t.Cleanup(func() { nav.SetNavService(nil) }) + + deps := common.Deps{ + TenantID: "t1", + DatasetID: "kb1", + Chat: proseChat{}, + Embed: mockEmbedder{dim: 8}, + } + products := []common.Product{ + {Content: "section one body", Meta: map[string]any{"kind": "summary", "level": 0}}, + {Content: "overall doc theme root summary", Vector: []float32{0.1, 0.2}, Meta: map[string]any{"kind": "root", "level": -1}}, + } + upsertTreeNav(context.Background(), deps, "d1", products) + + fake.mu.Lock() + defer fake.mu.Unlock() + if !fake.called { + t.Fatal("upsertTreeNav did not call NavService.UpsertDoc") + } + if fake.docID != "d1" { + t.Errorf("doc_id = %q, want d1", fake.docID) + } + if fake.kbID != "kb1" { + t.Errorf("kb_id = %q, want kb1", fake.kbID) + } + if !strings.Contains(fake.summary, "overall doc theme root summary") { + t.Errorf("summary = %q, want the tree root summary", fake.summary) + } +} + +// TestKnowledgeCompiler_TreeNavByProduct_NilServiceDoesNotAbort asserts the +// by-product hook tolerates a missing NavService without erroring (nav is a +// derived artifact; its absence must never abort compilation). +func TestKnowledgeCompiler_TreeNavByProduct_NilServiceDoesNotAbort(t *testing.T) { + nav.SetNavService(nil) + t.Cleanup(func() { nav.SetNavService(nil) }) + // Should not panic and must return normally. + upsertTreeNav(context.Background(), common.Deps{}, "d1", + []common.Product{{Content: "x", Meta: map[string]any{"kind": "root"}}}) +} + +// TestKnowledgeCompiler_StructureNavByProduct asserts the structure/page_index +// by-product hook summarizes the graph product entities and feeds NavService. +func TestKnowledgeCompiler_StructureNavByProduct(t *testing.T) { + fake := &fakeNavSvc{} + nav.SetNavService(fake) + t.Cleanup(func() { nav.SetNavService(nil) }) + + graphJSON := `{"entities":[{"name":"Engine","description":"a propulsion device"},{"name":"Fuel","description":"combustion source"}]}` + products := []common.Product{ + {Content: graphJSON, Vector: []float32{0.1, 0.2}, Meta: map[string]any{"kind": "graph", "compile_kwd": "page_index"}}, + } + // The by-product embeds the SUMMARY (not the graph JSON vector), so an + // embedder must be present. + upsertStructureNav(context.Background(), common.Deps{TenantID: "t1", DatasetID: "kb1", Embed: mockEmbedder{dim: 8}}, "d1", products) + + fake.mu.Lock() + defer fake.mu.Unlock() + if !fake.called { + t.Fatal("upsertStructureNav did not call NavService.UpsertDoc") + } + if fake.docID != "d1" { + t.Errorf("doc_id = %q, want d1", fake.docID) + } + if !strings.Contains(fake.summary, "Engine: a propulsion device") { + t.Errorf("summary = %q, want entity descriptions", fake.summary) + } +} + +// TestPageIndexSummary asserts entity descriptions are concatenated into a +// document summary. +func TestPageIndexSummary(t *testing.T) { + got := pageIndexSummary(`{"entities":[{"name":"A","description":"one thing"},{"name":"B","description":""}]}`) + if !strings.Contains(got, "A: one thing") { + t.Errorf("summary = %q, want entity A", got) + } + if strings.Contains(got, "B") { + t.Errorf("summary should skip empty-description entity, got %q", got) + } +} + func TestKnowledgeCompiler_Alias_Mindmap(t *testing.T) { installMockDeps(t) // "mind_map" is the deprecated alias for "mindmap"; both resolve to the @@ -388,20 +521,6 @@ func TestKnowledgeCompiler_Mindmap_EndToEnd(t *testing.T) { } } -func TestKnowledgeCompiler_Datasetnav_EndToEnd(t *testing.T) { - installProseDeps(t) - chunks := runVariant(t, "datasetnav", nil) - foundRoot := false - for _, c := range chunks { - if kind, _ := c["kc_kind"].(string); kind == "root" { - foundRoot = true - } - } - if !foundRoot { - t.Fatalf("datasetnav: no 'root' chunk; got %d chunks", len(chunks)) - } -} - // TestKnowledgeCompiler_EmitsChunks verifies that after compiling, the // component returns the knowledge units merged into the chunk stream (no // separate products/writer seam). The output shape is the chunker's @@ -786,140 +905,6 @@ func TestKnowledgeCompiler_Wiki_HistoricalDedupScopedByDataset(t *testing.T) { } } -// fakeDatasetnavLock records the key it was asked to acquire, so a test can -// assert the datasetnav rebuild lock is scoped to the dataset, not the document. -type fakeDatasetnavLock struct { - mu sync.Mutex - lastKey string -} - -func (f *fakeDatasetnavLock) Acquire(_ context.Context, key string) (bool, error) { - f.mu.Lock() - defer f.mu.Unlock() - f.lastKey = key - return true, nil -} -func (f *fakeDatasetnavLock) Release(_ context.Context, _ string) error { return nil } - -// TestKnowledgeCompiler_Datasetnav_LockKeyedByDataset is a regression for the -// Medium bug: the distributed rebuild lock used to be keyed by doc_id, so two -// runs against the same dataset (with different per-document doc_id values) -// took different locks and could interleave. This test supplies both doc_id -// ("d1") and dataset_id ("ds1") and asserts the lock key uses the dataset. -func TestKnowledgeCompiler_Datasetnav_LockKeyedByDataset(t *testing.T) { - lk := &fakeDatasetnavLock{} - common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { - return common.Deps{ - Chat: proseChat{}, - Embed: mockEmbedder{dim: 8}, - TenantID: tenantID, - Redis: lk, - }, nil - }) - t.Cleanup(func() { common.SetDepsResolver(nil) }) - - installVariantTemplateResolver(t, "datasetnav") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ - "compilation_template_id": "tpl-datasetnav", "llm_id": "llm1", "embedding_model": "emb1", - }) - if err != nil { - t.Fatalf("NewKnowledgeCompilerComponent: %v", err) - } - if _, err := c.Invoke(context.Background(), nil, map[string]any{ - "chunks": []any{ - map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog near the river bank."}, - map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river at dawn."}, - }, - "doc_id": "d1", - "dataset_id": "ds1", - "tenant_id": "t1", - }); err != nil { - t.Fatalf("Invoke: %v", err) - } - lk.mu.Lock() - key := lk.lastKey - lk.mu.Unlock() - want := "datasetnav:t1:ds1" - if key != want { - t.Fatalf("datasetnav lock key = %q, want %q (lock must be scoped to the dataset, not the document)", key, want) - } -} - -// recordingNavChat echoes each nav-group summary back verbatim (so the child -// summaries are identifiable) and records the root-synthesis user prompt so a -// test can assert which child summaries reached the root overview. -type recordingNavChat struct { - mu sync.Mutex - rootPrompt string -} - -func (r *recordingNavChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { - if strings.HasPrefix(req.UserPrompt, "Compose a navigation overview") { - r.mu.Lock() - r.rootPrompt = req.UserPrompt - r.mu.Unlock() - return &common.ChatResponse{Content: "root overview"}, nil - } - // Nav-group summary: echo the group text so the marker survives into the - // root prompt when the root is built from all child summaries. - return &common.ChatResponse{Content: strings.TrimSpace(req.UserPrompt)}, nil -} - -// TestKnowledgeCompiler_Datasetnav_RootIncludesAllSummaries is a regression for -// root-overview completeness: the root synthesis must be built from ALL child -// summaries, not from a partial buffer. The Go implementation buffers every -// product in Outputs.Products (no streaming sink), so the root always sees the -// full child set even when the result set is large. Here nav_radius>1 forces -// each chunk into its own nav node, and we assert every child marker appears in -// the recorded root-synthesis prompt. -func TestKnowledgeCompiler_Datasetnav_RootIncludesAllSummaries(t *testing.T) { - chat := &recordingNavChat{} - common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { - return common.Deps{ - Chat: chat, - Embed: mockEmbedder{dim: 8}, - TenantID: tenantID, - }, nil - }) - t.Cleanup(func() { common.SetDepsResolver(nil) }) - - installVariantTemplateResolver(t, "datasetnav") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ - "compilation_template_id": "tpl-datasetnav", "llm_id": "llm1", "embedding_model": "emb1", - // nav_radius > 1 guarantees no two chunks group together (cosine <= 1), - // so each chunk becomes its own nav node. - "extra": map[string]any{"nav_radius": 1.1}, - }) - if err != nil { - t.Fatalf("NewKnowledgeCompilerComponent: %v", err) - } - - markers := []string{"NAVMARKA", "NAVMARKB", "NAVMARKC", "NAVMARKD", "NAVMARKE"} - chunks := make([]any, len(markers)) - for i, m := range markers { - chunks[i] = map[string]any{"id": m, "text": m + " unique section body text"} - } - if _, err := c.Invoke(context.Background(), nil, map[string]any{ - "chunks": chunks, - "doc_id": "d1", - "tenant_id": "t1", - }); err != nil { - t.Fatalf("Invoke: %v", err) - } - - chat.mu.Lock() - root := chat.rootPrompt - chat.mu.Unlock() - if root == "" { - t.Fatalf("root synthesis prompt was never recorded (root node not built)") - } - for _, m := range markers { - if !strings.Contains(root, m) { - t.Fatalf("root overview is missing child summary %q; the root must be built from ALL child summaries.\nroot prompt:\n%s", m, root) - } - } -} - // fencedChat wraps otherwise-valid extraction JSON in a ```json ... ``` fence, // the most common way LLMs wrap JSON even when JSONMode is requested. type fencedChat struct{} diff --git a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go deleted file mode 100644 index f604e7be4b..0000000000 --- a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go +++ /dev/null @@ -1,776 +0,0 @@ -// Package datasetnav implements the "datasetnav" variant of KnowledgeCompiler, -// mirroring Python's dataset_nav.py incremental clustering: each input chunk -// (the in-run equivalent of a document leaf) is embedded and placed into the -// nearest nav_cluster via layered KNN descent — merged when the similarity -// clears _MERGE_THRESHOLD, given a fresh sibling cluster above _MIN_SIM, or -// given a new root-level cluster otherwise. Clusters that exceed the fanout -// or doc-count caps split via the same 2-means procedure as Python. -// -// The Go port keeps the whole tree in memory for one Invoke (no ES reads or -// writes): cross-run incremental upsert/removal, which in Python is an -// ES-backed read-modify-write, is the caller's concern. The per-run algorithm -// (thresholds, LLM merge/summary prompts, readable cluster names, split) is -// faithful to the Python original. -package datasetnav - -import ( - "context" - "encoding/json" - "fmt" - "math" - "strings" - - "ragflow/internal/ingestion/component/knowledge_compiler/common" -) - -// Thresholds and caps, mirrored from dataset_nav.py. -const ( - mergeThresholdDefault = 0.80 // merge doc into cluster - recurseThreshold = 0.65 // continue descending into children - minSim = 0.50 // minimum similarity to be considered related - maxFanout = 64 // max child count before triggering rebalance - maxDocsPerCluster = 50 // max docs per leaf cluster before triggering split -) - -// datasetnavLock is the minimal locking surface used to serialize dataset-level -// rebuilds. The production Redis client implements this via SET NX + TTL; tests -// inject an in-memory lock (or nil for no-op). -type datasetnavLock interface { - Acquire(ctx context.Context, key string) (bool, error) - Release(ctx context.Context, key string) error -} - -// navCluster is one internal tree node (Python's nav_cluster row). -type navCluster struct { - Name string - Desc string - Parent string // parent cluster name; "root" for depth-0 clusters - Depth int - DocIDs []string - Vector []float32 -} - -// navDoc is one document leaf (Python's nav_doc row). -type navDoc struct { - ChunkID string - Text string - Parent string // owning cluster name - Depth int - Vector []float32 -} - -// navChild is one child reference under a cluster. -type navChild struct { - name string // cluster name or doc chunk id - isDoc bool - vec []float32 -} - -// navTree holds the in-memory tree built over one Invoke. -type navTree struct { - clusters map[string]*navCluster - order []string // cluster names in creation order - docs []*navDoc - children map[string][]navChild // parent cluster name → children -} - -func newNavTree() *navTree { - return &navTree{clusters: map[string]*navCluster{}, children: map[string][]navChild{}} -} - -func (t *navTree) addCluster(c *navCluster) { - t.clusters[c.Name] = c - t.order = append(t.order, c.Name) -} - -// Run executes the datasetnav variant. -func Run(ctx context.Context, deps common.Deps, param common.Param, inputs common.Inputs) (common.Outputs, error) { - if deps.Embed == nil { - return common.Outputs{}, fmt.Errorf("datasetnav: embedder required") - } - docID := firstNonEmpty(inputs.DocID, deps.DatasetID) - if docID == "" { - docID = "unknown" - } - llmID := firstNonEmpty(param.LLMID, inputs.LLMID) - tenantID := deps.TenantID - - // The rebuild lock is scoped to the dataset, not the document (mirrors - // Python's _nav_lock_key(kb_id)): concurrent runs against the same dataset - // must share one lock. Fall back to docID only when no dataset id is known. - datasetID := firstNonEmpty(deps.DatasetID, docID) - - texts, chunkIDs := chunkTexts(inputs.Chunks) - if len(texts) == 0 { - return common.Outputs{}, nil - } - - if l, ok := deps.Redis.(datasetnavLock); ok && l != nil { - lockKey := "datasetnav:" + tenantID + ":" + datasetID - acquired, err := l.Acquire(ctx, lockKey) - if err != nil { - return common.Outputs{}, err - } - if !acquired { - return common.Outputs{}, fmt.Errorf("datasetnav: lock not acquired for %s", lockKey) - } - defer l.Release(ctx, lockKey) - } - - // Embed every chunk text once (Python embeds each doc summary on upsert). - vectors, err := deps.Embed.Encode(ctx, texts) - if err != nil { - return common.Outputs{}, err - } - if len(vectors) != len(texts) { - return common.Outputs{}, fmt.Errorf("datasetnav: embedding count mismatch (%d vs %d)", len(vectors), len(texts)) - } - - // Sequential placement, mirroring the lock-serialized per-document upserts - // in Python (also what makes the in-run tree deterministic). - tree := newNavTree() - var clusterDescs []string - mergeThreshold := mergeThresholdFor(param) - for i, text := range texts { - if err := ctx.Err(); err != nil { - return common.Outputs{}, err - } - desc, err := tree.place(ctx, deps, llmID, chunkIDs[i], text, vectors[i], mergeThreshold) - if err != nil { - return common.Outputs{}, err - } - clusterDescs = append(clusterDescs, desc...) - } - - // Emit cluster rows in creation order, then nav_doc leaves in placement - // order, then the root overview node (a Go-side convenience: Python's tree - // has no root row, but downstream consumers expect one overview product). - var products []common.Product - clusterProductID := map[string]string{} - // Pre-compute every cluster id (deterministic from StableRowID) before - // resolving parent edges, so a child emitted before its (later-created) - // parent still gets the correct ParentID instead of "" (reparenting bug). - // Cluster ids are tenant-scoped (as the nav_doc ids already are) so two - // tenants whose datasetID falls back to the same docID cannot collide. - for _, name := range tree.order { - clusterProductID[name] = common.StableRowID("dataset_nav", tenantID, datasetID, "cluster", name) - } - for _, name := range tree.order { - c := tree.clusters[name] - pid := "" - if c.Parent != "" && c.Parent != "root" { - pid = clusterProductID[c.Parent] - } - id := clusterProductID[c.Name] - products = append(products, common.Product{ - ID: id, - DocID: docID, - TenantID: tenantID, - Variant: common.VariantDatasetnav, - Content: payloadJSON(map[string]any{"type": "nav_cluster", "description": c.Desc}), - Vector: c.Vector, - ParentID: pid, - Meta: map[string]any{ - "kind": "nav_cluster", - "type": "nav_cluster", - "name": c.Name, - "parent_name": c.Parent, - "depth": c.Depth, - "doc_ids": append([]string{}, c.DocIDs...), - "size": len(c.DocIDs), - }, - }) - } - for _, d := range tree.docs { - products = append(products, common.Product{ - // Scope nav_doc ids by tenant and dataset (not just chunk id) so - // synthetic/positional chunk ids from different docs/datasets do not - // collide and overwrite each other across tenants (M5). - ID: common.StableRowID("dataset_nav", tenantID, datasetID, "doc", d.ChunkID), - DocID: docID, - TenantID: tenantID, - Variant: common.VariantDatasetnav, - Content: payloadJSON(map[string]any{"type": "nav_doc", "description": d.Text}), - Vector: d.Vector, - ParentID: clusterProductID[d.Parent], - Meta: map[string]any{ - "kind": "nav_doc", - "type": "nav_doc", - "name": navDocName(d.Parent, d.Text), - "parent_name": d.Parent, - "depth": d.Depth, - "doc_ids": []string{d.ChunkID}, - }, - }) - } - - // Root overview built from EVERY cluster description (all products are - // buffered in the same slice, so every node contributes to the root). - rootSummary, err := summarize(ctx, deps, llmID, "Compose a navigation overview from these section summaries:\n\n"+formatNavSummaries(clusterDescs)) - if err == nil && rootSummary != "" { - emb, e2 := deps.Embed.Encode(ctx, []string{rootSummary}) - if e2 == nil && len(emb) > 0 { - products = append(products, common.Product{ - ID: common.StableRowID(tenantID, docID, string(common.VariantDatasetnav), "root"), - DocID: docID, - TenantID: tenantID, - Variant: common.VariantDatasetnav, - Content: rootSummary, - Vector: emb[0], - Meta: map[string]any{ - "kind": "root", - "type": "nav_cluster", - "name": "root", - "depth": 0, - "size": len(tree.docs), - }, - }) - } - } - - out := common.Outputs{ - Products: products, - } - return out, nil -} - -// place inserts one document (chunk) into the tree, mirroring -// upsert_dataset_nav_doc's placement logic. It returns the descriptions of -// any clusters created for this document (for the root overview). -func (t *navTree) place(ctx context.Context, deps common.Deps, llmID, chunkID, text string, vec []float32, mergeThreshold float64) ([]string, error) { - bestName, bestParent, sim := t.findBestCluster(vec) - - switch { - case bestName != "" && sim >= mergeThreshold: - // ── Merge into the best cluster (mirrors _llm_merge + re-embed) ── - c := t.clusters[bestName] - newDesc, err := llmMerge(ctx, deps, llmID, c.Desc, text) - if err != nil { - return nil, err - } - if newDesc != c.Desc { - c.Desc = newDesc - emb, err := deps.Embed.Encode(ctx, []string{newDesc}) - if err != nil { - return nil, err - } - if len(emb) > 0 { - c.Vector = emb[0] - } - } - if !containsString(c.DocIDs, chunkID) { - c.DocIDs = append(c.DocIDs, chunkID) - } - t.addDoc(c, chunkID, text, vec, c.Depth+1) - if err := t.maybeSplit(ctx, deps, llmID, bestName); err != nil { - return nil, err - } - return nil, nil - - case bestName != "" && sim >= minSim: - // ── Create a sibling/child cluster (mirrors the _MIN_SIM branch) ── - parentForNew := bestParent - if parentForNew == "" { - parentForNew = bestName - } - parentDepth := 1 // Python's default when the parent row is absent - if p, ok := t.clusters[parentForNew]; ok { - parentDepth = p.Depth - } - title, desc, err := llmCreateSummary(ctx, deps, llmID, []string{text}) - if err != nil { - return nil, err - } - name := readableClusterName(title, text) - nc := &navCluster{ - Name: name, - Desc: desc, - Parent: parentForNew, - // Python stamps the new cluster at the PARENT's depth (not +1); - // the nav_doc goes one deeper. Mirrored, quirk included. - Depth: parentDepth, - DocIDs: []string{chunkID}, - Vector: vec, - } - t.addCluster(nc) - t.children[parentForNew] = append(t.children[parentForNew], navChild{name: name, vec: vec}) - t.addDoc(nc, chunkID, text, vec, parentDepth+1) - return []string{desc}, nil - - default: - // ── Create a root-level cluster (mirrors the else branch) ── - title, desc, err := llmCreateSummary(ctx, deps, llmID, []string{text}) - if err != nil { - return nil, err - } - name := readableClusterName(title, text) - nc := &navCluster{ - Name: name, - Desc: desc, - Parent: "root", - Depth: 0, - DocIDs: []string{chunkID}, - Vector: vec, - } - t.addCluster(nc) - t.addDoc(nc, chunkID, text, vec, 1) - return []string{desc}, nil - } -} - -func (t *navTree) addDoc(c *navCluster, chunkID, text string, vec []float32, depth int) { - t.docs = append(t.docs, &navDoc{ChunkID: chunkID, Text: text, Parent: c.Name, Depth: depth, Vector: vec}) - t.children[c.Name] = append(t.children[c.Name], navChild{name: chunkID, isDoc: true, vec: vec}) -} - -// findBestCluster mirrors _find_best_cluster: top-1 root cluster by cosine, -// then descend into the best-matching child while similarity stays above -// _RECURSE_THRESHOLD. Returns (best cluster name, its parent name, sim). -func (t *navTree) findBestCluster(vec []float32) (string, string, float64) { - var best *navCluster - bestSim := 0.0 - for _, c := range t.clusters { - if c.Depth != 0 { - continue - } - s := cosine(vec, c.Vector) - if best == nil || s > bestSim { - best, bestSim = c, s - } - } - if best == nil { - return "", "", 0.0 - } - bestName, bestParent, sim := best.Name, best.Parent, bestSim - visited := map[string]bool{bestName: true} - for sim >= recurseThreshold { - var child *navCluster - childSim := 0.0 - for _, k := range t.children[bestName] { - if k.isDoc || visited[k.name] { - continue - } - c := t.clusters[k.name] - if c == nil { - continue - } - if s := cosine(vec, c.Vector); child == nil || s > childSim { - child, childSim = c, s - } - } - if child == nil || childSim < recurseThreshold { - break - } - bestParent = best.Parent // the (old) best's parent, mirroring Python - bestName = child.Name - sim = childSim - best = child - visited[bestName] = true - } - return bestName, bestParent, sim -} - -// maybeSplit mirrors _maybe_split_cluster: when a cluster exceeds the fanout -// or doc-count caps, split its children into two 2-means groups and reparent -// each group under a fresh sub-cluster. -func (t *navTree) maybeSplit(ctx context.Context, deps common.Deps, llmID, clusterName string) error { - kids := t.children[clusterName] - var clusterKids, docKids []navChild - for _, k := range kids { - if k.isDoc { - docKids = append(docKids, k) - } else { - clusterKids = append(clusterKids, k) - } - } - if len(clusterKids)+len(docKids) <= maxFanout && len(docKids) <= maxDocsPerCluster { - return nil - } - if len(kids) < 4 { - return nil // Python requires >= 4 embeddings to attempt a split - } - - embs := make([][]float32, len(kids)) - for i, k := range kids { - embs[i] = k.vec - } - labels := twoMeans(embs) - - parent := t.clusters[clusterName] - depth := 0 - if parent != nil { - depth = parent.Depth + 1 - } - // navChild collects the new sub-cluster names so the split parent keeps a - // children entry pointing at them; without it findBestCluster can no longer - // descend into the split subtree (M3). - var subChildren []navChild - for gi := 0; gi < 2; gi++ { - var kidIdx []int - for i, lb := range labels { - if lb == gi { - kidIdx = append(kidIdx, i) - } - } - if len(kidIdx) == 0 { - continue - } - var docIDs, descs []string - for _, i := range kidIdx { - k := kids[i] - if k.isDoc { - if !containsString(docIDs, k.name) { - docIDs = append(docIDs, k.name) - } - for _, d := range t.docs { - if d.ChunkID == k.name { - descs = append(descs, d.Text) - } - } - continue - } - if c := t.clusters[k.name]; c != nil { - for _, d := range c.DocIDs { - if !containsString(docIDs, d) { - docIDs = append(docIDs, d) - } - } - descs = append(descs, c.Desc) - } - } - var title, desc string - if len(descs) > 0 { - var err error - title, desc, err = llmCreateSummary(ctx, deps, llmID, descs) - if err != nil { - return err - } - } else { - title, desc = fmt.Sprintf("Group %d", gi+1), fmt.Sprintf("Group %d", gi+1) - } - gname := readableClusterName(title, desc) - emb, err := deps.Embed.Encode(ctx, []string{desc}) - if err != nil { - return err - } - var gvec []float32 - if len(emb) > 0 { - gvec = emb[0] - } - nc := &navCluster{Name: gname, Desc: desc, Parent: clusterName, Depth: depth, DocIDs: docIDs, Vector: gvec} - t.addCluster(nc) - subChildren = append(subChildren, navChild{name: gname}) - - // Reparent the group's children to the new sub-cluster. - for _, i := range kidIdx { - k := kids[i] - if k.isDoc { - for _, d := range t.docs { - if d.ChunkID == k.name { - d.Parent = gname - d.Depth = depth + 1 - } - } - } else if c := t.clusters[k.name]; c != nil { - c.Parent = gname - c.Depth = depth + 1 - } - t.children[gname] = append(t.children[gname], k) - } - } - // The split parent's children become the new sub-clusters so findBestCluster - // can still descend into the split subtree (M3). When no sub-cluster was - // actually created the parent keeps an empty children list. - t.children[clusterName] = subChildren - return nil -} - -// twoMeans mirrors the runtime k-means-like split in _maybe_split_cluster: -// centroids seeded with the first and middle embeddings, 10 refinement -// rounds, squared-euclidean assignment, then a final relabel pass. -func twoMeans(embs [][]float32) []int { - n := len(embs) - if n == 0 { - return nil - } - centroids := [][]float32{append([]float32{}, embs[0]...), append([]float32{}, embs[n/2]...)} - assign := func(e []float32) int { - if sqDist(e, centroids[0]) < sqDist(e, centroids[1]) { - return 0 - } - return 1 - } - for iter := 0; iter < 10; iter++ { - var groups [2][][]float32 - for _, e := range embs { - g := assign(e) - groups[g] = append(groups[g], e) - } - for gi := 0; gi < 2; gi++ { - if len(groups[gi]) == 0 { - continue - } - dim := len(groups[gi][0]) - avg := make([]float32, dim) - for _, e := range groups[gi] { - for d := 0; d < dim; d++ { - avg[d] += e[d] - } - } - for d := range avg { - avg[d] /= float32(len(groups[gi])) - } - centroids[gi] = avg - } - } - labels := make([]int, n) - for i, e := range embs { - labels[i] = assign(e) - } - return labels -} - -func sqDist(a, b []float32) float64 { - var s float64 - for i := 0; i < len(a) && i < len(b); i++ { - d := float64(a[i] - b[i]) - s += d * d - } - return s -} - -func cosine(a, b []float32) float64 { - if len(a) == 0 || len(b) == 0 || len(a) != len(b) { - return 0 - } - var dot, na, nb float64 - for i := range a { - dot += float64(a[i]) * float64(b[i]) - na += float64(a[i]) * float64(a[i]) - nb += float64(b[i]) * float64(b[i]) - } - if na == 0 || nb == 0 { - return 0 - } - return dot / (math.Sqrt(na) * math.Sqrt(nb)) -} - -// ---- LLM helpers (prompts mirrored verbatim from dataset_nav.py) ---- - -var navLLMTemperature = 0.1 - -// llmMerge mirrors _llm_merge: fuse the existing cluster description with the -// new doc summary. The reply may be a JSON object ({"merged"|"result"}) or -// bare text; anything unusable keeps the old description. -func llmMerge(ctx context.Context, deps common.Deps, llmID, clusterDesc, docSummary string) (string, error) { - if deps.Chat == nil { - return clusterDesc, nil - } - prompt := "Merge the following two descriptions of the same topic into " + - "a single concise summary (1-3 sentences):\n\n" + - "Existing: " + clusterDesc + "\n\n" + - "New: " + docSummary + "\n\n" + - "Return ONLY the merged text, no commentary." - resp, err := deps.Chat.Chat(ctx, common.ChatRequest{LLMID: llmID, UserPrompt: prompt, Temperature: &navLLMTemperature}) - if err != nil { - return "", err - } - text := strings.TrimSpace(resp.Content) - if m := parseWholeJSON(text); m != nil { - if s, ok := m["merged"].(string); ok && strings.TrimSpace(s) != "" { - return s, nil - } - if s, ok := m["result"].(string); ok && strings.TrimSpace(s) != "" { - return s, nil - } - return clusterDesc, nil - } - if text != "" { - return text, nil - } - return clusterDesc, nil -} - -// llmCreateSummary mirrors _llm_create_summary: derive a readable (name, -// summary) pair from doc summaries, with the same fallbacks. -func llmCreateSummary(ctx context.Context, deps common.Deps, llmID string, docSummaries []string) (string, string, error) { - fallbackSummary := "" - if len(docSummaries) > 0 { - fallbackSummary = docSummaries[0] - } - if deps.Chat == nil { - return fallbackTitle(fallbackSummary), fallbackSummary, nil - } - prompt := "Given the document excerpts below, produce a short human-readable topic " + - "name and a concise description of their common topic.\n\n" + - strings.Join(docSummaries, "\n---\n") + "\n\n" + - `Return ONLY JSON: {"name": "<2-6 word topic title>", "summary": "<1-3 sentence description>"}` - resp, err := deps.Chat.Chat(ctx, common.ChatRequest{LLMID: llmID, UserPrompt: prompt, Temperature: &navLLMTemperature}) - if err != nil { - return "", "", err - } - text := strings.TrimSpace(resp.Content) - if m := parseWholeJSON(text); m != nil { - summary := fallbackSummary - if s, ok := m["summary"].(string); ok && strings.TrimSpace(s) != "" { - summary = strings.TrimSpace(s) - } else if s, ok := m["result"].(string); ok && strings.TrimSpace(s) != "" { - summary = strings.TrimSpace(s) - } - name := cleanTitle(firstStringOf(m["name"])) - if name == "" { - name = fallbackTitle(summary) - } - return name, summary, nil - } - if text != "" { - return fallbackTitle(text), text, nil - } - return fallbackTitle(fallbackSummary), fallbackSummary, nil -} - -// parseWholeJSON parses the response ONLY when it is a complete JSON object -// (mirrors gen_json's dict path; bare text takes the string fallback path). -func parseWholeJSON(s string) map[string]any { - if !strings.HasPrefix(s, "{") { - return nil - } - var m map[string]any - if err := json.Unmarshal([]byte(s), &m); err != nil { - return nil - } - return m -} - -// cleanTitle mirrors _clean_title: whitespace-normalized, capped at 48 chars. -func cleanTitle(title string) string { - return truncateRunes(strings.Join(strings.Fields(title), " "), 48) -} - -// fallbackTitle mirrors _fallback_title: first 6 words, else "Cluster". -func fallbackTitle(summary string) string { - words := strings.Fields(summary) - if len(words) > 6 { - words = words[:6] - } - if t := strings.Join(words, " "); t != "" { - return t - } - return "Cluster" -} - -// readableClusterName mirrors _readable_cluster_name: " <8-hex>". -func readableClusterName(title, seed string) string { - t := cleanTitle(title) - if t == "" { - t = "Cluster" - } - return t + " " + common.ContentHash(seed)[:8] -} - -// navDocName mirrors _make_nav_doc_row's name field: -// f"{parent_kwd}_{xxh64(summary)[:12]}". -func navDocName(parent, summary string) string { - return parent + "_" + common.ContentHash(summary)[:12] -} - -// summarize is the plain-text LLM helper used for the root overview. -func summarize(ctx context.Context, deps common.Deps, llmID, text string) (string, error) { - if deps.Chat == nil { - return "", nil - } - resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ - LLMID: llmID, - SystemPrompt: navSystemPrompt, - UserPrompt: text, - }) - if err != nil { - return "", err - } - return strings.TrimSpace(resp.Content), nil -} - -// formatNavSummaries renders cluster descriptions as a bullet list for the -// root overview prompt. -func formatNavSummaries(summaries []string) string { - var b strings.Builder - for _, s := range summaries { - b.WriteString("- ") - b.WriteString(s) - b.WriteString("\n") - } - return b.String() -} - -// mergeThresholdFor returns the merge threshold (nav_radius override kept for -// tuning/tests; default _MERGE_THRESHOLD). -func mergeThresholdFor(param common.Param) float64 { - if t, ok := param.Extra["nav_radius"].(float64); ok && t > 0 { - return t - } - return mergeThresholdDefault -} - -// chunkTexts returns the non-empty chunk texts and their ids in parallel -// order (Python skips docs without a summary; we skip empty chunks). -func chunkTexts(chunks []common.Chunk) ([]string, []string) { - var texts, ids []string - for i, c := range chunks { - t := firstNonEmpty(c.Text, c.Content) - if strings.TrimSpace(t) == "" { - continue - } - id := c.ID - if id == "" { - id = fmt.Sprintf("chunk-%d", i) - } - texts = append(texts, t) - ids = append(ids, id) - } - return texts, ids -} - -func payloadJSON(v map[string]any) string { - var b strings.Builder - enc := json.NewEncoder(&b) - enc.SetEscapeHTML(false) - if err := enc.Encode(v); err != nil { - return "{}" - } - return strings.TrimSpace(b.String()) -} - -func truncateRunes(s string, n int) string { - r := []rune(s) - if len(r) <= n { - return s - } - return string(r[:n]) -} - -func firstStringOf(v any) string { - if s, ok := v.(string); ok { - return s - } - return "" -} - -func containsString(haystack []string, needle string) bool { - for _, s := range haystack { - if s == needle { - return true - } - } - return false -} - -func firstNonEmpty(vals ...string) string { - for _, v := range vals { - if v != "" { - return v - } - } - return "" -} - -const navSystemPrompt = `You are a navigation assistant. Summarize the provided text into a concise label and overview that helps a user navigate to the relevant content. Output the summary only.` diff --git a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav_test.go b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav_test.go deleted file mode 100644 index 848e45f607..0000000000 --- a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav_test.go +++ /dev/null @@ -1,289 +0,0 @@ -package datasetnav - -import ( - "context" - "fmt" - "strings" - "testing" - - "ragflow/internal/ingestion/component/knowledge_compiler/common" -) - -// scriptedEmbedder maps texts to fixed vectors (fallback: halves). -type scriptedEmbedder struct { - vecs map[string][]float32 -} - -func (s scriptedEmbedder) Dimensions() int { return 2 } -func (s scriptedEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { - out := make([][]float32, len(texts)) - for i, t := range texts { - if v, ok := s.vecs[t]; ok { - out[i] = v - } else { - out[i] = []float32{0.5, 0.5} - } - } - return out, nil -} - -// navFakeChat answers the three nav LLM calls deterministically: merge keeps -// the existing description, create-summary returns canned JSON, the root -// overview is a fixed string. -type navFakeChat struct { - mergeCalls int - summaryCalls int -} - -func (m *navFakeChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { - switch { - case strings.Contains(req.UserPrompt, "Merge the following two descriptions"): - m.mergeCalls++ - existing := "" - if rest, ok := strings.CutPrefix(req.UserPrompt, "Merge the following two descriptions of the same topic into a single concise summary (1-3 sentences):\n\nExisting: "); ok { - existing, _, _ = strings.Cut(rest, "\n\nNew: ") - } - return &common.ChatResponse{Content: fmt.Sprintf(`{"merged": %q}`, existing)}, nil - case strings.Contains(req.UserPrompt, "Given the document excerpts below"): - m.summaryCalls++ - // Distinct summaries per call: readable cluster names hash the - // description, so identical canned summaries would collide (the same - // name-overwrite hazard exists in Python's _nav_cluster_id). - return &common.ChatResponse{Content: fmt.Sprintf(`{"name": "Topic", "summary": "a topic summary %d"}`, m.summaryCalls)}, nil - case strings.HasPrefix(req.UserPrompt, "Compose a navigation overview"): - return &common.ChatResponse{Content: "root overview"}, nil - default: - return &common.ChatResponse{Content: "ok"}, nil - } -} - -func TestTwoMeans(t *testing.T) { - embs := [][]float32{{1, 0}, {0.9, 0.1}, {0, 1}, {0.1, 0.9}} - labels := twoMeans(embs) - if len(labels) != 4 { - t.Fatalf("labels = %v", labels) - } - if labels[0] != labels[1] || labels[2] != labels[3] || labels[0] == labels[2] { - t.Fatalf("expected {x,x,y,y} grouping, got %v", labels) - } -} - -func TestReadableClusterName(t *testing.T) { - got := readableClusterName(" My Topic ", "seed text") - if !strings.HasPrefix(got, "My Topic ") { - t.Fatalf("readableClusterName = %q, want cleaned title + hash suffix", got) - } - if got := readableClusterName("", "seed"); !strings.HasPrefix(got, "Cluster ") { - t.Fatalf("empty title must fall back to Cluster: %q", got) - } -} - -func TestFallbackTitle(t *testing.T) { - if got := fallbackTitle("one two three four five six seven eight"); got != "one two three four five six" { - t.Fatalf("fallbackTitle = %q, want first 6 words", got) - } - if got := fallbackTitle(""); got != "Cluster" { - t.Fatalf("empty fallbackTitle = %q, want Cluster", got) - } -} - -func TestPlace_MergeSiblingRootBranches(t *testing.T) { - vecs := map[string][]float32{ - "A": {1, 0}, - "A2": {0.99, 0.1}, // cosine ≈ 0.995 with A → merge - "B": {0.7, 0.714}, // cosine ≈ 0.7 with A → sibling - "C": {0, 1}, // cosine 0 with A → new root - } - deps := common.Deps{Chat: &navFakeChat{}, Embed: scriptedEmbedder{vecs: vecs}} - tree := newNavTree() - - if _, err := tree.place(context.Background(), deps, "llm", "cA", "A", vecs["A"], mergeThresholdDefault); err != nil { - t.Fatal(err) - } - if _, err := tree.place(context.Background(), deps, "llm", "cA2", "A2", vecs["A2"], mergeThresholdDefault); err != nil { - t.Fatal(err) - } - if _, err := tree.place(context.Background(), deps, "llm", "cB", "B", vecs["B"], mergeThresholdDefault); err != nil { - t.Fatal(err) - } - if _, err := tree.place(context.Background(), deps, "llm", "cC", "C", vecs["C"], mergeThresholdDefault); err != nil { - t.Fatal(err) - } - - if len(tree.order) != 3 { - t.Fatalf("clusters = %d, want 3 (A-root, B-sibling, C-root): %v", len(tree.order), tree.order) - } - // A2 merged into A's cluster. - var clusterA *navCluster - for _, name := range tree.order { - if tree.clusters[name].Depth == 0 && containsString(tree.clusters[name].DocIDs, "cA") { - clusterA = tree.clusters[name] - } - } - if clusterA == nil || len(clusterA.DocIDs) != 2 || clusterA.DocIDs[1] != "cA2" { - t.Fatalf("A2 must merge into cluster A (doc_ids=%v)", clusterA) - } - // B's cluster: sibling under the virtual root (Python depth quirk: parent - // depth default 1 → cluster depth 1, parent "root"). - var clusterB *navCluster - for _, name := range tree.order { - c := tree.clusters[name] - if containsString(c.DocIDs, "cB") { - clusterB = c - } - } - if clusterB == nil || clusterB.Parent != "root" || clusterB.Depth != 1 { - t.Fatalf("B sibling cluster = %+v, want parent root depth 1", clusterB) - } - // C's cluster: a fresh root-level cluster. - var clusterC *navCluster - for _, name := range tree.order { - c := tree.clusters[name] - if containsString(c.DocIDs, "cC") { - clusterC = c - } - } - if clusterC == nil || clusterC.Parent != "root" || clusterC.Depth != 0 { - t.Fatalf("C root cluster = %+v, want parent root depth 0", clusterC) - } - if len(tree.docs) != 4 { - t.Fatalf("nav docs = %d, want 4", len(tree.docs)) - } -} - -func TestMaybeSplit(t *testing.T) { - vecs := map[string][]float32{} - deps := common.Deps{Chat: &navFakeChat{}, Embed: scriptedEmbedder{vecs: vecs}} - tree := newNavTree() - parent := &navCluster{Name: "P", Desc: "parent", Parent: "root", Depth: 0, Vector: []float32{1, 0}} - tree.addCluster(parent) - // 51 docs (> maxDocsPerCluster) alternating between two axes. - for i := 0; i < 51; i++ { - id := fmt.Sprintf("d%d", i) - v := []float32{1, 0} - if i%2 == 1 { - v = []float32{0, 1} - } - parent.DocIDs = append(parent.DocIDs, id) - tree.addDoc(parent, id, "text "+id, v, 1) - } - if err := tree.maybeSplit(context.Background(), deps, "llm", "P"); err != nil { - t.Fatal(err) - } - // Expect two new sub-clusters under P, children reparented, P left with - // no direct children. - if len(tree.order) != 3 { - t.Fatalf("clusters after split = %d, want 3 (P + 2 groups): %v", len(tree.order), tree.order) - } - // The split parent must now point at exactly the two new sub-clusters so - // findBestCluster can descend into the split subtree (M3). - if got := len(tree.children["P"]); got != 2 { - t.Fatalf("split parent children = %d, want 2 (the two new sub-clusters)", got) - } - subDocs := 0 - for _, name := range tree.order[1:] { - c := tree.clusters[name] - if c.Parent != "P" || c.Depth != 1 { - t.Errorf("sub-cluster %s parent/depth = %s/%d, want P/1", name, c.Parent, c.Depth) - } - subDocs += len(c.DocIDs) - } - if subDocs != 51 { - t.Errorf("sub-cluster doc_ids total = %d, want 51", subDocs) - } - for _, d := range tree.docs { - if d.Parent == "P" { - t.Errorf("doc %s not reparented off P", d.ChunkID) - } - if d.Depth != 2 { - t.Errorf("doc %s depth = %d, want 2", d.ChunkID, d.Depth) - } - } -} - -func TestMaybeSplit_BelowCaps(t *testing.T) { - tree := newNavTree() - parent := &navCluster{Name: "P", Desc: "p", Parent: "root", Depth: 0, Vector: []float32{1, 0}} - tree.addCluster(parent) - for i := 0; i < 3; i++ { - tree.addDoc(parent, fmt.Sprintf("d%d", i), "t", []float32{1, 0}, 1) - } - if err := tree.maybeSplit(context.Background(), common.Deps{}, "llm", "P"); err != nil { - t.Fatal(err) - } - if len(tree.order) != 1 { - t.Fatalf("below caps must not split: %v", tree.order) - } -} - -func TestRun_EndToEnd(t *testing.T) { - vecs := map[string][]float32{ - "alpha one": {1, 0}, - "alpha two": {0.99, 0.1}, - "beta one": {0, 1}, - } - deps := common.Deps{Chat: &navFakeChat{}, Embed: scriptedEmbedder{vecs: vecs}, TenantID: "t1", DatasetID: "ds1"} - p := common.Param{}.Defaults() - p.Variant = common.VariantDatasetnav - inputs := common.Inputs{ - DocID: "d1", - Chunks: []common.Chunk{ - {ID: "c1", Text: "alpha one"}, - {ID: "c2", Text: "alpha two"}, - {ID: "c3", Text: "beta one"}, - }, - } - out, err := Run(context.Background(), deps, p, inputs) - if err != nil { - t.Fatalf("Run: %v", err) - } - var clusters, docs, roots int - parentByID := map[string]string{} - for _, p := range out.Products { - switch p.Meta["kind"] { - case "nav_cluster": - clusters++ - parentByID[p.Meta["name"].(string)] = p.Meta["parent_name"].(string) - case "nav_doc": - docs++ - if p.ParentID == "" { - t.Errorf("nav_doc must link to its cluster product") - } - case "root": - roots++ - } - } - if clusters != 2 || docs != 3 || roots != 1 { - t.Fatalf("products = %d clusters + %d docs + %d roots, want 2+3+1", clusters, docs, roots) - } -} - -func TestLLMMergeFallbacks(t *testing.T) { - // Non-JSON text reply is used verbatim (mirrors gen_json's str path). - deps := common.Deps{Chat: chatFunc(func(req common.ChatRequest) (*common.ChatResponse, error) { - return &common.ChatResponse{Content: "merged text"}, nil - })} - got, err := llmMerge(context.Background(), deps, "llm", "old", "new") - if err != nil || got != "merged text" { - t.Fatalf("llmMerge text path = %q, %v", got, err) - } - // Empty reply keeps the old description. - deps2 := common.Deps{Chat: chatFunc(func(req common.ChatRequest) (*common.ChatResponse, error) { - return &common.ChatResponse{Content: ""}, nil - })} - got, _ = llmMerge(context.Background(), deps2, "llm", "old", "new") - if got != "old" { - t.Fatalf("empty reply must keep old desc, got %q", got) - } - // Nil chat keeps the old description without a call. - got, _ = llmMerge(context.Background(), common.Deps{}, "llm", "old", "new") - if got != "old" { - t.Fatalf("nil chat must keep old desc, got %q", got) - } -} - -type chatFunc func(common.ChatRequest) (*common.ChatResponse, error) - -func (f chatFunc) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { - return f(req) -} diff --git a/internal/service/dataset_artifact_nav_test.go b/internal/service/dataset_artifact_nav_test.go new file mode 100644 index 0000000000..88c1746f41 --- /dev/null +++ b/internal/service/dataset_artifact_nav_test.go @@ -0,0 +1,92 @@ +package service + +import ( + "context" + "testing" + + "ragflow/internal/service/nav" +) + +// fakeArtifactNav is an in-memory nav.NavService used to lock the limited +// DeleteNav/DeleteNavNode semantics (direct doc children only; no sub-cluster +// or cascade deletion). +type fakeArtifactNav struct { + clusters []nav.NavNode + children map[string][]nav.NavNode + removedDocs []string +} + +func (f *fakeArtifactNav) UpsertDoc(context.Context, nav.UpsertDocInput) error { return nil } +func (f *fakeArtifactNav) RemoveDoc(_ context.Context, _, _, docID string) error { + f.removedDocs = append(f.removedDocs, docID) + return nil +} +func (f *fakeArtifactNav) Search(context.Context, string, string, string, []float32, int) ([]nav.NavHit, error) { + return nil, nil +} +func (f *fakeArtifactNav) ListClusters(context.Context, string, string, int, int) ([]nav.NavNode, int64, error) { + return f.clusters, int64(len(f.clusters)), nil +} +func (f *fakeArtifactNav) ListChildren(_ context.Context, _, _, name string, _, _ int) ([]nav.NavNode, int64, error) { + return f.children[name], int64(len(f.children[name])), nil +} + +// TestDeleteNav_RemovesOnlyDirectDocChildren documents the limited semantic of +// the deprecated DeleteNav: it removes the nav_doc rows directly under root +// clusters, and does NOT touch sub-clusters (Python's full cascade is not +// implemented in the minimal loop). +func TestDeleteNav_RemovesOnlyDirectDocChildren(t *testing.T) { + fake := &fakeArtifactNav{ + clusters: []nav.NavNode{{Name: "C1", Description: "root"}}, + children: map[string][]nav.NavNode{ + "C1": {{Name: "Sub", Type: "cluster"}, {Name: "DocA", Type: "doc", DocID: "d1"}}, + "Sub": {{Name: "DocB", Type: "doc", DocID: "d2"}}, + }, + } + prev := nav.GetNavService() + nav.SetNavService(fake) + defer func() { nav.SetNavService(prev) }() + + svc := NewDatasetArtifactService() + n, err := svc.DeleteNav(context.Background(), "t1", "kb1") + if err != nil { + t.Fatalf("DeleteNav: %v", err) + } + // Only the direct doc child (d1) under the root cluster is removed; the + // sub-cluster's doc (d2) is NOT reached (no subtree traversal). + if n != 1 { + t.Errorf("deleted = %d, want 1 (direct doc child only)", n) + } + if len(fake.removedDocs) != 1 || fake.removedDocs[0] != "d1" { + t.Errorf("removed docs = %v, want [d1] only (no cascade into Sub)", fake.removedDocs) + } +} + +// TestDeleteNavNode_RemovesDirectChildren documents that DeleteNavNode drains +// only the immediate doc children of a named cluster. +func TestDeleteNavNode_RemovesDirectChildren(t *testing.T) { + fake := &fakeArtifactNav{ + children: map[string][]nav.NavNode{ + "C1": { + {Name: "Sub2", Type: "cluster"}, + {Name: "DocA", Type: "doc", DocID: "d1"}, + {Name: "DocB", Type: "doc", DocID: "d2"}, + }, + }, + } + prev := nav.GetNavService() + nav.SetNavService(fake) + defer func() { nav.SetNavService(prev) }() + + svc := NewDatasetArtifactService() + n, err := svc.DeleteNavNode(context.Background(), "t1", "kb1", "C1") + if err != nil { + t.Fatalf("DeleteNavNode: %v", err) + } + if n != 2 { + t.Errorf("deleted = %d, want 2 (direct doc children d1,d2)", n) + } + if len(fake.removedDocs) != 2 { + t.Errorf("removed docs = %v, want [d1 d2]", fake.removedDocs) + } +} diff --git a/internal/service/dataset_artifact_service.go b/internal/service/dataset_artifact_service.go index 38a524fcf0..af9a210a50 100644 --- a/internal/service/dataset_artifact_service.go +++ b/internal/service/dataset_artifact_service.go @@ -21,6 +21,7 @@ import ( "ragflow/internal/engine" "ragflow/internal/engine/types" + "ragflow/internal/service/nav" ) // Compile keyword constants used by the knowledge-compilation artifacts stored @@ -601,115 +602,120 @@ func (s *DatasetArtifactService) DeleteDocumentGraph(ctx context.Context, tenant return len(ids), nil } -// NavigationItem is a single navigation cluster. +// NavigationItem is a single navigation cluster (REST response shape, kept +// stable for frontend compatibility). type NavigationItem struct { Name string `json:"name"` Title string `json:"title"` Count int `json:"count"` } -// ListNavClusters returns the navigation clusters of a dataset. -func (s *DatasetArtifactService) ListNavClusters(ctx context.Context, tenantID, datasetID string) ([]NavigationItem, int64, error) { - chunks, total, err := s.searchCompiled(ctx, tenantID, datasetID, - map[string]interface{}{"compile_kwd": []string{CompileKwdDatasetNav}}, - []string{"nav_cluster_kwd", "title_kwd", "count_int"}, 0, 10000, nil) - if err != nil { - return nil, 0, err - } - items := make([]NavigationItem, 0, len(chunks)) - for _, c := range chunks { - count := intValue(c["count_int"]) - items = append(items, NavigationItem{ - Name: firstStringValue(c["nav_cluster_kwd"]), - Title: firstStringValue(c["title_kwd"]), - Count: count, - }) - } - return items, total, nil -} - -// DeleteNav deletes all navigation clusters of a dataset. -func (s *DatasetArtifactService) DeleteNav(ctx context.Context, tenantID, datasetID string) (int, error) { - docEngine := engine.Get() - if docEngine == nil { - return 0, fmt.Errorf("document engine is not initialized") - } - chunks, _, err := s.searchCompiled(ctx, tenantID, datasetID, - map[string]interface{}{"compile_kwd": []string{CompileKwdDatasetNav}}, []string{"id"}, 0, 10000, nil) - if err != nil { - return 0, err - } - if len(chunks) == 0 { - return 0, nil - } - ids := make([]string, 0, len(chunks)) - for _, c := range chunks { - if id, ok := c["id"].(string); ok { - ids = append(ids, id) - } - } - cond := map[string]interface{}{"id": ids, "kb_id": datasetID} - if _, err := docEngine.DeleteChunks(ctx, cond, wikiIndexName(tenantID), datasetID); err != nil { - return 0, err - } - return len(ids), nil -} - -// DeleteNavNode deletes a single navigation cluster by name. -func (s *DatasetArtifactService) DeleteNavNode(ctx context.Context, tenantID, datasetID, name string) (int, error) { - docEngine := engine.Get() - if docEngine == nil { - return 0, fmt.Errorf("document engine is not initialized") - } - chunks, _, err := s.searchCompiled(ctx, tenantID, datasetID, - map[string]interface{}{"compile_kwd": []string{CompileKwdDatasetNav}, "nav_cluster_kwd": []string{name}}, - []string{"id"}, 0, 10000, nil) - if err != nil { - return 0, err - } - if len(chunks) == 0 { - return 0, nil - } - ids := make([]string, 0, len(chunks)) - for _, c := range chunks { - if id, ok := c["id"].(string); ok { - ids = append(ids, id) - } - } - cond := map[string]interface{}{"id": ids, "kb_id": datasetID} - if _, err := docEngine.DeleteChunks(ctx, cond, wikiIndexName(tenantID), datasetID); err != nil { - return 0, err - } - return len(ids), nil -} - -// NavChildItem is a single child entry under a navigation cluster. +// NavChildItem is a single child entry under a navigation cluster (REST +// response shape, kept stable for frontend compatibility). type NavChildItem struct { Name string `json:"name"` Title string `json:"title"` Count int `json:"count"` } -// ListNavChildren returns the children of a navigation cluster. -func (s *DatasetArtifactService) ListNavChildren(ctx context.Context, tenantID, datasetID, name string) ([]NavChildItem, int64, error) { - chunks, total, err := s.searchCompiled(ctx, tenantID, datasetID, - map[string]interface{}{"compile_kwd": []string{CompileKwdDatasetNav}, "nav_cluster_kwd": []string{name}}, - []string{"nav_child_kwd", "title_kwd", "count_int"}, 0, 10000, nil) +// ListNavClusters returns the navigation clusters of a dataset. It is DEPRECATED +// and now delegates to the ES-backed NavService (internal/service datasetnav): +// the previous implementation queried nav_cluster_kwd/count_int fields that +// Python never writes, so it could never read the real nav tree. Do not add +// field-level patches here — route everything through NavService. +func (s *DatasetArtifactService) ListNavClusters(ctx context.Context, tenantID, datasetID string) ([]NavigationItem, int64, error) { + ns := nav.GetNavService() + if ns == nil { + return nil, 0, fmt.Errorf("datasetnav: NavService not initialized (SetNavService must be called at bootstrap)") + } + nodes, total, err := ns.ListClusters(ctx, tenantID, datasetID, 0, 10000) if err != nil { return nil, 0, err } - items := make([]NavChildItem, 0, len(chunks)) - for _, c := range chunks { - count := intValue(c["count_int"]) - items = append(items, NavChildItem{ - Name: firstStringValue(c["nav_child_kwd"]), - Title: firstStringValue(c["title_kwd"]), - Count: count, - }) + items := make([]NavigationItem, 0, len(nodes)) + for _, n := range nodes { + items = append(items, NavigationItem{Name: n.Name, Title: n.Description, Count: n.DocCount}) } return items, total, nil } +// ListNavChildren returns the children of a navigation cluster. DEPRECATED — +// delegates to NavService.ListChildren. +func (s *DatasetArtifactService) ListNavChildren(ctx context.Context, tenantID, datasetID, name string) ([]NavChildItem, int64, error) { + ns := nav.GetNavService() + if ns == nil { + return nil, 0, fmt.Errorf("datasetnav: NavService not initialized (SetNavService must be called at bootstrap)") + } + nodes, total, err := ns.ListChildren(ctx, tenantID, datasetID, name, 0, 10000) + if err != nil { + return nil, 0, err + } + items := make([]NavChildItem, 0, len(nodes)) + for _, n := range nodes { + items = append(items, NavChildItem{Name: n.Name, Title: n.Description, Count: n.DocCount}) + } + return items, total, nil +} + +// DeleteNav removes the direct nav_doc children of every root cluster of a +// dataset. DEPRECATED — this is the minimal-loop approximation of Python +// delete_nav: it drains only the immediate nav_doc rows under root clusters and +// does NOT implement Python's full subtree traversal or empty-cluster cascade +// cleanup. Prefer the NavService (future work) for a complete delete. Returns +// the number of nav_doc rows removed. +func (s *DatasetArtifactService) DeleteNav(ctx context.Context, tenantID, datasetID string) (int, error) { + ns := nav.GetNavService() + if ns == nil { + return 0, fmt.Errorf("datasetnav: NavService not initialized (SetNavService must be called at bootstrap)") + } + clusters, _, err := ns.ListClusters(ctx, tenantID, datasetID, 0, 10000) + if err != nil { + return 0, err + } + deleted := 0 + for _, c := range clusters { + children, _, err := ns.ListChildren(ctx, tenantID, datasetID, c.Name, 0, 10000) + if err != nil { + continue + } + for _, ch := range children { + if ch.DocID != "" { + if err := ns.RemoveDoc(ctx, tenantID, datasetID, ch.DocID); err != nil { + return deleted, err + } + deleted++ + } + } + } + return deleted, nil +} + +// DeleteNavNode deletes the direct nav_doc children of a named cluster. +// DEPRECATED — the minimal loop only drains immediate doc children (returns the +// count); it does NOT delete sub-clusters recursively nor perform Python's +// empty-cluster cascade. A full tree-node delete is future NavService work. +func (s *DatasetArtifactService) DeleteNavNode(ctx context.Context, tenantID, datasetID, name string) (int, error) { + ns := nav.GetNavService() + if ns == nil { + return 0, fmt.Errorf("datasetnav: NavService not initialized (SetNavService must be called at bootstrap)") + } + // Minimal loop has no per-node delete; drain direct children's docs. + children, _, err := ns.ListChildren(ctx, tenantID, datasetID, name, 0, 10000) + if err != nil { + return 0, err + } + deleted := 0 + for _, ch := range children { + if ch.DocID != "" { + if err := ns.RemoveDoc(ctx, tenantID, datasetID, ch.DocID); err != nil { + return deleted, err + } + deleted++ + } + } + return deleted, nil +} + // SkillTreeItem is a single skill-tree page summary. type SkillTreeItem struct { Kwd string `json:"kwd"` diff --git a/internal/service/nav/nav.go b/internal/service/nav/nav.go new file mode 100644 index 0000000000..9b590ac2c3 --- /dev/null +++ b/internal/service/nav/nav.go @@ -0,0 +1,96 @@ +// +// 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 nav defines the dataset-navigation service interface. It is kept as a +// dependency-light leaf package so that both the agent tool layer +// (internal/agent/tool) and the concrete implementation +// (internal/service/nlp) can depend on it without creating an import cycle — +// mirroring how RetrievalService lives in the agent tool package. See +// tasks/agentic_search_port_plan.md. +package nav + +import ( + "context" + "sync" +) + +// NavNode mirrors Python's _nav_item (dataset_api_service.py). +type NavNode struct { + Name string // name + Description string // content_with_weight payload description + DocCount int // doc_count_int (cluster) or 1 (leaf) + Type string // "cluster" | "doc" + DocID string // leaf doc_id; empty for cluster + HasChildren bool // is_cluster +} + +// NavHit is one KNN hit on a nav row. +type NavHit struct { + Type string // "nav_doc" | "nav_cluster" + DocID string + DocIDs []string + Name string + Score float64 +} + +// UpsertDocInput carries the per-document summary to place into the nav tree. +type UpsertDocInput struct { + TenantID string + KbID string + DocID string + Summary string // document summary text (tree product or page_index summary) + Embedd []float32 // optional precomputed embedding +} + +// NavService is the single read/write entrypoint for a dataset's navigation +// tree. It is the only nav consumer used by agent tools, REST handlers and the +// tree/structure compile-complete hooks. +type NavService interface { + // UpsertDoc places one document summary into the nav tree (incremental, + // ES-backed read-modify-write; deterministic placement in the minimal loop). + UpsertDoc(ctx context.Context, in UpsertDocInput) error + // RemoveDoc removes a document's nav rows for the given doc. The minimal-loop + // implementation deletes the nav_doc row(s) for the doc; empty-cluster + // cascade cleanup is NOT yet implemented. + RemoveDoc(ctx context.Context, tenantID, kbID, docID string) error + + // Search runs query KNN over nav rows and returns the routed doc ids. + Search(ctx context.Context, tenantID, kbID, query string, embd []float32, topK int) ([]NavHit, error) + // ListClusters returns the depth-0 clusters (parent_kwd=root). + ListClusters(ctx context.Context, tenantID, kbID string, page, pageSize int) ([]NavNode, int64, error) + // ListChildren returns the direct children of a cluster (parent_kwd=name). + ListChildren(ctx context.Context, tenantID, kbID, name string, page, pageSize int) ([]NavNode, int64, error) +} + +var ( + svcMu sync.RWMutex + svcInst NavService +) + +// SetNavService installs the (production or test) NavService singleton. +func SetNavService(s NavService) { + svcMu.Lock() + defer svcMu.Unlock() + svcInst = s +} + +// GetNavService returns the installed NavService. It may return nil until +// SetNavService is called during server bootstrap. +func GetNavService() NavService { + svcMu.RLock() + defer svcMu.RUnlock() + return svcInst +} diff --git a/internal/service/nav_embedder.go b/internal/service/nav_embedder.go new file mode 100644 index 0000000000..6598695291 --- /dev/null +++ b/internal/service/nav_embedder.go @@ -0,0 +1,80 @@ +// +// 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 service + +import ( + "context" + "fmt" + "strings" +) + +// NavEmbedder is the production implementation of nlp.NavEmbedder. It resolves +// the tenant's embedding model on each call and returns float32 vectors (the +// dataset-nav index stores q_<dim>_vec as float). It lives in the service +// package (not nlp) so it can import model_service without an import cycle. +type NavEmbedder struct { + modelSvc *ModelProviderService + // embdModelName is the composite embedding model name (e.g. + // "embedding_model@..." ). Empty falls back to resolving the tenant default. + embdModelName string +} + +// NewNavEmbedder builds the production embedder used by NavService. +func NewNavEmbedder(modelSvc *ModelProviderService, embdModelName string) *NavEmbedder { + return &NavEmbedder{modelSvc: modelSvc, embdModelName: embdModelName} +} + +// Encode embeds texts for the tenant and returns float32 vectors. +func (e *NavEmbedder) Encode(ctx context.Context, tenantID string, texts []string) ([][]float32, error) { + if e.modelSvc == nil { + return nil, fmt.Errorf("datasetnav: embedding model service not initialized") + } + name := e.embdModelName + if name == "" { + name = tenantID // composite name falls back to tenant default resolution + } + model, err := e.modelSvc.GetEmbeddingModel(ctx, tenantID, name) + if err != nil { + return nil, fmt.Errorf("datasetnav: resolve embedding model for tenant %s: %w", tenantID, err) + } + nonEmpty := make([]string, 0, len(texts)) + for _, t := range texts { + if strings.TrimSpace(t) != "" { + nonEmpty = append(nonEmpty, t) + } + } + if len(nonEmpty) == 0 { + return nil, nil + } + embeds, err := model.ModelDriver.Embed(ctx, model.ModelName, nonEmpty, model.APIConfig, nil, nil) + if err != nil { + return nil, err + } + out := make([][]float32, 0, len(embeds)) + for _, e := range embeds { + out = append(out, toF32(e.Embedding)) + } + return out, nil +} + +func toF32(v []float64) []float32 { + out := make([]float32, len(v)) + for i, x := range v { + out[i] = float32(x) + } + return out +} diff --git a/internal/service/nlp/datasetnav.go b/internal/service/nlp/datasetnav.go new file mode 100644 index 0000000000..2f8c133f47 --- /dev/null +++ b/internal/service/nlp/datasetnav.go @@ -0,0 +1,574 @@ +// +// 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 nlp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "ragflow/internal/engine" + "ragflow/internal/engine/types" + "ragflow/internal/service/nav" +) + +// Dataset-nav constants. These mirror Python's dataset_nav.py: nav rows live in +// the document index as compile_kwd="dataset_nav" rows with available_int=0 +// (invisible to the default retriever, which filters available_int=1), and the +// tree is threaded through parent_kwd ("root" for depth-0 clusters). +const ( + navCompileKwd = "dataset_nav" + navRootParent = "root" + + navMergeThreshold = 0.80 // sim >= this -> merge doc into cluster + navRecurse = 0.65 // descend while sim >= this + navMinSim = 0.50 // sim >= this -> new sibling cluster + navMaxDepth = 6 // max descent levels during best-cluster search +) + +// NavEmbedder embeds text. tenantID lets a production implementation resolve +// the tenant's embedding model. Kept as an interface so tests inject a stub. +type NavEmbedder interface { + Encode(ctx context.Context, tenantID string, texts []string) ([][]float32, error) +} + +// NavService is the concrete, ES-backed implementation of nav.NavService. +type NavService struct { + embed NavEmbedder + engine engine.DocEngine // optional; falls back to engine.Get() when nil +} + +// NewNavService builds the ES-backed NavService. embed may be nil when the +// caller guarantees every UpsertDocInput carries a precomputed Embedd. +func NewNavService(embed NavEmbedder) *NavService { + return &NavService{embed: embed} +} + +func (s *NavService) docEngine() (engine.DocEngine, error) { + if s.engine != nil { + return s.engine, nil + } + de := engine.Get() + if de == nil { + return nil, fmt.Errorf("document engine is not initialized") + } + return de, nil +} + +// navIndexName returns the tenant document index name (ragflow_<tenantID>). +func (s *NavService) navIndexName(tenantID string) string { + return fmt.Sprintf("ragflow_%s", tenantID) +} + +// navFilter builds the common nav filter that pins compile_kwd. +func navFilter(extra map[string]interface{}) map[string]interface{} { + f := map[string]interface{}{"compile_kwd": []string{navCompileKwd}} + for k, v := range extra { + f[k] = v + } + return f +} + +// navSearch runs a filtered read over the tenant index for the dataset. +func (s *NavService) navSearch(ctx context.Context, tenantID, kbID string, filter map[string]interface{}, selectFields []string, offset, limit int, matchExprs []interface{}) ([]map[string]interface{}, int64, error) { + de, err := s.docEngine() + if err != nil { + return nil, 0, err + } + merged := make(map[string]interface{}, len(filter)+1) + for k, v := range filter { + merged[k] = v + } + merged["kb_id"] = []string{kbID} + req := &types.SearchRequest{ + IndexNames: []string{s.navIndexName(tenantID)}, + KbIDs: []string{kbID}, + Offset: offset, + Limit: limit, + SelectFields: selectFields, + Filter: merged, + MatchExprs: matchExprs, + } + res, err := de.Search(ctx, req) + if err != nil { + return nil, 0, err + } + if res == nil { + return nil, 0, nil + } + return res.Chunks, res.Total, nil +} + +// ListClusters returns the depth-0 clusters (parent_kwd=root). +func (s *NavService) ListClusters(ctx context.Context, tenantID, kbID string, page, pageSize int) ([]nav.NavNode, int64, error) { + if pageSize <= 0 { + pageSize = 100 + } + offset := page * pageSize + chunks, total, err := s.navSearch(ctx, tenantID, kbID, + navFilter(map[string]interface{}{ + "type_kwd": []string{"nav_cluster"}, + "parent_kwd": []string{navRootParent}, + }), + []string{"title_kwd", "content_with_weight", "doc_count_int", "type_kwd"}, offset, pageSize, nil) + if err != nil { + return nil, 0, err + } + nodes := make([]nav.NavNode, 0, len(chunks)) + for _, c := range chunks { + nodes = append(nodes, s.nodeFromRow(c, "cluster")) + } + return nodes, total, nil +} + +// ListChildren returns the direct children of a cluster (parent_kwd=name). +func (s *NavService) ListChildren(ctx context.Context, tenantID, kbID, name string, page, pageSize int) ([]nav.NavNode, int64, error) { + if pageSize <= 0 { + pageSize = 100 + } + offset := page * pageSize + chunks, total, err := s.navSearch(ctx, tenantID, kbID, + navFilter(map[string]interface{}{"parent_kwd": []string{name}}), + []string{"title_kwd", "content_with_weight", "doc_count_int", "type_kwd", "doc_id"}, offset, pageSize, nil) + if err != nil { + return nil, 0, err + } + nodes := make([]nav.NavNode, 0, len(chunks)) + for _, c := range chunks { + typ := firstStringValue(c["type_kwd"]) + nodeType := "doc" + if typ == "nav_cluster" { + nodeType = "cluster" + } + nodes = append(nodes, s.nodeFromRow(c, nodeType)) + } + return nodes, total, nil +} + +// nodeFromRow converts an engine row into a NavNode. +func (s *NavService) nodeFromRow(row map[string]interface{}, fallbackType string) nav.NavNode { + node := nav.NavNode{ + Name: firstStringValue(row["title_kwd"]), + DocCount: intValue(row["doc_count_int"]), + Type: fallbackType, + DocID: firstStringValue(row["doc_id"]), + } + if t := firstStringValue(row["type_kwd"]); t != "" { + if t == "nav_cluster" { + node.Type = "cluster" + } else { + node.Type = "doc" + } + } + if payload, ok := row["content_with_weight"].(string); ok { + var m map[string]interface{} + if err := json.Unmarshal([]byte(payload), &m); err == nil { + if d, ok := m["description"].(string); ok { + node.Description = d + } + } + } + node.HasChildren = node.Type == "cluster" + if node.DocCount <= 0 { + node.DocCount = 1 + } + return node +} + +// Search runs query KNN over nav rows and returns routed doc ids. +func (s *NavService) Search(ctx context.Context, tenantID, kbID, query string, embd []float32, topK int) ([]nav.NavHit, error) { + if topK <= 0 { + topK = 8 + } + vec := embd + if len(vec) == 0 { + if s.embed == nil { + return nil, fmt.Errorf("datasetnav: no embedding available for Search") + } + embeddings, err := s.embed.Encode(ctx, tenantID, []string{query}) + if err != nil { + return nil, err + } + if len(embeddings) == 0 { + return nil, fmt.Errorf("datasetnav: embedding produced no vector") + } + vec = embeddings[0] + } + f64 := f32ToF64Slice(vec) + chunks, _, err := s.navSearch(ctx, tenantID, kbID, + navFilter(nil), + []string{"type_kwd", "title_kwd", "doc_id", "doc_ids_kwd", "_score"}, 0, topK, + []interface{}{&types.MatchDenseExpr{ + VectorColumnName: fmt.Sprintf("q_%d_vec", len(f64)), + EmbeddingData: f64, + EmbeddingDataType: "float", + DistanceType: "cosine", + TopN: topK, + ExtraOptions: map[string]interface{}{"similarity": 0.0}, + }}) + if err != nil { + return nil, err + } + hits := make([]nav.NavHit, 0, len(chunks)) + for _, c := range chunks { + h := nav.NavHit{ + Type: firstStringValue(c["type_kwd"]), + Name: firstStringValue(c["title_kwd"]), + DocID: firstStringValue(c["doc_id"]), + } + if sc, ok := c["_score"].(float64); ok { + h.Score = sc + } else if sc, ok := c["_score"].(float32); ok { + h.Score = float64(sc) + } + if ds, ok := c["doc_ids_kwd"].([]interface{}); ok { + for _, d := range ds { + if dd, ok := d.(string); ok { + h.DocIDs = append(h.DocIDs, dd) + } + } + } + hits = append(hits, h) + } + return hits, nil +} + +// UpsertDoc places one document summary into the nav tree. Minimal closed loop: +// deterministic placement (KNN find best cluster -> merge if sim>=0.80, else a +// new root-level cluster). No LLM, no split/rebalance, no cascade cleanup. +func (s *NavService) UpsertDoc(ctx context.Context, in nav.UpsertDocInput) error { + if strings.TrimSpace(in.Summary) == "" { + return nil + } + de, err := s.docEngine() + if err != nil { + return err + } + if s.embed == nil && len(in.Embedd) == 0 { + return fmt.Errorf("datasetnav: embedder required for UpsertDoc") + } + vec := in.Embedd + if len(vec) == 0 { + embeddings, err := s.embed.Encode(ctx, in.TenantID, []string{in.Summary}) + if err != nil { + return err + } + if len(embeddings) == 0 { + return nil + } + vec = embeddings[0] + } + + // storeGet: skip if a nav_doc for this doc already exists with same summary. + existing, _, err := s.navSearch(ctx, in.TenantID, in.KbID, + navFilter(map[string]interface{}{"doc_id": []string{in.DocID}}), + []string{"content_with_weight"}, 0, 1, nil) + if err != nil { + return err + } + if len(existing) > 0 { + if payload, ok := existing[0]["content_with_weight"].(string); ok { + var m map[string]interface{} + if err := json.Unmarshal([]byte(payload), &m); err == nil { + if d, _ := m["description"].(string); d == in.Summary { + return nil // unchanged + } + } + } + // Changed summary: remove the old nav_doc first (no cascade in minimal loop). + if _, err := s.deleteNavDoc(ctx, in.TenantID, in.KbID, in.DocID); err != nil { + return err + } + } + + bestName, sim, bestDepth, err := s.findBestCluster(ctx, in.TenantID, in.KbID, vec) + if err != nil { + return err + } + idx := s.navIndexName(in.TenantID) + + if bestName != "" && sim >= navMergeThreshold { + parent := bestName + if err := s.appendDocToCluster(ctx, de, in.TenantID, in.KbID, bestName, in.DocID); err != nil { + return err + } + _, err = de.InsertChunks(ctx, []map[string]interface{}{{ + "compile_kwd": navCompileKwd, + "available_int": 0, + "type_kwd": "nav_doc", + "title_kwd": in.DocID, + "parent_kwd": parent, + // The nav_doc sits one level below its (possibly nested) parent + // cluster, so its depth is parentDepth+1 — not a hard-coded 1. + "depth_int": bestDepth + 1, + "doc_id": in.DocID, + "doc_count_int": 1, + "content_with_weight": payloadJSONNav(map[string]interface{}{"type": "nav_doc", "description": in.Summary}), + "q_" + fmt.Sprintf("%d", len(vec)) + "_vec": f32ToF64Slice(vec), + }}, idx, in.KbID) + return err + } + + // A similar-but-not-mergeable cluster creates a sibling sub-cluster (Python + // _MIN_SIM=0.50); otherwise a fresh root cluster. This keeps the nav tree + // from degrading into one root per document. + parent := navRootParent + depth := 0 + if bestName != "" && sim >= navMinSim { + parent = bestName + // A sibling of the (possibly nested) best cluster is one level deeper + // than it, so depth = parentDepth+1 — not a hard-coded 1. + depth = bestDepth + 1 + } + name := navDocName(in.DocID, in.Summary) + _, err = de.InsertChunks(ctx, []map[string]interface{}{{ + "compile_kwd": navCompileKwd, + "available_int": 0, + "type_kwd": "nav_cluster", + "title_kwd": name, + "parent_kwd": parent, + "depth_int": depth, + "doc_count_int": 1, + "doc_ids_kwd": []string{in.DocID}, + "content_with_weight": payloadJSONNav(map[string]interface{}{"type": "nav_cluster", "description": in.Summary}), + "q_" + fmt.Sprintf("%d", len(vec)) + "_vec": f32ToF64Slice(vec), + }}, idx, in.KbID) + return err +} + +// findBestCluster finds the best-matching cluster via level-by-level descent +// (mirroring Python _find_best_cluster). It KNNs the current level's clusters +// and, when the best match is >= recurse threshold, descends into that cluster's +// children. Returns the best cluster name, similarity, and its depth (0 = root) +// so callers can assign consistent child depth_int values. +func (s *NavService) findBestCluster(ctx context.Context, tenantID, kbID string, vec []float32) (string, float64, int, error) { + f64 := f32ToF64Slice(vec) + parent := navRootParent + bestName := "" + bestSim := 0.0 + bestDepth := 0 + for level := 0; level < navMaxDepth; level++ { + // KNN among clusters whose parent is the current level. + chunks, _, err := s.navSearch(ctx, tenantID, kbID, + navFilter(map[string]interface{}{ + "type_kwd": []string{"nav_cluster"}, + "parent_kwd": []string{parent}, + }), + []string{"title_kwd", "_score"}, 0, 1, + []interface{}{&types.MatchDenseExpr{ + VectorColumnName: fmt.Sprintf("q_%d_vec", len(f64)), + EmbeddingData: f64, + EmbeddingDataType: "float", + DistanceType: "cosine", + TopN: 1, + ExtraOptions: map[string]interface{}{"similarity": 0.0}, + }}) + if err != nil { + return "", 0, 0, err + } + if len(chunks) == 0 { + break + } + name := firstStringValue(chunks[0]["title_kwd"]) + sim := rowScore(chunks[0]) + // Keep the STRONGEST match seen so far across all levels, so a strong + // ancestor is never displaced by a weaker descendant. Record its depth + // so the caller can set consistent child depth_int values. + if sim > bestSim { + bestName, bestSim, bestDepth = name, sim, level + } + // Descend only while the current match is strong enough that a deeper + // child could be a better target. + if sim < navRecurse { + break + } + parent = name + } + return bestName, bestSim, bestDepth, nil +} + +// rowScore extracts the engine's _score field as float64. +func rowScore(row map[string]interface{}) float64 { + if sc, ok := row["_score"].(float64); ok { + return sc + } + if sc, ok := row["_score"].(float32); ok { + return float64(sc) + } + return 0 +} + +// appendDocToCluster appends a doc id to a cluster's doc_ids_kwd and bumps its +// doc_count_int. Implemented as a read-modify-write. +func (s *NavService) appendDocToCluster(ctx context.Context, de engine.DocEngine, tenantID, kbID, clusterName, docID string) error { + chunks, _, err := s.navSearch(ctx, tenantID, kbID, + navFilter(map[string]interface{}{"type_kwd": []string{"nav_cluster"}, "title_kwd": []string{clusterName}}), + []string{"doc_ids_kwd", "doc_count_int"}, 0, 1, nil) + if err != nil { + return err + } + if len(chunks) == 0 { + return nil + } + ids := []string{} + if raw, ok := chunks[0]["doc_ids_kwd"].([]interface{}); ok { + for _, d := range raw { + if dd, ok := d.(string); ok { + ids = append(ids, dd) + } + } + } + found := false + for _, id := range ids { + if id == docID { + found = true + break + } + } + if !found { + ids = append(ids, docID) + } + count := intValue(chunks[0]["doc_count_int"]) + if !found { + count++ + } + // Pin the update to the nav_cluster row only: a regular chunk sharing the + // same title_kwd must never be clobbered. The read-modify-write here is + // expected to run under a per-dataset lock held by the UpsertDoc caller; + // without it, concurrent appends to the same cluster can lose updates. + return de.UpdateChunks(ctx, + map[string]interface{}{ + "compile_kwd": []string{navCompileKwd}, + "type_kwd": []string{"nav_cluster"}, + "title_kwd": []string{clusterName}, + "kb_id": kbID, + }, + map[string]interface{}{"doc_ids_kwd": ids, "doc_count_int": count}, + s.navIndexName(tenantID), kbID) +} + +// deleteNavDoc deletes a nav_doc row by doc_id. +func (s *NavService) deleteNavDoc(ctx context.Context, tenantID, kbID, docID string) (int64, error) { + de, err := s.docEngine() + if err != nil { + return 0, err + } + chunks, _, err := s.navSearch(ctx, tenantID, kbID, + navFilter(map[string]interface{}{"doc_id": []string{docID}}), + []string{"id"}, 0, 100, nil) + if err != nil { + return 0, err + } + ids := make([]string, 0, len(chunks)) + for _, c := range chunks { + if id, ok := c["id"].(string); ok { + ids = append(ids, id) + } + } + if len(ids) == 0 { + return 0, nil + } + return de.DeleteChunks(ctx, map[string]interface{}{"id": ids, "kb_id": kbID}, s.navIndexName(tenantID), kbID) +} + +// RemoveDoc removes a document's nav_doc (no cascade cleanup in minimal loop). +func (s *NavService) RemoveDoc(ctx context.Context, tenantID, kbID, docID string) error { + _, err := s.deleteNavDoc(ctx, tenantID, kbID, docID) + return err +} + +// f32ToF64Slice converts a float32 vector to float64. +func f32ToF64Slice(v []float32) []float64 { + out := make([]float64, len(v)) + for i, x := range v { + out[i] = float64(x) + } + return out +} + +// payloadJSONNav marshals a nav payload map into the content_with_weight JSON. +func payloadJSONNav(v map[string]interface{}) string { + b, err := json.Marshal(v) + if err != nil { + return "{}" + } + return string(b) +} + +// navDocName builds a deterministic cluster name for the minimal loop. +func navDocName(docID, summary string) string { + return docID + "_" + contentHash8(summary) +} + +// contentHash8 is a stable 8-char hash. +func contentHash8(s string) string { + h := uint32(2166136261) + for i := 0; i < len(s); i++ { + h ^= uint32(s[i]) + h *= 16777619 + } + return fmt.Sprintf("%08x", h) +} + +// firstStringValue returns the first string value of a (possibly list-wrapped) +// engine field. +func firstStringValue(v interface{}) string { + switch tv := v.(type) { + case string: + return tv + case []string: + if len(tv) > 0 { + return tv[0] + } + case []interface{}: + if len(tv) > 0 { + if s, ok := tv[0].(string); ok { + return s + } + } + } + return "" +} + +// intValue returns the integer value of an engine field. +func intValue(v interface{}) int { + switch tv := v.(type) { + case float64: + return int(tv) + case float32: + return int(tv) + case int: + return tv + case int64: + return int(tv) + case []float64: + if len(tv) > 0 { + return int(tv[0]) + } + case []interface{}: + if len(tv) > 0 { + switch n := tv[0].(type) { + case float64: + return int(n) + case int: + return n + } + } + } + return 0 +} diff --git a/internal/service/nlp/datasetnav_integration_test.go b/internal/service/nlp/datasetnav_integration_test.go new file mode 100644 index 0000000000..194bf36335 --- /dev/null +++ b/internal/service/nlp/datasetnav_integration_test.go @@ -0,0 +1,139 @@ +//go:build integration +// +build integration + +package nlp + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "ragflow/internal/engine" + "ragflow/internal/engine/types" + "ragflow/internal/server" + "ragflow/internal/service/nav" + + "go.uber.org/zap" +) + +// repoRootOf walks up from the package directory to the repository root (the +// dir containing go.mod). Kept local to this test package so it needs no shared +// helper from another package. +func repoRootOf(t *testing.T) string { + t.Helper() + dir, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + for { + if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil { + return dir + } + parent := filepath.Dir(dir) + if parent == dir { + t.Fatal("repository root (go.mod) not found above cwd") + } + dir = parent + } +} + +// findNavRow reads the nav row for a doc directly from the document engine. +func findNavRow(t *testing.T, tenantID, kbID, docID string) map[string]interface{} { + t.Helper() + de := engine.Get() + if de == nil { + t.Skip("no live document engine") + } + idx := "ragflow_" + tenantID + req := &types.SearchRequest{ + IndexNames: []string{idx}, + Filter: map[string]interface{}{"doc_id": []string{docID}, "compile_kwd": []string{"dataset_nav"}}, + SelectFields: []string{"available_int", "compile_kwd", "type_kwd"}, + Limit: 10, + } + res, err := de.Search(context.Background(), req) + if err != nil { + t.Fatalf("search nav row: %v", err) + } + for _, row := range res.Chunks { + return row + } + t.Fatalf("no nav row found for doc %s", docID) + return nil +} + +// TestDatasetNav_AvailableIntZero_Isolation is an integration test against a real +// ES/Infinity backend (requires conf/service_conf.yaml + a live document store). +// It verifies acceptance criterion #6: a nav row written with available_int=0 is +// invisible to the default retriever (which filters available_int=1) but IS +// reachable through NavService.Search. +// +// Run with: bash build.sh --test-integration ./internal/service/nlp/... +func TestDatasetNav_AvailableIntZero_Isolation(t *testing.T) { + server.SetLogger(zap.NewNop()) + configPath := filepath.Join(repoRootOf(t), "conf", "service_conf.yaml") + if err := server.Init(configPath); err != nil { + t.Fatalf("init service config: %v", err) + } + if err := engine.Init(); err != nil { + t.Fatalf("init document engine: %v", err) + } + if engine.Get() == nil { + t.Skip("no live document engine configured") + } + + tenantID := "navint_t1" + kbID := "navint_kb1" + docID := "navint_doc1" + + ns := NewNavService(stubNavEmbedder{}) + if err := ns.UpsertDoc(context.Background(), nav.UpsertDocInput{ + TenantID: tenantID, KbID: kbID, DocID: docID, Summary: "rocket propulsion integration evidence", + }); err != nil { + t.Fatalf("upsert nav doc: %v", err) + } + t.Cleanup(func() { _ = ns.RemoveDoc(context.Background(), tenantID, kbID, docID) }) + + // NavService.Search must find the nav row (reads nav rows directly). + hits, err := ns.Search(context.Background(), tenantID, kbID, "rocket propulsion", nil, 5) + if err != nil { + t.Fatalf("nav search: %v", err) + } + if len(hits) == 0 { + t.Fatal("NavService.Search returned no hits; nav row is not reachable") + } + + // The written nav row must carry compile_kwd=dataset_nav and available_int=0, + // so the default retriever (available_int=1 filter) will not surface it. + row := findNavRow(t, tenantID, kbID, docID) + // compile_kwd may come back list-wrapped by the engine, so use firstStrOrSlice. + if ck := firstStrOrSlice(row["compile_kwd"]); !strings.Contains(ck, "dataset_nav") { + t.Errorf("nav row compile_kwd = %q, want dataset_nav", ck) + } + avail := intValue(row["available_int"]) + if avail != 0 { + t.Errorf("nav row available_int = %d, want 0 (so it is hidden from the default retriever)", avail) + } +} + +// firstStrOrSlice returns the first string of a value that may be a plain string +// or a list-wrapped string (engine fields are often returned as []interface{}). +func firstStrOrSlice(v interface{}) string { + switch tv := v.(type) { + case string: + return tv + case []string: + if len(tv) > 0 { + return tv[0] + } + case []interface{}: + if len(tv) > 0 { + if s, ok := tv[0].(string); ok { + return s + } + } + } + return "" +} diff --git a/internal/service/nlp/datasetnav_test.go b/internal/service/nlp/datasetnav_test.go new file mode 100644 index 0000000000..bc4c52a548 --- /dev/null +++ b/internal/service/nlp/datasetnav_test.go @@ -0,0 +1,390 @@ +package nlp + +import ( + "context" + "math" + "testing" + + "ragflow/internal/engine/types" + "ragflow/internal/service/nav" + + "gorm.io/gorm" +) + +// memNavEngine is an in-memory DocEngine double sufficient for the datasetnav +// minimal closed loop. It stores rows in a slice and supports the filtered +// search / insert / update / delete operations NavService uses. Dense-vector +// KNN is approximated by a deterministic cosine over a synthetic q_<dim>_vec. +type memNavEngine struct { + rows []map[string]interface{} + nextID int +} + +func newMemNavEngine() *memNavEngine { return &memNavEngine{} } + +func (m *memNavEngine) InsertChunks(_ context.Context, chunks []map[string]interface{}, _ string, datasetID string) ([]string, error) { + ids := make([]string, 0, len(chunks)) + for _, c := range chunks { + m.nextID++ + id := "nav" + strconvItoa(m.nextID) + cp := make(map[string]interface{}, len(c)+2) + for k, v := range c { + cp[k] = v + } + cp["id"] = id + cp["kb_id"] = datasetID + m.rows = append(m.rows, cp) + ids = append(ids, id) + } + return ids, nil +} + +func (m *memNavEngine) UpdateChunks(_ context.Context, cond map[string]interface{}, newValue map[string]interface{}, _ string, _ string) error { + for _, r := range m.rows { + if matchNavRow(r, cond) { + for k, v := range newValue { + r[k] = v + } + } + } + return nil +} + +func (m *memNavEngine) DeleteChunks(_ context.Context, cond map[string]interface{}, _ string, _ string) (int64, error) { + out := m.rows[:0] + var deleted int64 + for _, r := range m.rows { + if matchNavRow(r, cond) { + deleted++ + continue + } + out = append(out, r) + } + m.rows = out + return deleted, nil +} + +func (m *memNavEngine) Search(_ context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + var matched []map[string]interface{} + hasDense := len(req.MatchExprs) > 0 + var queryVec []float64 + if hasDense { + if de, ok := req.MatchExprs[0].(*types.MatchDenseExpr); ok { + queryVec = de.EmbeddingData + } + } + for _, r := range m.rows { + if !matchNavRow(r, req.Filter) { + continue + } + if hasDense { + col := "q_" + strconvItoa(len(queryVec)) + "_vec" + rv, ok := r[col].([]float64) + if !ok { + continue + } + r["_score"] = cosineNav(rv, queryVec) + } + matched = append(matched, r) + } + if hasDense { + for i := 1; i < len(matched); i++ { + for j := i; j > 0 && scoreNavOf(matched[j]) > scoreNavOf(matched[j-1]); j-- { + matched[j], matched[j-1] = matched[j-1], matched[j] + } + } + } + offset, limit := req.Offset, req.Limit + if offset > len(matched) { + offset = len(matched) + } + end := offset + limit + if limit <= 0 || end > len(matched) { + end = len(matched) + } + return &types.SearchResult{Chunks: matched[offset:end], Total: int64(len(matched))}, nil +} + +func (m *memNavEngine) DropChunkStore(context.Context, string, string) error { return nil } +func (m *memNavEngine) ChunkStoreExists(context.Context, string, string) (bool, error) { + return true, nil +} +func (m *memNavEngine) Close() error { return nil } +func (m *memNavEngine) Ping(context.Context) error { return nil } +func (m *memNavEngine) GetType() string { return "mem" } +func (m *memNavEngine) SupportsPageRank() bool { return false } +func (m *memNavEngine) CreateChunkStore(context.Context, string, string, int, string) error { + return nil +} +func (m *memNavEngine) GetChunk(context.Context, string, string, []string) (interface{}, error) { + return nil, nil +} +func (m *memNavEngine) CreateMetadataStore(context.Context, string) error { return nil } +func (m *memNavEngine) InsertMetadata(context.Context, []map[string]interface{}, string) ([]string, error) { + return nil, nil +} +func (m *memNavEngine) UpdateMetadata(context.Context, string, string, map[string]interface{}, string) error { + return nil +} +func (m *memNavEngine) DeleteMetadata(context.Context, map[string]interface{}, string) (int64, error) { + return 0, nil +} +func (m *memNavEngine) DeleteMetadataKeys(context.Context, string, string, []string, string) error { + return nil +} +func (m *memNavEngine) DropMetadataStore(context.Context, string) error { return nil } +func (m *memNavEngine) MetadataStoreExists(context.Context, string) (bool, error) { + return true, nil +} +func (m *memNavEngine) SearchMetadata(context.Context, *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) { + return nil, nil +} +func (m *memNavEngine) IndexDocument(context.Context, string, string, interface{}) error { + return nil +} +func (m *memNavEngine) DeleteDocument(context.Context, string, string) error { return nil } +func (m *memNavEngine) BulkIndex(context.Context, string, []interface{}) (interface{}, error) { + return nil, nil +} +func (m *memNavEngine) GetFields([]map[string]interface{}, []string) map[string]map[string]interface{} { + return nil +} +func (m *memNavEngine) GetAggregation([]map[string]interface{}, string) []map[string]interface{} { + return nil +} +func (m *memNavEngine) GetHighlight([]map[string]interface{}, []string, string) map[string]string { + return nil +} +func (m *memNavEngine) RunSQL(context.Context, string, string, []string, string) ([]map[string]interface{}, error) { + return nil, nil +} +func (m *memNavEngine) GetChunkIDs([]map[string]interface{}) []string { return nil } +func (m *memNavEngine) KNNScores(context.Context, []map[string]interface{}, []float64, int) (map[string]interface{}, error) { + return nil, nil +} +func (m *memNavEngine) GetScores(map[string]interface{}) map[string]float64 { return nil } +func (m *memNavEngine) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { + return nil +} + +func matchNavRow(row map[string]interface{}, cond map[string]interface{}) bool { + for k, v := range cond { + rv, ok := row[k] + if !ok { + return false + } + switch want := v.(type) { + case []string: + ok = false + for _, w := range want { + if rv == w { + ok = true + break + } + } + if !ok { + return false + } + default: + if rv != v { + return false + } + } + } + return true +} + +func scoreNavOf(r map[string]interface{}) float64 { + switch s := r["_score"].(type) { + case float64: + return s + case float32: + return float64(s) + } + return 0 +} + +func cosineNav(a, b []float64) float64 { + if len(a) == 0 || len(b) == 0 || len(a) != len(b) { + return 0 + } + var dot, na, nb float64 + for i := range a { + dot += a[i] * b[i] + na += a[i] * a[i] + nb += b[i] * b[i] + } + if na == 0 || nb == 0 { + return 0 + } + return dot / (math.Sqrt(na) * math.Sqrt(nb)) +} + +func strconvItoa(n int) string { + if n == 0 { + return "0" + } + digits := []byte{} + for n > 0 { + digits = append([]byte{byte('0' + n%10)}, digits...) + n /= 10 + } + return string(digits) +} + +// stubNavEmbedder returns a fixed deterministic vector per distinct text. +type stubNavEmbedder struct{} + +func (stubNavEmbedder) Encode(_ context.Context, _ string, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, t := range texts { + dim := 8 + v := make([]float32, dim) + for d := 0; d < dim; d++ { + v[d] = float32(int(t[0]) + d) + } + out[i] = v + } + return out, nil +} + +func newTestNav(eng *memNavEngine) *NavService { + ns := NewNavService(stubNavEmbedder{}) + ns.engine = eng + return ns +} + +// TestNavService_UpsertDoc_WritesNavRow asserts acceptance #1: after UpsertDoc +// the row carries both compile_kwd=dataset_nav and available_int=0. +func TestNavService_UpsertDoc_WritesNavRow(t *testing.T) { + eng := newMemNavEngine() + ns := newTestNav(eng) + if err := ns.UpsertDoc(context.Background(), navUpsertInput("t1", "kb1", "d1", "alpha")); err != nil { + t.Fatalf("UpsertDoc: %v", err) + } + if len(eng.rows) == 0 { + t.Fatal("expected at least one nav row") + } + row := eng.rows[0] + if row["compile_kwd"] != "dataset_nav" { + t.Errorf("compile_kwd = %v, want dataset_nav", row["compile_kwd"]) + } + if row["available_int"] != 0 { + t.Errorf("available_int = %v, want 0", row["available_int"]) + } + if row["type_kwd"] != "nav_cluster" { + t.Errorf("type_kwd = %v, want nav_cluster", row["type_kwd"]) + } + if row["parent_kwd"] != "root" { + t.Errorf("parent_kwd = %v, want root", row["parent_kwd"]) + } +} + +// TestNavService_ListClusters_FiltersRoot asserts acceptance #3. +func TestNavService_ListClusters_FiltersRoot(t *testing.T) { + eng := newMemNavEngine() + ns := newTestNav(eng) + if err := ns.UpsertDoc(context.Background(), navUpsertInput("t1", "kb1", "d1", "aaa")); err != nil { + t.Fatal(err) + } + clusters, total, err := ns.ListClusters(context.Background(), "t1", "kb1", 0, 10) + if err != nil { + t.Fatal(err) + } + if total != 1 || len(clusters) != 1 { + t.Fatalf("expected 1 root cluster, got total=%d len=%d", total, len(clusters)) + } + if clusters[0].Type != "cluster" { + t.Errorf("cluster type = %s, want cluster", clusters[0].Type) + } + if clusters[0].DocCount < 1 { + t.Errorf("cluster doc_count = %d, want >=1", clusters[0].DocCount) + } +} + +// TestNavService_Search_ReturnsHit asserts acceptance #5. +func TestNavService_Search_ReturnsHit(t *testing.T) { + eng := newMemNavEngine() + ns := newTestNav(eng) + if err := ns.UpsertDoc(context.Background(), navUpsertInput("t1", "kb1", "d1", "aaa")); err != nil { + t.Fatal(err) + } + hits, err := ns.Search(context.Background(), "t1", "kb1", "aaa", nil, 5) + if err != nil { + t.Fatal(err) + } + if len(hits) == 0 { + t.Fatal("expected hits") + } + if hits[0].Name == "" { + t.Error("hit name empty") + } +} + +// TestNavService_Acceptance4_ListChildren asserts acceptance #4: ListChildren +// returns only the rows whose parent_kwd=name. Two docs with identical stub +// vectors merge into one root cluster: the first becomes the cluster itself, +// the second merges in as a nav_doc (parent_kwd=clusterName). So the cluster +// doc_count reflects both docs, and exactly one nav_doc sits under it. +func TestNavService_Acceptance4_ListChildren(t *testing.T) { + eng := newMemNavEngine() + ns := newTestNav(eng) + // Two docs that merge into one root cluster (same first char -> identical + // stub vectors -> sim=1.0 >= merge threshold). + if err := ns.UpsertDoc(context.Background(), navUpsertInput("t1", "kb1", "d1", "aaa one")); err != nil { + t.Fatal(err) + } + if err := ns.UpsertDoc(context.Background(), navUpsertInput("t1", "kb1", "d2", "aaa two")); err != nil { + t.Fatal(err) + } + clusters, _, err := ns.ListClusters(context.Background(), "t1", "kb1", 0, 10) + if err != nil { + t.Fatal(err) + } + if len(clusters) != 1 { + t.Fatalf("expected 1 cluster, got %d", len(clusters)) + } + if clusters[0].DocCount != 2 { + t.Fatalf("cluster doc_count = %d, want 2 (both docs merged into the cluster)", clusters[0].DocCount) + } + name := clusters[0].Name + children, total, err := ns.ListChildren(context.Background(), "t1", "kb1", name, 0, 10) + if err != nil { + t.Fatal(err) + } + // Exactly one nav_doc (for d2) sits under the cluster; d1 is the cluster. + if total != 1 || len(children) != 1 { + t.Fatalf("expected 1 child under cluster, got total=%d len=%d", total, len(children)) + } + if children[0].DocID != "d2" { + t.Errorf("child doc_id = %q, want d2", children[0].DocID) + } + if children[0].Type != "doc" { + t.Errorf("child type = %q, want doc", children[0].Type) + } +} + +// TestNavService_NavDocDepth asserts a nav_doc merged under a root cluster +// (depth 0) gets depth_int = parentDepth+1 = 1, not a hard-coded value. +func TestNavService_NavDocDepth(t *testing.T) { + eng := newMemNavEngine() + ns := newTestNav(eng) + if err := ns.UpsertDoc(context.Background(), navUpsertInput("t1", "kb1", "d1", "aaa one")); err != nil { + t.Fatal(err) + } + if err := ns.UpsertDoc(context.Background(), navUpsertInput("t1", "kb1", "d2", "aaa two")); err != nil { + t.Fatal(err) + } + // The nav_doc for d2 sits under the root cluster; its depth_int must be 1. + for _, row := range eng.rows { + if row["doc_id"] == "d2" { + if d, ok := row["depth_int"].(int); !ok || d != 1 { + t.Errorf("nav_doc d2 depth_int = %v, want 1 (parentDepth 0 + 1)", row["depth_int"]) + } + } + } +} + +func navUpsertInput(tenant, kb, doc, summary string) nav.UpsertDocInput { + return nav.UpsertDocInput{TenantID: tenant, KbID: kb, DocID: doc, Summary: summary} +}