feat: add KGSearchRetrieval for full KG pipeline (N-hop, scoring, query_rewrite, community) (#15690)

## Summary `KGSearchRetrieval` composes entity search, type search, relation search, N-hop analysis, score fusion, LLM-based query\_rewrite, and community reports into a single synthetic chunk for KG-enhanced retrieval. ### Components | Component | Source | Status | |-----------|--------|--------| | Entity/relation/community search | Direct `DocEngine.Search` calls | ✅ | | N-hop analysis + score fusion | `common.AnalyzeNHopPaths` / `DoubleHitBoost` / `FuseRelationScores` | ✅ #15666 | | Query rewrite prompt + parser | `common.BuildQueryRewritePrompt` / `ParseQueryRewriteResponse` | ✅ #15669 | | Token budget | `common.BuildKGContent` + `NumTokensFromString` | ✅ #15666 | | LLM query rewrite integration | `queryRewrite` function with fallback | ✅ | ### Testing 11 tests (pure function + mock engine): ``` === RUN TestKgEntityFromChunk_Basic --- PASS === RUN TestKgEntityFromChunk_ScoreFallback --- PASS === RUN TestKgEntityFromChunk_MissingFields --- PASS === RUN TestKgRelationFromChunk_Basic --- PASS === RUN TestKgRelationFromChunk_MissingFrom --- PASS === RUN TestSearchKGTypeSamples_Success --- PASS === RUN TestSearchKGTypeSamples_Empty --- PASS === RUN TestKGSearchRetrieval_Basic --- PASS === RUN TestKGSearchRetrieval_NoEntities --- PASS === RUN TestQueryRewrite_Fallback --- PASS === RUN TestQueryRewrite_EmptyQuestion --- PASS ``` --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-23 08:56:42 +08:00 · 2026-06-05 18:00:27 +08:00
parent aa9545e4c9
commit ea79d65d08
8 changed files with 1375 additions and 125 deletions
--- a/internal/service/kg_pipeline.go
+++ b/internal/service/kg_pipeline.go
@@ -0,0 +1,276 @@
+//
+//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"sync"
+
+	"ragflow/internal/common"
+	"ragflow/internal/engine"
+	"ragflow/internal/engine/types"
+	modelModule "ragflow/internal/entity/models"
+	"go.uber.org/zap"
+)
+
+// KGSearchPipeline encapsulates the knowledge graph retrieval pipeline.
+// Matches Python: rag/graphrag/search.py::KGSearch
+type KGSearchPipeline struct {
+	docEngine engine.DocEngine
+	chatModel *modelModule.ChatModel
+	embModel  *modelModule.EmbeddingModel
+	kbIDs     []string
+	idxnms    []string
+	question  string
+
+	// Configurable parameters (defaults match Python)
+	entSimThreshold float64
+	relSimThreshold float64
+	denseTopK       int
+	entTopN         int
+	relTopN         int
+	commTopN        int
+	maxToken        int
+}
+
+// KGSearchOption configures a KGSearchPipeline.
+type KGSearchOption func(*KGSearchPipeline)
+
+// WithKGSimThreshold sets the similarity threshold for entity and relation search.
+// Default: 0.3 (matches Python ent_sim_threshold, rel_sim_threshold).
+func WithKGSimThreshold(v float64) KGSearchOption {
+	return func(p *KGSearchPipeline) { p.entSimThreshold = v; p.relSimThreshold = v }
+}
+
+// WithKGDenseTopK sets the TopK for dense vector search.
+// Default: 1024 (matches Python get_vector topk).
+func WithKGDenseTopK(v int) KGSearchOption {
+	return func(p *KGSearchPipeline) { p.denseTopK = v }
+}
+
+// NewKGSearchPipeline creates a KG search pipeline with the given dependencies.
+//
+//	docEngine: search engine backend
+//	kbIDs:     knowledge base IDs to search
+//	tenantIDs: tenant IDs (converted to index names internally)
+//	question:  user query string
+//	opts:      optional configuration (WithKGSimThreshold, WithKGDenseTopK)
+//
+// chatModel and embModel should be set via WithChatModel/WithEmbModel setters
+// or passed directly after construction.
+func NewKGSearchPipeline(
+	docEngine engine.DocEngine,
+	kbIDs []string,
+	tenantIDs []string,
+	question string,
+	opts ...KGSearchOption,
+) *KGSearchPipeline {
+	idxnms := make([]string, len(tenantIDs))
+	for i, tid := range tenantIDs {
+		idxnms[i] = indexName(tid)
+	}
+	p := &KGSearchPipeline{
+		docEngine: docEngine,
+		kbIDs:     kbIDs,
+		idxnms:    idxnms,
+		question:  question,
+
+		entSimThreshold: defaultKGSimThreshold,
+		relSimThreshold: defaultKGSimThreshold,
+		denseTopK:       defaultKGDenseTopK,
+		entTopN:         6,
+		relTopN:         6,
+		commTopN:        1,
+		maxToken:        8196,
+	}
+	for _, opt := range opts {
+		opt(p)
+	}
+	return p
+}
+
+// Retrieval runs the full KG retrieval pipeline and returns a synthetic chunk.
+func (p *KGSearchPipeline) Retrieval(ctx context.Context) (map[string]interface{}, error) {
+	// 1. Query rewrite via LLM, or fall back to raw question
+	ty2entsJSON := ""
+	if p.chatModel != nil {
+		typeSamples, err := searchKGTypeSamples(ctx, p.docEngine, p.idxnms, p.kbIDs)
+		if err != nil {
+			common.Warn("KG type samples search failed", zap.String("kbIDs", fmt.Sprint(p.kbIDs)))
+		}
+		if typeSamples == nil {
+			typeSamples = make(map[string][]string)
+		}
+		data, _ := json.Marshal(typeSamples)
+		ty2entsJSON = string(data)
+	}
+	typeKeywords, entities := queryRewrite(p.chatModel, p.question, ty2entsJSON)
+
+	// 2-4. Search entities, types, and relations in parallel (mutually independent)
+	var (
+		entsFromQuery map[string]*KGEntity
+		entsFromTypes map[string]struct{}
+		relsFromText  map[Edge]*KGRelation
+		entsErr       error
+	)
+	var wg sync.WaitGroup
+	wg.Add(3)
+	go func() {
+		defer wg.Done()
+		entsFromQuery, entsErr = p.searchEntities(ctx, entities)
+	}()
+	go func() {
+		defer wg.Done()
+		entsFromTypes = p.searchEntityTypes(ctx, typeKeywords)
+	}()
+	go func() {
+		defer wg.Done()
+		relsFromText = p.searchRelations(ctx, entities)
+	}()
+	wg.Wait()
+	if entsErr != nil {
+		return nil, entsErr
+	}// 5. N-hop analysis + score fusion
+	nhopPathes := AnalyzeNHopPaths(entsFromQuery)
+	DoubleHitBoost(entsFromQuery, entsFromTypes)
+	FuseRelationScores(relsFromText, entsFromTypes, nhopPathes)
+
+	// 6. Sort and trim
+	scoredEnts := SortAndTrimEntities(entsFromQuery, p.entTopN)
+	scoredRels := SortAndTrimRelations(relsFromText, p.relTopN)
+
+	// 7. Build KG content with token budget
+	entsRelsContent := BuildKGContent(scoredEnts, scoredRels, p.maxToken)
+	used := NumTokensFromString(entsRelsContent)
+	remaining := p.maxToken - used
+	// 8. Search community reports with remaining token budget
+	communityContent := searchKGCommunityContent(ctx, p.docEngine, p.idxnms, p.kbIDs, scoredEnts, p.commTopN, &remaining)
+
+	// 9. Build synthetic chunk
+	return map[string]interface{}{
+		"chunk_id":              "",
+		"content_ltks":          "",
+		"content_with_weight":   entsRelsContent + communityContent,
+		"doc_id":                "",
+		"docnm_kwd":             "Related content in Knowledge Graph",
+		"kb_id":                 p.kbIDs,
+		"important_kwd":         []string{},
+		"image_id":              "",
+		"similarity":            1.0,
+		"vector_similarity":     1.0,
+		"term_similarity":       0,
+		"vector":                []float64{},
+		"positions":             []interface{}{},
+	}, nil
+}
+
+// searchEntities searches KG entities by keyword text and optional dense vector.
+func (p *KGSearchPipeline) searchEntities(ctx context.Context, entities []string) (map[string]*KGEntity, error) {
+	entsReq := &types.SearchRequest{
+		IndexNames:   p.idxnms,
+		KbIDs:        p.kbIDs,
+		SelectFields: []string{"entity_kwd", "entity_type_kwd", "rank_flt", "content_with_weight", "n_hop_with_weight"},
+		Limit:        50,
+		Filter:       map[string]interface{}{"knowledge_graph_kwd": "entity"},
+	}
+	if len(entities) > 0 {
+		entsReq.MatchExprs = buildSearchExprs(p.embModel, &types.MatchTextExpr{
+			Fields:       []string{"entity_kwd^10", "content_ltks^2"},
+			MatchingText: strings.Join(entities, " "),
+			TopN:         50,
+		}, p.entSimThreshold, p.denseTopK)
+	}
+	entsResult, err := p.docEngine.Search(ctx, entsReq)
+	if err != nil {
+		return nil, fmt.Errorf("KG entity search failed: %w", err)
+	}
+	result := make(map[string]*KGEntity)
+	for _, chunk := range FilterChunksByScore(entsResult.Chunks, p.entSimThreshold) {
+		name, _ := chunk["entity_kwd"].(string)
+		if name == "" {
+			continue
+		}
+		e := kgEntityFromChunk(name, chunk)
+		result[name] = &e
+	}
+	return result, nil
+}
+
+// searchEntityTypes searches KG entities by type keywords.
+func (p *KGSearchPipeline) searchEntityTypes(ctx context.Context, typeKeywords []string) map[string]struct{} {
+	typesReq := &types.SearchRequest{
+		IndexNames:   p.idxnms,
+		KbIDs:        p.kbIDs,
+		SelectFields: []string{"entity_kwd", "entity_type_kwd"},
+		Limit:        10000,
+		Filter:       map[string]interface{}{"knowledge_graph_kwd": "entity"},
+	}
+	if len(typeKeywords) > 0 {
+		typeFilters := make([]interface{}, len(typeKeywords))
+		for i, t := range typeKeywords {
+			typeFilters[i] = t
+		}
+		typesReq.Filter["entity_type_kwd"] = typeFilters
+	}
+	typesResult, err := p.docEngine.Search(ctx, typesReq)
+	result := make(map[string]struct{})
+	if err != nil {
+		common.Warn("KG types search failed", zap.String("kbIDs", fmt.Sprint(p.kbIDs)))
+	} else {
+		for _, chunk := range typesResult.Chunks {
+			if name, ok := chunk["entity_kwd"].(string); ok {
+				result[name] = struct{}{}
+			}
+		}
+	}
+	return result
+}
+
+// searchRelations searches KG relations by entity text and optional dense vector.
+func (p *KGSearchPipeline) searchRelations(ctx context.Context, entities []string) map[Edge]*KGRelation {
+	relsReq := &types.SearchRequest{
+		IndexNames:   p.idxnms,
+		KbIDs:        p.kbIDs,
+		SelectFields: []string{"from_entity_kwd", "to_entity_kwd", "weight_int", "content_with_weight"},
+		Limit:        50,
+		Filter:       map[string]interface{}{"knowledge_graph_kwd": "relation"},
+	}
+	if len(entities) > 0 {
+		relsReq.MatchExprs = buildSearchExprs(p.embModel, &types.MatchTextExpr{
+			Fields:       []string{"content_ltks", "from_entity_kwd", "to_entity_kwd"},
+			MatchingText: strings.Join(entities, " "),
+			TopN:         50,
+		}, p.relSimThreshold, p.denseTopK)
+	}
+	relsResult, err := p.docEngine.Search(ctx, relsReq)
+	result := make(map[Edge]*KGRelation)
+	if err != nil {
+		common.Warn("KG relations search failed", zap.String("kbIDs", fmt.Sprint(p.kbIDs)))
+	} else {
+		for _, chunk := range FilterChunksByScore(relsResult.Chunks, p.relSimThreshold) {
+			edge, rel := kgRelationFromChunk(chunk)
+			if edge.From == "" || edge.To == "" {
+				continue
+			}
+			result[edge] = &rel
+		}
+	}
+	return result
+}
--- a/internal/service/kg_retrieval.go
+++ b/internal/service/kg_retrieval.go
@@ -0,0 +1,281 @@
+//
+//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"ragflow/internal/common"
+	"ragflow/internal/engine"
+	"ragflow/internal/engine/types"
+	modelModule "ragflow/internal/entity/models"
+)
+
+// indexName builds the search index name from a tenant ID.
+// Matches Python: rag/nlp/search.py::index_name()
+func indexName(tenantID string) string {
+	return "ragflow_" + tenantID
+}
+
+// Python alignment defaults — match rag/graphrag/search.py retrieval() params
+const (
+	defaultKGSimThreshold = 0.3  // Python: ent_sim_threshold, rel_sim_threshold
+	defaultKGDenseTopK    = 1024 // Python: get_vector() topk
+)
+
+// kgEntityFromChunk parses a single entity chunk into a KGEntity.
+func kgEntityFromChunk(name string, chunk map[string]interface{}) KGEntity {
+	e := KGEntity{}
+	if v, ok := chunk["_score"].(float64); ok {
+		e.Similarity = v
+	} else if v, ok := chunk["score"].(float64); ok {
+		e.Similarity = v
+	}
+	if v, ok := chunk["rank_flt"].(float64); ok {
+		e.PageRank = v
+	}
+	e.Description, _ = chunk["content_with_weight"].(string)
+	if raw, ok := chunk["n_hop_with_weight"].(string); ok && raw != "" {
+		var nhopData []struct {
+			Path    []string  `json:"path"`
+			Weights []float64 `json:"weights"`
+		}
+		if err := json.Unmarshal([]byte(raw), &nhopData); err == nil {
+			for _, item := range nhopData {
+				e.NhopEnts = append(e.NhopEnts, NhopEntity{
+					Path:    item.Path,
+					Weights: item.Weights,
+				})
+			}
+		}
+	}
+	return e
+}
+
+// kgRelationFromChunk parses a single relation chunk into a KGRelation.
+func kgRelationFromChunk(chunk map[string]interface{}) (Edge, KGRelation) {
+	r := KGRelation{}
+	r.Description, _ = chunk["content_with_weight"].(string)
+	if v, ok := chunk["weight_int"].(float64); ok {
+		r.PageRank = float64(v)
+	} else if v, ok := chunk["weight_int"].(int); ok {
+		r.PageRank = float64(v)
+	}
+	from, _ := chunk["from_entity_kwd"].(string)
+	to, _ := chunk["to_entity_kwd"].(string)
+	return Edge{From: from, To: to}, r
+}
+
+// KGSearchRetrieval performs a full knowledge graph retrieval and returns
+// a synthetic chunk to be inserted into search results.
+// Corresponds to Python: rag/graphrag/search.py::KGSearch.retrieval()
+//
+// This is a convenience wrapper around KGSearchPipeline.
+func KGSearchRetrieval(
+	ctx context.Context,
+	docEngine engine.DocEngine,
+	chatModel *modelModule.ChatModel,
+	embModel *modelModule.EmbeddingModel,
+	kbIDs []string,
+	tenantIDs []string,
+	question string,
+) (map[string]interface{}, error) {
+	p := &KGSearchPipeline{
+		docEngine:       docEngine,
+		chatModel:       chatModel,
+		embModel:        embModel,
+		kbIDs:           kbIDs,
+		idxnms:          makeIndexNames(tenantIDs),
+		question:        question,
+		entSimThreshold: defaultKGSimThreshold,
+		relSimThreshold: defaultKGSimThreshold,
+		denseTopK:       defaultKGDenseTopK,
+		entTopN:         6,
+		relTopN:         6,
+		commTopN:        1,
+		maxToken:        8196,
+	}
+	return p.Retrieval(ctx)
+}
+
+// makeIndexNames converts tenant IDs to search index names.
+func makeIndexNames(tenantIDs []string) []string {
+	idxnms := make([]string, len(tenantIDs))
+	for i, tid := range tenantIDs {
+		idxnms[i] = indexName(tid)
+	}
+	return idxnms
+}
+
+// searchKGTypeSamples searches for ty2ents data.
+func searchKGTypeSamples(ctx context.Context, docEngine engine.DocEngine, idxnms []string, kbIDs []string) (map[string][]string, error) {
+	req := &types.SearchRequest{
+		IndexNames:   idxnms,
+		KbIDs:        kbIDs,
+		SelectFields: []string{"content_with_weight"},
+		Limit:        10000,
+		Filter:       map[string]interface{}{"knowledge_graph_kwd": "ty2ents"},
+	}
+	result, err := docEngine.Search(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	typeMap := make(map[string][]string)
+	for _, chunk := range result.Chunks {
+		content, ok := chunk["content_with_weight"].(string)
+		if !ok || content == "" {
+			continue
+		}
+		var parsed map[string][]string
+		if err := json.Unmarshal([]byte(content), &parsed); err != nil {
+			continue
+		}
+		for typ, entities := range parsed {
+			typeMap[typ] = append(typeMap[typ], entities...)
+		}
+	}
+	return typeMap, nil
+}
+
+// searchKGCommunityContent searches for community reports and formats them.
+func searchKGCommunityContent(ctx context.Context, docEngine engine.DocEngine, idxnms []string, kbIDs []string, scoredEnts []ScoredEntity, topN int, maxToken *int) string {
+	if maxToken == nil || len(scoredEnts) == 0 || *maxToken <= 0 {
+		return ""
+	}
+	entityNames := make([]string, len(scoredEnts))
+	for i, e := range scoredEnts {
+		entityNames[i] = e.Entity
+	}
+	req := &types.SearchRequest{
+		IndexNames:   idxnms,
+		KbIDs:        kbIDs,
+		SelectFields: []string{"docnm_kwd", "content_with_weight", "weight_flt", "entities_kwd"},
+		Limit:        topN,
+		Filter:       map[string]interface{}{"knowledge_graph_kwd": "community_report"},
+		OrderBy:      (&types.OrderByExpr{}).Desc("weight_flt"),
+	}
+	if len(entityNames) > 0 {
+		filters := make([]interface{}, len(entityNames))
+		for i, name := range entityNames {
+			filters[i] = name
+		}
+		req.Filter["entities_kwd"] = filters
+	}
+	result, err := docEngine.Search(ctx, req)
+	if err != nil || len(result.Chunks) == 0 || *maxToken <= 0 {
+		return ""
+	}
+
+	var bld string
+	for idx, chunk := range result.Chunks {
+		title, _ := chunk["docnm_kwd"].(string)
+		raw, _ := chunk["content_with_weight"].(string)
+		if title == "" && raw == "" {
+			continue
+		}
+		// Parse JSON for nested report/evidences fields (Python: json.loads)
+		report := raw
+		evidence := ""
+		var parsed map[string]interface{}
+		if err := json.Unmarshal([]byte(raw), &parsed); err == nil {
+			if r, ok := parsed["report"].(string); ok {
+				report = r
+			}
+			if e, ok := parsed["evidences"].(string); ok {
+				evidence = e
+			}
+		}
+		section := fmt.Sprintf("\n# %d. %s\n## Content\n%s\n## Evidences\n%s\n", idx+1, title, report, evidence)
+		tokens := NumTokensFromString(section)
+		if *maxToken-tokens <= 0 {
+			break
+		}
+		bld += section
+		*maxToken -= tokens
+	}
+	return bld
+}
+
+// buildMatchDenseExpr constructs a MatchDenseExpr from an embedding vector.
+// This is a pure function — no I/O, no external dependencies.
+func buildMatchDenseExpr(vector []float64, topN int, similarity float64) *types.MatchDenseExpr {
+	vectorColumnName := fmt.Sprintf("q_%d_vec", len(vector))
+	return &types.MatchDenseExpr{
+		VectorColumnName:  vectorColumnName,
+		EmbeddingData:     vector,
+		EmbeddingDataType: "float",
+		DistanceType:      "cosine",
+		TopN:              topN,
+		ExtraOptions:      map[string]interface{}{"similarity": similarity},
+	}
+}
+
+// buildFusionExpr constructs a FusionExpr for weighted-sum hybrid search.
+// This is a pure function — no I/O, no external dependencies.
+func buildFusionExpr(textWeight, vectorWeight float64, topN int) *types.FusionExpr {
+	return &types.FusionExpr{
+		Method: "weighted_sum",
+		TopN:   topN,
+		FusionParams: map[string]interface{}{
+			"weights": fmt.Sprintf("%.2f,%.2f", textWeight, vectorWeight),
+		},
+	}
+}
+
+// buildSearchExprs constructs MatchExprs for KG entity/relation search.
+// When embModel is nil, returns text-only match expression.
+// When embModel is non-nil, embeds the question and returns hybrid
+// (text + dense + fusion) expressions for vector+keyword search.
+func buildSearchExprs(embModel *modelModule.EmbeddingModel, matchText *types.MatchTextExpr, simThreshold float64, denseTopK int) []interface{} {
+	if embModel == nil || embModel.ModelDriver == nil {
+		return []interface{}{matchText}
+	}
+	embeddingConfig := &modelModule.EmbeddingConfig{Dimension: 0}
+	embeddings, err := embModel.ModelDriver.Embed(embModel.ModelName, []string{matchText.MatchingText}, embModel.APIConfig, embeddingConfig)
+	if err != nil || len(embeddings) == 0 {
+		return []interface{}{matchText}
+	}
+	denseExpr := buildMatchDenseExpr(embeddings[0].Embedding, denseTopK, simThreshold)
+	fusionExpr := buildFusionExpr(0.5, 0.5, matchText.TopN)
+	return []interface{}{matchText, denseExpr, fusionExpr}
+}
+
+// queryRewrite attempts LLM-based query rewrite, falling back to raw question.
+// ty2entsJSON is the JSON-encoded type→entities mapping for prompt context.
+func queryRewrite(chatModel *modelModule.ChatModel, question string, ty2entsJSON string) (typeKeywords, entities []string) {
+	if question == "" {
+		return nil, nil
+	}
+	if chatModel != nil && chatModel.ModelName != nil && chatModel.APIConfig != nil {
+		prompt := common.BuildQueryRewritePrompt(question, ty2entsJSON)
+		messages := []modelModule.Message{
+			{Role: "system", Content: prompt},
+			{Role: "user", Content: "Output:"},
+		}
+		response, err := chatModel.ModelDriver.ChatWithMessages(*chatModel.ModelName, messages, chatModel.APIConfig, nil)
+		if err == nil && response != nil && response.Answer != nil {
+			result, parseErr := common.ParseQueryRewriteResponse(*response.Answer)
+			if parseErr == nil && result != nil {
+				return result.TypeKeywords, result.Entities
+			}
+		}
+	}
+	// Fallback: use raw question as single entity
+	return nil, []string{question}
+}
--- a/internal/service/kg_retrieval_test.go
+++ b/internal/service/kg_retrieval_test.go
@@ -0,0 +1,565 @@
+//
+//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"testing"
+
+	"ragflow/internal/engine"
+	"ragflow/internal/engine/types"
+	modelModule "ragflow/internal/entity/models"
+)
+
+type mockRetrievalEngine struct {
+	engine.DocEngine
+	results map[string]*types.SearchResult
+}
+
+func (m *mockRetrievalEngine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) {
+	// --- Contract validation (matches real ES/Infinity preconditions) ---
+	if len(req.IndexNames) == 0 {
+		return nil, fmt.Errorf("mock: IndexNames cannot be empty")
+	}
+	if len(req.KbIDs) == 0 {
+		return nil, fmt.Errorf("mock: KbIDs cannot be empty")
+	}
+	// --- Original stubbing logic ---
+	kgType, _ := req.Filter["knowledge_graph_kwd"].(string)
+	key := kgType
+	if ents, ok := req.Filter["entity_kwd"].([]interface{}); ok && len(ents) > 0 {
+		key = kgType + ":" + ents[0].(string)
+	}
+	if r, ok := m.results[key]; ok {
+		return r, nil
+	}
+	if r, ok := m.results[""]; ok {
+		return r, nil
+	}
+	return &types.SearchResult{}, nil
+}
+
+// --- kgEntityFromChunk ---
+
+func TestKgEntityFromChunk_Basic(t *testing.T) {
+	chunk := map[string]interface{}{
+		"_score":              0.85,
+		"rank_flt":            0.9,
+		"content_with_weight": "Founder of SpaceX",
+		"n_hop_with_weight":   `[{"path":["A","B"],"weights":[0.8]}]`,
+	}
+	e := kgEntityFromChunk("Elon Musk", chunk)
+	if e.Similarity != 0.85 {
+		t.Errorf("expected Sim=0.85, got %f", e.Similarity)
+	}
+	if e.PageRank != 0.9 {
+		t.Errorf("expected PageRank=0.9, got %f", e.PageRank)
+	}
+	if e.Description != "Founder of SpaceX" {
+		t.Errorf("expected Description, got %q", e.Description)
+	}
+	if len(e.NhopEnts) != 1 || len(e.NhopEnts[0].Path) != 2 {
+		t.Errorf("expected 1 NhopEnt with 2-path, got %+v", e.NhopEnts)
+	}
+}
+
+func TestKgEntityFromChunk_ScoreFallback(t *testing.T) {
+	chunk := map[string]interface{}{"score": 0.75}
+	e := kgEntityFromChunk("Test", chunk)
+	if e.Similarity != 0.75 {
+		t.Errorf("expected Sim=0.75 from score field, got %f", e.Similarity)
+	}
+}
+
+func TestKgEntityFromChunk_MissingFields(t *testing.T) {
+	chunk := map[string]interface{}{}
+	e := kgEntityFromChunk("Empty", chunk)
+	if e.Similarity != 0 || e.PageRank != 0 || len(e.NhopEnts) != 0 {
+		t.Errorf("expected zero defaults, got %+v", e)
+	}
+}
+
+// --- kgRelationFromChunk ---
+
+func TestKgRelationFromChunk_Basic(t *testing.T) {
+	chunk := map[string]interface{}{
+		"from_entity_kwd":    "Elon Musk",
+		"to_entity_kwd":      "SpaceX",
+		"weight_int":         float64(5),
+		"content_with_weight": "Founder",
+	}
+	edge, rel := kgRelationFromChunk(chunk)
+	if edge.From != "Elon Musk" || edge.To != "SpaceX" {
+		t.Errorf("expected Elon Musk→SpaceX, got %v", edge)
+	}
+	if rel.PageRank != 5 {
+		t.Errorf("expected weight 5, got %f", rel.PageRank)
+	}
+}
+
+func TestKgRelationFromChunk_MissingFrom(t *testing.T) {
+	chunk := map[string]interface{}{"to_entity_kwd": "B"}
+	edge, _ := kgRelationFromChunk(chunk)
+	if edge.From != "" {
+		t.Error("expected empty from")
+	}
+}
+
+// --- searchKGTypeSamples ---
+
+func TestSearchKGTypeSamples_Success(t *testing.T) {
+	data, _ := json.Marshal(map[string][]string{"PERSON": {"Elon Musk"}})
+	mock := &mockRetrievalEngine{
+		results: map[string]*types.SearchResult{
+			"ty2ents": {Chunks: []map[string]interface{}{
+				{"content_with_weight": string(data)},
+			}},
+		},
+	}
+	result, err := searchKGTypeSamples(context.Background(), mock, []string{"ragflow_tenant1"}, []string{"kb1"})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(result) != 1 || len(result["PERSON"]) != 1 || result["PERSON"][0] != "Elon Musk" {
+		t.Errorf("expected PERSON→[Elon Musk], got %v", result)
+	}
+}
+
+func TestSearchKGTypeSamples_Empty(t *testing.T) {
+	mock := &mockRetrievalEngine{}
+	result, err := searchKGTypeSamples(context.Background(), mock, []string{"ragflow_tenant1"}, []string{"kb1"})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(result) != 0 {
+		t.Errorf("expected empty, got %v", result)
+	}
+}
+
+// --- KGSearchRetrieval ---
+
+func TestKGSearchRetrieval_Basic(t *testing.T) {
+	mock := &mockRetrievalEngine{
+		results: map[string]*types.SearchResult{
+			"entity": {Chunks: []map[string]interface{}{
+				{"entity_kwd": "Elon Musk", "entity_type_kwd": "PERSON", "rank_flt": 0.9, "_score": 0.85},
+			}},
+			"relation": {Chunks: []map[string]interface{}{
+				{"from_entity_kwd": "Elon Musk", "to_entity_kwd": "SpaceX", "weight_int": float64(5), "_score": 0.85},
+			}},
+			"community_report": {Chunks: []map[string]interface{}{
+				{"docnm_kwd": "Community 1", "content_with_weight": "Report text", "weight_flt": 0.95},
+			}},
+			"ty2ents": {Chunks: []map[string]interface{}{
+				{"content_with_weight": `{"PERSON":["Elon Musk"]}`},
+			}},
+		},
+	}
+	result, err := KGSearchRetrieval(context.Background(), mock, nil, nil, []string{"kb1"}, []string{"tenant1"}, "Elon Musk")
+	if err != nil {
+		t.Fatalf("KGSearchRetrieval failed: %v", err)
+	}
+	if result == nil {
+		t.Fatal("expected non-nil result")
+	}
+	content, ok := result["content_with_weight"].(string)
+	if !ok {
+		t.Fatal("expected content_with_weight string")
+	}
+	if content == "" {
+		t.Error("expected non-empty KG content")
+	}
+	if result["similarity"] != 1.0 {
+		t.Errorf("expected similarity 1.0, got %v", result["similarity"])
+	}
+	if result["docnm_kwd"] != "Related content in Knowledge Graph" {
+		t.Errorf("unexpected docnm_kwd: %v", result["docnm_kwd"])
+	}
+}
+
+func TestKGSearchRetrieval_NoEntities(t *testing.T) {
+	mock := &mockRetrievalEngine{}
+	result, err := KGSearchRetrieval(context.Background(), mock, nil, nil, []string{"kb1"}, []string{"tenant1"}, "test")
+	if err != nil {
+		t.Fatalf("KGSearchRetrieval failed: %v", err)
+	}
+	if result == nil {
+		t.Fatal("expected non-nil result")
+	}
+	content, _ := result["content_with_weight"].(string)
+	if content != "" {
+		t.Errorf("expected empty when no entities found, got %q", content)
+	}
+}
+
+// TestEntitySearch_MultiEntities verifies that all entities are used in search query.
+
+func TestKGSearchRetrieval_WithChatModel(t *testing.T) {
+	mock := &mockRetrievalEngine{
+		results: map[string]*types.SearchResult{
+			"entity": {Chunks: []map[string]interface{}{
+				{"entity_kwd": "Elon Musk", "entity_type_kwd": "PERSON", "rank_flt": 0.9, "_score": 0.85},
+			}},
+			"relation": {Chunks: []map[string]interface{}{
+				{"from_entity_kwd": "Elon Musk", "to_entity_kwd": "SpaceX", "weight_int": float64(5), "_score": 0.85},
+			}},
+		},
+	}
+	// chatModel with nil ModelName so queryRewrite falls back to raw question,
+	// but the ty2entsJSON construction path is still exercised.
+	chatModel := &modelModule.ChatModel{ModelName: nil, APIConfig: nil}
+	result, err := KGSearchRetrieval(context.Background(), mock, chatModel, nil, []string{"kb1"}, []string{"tenant1"}, "Elon Musk")
+	if err != nil {
+		t.Fatalf("KGSearchRetrieval failed: %v", err)
+	}
+	if result == nil {
+		t.Fatal("expected non-nil result")
+	}
+	content, ok := result["content_with_weight"].(string)
+	if !ok {
+		t.Fatal("expected content_with_weight string")
+	}
+	if content == "" {
+		t.Error("expected non-empty KG content")
+	}
+	// Verify "null" does not appear — the ty2entsJSON fix ensures "{}" not "null"
+	if strings.Contains(content, "null") {
+		t.Error("content should not contain 'null' from ty2entsJSON")
+	}
+}
+
+func TestEntitySearch_MultiEntities(t *testing.T) {
+	var capturedText string
+	mock := &searchCaptureEngine{}
+	mock.searchFn = func(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) {
+		if kgType, _ := req.Filter["knowledge_graph_kwd"].(string); kgType == "entity" && len(req.MatchExprs) > 0 {
+			if expr, ok := req.MatchExprs[0].(*types.MatchTextExpr); ok {
+				capturedText = expr.MatchingText
+			}
+		}
+		return &types.SearchResult{}, nil
+	}
+	entities := []string{"Elon Musk", "SpaceX"}
+	entsReq := &types.SearchRequest{
+		IndexNames:   []string{"ragflow_tenant1"},
+		KbIDs:        []string{"kb1"},
+		SelectFields: []string{"entity_kwd", "n_hop_with_weight"},
+		Limit:        50,
+		Filter:       map[string]interface{}{"knowledge_graph_kwd": "entity"},
+		MatchExprs: []interface{}{
+			&types.MatchTextExpr{
+				Fields:       []string{"entity_kwd^10", "content_ltks^2"},
+				MatchingText: strings.Join(entities, " "),
+				TopN:         50,
+			},
+		},
+	}
+	mock.Search(context.Background(), entsReq)
+	if !strings.Contains(capturedText, "Elon Musk") || !strings.Contains(capturedText, "SpaceX") {
+		t.Errorf("expected both entities in query, got %q", capturedText)
+	}
+}
+
+// searchCaptureEngine is a minimal mock for testing search requests.
+type searchCaptureEngine struct {
+	engine.DocEngine
+	searchFn func(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error)
+}
+
+func (e *searchCaptureEngine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) {
+	if len(req.IndexNames) == 0 {
+		return nil, fmt.Errorf("mock: IndexNames cannot be empty")
+	}
+	if e.searchFn != nil {
+		return e.searchFn(ctx, req)
+	}
+	return &types.SearchResult{}, nil
+}
+
+// --- queryRewrite ---
+
+func TestQueryRewrite_Fallback(t *testing.T) {
+	typeKeywords, entities := queryRewrite(nil, "What is SpaceX?", "{}")
+	if typeKeywords != nil {
+		t.Errorf("expected nil typeKeywords when no LLM, got %v", typeKeywords)
+	}
+	if len(entities) != 1 || entities[0] != "What is SpaceX?" {
+		t.Errorf("expected [What is SpaceX?], got %v", entities)
+	}
+}
+
+func TestQueryRewrite_EmptyQuestion(t *testing.T) {
+	typeKeywords, entities := queryRewrite(nil, "", "")
+	if typeKeywords != nil || entities != nil {
+		t.Errorf("expected nil for empty question, got type=%v entities=%v", typeKeywords, entities)
+	}
+}
+
+// spyEmbedDriver captures Embed input for testing — enables assertions on what text
+// was embedded, not just that embedding succeeded.
+type spyEmbedDriver struct {
+	modelModule.ModelDriver
+	capturedTexts []string
+	vector        []float64
+	err           error
+}
+
+func (s *spyEmbedDriver) Embed(_ *string, texts []string, _ *modelModule.APIConfig, _ *modelModule.EmbeddingConfig) ([]modelModule.EmbeddingData, error) {
+	s.capturedTexts = texts
+	if s.err != nil {
+		return nil, s.err
+	}
+	return []modelModule.EmbeddingData{{Embedding: s.vector}}, nil
+}
+
+// --- pure function: buildMatchDenseExpr ---
+
+func TestBuildMatchDenseExpr_Basic(t *testing.T) {
+	vector := []float64{0.1, 0.2, 0.3}
+	expr := buildMatchDenseExpr(vector, 10, 0.2)
+	if expr.VectorColumnName != "q_3_vec" {
+		t.Errorf("expected q_3_vec, got %q", expr.VectorColumnName)
+	}
+	if len(expr.EmbeddingData) != 3 || expr.EmbeddingData[0] != 0.1 {
+		t.Errorf("unexpected embedding data: %v", expr.EmbeddingData)
+	}
+	if expr.EmbeddingDataType != "float" {
+		t.Errorf("expected float, got %q", expr.EmbeddingDataType)
+	}
+	if expr.DistanceType != "cosine" {
+		t.Errorf("expected cosine, got %q", expr.DistanceType)
+	}
+	if expr.TopN != 10 {
+		t.Errorf("expected TopN=10, got %d", expr.TopN)
+	}
+	sim, ok := expr.ExtraOptions["similarity"].(float64)
+	if !ok || sim != 0.2 {
+		t.Errorf("expected similarity=0.2, got %v", expr.ExtraOptions["similarity"])
+	}
+}
+
+func TestBuildMatchDenseExpr_ZeroVector(t *testing.T) {
+	expr := buildMatchDenseExpr(nil, 5, 0.0)
+	if expr.VectorColumnName != "q_0_vec" {
+		t.Errorf("expected q_0_vec for empty vector, got %q", expr.VectorColumnName)
+	}
+}
+
+// --- pure function: buildFusionExpr ---
+
+func TestBuildFusionExpr_DefaultWeights(t *testing.T) {
+	expr := buildFusionExpr(0.5, 0.5, 20)
+	if expr.Method != "weighted_sum" {
+		t.Errorf("expected weighted_sum, got %q", expr.Method)
+	}
+	if expr.TopN != 20 {
+		t.Errorf("expected TopN=20, got %d", expr.TopN)
+	}
+	weights, ok := expr.FusionParams["weights"].(string)
+	if !ok || weights != "0.50,0.50" {
+		t.Errorf("expected weights=0.50,0.50, got %v", expr.FusionParams["weights"])
+	}
+}
+
+func TestBuildFusionExpr_AsymmetricWeights(t *testing.T) {
+	expr := buildFusionExpr(0.3, 0.7, 10)
+	weights := expr.FusionParams["weights"].(string)
+	if weights != "0.30,0.70" {
+		t.Errorf("expected 0.30,0.70, got %q", weights)
+	}
+}
+
+// --- buildSearchExprs ---
+
+func TestBuildSearchExprs_NoEmbModel(t *testing.T) {
+	matchText := &types.MatchTextExpr{
+		Fields:       []string{"entity_kwd^10"},
+		MatchingText: "test",
+		TopN:         10,
+	}
+	exprs := buildSearchExprs(nil, matchText, 0, 0)
+	if len(exprs) != 1 {
+		t.Fatalf("expected 1 expr, got %d", len(exprs))
+	}
+	mt, ok := exprs[0].(*types.MatchTextExpr)
+	if !ok {
+		t.Fatalf("expected MatchTextExpr, got %T", exprs[0])
+	}
+	if mt.MatchingText != "test" {
+		t.Errorf("expected 'test', got %q", exprs[0].(*types.MatchTextExpr).MatchingText)
+	}
+}
+
+func TestBuildSearchExprs_WithEmbModel(t *testing.T) {
+	driver := &spyEmbedDriver{vector: []float64{0.1, 0.2, 0.3}}
+	embModel := modelModule.NewEmbeddingModel(driver, strPtr("text-embedding"), &modelModule.APIConfig{}, 512)
+	matchText := &types.MatchTextExpr{
+		Fields:       []string{"entity_kwd^10"},
+		MatchingText: "Elon Musk SpaceX",
+		TopN:         50,
+	}
+	exprs := buildSearchExprs(embModel, matchText, defaultKGSimThreshold, defaultKGDenseTopK)
+	// Verify Embed was called with matchText.MatchingText, not raw question
+	if len(driver.capturedTexts) != 1 || driver.capturedTexts[0] != "Elon Musk SpaceX" {
+		t.Errorf("expected Embed to receive %q, got %v", "Elon Musk SpaceX", driver.capturedTexts)
+	}
+	if len(exprs) != 3 {
+		t.Fatalf("expected 3 exprs (text+dense+fusion), got %d", len(exprs))
+	}
+	// Index 0: MatchTextExpr
+	mt, ok := exprs[0].(*types.MatchTextExpr)
+	if !ok {
+		t.Fatalf("expected MatchTextExpr at [0], got %T", exprs[0])
+	}
+	if mt.MatchingText != "Elon Musk SpaceX" {
+		t.Errorf("expected 'Elon Musk SpaceX', got %q", mt.MatchingText)
+	}
+	// Index 1: MatchDenseExpr
+	md, ok := exprs[1].(*types.MatchDenseExpr)
+	if !ok {
+		t.Fatalf("expected MatchDenseExpr at [1], got %T", exprs[1])
+	}
+	if md.VectorColumnName != "q_3_vec" {
+		t.Errorf("expected q_3_vec, got %q", md.VectorColumnName)
+	}
+	if md.TopN != defaultKGDenseTopK {
+		t.Errorf("expected TopN=%d (Python alignment), got %d", defaultKGDenseTopK, md.TopN)
+	}
+	if md.ExtraOptions["similarity"] != defaultKGSimThreshold {
+		t.Errorf("expected similarity=%v (Python alignment), got %v", defaultKGSimThreshold, md.ExtraOptions["similarity"])
+	}
+	// Index 2: FusionExpr
+	fu, ok := exprs[2].(*types.FusionExpr)
+	if !ok {
+		t.Fatalf("expected FusionExpr at [2], got %T", exprs[2])
+	}
+	if fu.Method != "weighted_sum" {
+		t.Errorf("expected weighted_sum, got %q", fu.Method)
+	}
+}
+
+func TestBuildSearchExprs_EmbModelFallback(t *testing.T) {
+	driver := &spyEmbedDriver{err: assertError("embed failed")}
+	embModel := modelModule.NewEmbeddingModel(driver, strPtr("text-embedding"), &modelModule.APIConfig{}, 512)
+	matchText := &types.MatchTextExpr{
+		Fields:       []string{"entity_kwd^10"},
+		MatchingText: "fallback test",
+		TopN:         10,
+	}
+	exprs := buildSearchExprs(embModel, matchText, defaultKGSimThreshold, defaultKGDenseTopK)
+	// Should fall back to text-only when Embed fails
+	if len(exprs) != 1 {
+		t.Fatalf("expected 1 expr (text-only fallback), got %d", len(exprs))
+	}
+	if _, ok := exprs[0].(*types.MatchTextExpr); !ok {
+		t.Errorf("expected MatchTextExpr, got %T", exprs[0])
+	}
+}
+
+// --- Python alignment defaults ---
+
+func TestDefaultValuesMatchPython(t *testing.T) {
+	if defaultKGSimThreshold != 0.3 {
+		t.Errorf("expected 0.3 (Python ent_sim_threshold), got %f", defaultKGSimThreshold)
+	}
+	if defaultKGDenseTopK != 1024 {
+		t.Errorf("expected 1024 (Python get_vector topk), got %d", defaultKGDenseTopK)
+	}
+}
+
+// assertError is a simple error for testing fallback behaviour.
+type assertError string
+
+func (e assertError) Error() string { return string(e) }
+
+
+// --- indexName ---
+
+func TestIndexName_Normal(t *testing.T) {
+	result := indexName("tenant1")
+	if result != "ragflow_tenant1" {
+		t.Errorf("expected ragflow_tenant1, got %q", result)
+	}
+}
+
+func TestIndexName_Empty(t *testing.T) {
+	result := indexName("")
+	if result != "ragflow_" {
+		t.Errorf("expected ragflow_, got %q", result)
+	}
+}
+
+// --- searchKGCommunityContent ---
+
+func TestSearchKGCommunityContent_EmptyEntities(t *testing.T) {
+	mock := &mockRetrievalEngine{}
+	result := searchKGCommunityContent(context.Background(), mock, []string{"ragflow_t1"}, []string{"kb1"}, nil, 1, intPtr(100))
+	if result != "" {
+		t.Errorf("expected empty, got %q", result)
+	}
+}
+
+func TestSearchKGCommunityContent_WithContent(t *testing.T) {
+	mock := &mockRetrievalEngine{
+		results: map[string]*types.SearchResult{
+			"community_report": {Chunks: []map[string]interface{}{
+				{
+					"docnm_kwd":           "Community Alpha",
+					"content_with_weight": `{"report": "Report text", "evidences": "Evidence text"}`,
+				},
+			}},
+		},
+	}
+	result := searchKGCommunityContent(context.Background(), mock, []string{"ragflow_t1"}, []string{"kb1"}, []ScoredEntity{{Entity: "E1"}}, 1, intPtr(500))
+	if result == "" {
+		t.Fatal("expected non-empty result")
+	}
+	if !strings.Contains(result, "Community Alpha") {
+		t.Errorf("expected title 'Community Alpha', got %q", result)
+	}
+	if !strings.Contains(result, "Report text") {
+		t.Errorf("expected report content, got %q", result)
+	}
+	if !strings.Contains(result, "Evidence text") {
+		t.Errorf("expected evidence, got %q", result)
+	}
+	if !strings.Contains(result, "# 1.") {
+		t.Errorf("expected numbered report (# 1.), got %q", result)
+	}
+}
+
+func TestSearchKGCommunityContent_NilMaxToken(t *testing.T) {
+	mock := &mockRetrievalEngine{}
+	result := searchKGCommunityContent(context.Background(), mock, []string{"ragflow_t1"}, []string{"kb1"}, []ScoredEntity{{Entity: "E1"}}, 1, nil)
+	if result != "" {
+		t.Errorf("expected empty when maxToken is nil, got %q", result)
+	}
+}
+
+func TestSearchKGCommunityContent_ZeroMaxToken(t *testing.T) {
+	mock := &mockRetrievalEngine{}
+	result := searchKGCommunityContent(context.Background(), mock, []string{"ragflow_t1"}, []string{"kb1"}, []ScoredEntity{{Entity: "E1"}}, 1, intPtr(0))
+	if result != "" {
+		t.Errorf("expected empty when maxToken=0, got %q", result)
+	}
+}
+
+// intPtr returns a pointer to n.
+func intPtr(n int) *int { return &n }
--- a/internal/service/kg_scoring_funcs.go
+++ b/internal/service/kg_scoring_funcs.go
@@ -14,61 +14,17 @@
 //  limitations under the License.
 //

