mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-07 08:01:13 +08:00
Port Python agentic search to Go (nav service, harness, tools) (#17702)
Port Python rag/advanced_rag agentic search to Go: ES-backed dataset-nav service, agentic-search harness, and agent tools. Includes agentic-search port plan and self-review docs.
This commit is contained in:
Zhichang Yu
110
internal/agent/chat/chat.go
Normal file
110
internal/agent/chat/chat.go
Normal file
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
107
internal/agent/harness/agentic_rag.go
Normal file
107
internal/agent/harness/agentic_rag.go
Normal file
@@ -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
|
||||
}
|
||||
124
internal/agent/harness/answer.go
Normal file
124
internal/agent/harness/answer.go
Normal file
@@ -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 ""
|
||||
}
|
||||
71
internal/agent/harness/answer_test.go
Normal file
71
internal/agent/harness/answer_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
225
internal/agent/harness/datasetnav.go
Normal file
225
internal/agent/harness/datasetnav.go
Normal file
@@ -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": [<index>, ...]}`
|
||||
|
||||
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
|
||||
}
|
||||
102
internal/agent/harness/datasetnav_test.go
Normal file
102
internal/agent/harness/datasetnav_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
281
internal/agent/harness/navigation.go
Normal file
281
internal/agent/harness/navigation.go
Normal file
@@ -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": "<answer, or empty>", "relevant_entities": ["<entity name>", ...]}`
|
||||
|
||||
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
|
||||
}
|
||||
65
internal/agent/harness/navigation_test.go
Normal file
65
internal/agent/harness/navigation_test.go
Normal file
@@ -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
|
||||
}
|
||||
247
internal/agent/harness/orchestrator.go
Normal file
247
internal/agent/harness/orchestrator.go
Normal file
@@ -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
|
||||
}
|
||||
135
internal/agent/harness/orchestrator_test.go
Normal file
135
internal/agent/harness/orchestrator_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
164
internal/agent/harness/planner.go
Normal file
164
internal/agent/harness/planner.go
Normal file
@@ -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
|
||||
}
|
||||
147
internal/agent/harness/production.go
Normal file
147
internal/agent/harness/production.go
Normal file
@@ -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)
|
||||
}
|
||||
144
internal/agent/harness/production_test.go
Normal file
144
internal/agent/harness/production_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
133
internal/agent/harness/route.go
Normal file
133
internal/agent/harness/route.go
Normal file
@@ -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)^.*</think>`)
|
||||
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)
|
||||
}
|
||||
140
internal/agent/harness/route_test.go
Normal file
140
internal/agent/harness/route_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
250
internal/agent/harness/sufficiency.go
Normal file
250
internal/agent/harness/sufficiency.go
Normal file
@@ -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
|
||||
}
|
||||
}
|
||||
151
internal/agent/harness/types.go
Normal file
151
internal/agent/harness/types.go
Normal file
@@ -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
|
||||
}
|
||||
345
internal/agent/tool/agentic_search.go
Normal file
345
internal/agent/tool/agentic_search.go
Normal file
@@ -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 ".<digit>" 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 <em>.
|
||||
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, "<em>${1}</em>")
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
78
internal/agent/tool/agentic_search_test.go
Normal file
78
internal/agent/tool/agentic_search_test.go
Normal file
@@ -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, "<em>rocket</em>") {
|
||||
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)
|
||||
}
|
||||
}
|
||||
59
internal/agent/tool/canvas_ctx.go
Normal file
59
internal/agent/tool/canvas_ctx.go
Normal file
@@ -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
|
||||
}
|
||||
220
internal/agent/tool/dataset_navigation.go
Normal file
220
internal/agent/tool/dataset_navigation.go
Normal file
@@ -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)
|
||||
}
|
||||
134
internal/agent/tool/dataset_navigation_test.go
Normal file
134
internal/agent/tool/dataset_navigation_test.go
Normal file
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user