Skip to content

Commit e3ddf7d

Browse files
xugangqiangclaude
andcommitted
refactor: KG search pipeline + tests + fix nil DAO panic
- Extract KGSearchPipeline struct from monolithic KGSearchRetrieval - Parallelize three independent searches (entity/types/relations) with sync.WaitGroup - Move scoring/formatting/types from internal/common/ to internal/service/ - Fix nil ModelDriver panic in buildSearchExprs - Fix searchKGTypeSamples error discarded (add common.Warn) - Add FilterChunksByScore for post-search _score filtering (Python alignment) - Fix Embed input to use matchText.MatchingText instead of question - Add Python alignment constants (simThreshold=0.3, denseTopK=1024) - Fix community report format (JSON parse, Evidences section, numbering) - Add contract validation to mocks (IndexNames, KbIDs) - Convert fakeEmbedDriver stub to spyEmbedDriver with captured inputs - Fix relation test chunks missing _score - Fix bare type assertion in TestBuildSearchExprs_NoEmbModel - Add searchCaptureEngine IndexNames validation - Fix nil userTenantDAO panic in RemoveMember/AcceptInvite - Recover accidentally truncated task_test.go - 80+ tests covering all paths Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 8ba64b2 commit e3ddf7d

8 files changed

Lines changed: 932 additions & 237 deletions

File tree

internal/service/kg_pipeline.go