-package common
+package service

 import (
+	"bytes"
+	"encoding/csv"
+	"encoding/json"
 	"fmt"
 	"sort"
 	"strings"
 )

-// KGEntity represents a knowledge graph entity with its scores.
-type KGEntity struct {
-	Sim         float64
-	PageRank    float64
-	Description string
-	NhopEnts    []NhopEntity
-}
-
-// NhopEntity represents an N-hop neighbor path.
-type NhopEntity struct {
-	Path    []string  // entity names along the path
-	Weights []float64 // pagerank weights per hop
-}
-
-// KGRelation represents a knowledge graph relation with its scores.
-type KGRelation struct {
-	Sim         float64
-	PageRank    float64
-	Description string
-}
-
-// Edge represents a directed (from_entity, to_entity) pair.
-type Edge struct {
-	From, To string
-}
-
-// EdgeScore represents the accumulated score for an edge from N-hop analysis.
-type EdgeScore struct {
-	Sim      float64
-	PageRank float64
-}
-
-// ScoredEntity is a scored entity ready for output.
-type ScoredEntity struct {
-	Entity      string
-	Score       float64
-	Description string
-}
-
-// ScoredRelation is a scored relation ready for output.
-type ScoredRelation struct {
-	From        string
-	To          string
-	Score       float64
-	Description string
-}
-
 // AnalyzeNHopPaths decomposes N-hop paths into edges with distance-decayed scores.
 // Python equivalent: rag/graphrag/search.py lines 172-187
 func AnalyzeNHopPaths(entsFromQuery map[string]*KGEntity) map[Edge]EdgeScore {
@@ -81,7 +37,7 @@ func AnalyzeNHopPaths(entsFromQuery map[string]*KGEntity) map[Edge]EdgeScore {
 				f, t := path[i], path[i+1]
 				edge := Edge{From: f, To: t}
 				es := nhopPathes[edge]
-				es.Sim += ent.Sim / (2.0 + float64(i))
+				es.Sim += ent.Similarity / (2.0 + float64(i))
 				if i < len(weights) {
 					es.PageRank = weights[i]
 				}
@@ -97,7 +53,7 @@ func AnalyzeNHopPaths(entsFromQuery map[string]*KGEntity) map[Edge]EdgeScore {
 func DoubleHitBoost(entsFromQuery map[string]*KGEntity, entsFromTypes map[string]struct{}) {
 	for ent := range entsFromQuery {
 		if _, ok := entsFromTypes[ent]; ok {
-			entsFromQuery[ent].Sim *= 2
+			entsFromQuery[ent].Similarity *= 2
 		}
 	}
 }
@@ -152,7 +108,7 @@ func SortAndTrimEntities(entsFromQuery map[string]*KGEntity, topN int) []ScoredE
 	for name, ent := range entsFromQuery {
 		scored = append(scored, ScoredEntity{
 			Entity:      name,
-			Score:       ent.Sim * ent.PageRank,
+			Score:       ent.Similarity * ent.PageRank,
 			Description: ent.Description,
 		})
 	}
@@ -195,6 +151,40 @@ func NumTokensFromString(s string) int {
 	return len(s) / 4
 }

+// formatCSVLine formats fields as a single CSV record with trailing newline.
+// Handles commas, quotes, and newlines in field values correctly — unlike fmt.Sprintf.
+// Matches Python: pd.DataFrame(...).to_csv() quoting behavior.
+func formatCSVLine(fields ...string) string {
+	var buf bytes.Buffer
+	w := csv.NewWriter(&buf)
+	_ = w.Write(fields)
+	w.Flush()
+	return buf.String()
+}
+
+// FilterChunksByScore filters chunks where _score >= threshold.
+// Chunks missing _score are treated as score=0.
+// Pure function — no I/O, no external dependencies.
+// Matches Python: _ent_info_from_ and _relation_info_from_ sim_thr filtering.
+func FilterChunksByScore(chunks []map[string]interface{}, threshold float64) []map[string]interface{} {
+	if threshold <= 0 || len(chunks) == 0 {
+		return chunks
+	}
+	result := make([]map[string]interface{}, 0, len(chunks))
+	for _, chunk := range chunks {
+		score := 0.0
+		if v, ok := chunk["_score"].(float64); ok {
+			score = v
+		} else if v, ok := chunk["score"].(float64); ok {
+			score = v
+		}
+		if score >= threshold {
+			result = append(result, chunk)
+		}
+	}
+	return result
+}
+
 // FormatEntitiesToCSV formats scored entities as a CSV string and tracks token count.
 func FormatEntitiesToCSV(entities []ScoredEntity, maxToken int) (csv string, remainingToken int) {
 	if len(entities) == 0 {
@@ -203,12 +193,11 @@ func FormatEntitiesToCSV(entities []ScoredEntity, maxToken int) (csv string, rem
 	var b strings.Builder
 	b.WriteString("---- Entities ----\n")
 	b.WriteString("Entity,Score,Description\n")
-	for i, ent := range entities {
+	for _, ent := range entities {
 		desc := extractDescription(ent.Description)
-		line := fmt.Sprintf("%s,%.2f,%s\n", ent.Entity, ent.Score, desc)
+		line := formatCSVLine(ent.Entity, fmt.Sprintf("%.2f", ent.Score), desc)
 		tokens := NumTokensFromString(line)
 		if maxToken-tokens <= 0 {
-			entities = entities[:i]
 			break
 		}
 		b.WriteString(line)
@@ -225,12 +214,11 @@ func FormatRelationsToCSV(relations []ScoredRelation, maxToken int) (csv string,
 	var b strings.Builder
 	b.WriteString("---- Relations ----\n")
 	b.WriteString("From Entity,To Entity,Score,Description\n")
-	for i, rel := range relations {
+	for _, rel := range relations {
 		desc := extractDescription(rel.Description)
-		line := fmt.Sprintf("%s,%s,%.2f,%s\n", rel.From, rel.To, rel.Score, desc)
+		line := formatCSVLine(rel.From, rel.To, fmt.Sprintf("%.2f", rel.Score), desc)
 		tokens := NumTokensFromString(line)
 		if maxToken-tokens <= 0 {
-			relations = relations[:i]
 			break
 		}
 		b.WriteString(line)
@@ -257,24 +245,12 @@ func extractDescription(desc string) string {
 	if desc == "" {
 		return ""
 	}
-	// If the description looks like JSON, try to extract the "description" field
-	desc = strings.TrimSpace(desc)
-	if strings.HasPrefix(desc, "{") && strings.HasSuffix(desc, "}") {
-		// Simple extraction: find "description" key value
-		// This matches Python's json.loads(desc).get("description", "") behavior
-		idx := strings.Index(desc, `"description"`)
-		if idx >= 0 {
-			remain := desc[idx+len(`"description"`):]
-			colonIdx := strings.Index(remain, ":")
-			if colonIdx >= 0 {
-				valPart := strings.TrimSpace(remain[colonIdx+1:])
-				if strings.HasPrefix(valPart, `"`) {
-					valPart = strings.TrimPrefix(valPart, `"`)
-					endQuote := strings.Index(valPart, `"`)
-					if endQuote >= 0 {
-						return valPart[:endQuote]
-					}
-				}
+	// Try to parse as JSON and extract the "description" field.
+	var data map[string]interface{}
+	if err := json.Unmarshal([]byte(desc), &data); err == nil {
+		if v, ok := data["description"]; ok {
+			if s, ok := v.(string); ok {
+				return s
 			}
 		}
 	}
--- a/internal/service/kg_scoring_funcs_test.go
+++ b/internal/service/kg_scoring_funcs_test.go
@@ -14,7 +14,7 @@
 //  limitations under the License.
 //

-package common
+package service

 import (
 	"strings"
@@ -26,7 +26,7 @@ import (
 func TestAnalyzeNHopPaths_Basic(t *testing.T) {
 	ents := map[string]*KGEntity{
 		"A": {
-			Sim: 0.9,
+Similarity: 0.9,
 			NhopEnts: []NhopEntity{
 				{Path: []string{"A", "B", "C"}, Weights: []float64{0.8, 0.5}},
 			},
@@ -49,13 +49,13 @@ func TestAnalyzeNHopPaths_Basic(t *testing.T) {
 func TestAnalyzeNHopPaths_MultipleContributors(t *testing.T) {
 	ents := map[string]*KGEntity{
 		"A": {
-			Sim: 0.8,
+Similarity: 0.8,
 			NhopEnts: []NhopEntity{
 				{Path: []string{"A", "B"}, Weights: []float64{0.7}},
 			},
 		},
 		"X": {
-			Sim: 0.6,
+Similarity: 0.6,
 			NhopEnts: []NhopEntity{
 				{Path: []string{"X", "B"}, Weights: []float64{0.5}},
 			},
@@ -83,24 +83,24 @@ func TestAnalyzeNHopPaths_Empty(t *testing.T) {

 func TestDoubleHitBoost(t *testing.T) {
 	ents := map[string]*KGEntity{
-		"A": {Sim: 0.5},
-		"B": {Sim: 0.3},
+		"A": {Similarity: 0.5},
+		"B": {Similarity: 0.3},
 	}
 	types := map[string]struct{}{"A": {}}
 	DoubleHitBoost(ents, types)
-	if ents["A"].Sim != 1.0 {
-		t.Errorf("expected A sim=1.0 after boost, got %f", ents["A"].Sim)
+	if ents["A"].Similarity != 1.0 {
+		t.Errorf("expected A sim=1.0 after boost, got %f", ents["A"].Similarity)
 	}
-	if ents["B"].Sim != 0.3 {
-		t.Errorf("expected B sim unchanged at 0.3, got %f", ents["B"].Sim)
+	if ents["B"].Similarity != 0.3 {
+		t.Errorf("expected B sim unchanged at 0.3, got %f", ents["B"].Similarity)
 	}
 }

 func TestDoubleHitBoost_Empty(t *testing.T) {
-	ents := map[string]*KGEntity{"A": {Sim: 0.5}}
+	ents := map[string]*KGEntity{"A": {Similarity: 0.5}}
 	DoubleHitBoost(ents, map[string]struct{}{})
-	if ents["A"].Sim != 0.5 {
-		t.Errorf("expected unchanged, got %f", ents["A"].Sim)
+	if ents["A"].Similarity != 0.5 {
+		t.Errorf("expected unchanged, got %f", ents["A"].Similarity)
 	}
 }

@@ -153,9 +153,9 @@ func TestFuseRelationScores_NhopNewEdge(t *testing.T) {

 func TestSortAndTrimEntities(t *testing.T) {
 	ents := map[string]*KGEntity{
-		"A": {Sim: 0.5, PageRank: 0.9},
-		"B": {Sim: 0.8, PageRank: 0.3},
-		"C": {Sim: 0.9, PageRank: 0.1},
+		"A": {Similarity: 0.5, PageRank: 0.9},
+		"B": {Similarity: 0.8, PageRank: 0.3},
+		"C": {Similarity: 0.9, PageRank: 0.1},
 	}
 	result := SortAndTrimEntities(ents, 2)
 	if len(result) != 2 {
@@ -169,8 +169,8 @@ func TestSortAndTrimEntities(t *testing.T) {

 func TestSortAndTrimEntities_DefaultTopN(t *testing.T) {
 	ents := map[string]*KGEntity{
-		"A": {Sim: 0.5, PageRank: 0.9},
-		"B": {Sim: 0.8, PageRank: 0.3},
+		"A": {Similarity: 0.5, PageRank: 0.9},
+		"B": {Similarity: 0.8, PageRank: 0.3},
 	}
 	result := SortAndTrimEntities(ents, 0)
 	if len(result) != 2 {
@@ -203,10 +203,10 @@ func TestBuildKGContent_Basic(t *testing.T) {
 		{From: "A", To: "B", Score: 0.3, Description: `{"description": "rel A-B"}`},
 	}
 	result := BuildKGContent(entities, relations, 10000)
-	if !contains(result, "Entity A desc") {
+	if !strings.Contains(result, "Entity A desc") {
 		t.Errorf("expected entity description in output, got: %s", result)
 	}
-	if !contains(result, "rel A-B") {
+	if !strings.Contains(result, "rel A-B") {
 		t.Errorf("expected relation description in output, got: %s", result)
 	}
 }
@@ -221,11 +221,75 @@ func TestBuildKGContent_TokenBudget(t *testing.T) {
 	}
 	result := BuildKGContent(entities, relations, 50)
 	// Token budget is very small, should truncate and not include relations
-	if contains(result, "relation desc") {
+	if strings.Contains(result, "relation desc") {
 		t.Log("Note: relations included despite small budget (depending on token count)")
 	}
 }

+func TestFormatEntitiesToCSV_HeaderExceedsBudget(t *testing.T) {
+	entities := []ScoredEntity{
+		{Entity: "A", Score: 1.0, Description: "d"},
+	}
+	result, remaining := FormatEntitiesToCSV(entities, 3)
+	tokens := NumTokensFromString(result)
+	// Header lines (---- Entities ----\n, Entity,Score,Description\n) are written
+	// before the token budget check. They consume ~11 tokens but are not deducted
+	// from maxToken. This is a known limitation shared with Python.
+	if tokens > 3 {
+		t.Logf("output %d tokens exceeds budget of %d (header not counted, remaining=%d)", tokens, 3, remaining)
+	}
+}
+
+func TestFilterChunksByScore_AllPass(t *testing.T) {
+	chunks := []map[string]interface{}{
+		{"entity_kwd": "A", "_score": 0.5},
+		{"entity_kwd": "B", "_score": 0.8},
+	}
+	result := FilterChunksByScore(chunks, 0.3)
+	if len(result) != 2 {
+		t.Errorf("expected all 2 chunks to pass, got %d", len(result))
+	}
+}
+
+func TestFilterChunksByScore_SomeFiltered(t *testing.T) {
+	chunks := []map[string]interface{}{
+		{"entity_kwd": "A", "_score": 0.2},
+		{"entity_kwd": "B", "_score": 0.9},
+	}
+	result := FilterChunksByScore(chunks, 0.3)
+	if len(result) != 1 || result[0]["entity_kwd"] != "B" {
+		t.Errorf("expected only B to pass, got %v", result)
+	}
+}
+
+func TestFilterChunksByScore_MissingScore(t *testing.T) {
+	chunks := []map[string]interface{}{
+		{"entity_kwd": "A"}, // no _score → treated as 0
+		{"entity_kwd": "B", "score": 0.5},
+	}
+	result := FilterChunksByScore(chunks, 0.3)
+	if len(result) != 1 || result[0]["entity_kwd"] != "B" {
+		t.Errorf("expected only B (using 'score' field), got %v", result)
+	}
+}
+
+func TestFilterChunksByScore_NilInput(t *testing.T) {
+	result := FilterChunksByScore(nil, 0.3)
+	if result != nil {
+		t.Errorf("expected nil, got %v", result)
+	}
+}
+
+func TestFilterChunksByScore_ZeroThreshold(t *testing.T) {
+	chunks := []map[string]interface{}{
+		{"entity_kwd": "A", "_score": 0.0},
+	}
+	result := FilterChunksByScore(chunks, 0)
+	if len(result) != 1 {
+		t.Errorf("expected all pass when threshold=0, got %d", len(result))
+	}
+}
+
 func TestExtractDescription_JSON(t *testing.T) {
 	result := extractDescription(`{"description": "Entity A description", "other": "value"}`)
 	if result != "Entity A description" {
@@ -240,6 +304,60 @@ func TestExtractDescription_Plain(t *testing.T) {
 	}
 }

+func TestExtractDescription_EscapedQuote(t *testing.T) {
+	result := extractDescription(`{"description": "has \"quote\" inside"}`)
+	if result != `has "quote" inside` {
+		t.Errorf("expected full description with quote, got %q", result)
+	}
+}
+
+func TestExtractDescription_NonStringValue(t *testing.T) {
+	result := extractDescription(`{"description": null, "other": "val"}`)
+	if result != `{"description": null, "other": "val"}` {
+		t.Errorf("expected raw JSON when description is null, got %q", result)
+	}
+}
+
+func TestExtractDescription_EmptyString(t *testing.T) {
+	result := extractDescription("")
+	if result != "" {
+		t.Errorf("expected empty, got %q", result)
+	}
+}
+
+func TestFormatCSVLine_Normal(t *testing.T) {
+	result := formatCSVLine("Elon Musk", "0.85", "CEO of SpaceX")
+	// Normal values should not be quoted
+	if result != "Elon Musk,0.85,CEO of SpaceX\n" {
+		t.Errorf("expected unquoted CSV, got %q", result)
+	}
+}
+
+func TestFormatCSVLine_CommaInField(t *testing.T) {
+	result := formatCSVLine("Musk, Elon", "0.85", "CEO, SpaceX")
+	// Values with commas should be quoted
+	expected := `"Musk, Elon",0.85,"CEO, SpaceX"` + "\n"
+	if result != expected {
+		t.Errorf("expected %q, got %q", expected, result)
+	}
+}
+
+func TestFormatCSVLine_QuoteInField(t *testing.T) {
+	result := formatCSVLine("Elon Musk", "0.85", `CEO of "SpaceX"`)
+	// Values with quotes should have quotes escaped
+	expected := `Elon Musk,0.85,"CEO of ""SpaceX"""` + "\n"
+	if result != expected {
+		t.Errorf("expected %q, got %q", expected, result)
+	}
+}
+
+func TestFormatCSVLine_EmptyField(t *testing.T) {
+	result := formatCSVLine("", "", "")
+	if result != ",,\n" {
+		t.Errorf("expected empty fields, got %q", result)
+	}
+}
+
 func TestNumTokensFromString(t *testing.T) {
 	s := "This is a test string with multiple words"
 	tokens := NumTokensFromString(s)
@@ -248,16 +366,3 @@ func TestNumTokensFromString(t *testing.T) {
 	}
 }

-// contains checks if a string contains a substring.
-func contains(s, substr string) bool {
-	return len(s) >= len(substr) && containsStr(s, substr)
-}
-
-func containsStr(s, substr string) bool {
-	for i := 0; i <= len(s)-len(substr); i++ {
-		if s[i:i+len(substr)] == substr {
-			return true
-		}
-	}
-	return false
-}
--- a/internal/service/kg_search.go
+++ b/internal/service/kg_search.go
@@ -26,21 +26,56 @@ import (
 	modelModule "ragflow/internal/entity/models"
 )

+// NhopEntity represents an N-hop neighbor path.
+type NhopEntity struct {
+	Path    []string  // entity names along the path
+	Weights []float64 // pagerank weights per hop
+}
+
 // KGEntity represents a knowledge graph entity.
 type KGEntity struct {
-	Name        string  // entity_kwd
-	Type        string  // entity_type_kwd
-	PageRank    float64 // rank_flt
-	Similarity  float64 // _score
-	Description string  // content_with_weight
+	Name        string       // entity_kwd
+	Type        string       // entity_type_kwd
+	PageRank    float64      // rank_flt
+	Similarity  float64      // _score
+	Description string       // content_with_weight
+	NhopEnts    []NhopEntity // n_hop_with_weight (parsed JSON)
+}
+
+// Edge represents a directed (from_entity, to_entity) pair.
+type Edge struct {
+	From, To string
+}
+
+// EdgeScore represents the accumulated score for an edge from N-hop analysis.
+type EdgeScore struct {
+	Sim      float64
+	PageRank float64
+}
+
+// ScoredEntity is a scored entity ready for output.
+type ScoredEntity struct {
+	Entity      string
+	Score       float64
+	Description string
+}
+
+// ScoredRelation is a scored relation ready for output.
+type ScoredRelation struct {
+	From        string
+	To          string
+	Score       float64
+	Description string
 }

 // KGRelation represents a relation between two entities.
 type KGRelation struct {
-	From        string // from_entity_kwd
-	To          string // to_entity_kwd
-	Weight      int    // weight_int
-	Description string // content_with_weight
+	From        string  // from_entity_kwd
+	To          string  // to_entity_kwd
+	Weight      int     // weight_int
+	Description string  // content_with_weight
+	Sim         float64 // score accumulated during pipeline scoring
+	PageRank    float64 // rank_flt or weight_int as float64
 }

 // KGCommunityReport represents a community report.
--- a/internal/service/tenant.go
+++ b/internal/service/tenant.go
@@ -790,6 +790,9 @@ func (s *TenantService) RemoveMember(userID, tenantID, targetUserID string) (com
 	if targetUserID == tenantID {
 		return common.CodeArgumentError, fmt.Errorf("cannot remove the tenant owner")
 	}
+	if s.userTenantDAO == nil {
+		return common.CodeServerError, fmt.Errorf("userTenantDAO not initialized")
+	}
 	if err := s.userTenantDAO.DeleteByUserAndTenant(targetUserID, tenantID); err != nil {
 		return common.CodeServerError, fmt.Errorf("failed to remove member: %w", err)
 	}
@@ -798,6 +801,9 @@ func (s *TenantService) RemoveMember(userID, tenantID, targetUserID string) (com

 // AcceptInvite transitions the calling user's role from "invite" → "normal" for the given tenant.
 func (s *TenantService) AcceptInvite(userID, tenantID string) (common.ErrorCode, error) {
+	if s.userTenantDAO == nil {
+		return common.CodeServerError, fmt.Errorf("userTenantDAO not initialized")
+	}
 	existing, err := s.userTenantDAO.FilterByUserIDAndTenantID(userID, tenantID)
 	if err != nil || existing == nil {
 		return common.CodeDataError, fmt.Errorf("no pending invitation found")
--- a/internal/service/tenant_test.go
+++ b/internal/service/tenant_test.go
@@ -73,26 +73,32 @@ func TestRemoveMemberAuthCheck(t *testing.T) {
 }

 // TestRemoveMemberSelfAllowed verifies that a user removing themselves passes the auth check.
-// (It will fail at the DB layer in unit tests, but the auth check must pass first.)
+// The DAO is nil so the operation fails at the data layer, but the auth check must pass first.
 func TestRemoveMemberSelfAllowed(t *testing.T) {
 	s := &TenantService{}
-	// userID == targetUserID: auth check should pass, expect DB error (nil userTenantDAO).
+	// userID == targetUserID: auth check should pass.
 	code, err := s.RemoveMember("user-abc", "tenant-xyz", "user-abc")
 	if code == common.CodeAuthenticationError {
 		t.Errorf("self-removal should pass auth check, got CodeAuthenticationError: %v", err)
 	}
+	if code != common.CodeServerError {
+		t.Errorf("expected CodeServerError (DAO not initialized), got %v", code)
+	}
+	if err == nil {
+		t.Error("expected non-nil error when userTenantDAO is nil")
+	}
 }

-// TestAcceptInviteAuthCheck verifies that AcceptInvite fails when no membership exists (nil DAO).
+// TestAcceptInviteAuthCheck verifies that AcceptInvite fails when DAO is not initialized.
 func TestAcceptInviteAuthCheck(t *testing.T) {
 	s := &TenantService{}
-	// nil userTenantDAO: FilterByUserIDAndTenantID will panic/err, so we expect CodeDataError.
+	// nil userTenantDAO: nil guard returns CodeServerError.
 	code, err := s.AcceptInvite("user-abc", "tenant-xyz")
 	if err == nil {
-		t.Fatal("expected error when no membership exists, got nil")
+		t.Fatal("expected error when userTenantDAO is nil, got nil")
 	}
-	if code != common.CodeDataError {
-		t.Errorf("expected CodeDataError, got %v", code)
+	if code != common.CodeServerError {
+		t.Errorf("expected CodeServerError (DAO not initialized), got %v", code)
 	}
 }