diff --git a/.env.example b/.env.example index c9fedbaf..ebfc80c0 100644 --- a/.env.example +++ b/.env.example @@ -293,6 +293,9 @@ RETRIEVE_DRIVER=postgres # 知识图谱唯一全局开关。非 true 则禁用图谱构建与检索（构建阶段需调用大模型，耗时较长）。 # 注：ENABLE_GRAPH_RAG 自 v0.1.6 起已被 NEO4J_ENABLE 取代，Go 主应用不再读取，已移除。 # NEO4J_ENABLE=false +# 快速问答自动图谱检索默认关闭；仅在实体查询、Rerank、Merge容量验收后显式启用。 +# 智能推理的按需图谱工具不受此开关影响。 +# QUICK_ANSWER_GRAPH_RETRIEVAL_ENABLE=false # Neo4j 连接 URI。bolt:// 直连单机（推荐，无路由开销）；neo4j:// 走集群路由发现（单机也能用但多一次探测）。 # WeKnora 默认部署单机 Neo4j，故用 bolt://。 # NEO4J_URI=bolt://neo4j:7687 diff --git a/internal/application/repository/retriever/neo4j/repository.go b/internal/application/repository/retriever/neo4j/repository.go index 0c563ff1..a23004fa 100644 --- a/internal/application/repository/retriever/neo4j/repository.go +++ b/internal/application/repository/retriever/neo4j/repository.go @@ -17,6 +17,44 @@ type Neo4jRepository struct { nodePrefix string } +const ( + maxGraphSearchEntities = 8 + maxGraphSearchRows = 100 +) + +func normalizeSearchNodes(nodes []string) []string { + result := make([]string, 0, min(len(nodes), maxGraphSearchEntities)) + seen := make(map[string]struct{}, len(nodes)) + for _, node := range nodes { + node = strings.TrimSpace(node) + if node == "" { + continue + } + key := strings.ToLower(node) + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + result = append(result, node) + if len(result) == maxGraphSearchEntities { + break + } + } + return result +} + +func buildSearchNodeQuery(labelExpr string) string { + return ` + MATCH (n:` + labelExpr + `)-[r]-(m:` + labelExpr + `) + WHERE ANY(nodeText IN $nodes WHERE n.name CONTAINS nodeText) + WITH n, r, m, + ANY(nodeText IN $nodes WHERE n.name = nodeText) AS exact_match + ORDER BY exact_match DESC, n.name ASC, type(r) ASC, m.name ASC + LIMIT $row_limit + RETURN n, r, m + ` +} + // NewNeo4jRepository creates a new Neo4j repository func NewNeo4jRepository(driver neo4j.Driver) interfaces.RetrieveGraphRepository { return &Neo4jRepository{driver: driver, nodePrefix: "ENTITY"} @@ -170,17 +208,20 @@ func (n *Neo4jRepository) SearchNode( logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH") return nil, nil } + nodes = normalizeSearchNodes(nodes) + if len(nodes) == 0 { + return &types.GraphData{}, nil + } session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeRead}) defer session.Close(ctx) result, err := session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) { labelExpr := n.Label(namespace) - query := ` - MATCH (n:` + labelExpr + `)-[r]-(m:` + labelExpr + `) - WHERE ANY(nodeText IN $nodes WHERE n.name CONTAINS nodeText) - RETURN n, r, m - ` - params := map[string]interface{}{"nodes": nodes} + query := buildSearchNodeQuery(labelExpr) + params := map[string]interface{}{ + "nodes": nodes, + "row_limit": maxGraphSearchRows, + } result, err := tx.Run(ctx, query, params) if err != nil { return nil, fmt.Errorf("failed to run query: %v", err) diff --git a/internal/application/repository/retriever/neo4j/repository_search_test.go b/internal/application/repository/retriever/neo4j/repository_search_test.go new file mode 100644 index 00000000..ffe6ceb8 --- /dev/null +++ b/internal/application/repository/retriever/neo4j/repository_search_test.go @@ -0,0 +1,31 @@ +package neo4j + +import ( + "reflect" + "strings" + "testing" +) + +func TestNormalizeSearchNodesDropsBlankDeduplicatesAndCaps(t *testing.T) { + got := normalizeSearchNodes([]string{ + " CRA ", "", "cra", "控制柜", "control", "CONTROL", "A", "B", "C", "D", "E", "F", + }) + want := []string{"CRA", "控制柜", "control", "A", "B", "C", "D", "E"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("normalizeSearchNodes() = %#v, want %#v", got, want) + } +} + +func TestBuildSearchNodeQueryIsDeterministicAndBounded(t *testing.T) { + query := buildSearchNodeQuery("ENTITY_KB") + for _, required := range []string{ + "exact_match", + "ORDER BY exact_match DESC", + "LIMIT $row_limit", + "RETURN n, r, m", + } { + if !strings.Contains(query, required) { + t.Fatalf("query must contain %q:\n%s", required, query) + } + } +} diff --git a/internal/application/service/chat_pipeline/extract_entity.go b/internal/application/service/chat_pipeline/extract_entity.go index 233779a6..92aa6f9f 100644 --- a/internal/application/service/chat_pipeline/extract_entity.go +++ b/internal/application/service/chat_pipeline/extract_entity.go @@ -26,6 +26,14 @@ type PluginExtractEntity struct { knowledgeRepo interfaces.KnowledgeRepository } +// quickAnswerGraphRetrievalEnabled requires an explicit opt-in in addition to +// the global Neo4j switch. Smart-reasoning uses a separate on-demand Agent tool +// path and is not gated by this function. +func quickAnswerGraphRetrievalEnabled() bool { + return strings.EqualFold(strings.TrimSpace(os.Getenv("NEO4J_ENABLE")), "true") && + strings.EqualFold(strings.TrimSpace(os.Getenv("QUICK_ANSWER_GRAPH_RETRIEVAL_ENABLE")), "true") +} + // NewPluginExtractEntity creates a new extract-entity plugin instance // Also registers the plugin with the event manager func NewPluginExtractEntity( @@ -57,8 +65,8 @@ func (p *PluginExtractEntity) ActivationEvents() []types.EventType { func (p *PluginExtractEntity) OnEvent(ctx context.Context, eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError, ) *PluginError { - if strings.ToLower(os.Getenv("NEO4J_ENABLE")) != "true" { - logger.Debugf(ctx, "skipping extract entity, neo4j is disabled") + if !quickAnswerGraphRetrievalEnabled() { + logger.Debugf(ctx, "skipping quick-answer entity extraction, graph retrieval is not explicitly enabled") return next() } diff --git a/internal/application/service/chat_pipeline/extract_entity_graph_gate_test.go b/internal/application/service/chat_pipeline/extract_entity_graph_gate_test.go new file mode 100644 index 00000000..d6e4565d --- /dev/null +++ b/internal/application/service/chat_pipeline/extract_entity_graph_gate_test.go @@ -0,0 +1,30 @@ +package chatpipeline + +import "testing" + +func TestQuickAnswerGraphRetrievalDisabledUnlessExplicitlyEnabled(t *testing.T) { + t.Setenv("NEO4J_ENABLE", "true") + t.Setenv("QUICK_ANSWER_GRAPH_RETRIEVAL_ENABLE", "") + + if quickAnswerGraphRetrievalEnabled() { + t.Fatal("quick-answer graph retrieval must be disabled unless explicitly enabled") + } +} + +func TestQuickAnswerGraphRetrievalRequiresBothFlags(t *testing.T) { + t.Run("both enabled", func(t *testing.T) { + t.Setenv("NEO4J_ENABLE", "true") + t.Setenv("QUICK_ANSWER_GRAPH_RETRIEVAL_ENABLE", "true") + if !quickAnswerGraphRetrievalEnabled() { + t.Fatal("quick-answer graph retrieval should be enabled when both flags are true") + } + }) + + t.Run("global graph disabled", func(t *testing.T) { + t.Setenv("NEO4J_ENABLE", "false") + t.Setenv("QUICK_ANSWER_GRAPH_RETRIEVAL_ENABLE", "true") + if quickAnswerGraphRetrievalEnabled() { + t.Fatal("quick-answer graph retrieval must not bypass the global Neo4j switch") + } + }) +} diff --git a/internal/application/service/chat_pipeline/merge.go b/internal/application/service/chat_pipeline/merge.go index 82a6506d..3989f3b4 100644 --- a/internal/application/service/chat_pipeline/merge.go +++ b/internal/application/service/chat_pipeline/merge.go @@ -105,10 +105,21 @@ func (p *PluginMerge) selectInputResults(ctx context.Context, chatManage *types. pipelineWarn(ctx, "Merge", "fallback", map[string]interface{}{ "reason": "empty_rerank_result", }) - result := chatManage.SearchResult - sort.Slice(result, func(i, j int) bool { - return result[i].Score > result[j].Score + result := append([]*types.SearchResult(nil), chatManage.SearchResult...) + sort.SliceStable(result, func(i, j int) bool { + if result[i].Score != result[j].Score { + return result[i].Score > result[j].Score + } + return result[i].ID < result[j].ID }) + limit := max(1, chatManage.RerankTopK) + if len(result) > limit { + pipelineWarn(ctx, "Merge", "fallback_bounded", map[string]interface{}{ + "before": len(result), + "after": limit, + }) + result = result[:limit] + } return result } diff --git a/internal/application/service/chat_pipeline/merge_fallback_bounds_test.go b/internal/application/service/chat_pipeline/merge_fallback_bounds_test.go new file mode 100644 index 00000000..ff592e93 --- /dev/null +++ b/internal/application/service/chat_pipeline/merge_fallback_bounds_test.go @@ -0,0 +1,27 @@ +package chatpipeline + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func TestMergeFallbackInputIsSortedAndBounded(t *testing.T) { + chatManage := &types.ChatManage{ + PipelineRequest: types.PipelineRequest{RerankTopK: 2}, + PipelineState: types.PipelineState{SearchResult: []*types.SearchResult{ + {ID: "low", Score: 0.1}, + {ID: "high", Score: 0.9}, + {ID: "middle", Score: 0.5}, + }}, + } + + got := (&PluginMerge{}).selectInputResults(context.Background(), chatManage) + if len(got) != 2 { + t.Fatalf("merge fallback returned %d candidates, want 2", len(got)) + } + if got[0].ID != "high" || got[1].ID != "middle" { + t.Fatalf("merge fallback order = [%s %s], want [high middle]", got[0].ID, got[1].ID) + } +} diff --git a/internal/application/service/chat_pipeline/rerank.go b/internal/application/service/chat_pipeline/rerank.go index 8887bc53..a618c3f0 100644 --- a/internal/application/service/chat_pipeline/rerank.go +++ b/internal/application/service/chat_pipeline/rerank.go @@ -20,6 +20,16 @@ type PluginRerank struct { modelService interfaces.ModelService // Service to access rerank models } +func boundedRerankFallback(candidates []*types.SearchResult, topK int) []*types.SearchResult { + if topK <= 0 { + topK = 1 + } + if len(candidates) > topK { + candidates = candidates[:topK] + } + return append([]*types.SearchResult(nil), candidates...) +} + // NewPluginRerank creates a new rerank plugin instance func NewPluginRerank(eventManager *EventManager, modelService interfaces.ModelService) *PluginRerank { res := &PluginRerank{ @@ -134,7 +144,7 @@ func (p *PluginRerank) OnEvent(ctx context.Context, "error": rerankErr.Error(), "candidate_cnt": len(candidatesToRerank), }) - chatManage.SearchResult = candidatesToRerank + chatManage.SearchResult = boundedRerankFallback(candidatesToRerank, chatManage.RerankTopK) spanOutput = map[string]interface{}{ "stage": "api_error_fallback", "candidate_count": len(candidatesToRerank), @@ -166,7 +176,7 @@ func (p *PluginRerank) OnEvent(ctx context.Context, "error": rerankErr.Error(), "candidate_cnt": len(candidatesToRerank), }) - chatManage.SearchResult = candidatesToRerank + chatManage.SearchResult = boundedRerankFallback(candidatesToRerank, chatManage.RerankTopK) spanOutput = map[string]interface{}{ "stage": "api_error_fallback", "candidate_count": len(candidatesToRerank), diff --git a/internal/application/service/chat_pipeline/rerank_fallback_bounds_test.go b/internal/application/service/chat_pipeline/rerank_fallback_bounds_test.go new file mode 100644 index 00000000..e0f15543 --- /dev/null +++ b/internal/application/service/chat_pipeline/rerank_fallback_bounds_test.go @@ -0,0 +1,22 @@ +package chatpipeline + +import ( + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func TestBoundedRerankFallbackCapsCandidates(t *testing.T) { + candidates := []*types.SearchResult{{ID: "first"}, {ID: "second"}, {ID: "third"}} + + got := boundedRerankFallback(candidates, 2) + want := []string{"first", "second"} + if len(got) != len(want) { + t.Fatalf("fallback returned %d results, want %d", len(got), len(want)) + } + for i, id := range want { + if got[i].ID != id { + t.Fatalf("result[%d] = %q, want %q", i, got[i].ID, id) + } + } +} diff --git a/internal/application/service/chat_pipeline/search_entity.go b/internal/application/service/chat_pipeline/search_entity.go index 90415e3a..a3fca302 100644 --- a/internal/application/service/chat_pipeline/search_entity.go +++ b/internal/application/service/chat_pipeline/search_entity.go @@ -17,6 +17,8 @@ type PluginSearchEntity struct { knowledgeRepo interfaces.KnowledgeRepository } +const maxGraphChunkResults = 10 + // NewPluginSearchEntity creates a new plugin search entity func NewPluginSearchEntity( eventManager *EventManager, @@ -200,6 +202,12 @@ func filterSeenChunk(ctx context.Context, graph *types.GraphData, searchResult [ } seen[chunkID] = true chunkIDs = append(chunkIDs, chunkID) + if len(chunkIDs) == maxGraphChunkResults { + break + } + } + if len(chunkIDs) == maxGraphChunkResults { + break } } logger.Infof(ctx, "filterSeenChunk: new chunkIDs count: %d", len(chunkIDs)) diff --git a/internal/application/service/chat_pipeline/search_entity_bounds_test.go b/internal/application/service/chat_pipeline/search_entity_bounds_test.go new file mode 100644 index 00000000..1cab2569 --- /dev/null +++ b/internal/application/service/chat_pipeline/search_entity_bounds_test.go @@ -0,0 +1,27 @@ +package chatpipeline + +import ( + "context" + "fmt" + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func TestFilterSeenChunkCapsGraphChunkIDs(t *testing.T) { + chunkIDs := make([]string, 12) + for i := range chunkIDs { + chunkIDs[i] = fmt.Sprintf("chunk-%02d", i) + } + graph := &types.GraphData{Node: []*types.GraphNode{{Name: "hub", Chunks: chunkIDs}}} + + got := filterSeenChunk(context.Background(), graph, nil) + if len(got) != 10 { + t.Fatalf("filterSeenChunk returned %d chunk IDs, want 10", len(got)) + } + for i, id := range got { + if id != chunkIDs[i] { + t.Fatalf("result[%d] = %q, want %q", i, id, chunkIDs[i]) + } + } +}