diff --git a/internal/service/kg_pipeline.go b/internal/service/kg_pipeline.go new file mode 100644 index 0000000000..fd2a01835f --- /dev/null +++ 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 +} diff --git a/internal/service/kg_retrieval.go b/internal/service/kg_retrieval.go new file mode 100644 index 0000000000..000f21a783 --- /dev/null +++ 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} +} diff --git a/internal/service/kg_retrieval_test.go b/internal/service/kg_retrieval_test.go new file mode 100644 index 0000000000..7e3989abd7 --- /dev/null +++ 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 } diff --git a/internal/common/kg_scoring.go b/internal/service/kg_scoring_funcs.go similarity index 72% rename from internal/common/kg_scoring.go rename to internal/service/kg_scoring_funcs.go index 977ad32644..f739e62fe6 100644 --- a/internal/common/kg_scoring.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 } } } diff --git a/internal/common/kg_scoring_test.go b/internal/service/kg_scoring_funcs_test.go similarity index 59% rename from internal/common/kg_scoring_test.go rename to internal/service/kg_scoring_funcs_test.go index 12da289004..62ed85d6b0 100644 --- a/internal/common/kg_scoring_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 -} diff --git a/internal/service/kg_search.go b/internal/service/kg_search.go index a5373eb650..3f93c98e6f 100644 --- 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. diff --git a/internal/service/tenant.go b/internal/service/tenant.go index c7bb26bd2a..6166559980 100644 --- 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") diff --git a/internal/service/tenant_test.go b/internal/service/tenant_test.go index c28095b3a1..eefefb9419 100644 --- 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) } }