Lines changed: 265 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,265 @@
1+
//
2+
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
//
16+
17+
package service
18+
19+
import (
20+
"context"
21+
"encoding/json"
22+
"fmt"
23+
"strings"
24+
"sync"
25+
26+
"ragflow/internal/common"
27+
"ragflow/internal/engine"
28+
"ragflow/internal/engine/types"
29+
modelModule "ragflow/internal/entity/models"
30+
"go.uber.org/zap"
31+
)
32+
33+
// KGSearchPipeline encapsulates the knowledge graph retrieval pipeline.
34+
// Matches Python: rag/graphrag/search.py::KGSearch
35+
type KGSearchPipeline struct {
36+
docEngine engine.DocEngine
37+
chatModel *modelModule.ChatModel
38+
embModel *modelModule.EmbeddingModel
39+
kbIDs []string
40+
idxnms []string
41+
question string
42+
43+
// Configurable parameters (defaults match Python)
44+
entSimThreshold float64
45+
relSimThreshold float64
46+
denseTopK int
47+
entTopN int
48+
relTopN int
49+
commTopN int
50+
maxToken int
51+
}
52+
53+
// KGSearchOption configures a KGSearchPipeline.
54+
type KGSearchOption func(*KGSearchPipeline)
55+
56+
// WithKGSimThreshold sets the similarity threshold for entity and relation search.
57+
// Default: 0.3 (matches Python ent_sim_threshold, rel_sim_threshold).
58+
func WithKGSimThreshold(v float64) KGSearchOption {
59+
return func(p *KGSearchPipeline) { p.entSimThreshold = v; p.relSimThreshold = v }
60+
}
61+
62+
// WithKGDenseTopK sets the TopK for dense vector search.
63+
// Default: 1024 (matches Python get_vector topk).
64+
func WithKGDenseTopK(v int) KGSearchOption {
65+
return func(p *KGSearchPipeline) { p.denseTopK = v }
66+
}
67+
68+
// NewKGSearchPipeline creates a KG search pipeline with the given dependencies.
69+
//
70+
// docEngine: search engine backend
71+
// kbIDs: knowledge base IDs to search
72+
// tenantIDs: tenant IDs (converted to index names internally)
73+
// question: user query string
74+
// opts: optional configuration (WithKGSimThreshold, WithKGDenseTopK)
75+
//
76+
// chatModel and embModel should be set via WithChatModel/WithEmbModel setters
77+
// or passed directly after construction.
78+
func NewKGSearchPipeline(
79+
docEngine engine.DocEngine,
80+
kbIDs []string,
81+
tenantIDs []string,
82+
question string,
83+
opts ...KGSearchOption,
84+
) *KGSearchPipeline {
85+
idxnms := make([]string, len(tenantIDs))
86+
for i, tid := range tenantIDs {
87+
idxnms[i] = indexName(tid)
88+
}
89+
p := &KGSearchPipeline{
90+
docEngine: docEngine,
91+
kbIDs: kbIDs,
92+
idxnms: idxnms,
93+
question: question,
94+
95+
entSimThreshold: defaultKGSimThreshold,
96+
relSimThreshold: defaultKGSimThreshold,
97+
denseTopK: defaultKGDenseTopK,
98+
entTopN: 6,
99+
relTopN: 6,
100+
commTopN: 1,
101+
maxToken: 8196,
102+
}
103+
for _, opt := range opts {
104+
opt(p)
105+
}
106+
return p
107+
}
108+
109+
// Retrieval runs the full KG retrieval pipeline and returns a synthetic chunk.
110+
func (p *KGSearchPipeline) Retrieval(ctx context.Context) (map[string]interface{}, error) {
111+
// 1. Query rewrite via LLM, or fall back to raw question
112+
ty2entsJSON := ""
113+
if p.chatModel != nil {
114+
typeSamples, err := searchKGTypeSamples(ctx, p.docEngine, p.idxnms, p.kbIDs)
115+
if err != nil {
116+
common.Warn("KG type samples search failed", zap.String("kbIDs", fmt.Sprint(p.kbIDs)))
117+
}
118+
if typeSamples == nil {
119+
typeSamples = make(map[string][]string)
120+
}
121+
data, _ := json.Marshal(typeSamples)
122+
ty2entsJSON = string(data)
123+
}
124+
typeKeywords, entities := queryRewrite(p.chatModel, p.question, ty2entsJSON)
125+
126+
// 2-4. Search entities, types, and relations in parallel
127+
// (mutually independent, can run concurrently for ~3x latency reduction)
128+
var (
129+
entsFromQuery map[string]*KGEntity
130+
entsFromTypes map[string]struct{}
131+
relsFromText map[Edge]*KGRelation
132+
entsErr error
133+
)
134+
var wg sync.WaitGroup
135+
wg.Add(3)
136+
go func() {
137+
defer wg.Done()
138+
entsReq := &types.SearchRequest{
139+
IndexNames: p.idxnms,
140+
KbIDs: p.kbIDs,
141+
SelectFields: []string{"entity_kwd", "entity_type_kwd", "rank_flt", "content_with_weight", "n_hop_with_weight"},
142+
Limit: 50,
143+
Filter: map[string]interface{}{"knowledge_graph_kwd": "entity"},
144+
}
145+
if len(entities) > 0 {
146+
entsReq.MatchExprs = buildSearchExprs(p.embModel, &types.MatchTextExpr{
147+
Fields: []string{"entity_kwd^10", "content_ltks^2"},
148+
MatchingText: strings.Join(entities, " "),
149+
TopN: 50,
150+
}, p.entSimThreshold, p.denseTopK)
151+
}
152+
entsResult, err := p.docEngine.Search(ctx, entsReq)
153+
if err != nil {
154+
entsErr = fmt.Errorf("KG entity search failed: %w", err)
155+
return
156+
}
157+
result := make(map[string]*KGEntity)
158+
for _, chunk := range FilterChunksByScore(entsResult.Chunks, p.entSimThreshold) {
159+
name, _ := chunk["entity_kwd"].(string)
160+
if name == "" {
161+
continue
162+
}
163+
e := kgEntityFromChunk(name, chunk)
164+
result[name] = &e
165+
}
166+
entsFromQuery = result
167+
}()
168+
go func() {
169+
defer wg.Done()
170+
typesReq := &types.SearchRequest{
171+
IndexNames: p.idxnms,
172+
KbIDs: p.kbIDs,
173+
SelectFields: []string{"entity_kwd", "entity_type_kwd"},
174+
Limit: 10000,
175+
Filter: map[string]interface{}{"knowledge_graph_kwd": "entity"},
176+
}
177+
if len(typeKeywords) > 0 {
178+
typeFilters := make([]interface{}, len(typeKeywords))
179+
for i, t := range typeKeywords {
180+
typeFilters[i] = t
181+
}
182+
typesReq.Filter["entity_type_kwd"] = typeFilters
183+
}
184+
typesResult, err := p.docEngine.Search(ctx, typesReq)
185+
result := make(map[string]struct{})
186+
if err != nil {
187+
common.Warn("KG types search failed", zap.String("kbIDs", fmt.Sprint(p.kbIDs)))
188+
} else {
189+
for _, chunk := range typesResult.Chunks {
190+
if name, ok := chunk["entity_kwd"].(string); ok {
191+
result[name] = struct{}{}
192+
}
193+
}
194+
}
195+
entsFromTypes = result
196+
}()
197+
go func() {
198+
defer wg.Done()
199+
relsReq := &types.SearchRequest{
200+
IndexNames: p.idxnms,
201+
KbIDs: p.kbIDs,
202+
SelectFields: []string{"from_entity_kwd", "to_entity_kwd", "weight_int", "content_with_weight"},
203+
Limit: 50,
204+
Filter: map[string]interface{}{"knowledge_graph_kwd": "relation"},
205+
}
206+
if len(entities) > 0 {
207+
relsReq.MatchExprs = buildSearchExprs(p.embModel, &types.MatchTextExpr{
208+
Fields: []string{"content_ltks", "from_entity_kwd", "to_entity_kwd"},
209+
MatchingText: strings.Join(entities, " "),
210+
TopN: 50,
211+
}, p.relSimThreshold, p.denseTopK)
212+
}
213+
relsResult, err := p.docEngine.Search(ctx, relsReq)
214+
result := make(map[Edge]*KGRelation)
215+
if err != nil {
216+
common.Warn("KG relations search failed", zap.String("kbIDs", fmt.Sprint(p.kbIDs)))
217+
} else {
218+
for _, chunk := range FilterChunksByScore(relsResult.Chunks, p.relSimThreshold) {
219+
edge, rel := kgRelationFromChunk(chunk)
220+
if edge.From == "" || edge.To == "" {
221+
continue
222+
}
223+
result[edge] = &rel
224+
}
225+
}
226+
relsFromText = result
227+
}()
228+
wg.Wait()
229+
if entsErr != nil {
230+
return nil, entsErr
231+
}
232+
233+
// 5. N-hop analysis + score fusion
234+
nhopPathes := AnalyzeNHopPaths(entsFromQuery)
235+
DoubleHitBoost(entsFromQuery, entsFromTypes)
236+
FuseRelationScores(relsFromText, entsFromTypes, nhopPathes)
237+
238+
// 6. Sort and trim
239+
scoredEnts := SortAndTrimEntities(entsFromQuery, p.entTopN)
240+
scoredRels := SortAndTrimRelations(relsFromText, p.relTopN)
241+
242+
// 7. Build KG content with token budget
243+
entsRelsContent := BuildKGContent(scoredEnts, scoredRels, p.maxToken)
244+
used := NumTokensFromString(entsRelsContent)
245+
remaining := p.maxToken - used
246+
// 8. Search community reports with remaining token budget
247+
communityContent := searchKGCommunityContent(ctx, p.docEngine, p.idxnms, p.kbIDs, scoredEnts, p.commTopN, &remaining)
248+
249+
// 9. Build synthetic chunk
250+
return map[string]interface{}{
251+
"chunk_id": "",
252+
"content_ltks": "",
253+
"content_with_weight": entsRelsContent + communityContent,
254+
"doc_id": "",
255+
"docnm_kwd": "Related content in Knowledge Graph",
256+
"kb_id": p.kbIDs,
257+
"important_kwd": []string{},
258+
"image_id": "",
259+
"similarity": 1.0,
260+
"vector_similarity": 1.0,
261+
"term_similarity": 0,
262+
"vector": []float64{},
263+
"positions": []interface{}{},
264+
}, nil
265+
}

0 commit comments

Comments
 (0)