From fc451df5d94c2e2d5465fdd040ebd2cdc46e3747 Mon Sep 17 00:00:00 2001 From: qinling0210 <88864212+qinling0210@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:27:52 +0800 Subject: [PATCH] =?UTF-8?q?fix(chunk):=20scope=20the=20chunk=20list=20and?= =?UTF-8?q?=20detail=20APIs=20to=20the=20route's=20datas=E2=80=A6=20(#2005?= =?UTF-8?q?9)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/dao/kb.go | 10 -- internal/engine/infinity/chunk.go | 11 ++ internal/handler/chunk.go | 10 +- internal/service/chunk/chunk.go | 132 +++++++------- internal/service/chunk/chunk_test.go | 257 +++++++++++++++++++++++++++ internal/service/chunk_types.go | 4 +- 6 files changed, 341 insertions(+), 83 deletions(-) diff --git a/internal/dao/kb.go b/internal/dao/kb.go index 92e4d4b804..cdbd3f4c59 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -467,16 +467,6 @@ func (dao *KnowledgebaseDAO) DecreaseDocumentNum(ctx context.Context, db *gorm.D }).Error } -// GetKBIDsByTenantID retrieves all knowledge base IDs for a tenant -// This matches the Python get_kb_ids method -func (dao *KnowledgebaseDAO) GetKBIDsByTenantID(ctx context.Context, db *gorm.DB, tenantID string) ([]string, error) { - var kbIDs []string - err := db.WithContext(ctx).Model(&entity.Knowledgebase{}). - Where("tenant_id = ? AND status = ?", tenantID, string(entity.StatusValid)). - Pluck("id", &kbIDs).Error - return kbIDs, err -} - // GetAllIDs retrieves all knowledge base IDs // This matches the Python get_all_ids method func (dao *KnowledgebaseDAO) GetAllIDs(ctx context.Context, db *gorm.DB) ([]string, error) { diff --git a/internal/engine/infinity/chunk.go b/internal/engine/infinity/chunk.go index ffd4865504..da6098b88d 100644 --- a/internal/engine/infinity/chunk.go +++ b/internal/engine/infinity/chunk.go @@ -1173,6 +1173,17 @@ func (e *Engine) Search(ctx context.Context, req *types.SearchRequest) (*types.S } } + // Filter-only queries (e.g. the management chunk list) are the ones + // that can silently span KBs, so echo per-table outcome at debug + // level: which table was queried with which filter, and how many + // rows came back. + if !hasTextMatch && !hasVectorMatch { + common.Debug("Infinity filter-only search", + zap.String("table", tableName), + zap.String("filter", filterStr), + zap.Int("rows", len(searchChunks))) + } + // Parse total_hits_count from ExtraInfo var tableTotal int64 if df.ExtraInfo != "" { diff --git a/internal/handler/chunk.go b/internal/handler/chunk.go index e6207c942e..483ff91565 100644 --- a/internal/handler/chunk.go +++ b/internal/handler/chunk.go @@ -151,14 +151,18 @@ func (h *ChunkHandler) Get(c *gin.Context) { return } + datasetID := c.Param("dataset_id") + documentID := c.Param("document_id") chunkID := c.Param("chunk_id") - if chunkID == "" { - common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "chunk_id is required") + if datasetID == "" || documentID == "" || chunkID == "" { + common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "dataset_id, document_id and chunk_id are required") return } req := &service.GetChunkRequest{ - ChunkID: chunkID, + DatasetID: datasetID, + DocumentID: documentID, + ChunkID: chunkID, } ctx := c.Request.Context() diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index af15ff71b4..61a71873ef 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -536,74 +536,70 @@ func (s *ChunkService) Get(ctx context.Context, req *service.GetChunkRequest, us return nil, fmt.Errorf("user has no accessible tenants") } - // Try each tenant to find the chunk - var chunk map[string]interface{} + // Find the tenant that owns this dataset + var targetTenantID string for _, tenant := range tenants { - // Get kbIDs for this tenant - kbIDs, err := s.kbDAO.GetKBIDsByTenantID(ctx, dao.DB, tenant.TenantID) - if err != nil { - continue - } - - indexName := fmt.Sprintf("ragflow_%s", tenant.TenantID) - - doc, err := s.docEngine.GetChunk(ctx, indexName, req.ChunkID, kbIDs) - if err != nil { - continue - } - - if doc != nil { - chunk, ok := doc.(map[string]interface{}) - if ok { - result := make(map[string]interface{}) - skipFields := map[string]bool{ - "id": true, "authors": true, "_score": true, "SCORE": true, - } - for k, v := range chunk { - if skipFields[k] || strings.HasSuffix(k, "_vec") || strings.Contains(k, "_sm_") || strings.HasSuffix(k, "_tks") || strings.HasSuffix(k, "_ltks") { - continue - } - switch k { - case "content": - result["content_with_weight"] = v - case "docnm": - result["docnm_kwd"] = v - case "important_keywords": - utility.SetFieldArray(result, "important_kwd", v) - case "questions": - utility.SetFieldArray(result, "question_kwd", v) - case "entities_kwd", "entity_kwd", "entity_type_kwd", "from_entity_kwd", - "name_kwd", "raptor_kwd", "removed_kwd", "source_id", "tag_kwd", - "to_entity_kwd", "toc_kwd", "authors_tks", "doc_type_kwd": - if utility.IsEmpty(v) { - result[k] = []interface{}{} - } else { - result[k] = v - } - case "tag_feas": - if utility.IsEmpty(v) { - result[k] = map[string]interface{}{} - } else { - result[k] = v - } - case "create_timestamp_flt", "rank_flt", "weight_flt": - if floatVal, ok := utility.ToFloat64(v); ok { - result[k] = utility.JSONFloat64(floatVal) - } - default: - result[k] = v - } - } - return &service.GetChunkResponse{Chunk: result}, nil - } + kb, err := s.kbDAO.GetByIDAndTenantID(ctx, dao.DB, req.DatasetID, tenant.TenantID) + if err == nil && kb != nil { + targetTenantID = tenant.TenantID + break } } + if targetTenantID == "" { + return nil, fmt.Errorf("user does not have access to this dataset") + } - if chunk == nil { + // Verify the document belongs to the dataset, mirroring Python's get_chunk + // (DocumentService.query(id=document_id, kb_id=dataset_id)). + doc, err := dao.NewDocumentDAO().GetByID(ctx, dao.DB, req.DocumentID) + if err != nil || doc == nil { + return nil, fmt.Errorf("document not found") + } + if doc.KbID != req.DatasetID { + return nil, fmt.Errorf("document does not belong to this dataset") + } + + // The lookup stays inside the dataset named in the route, and the row must + // resolve to the document named in the route: a chunk id alone is not + // authority to read another dataset's (or another document's) chunk. + indexName := fmt.Sprintf("ragflow_%s", targetTenantID) + rawChunk, err := s.docEngine.GetChunk(ctx, indexName, req.ChunkID, []string{req.DatasetID}) + if err != nil { + return nil, fmt.Errorf("chunk not found") + } + chunk, ok := rawChunk.(map[string]interface{}) + if !ok || chunk == nil { + return nil, fmt.Errorf("chunk not found") + } + if documentID, _ := chunk["doc_id"].(string); documentID != req.DocumentID { + return nil, fmt.Errorf("chunk not found") + } + if !utility.IsEmpty(chunk["compile_kwd"]) { return nil, fmt.Errorf("chunk not found") } - return &service.GetChunkResponse{Chunk: chunk}, nil + // Return the stored row with only the tokenized/vector runtime fields + // dropped, mirroring Python's _strip_chunk_runtime_fields. The frontend + // normalizes both this shape and the renamed one (mapChunkToLegacy), so + // returning the row verbatim is what Python does and costs no UI change. + result := make(map[string]interface{}, len(chunk)) + for k, v := range chunk { + if isChunkRuntimeField(k) { + continue + } + result[k] = v + } + return &service.GetChunkResponse{Chunk: result}, nil +} + +// isChunkRuntimeField reports the fields Python's _strip_chunk_runtime_fields +// removes before returning a chunk: the regex (_vec$|_sm_|_tks|_ltks) — _vec is +// anchored at the end, the other three are substring matches. +func isChunkRuntimeField(name string) bool { + return strings.HasSuffix(name, "_vec") || + strings.Contains(name, "_sm_") || + strings.Contains(name, "_tks") || + strings.Contains(name, "_ltks") } const ( @@ -832,12 +828,6 @@ func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest, return nil, fmt.Errorf("user does not have access to this document") } - // Get kbIDs for this tenant - kbIDs, err := s.kbDAO.GetKBIDsByTenantID(ctx, dao.DB, targetTenantID) - if err != nil { - return nil, fmt.Errorf("failed to get kb ids: %w", err) - } - indexName := fmt.Sprintf("ragflow_%s", targetTenantID) page := common.CoalesceInt(req.Page, 1) @@ -861,11 +851,15 @@ func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest, } } - // Build search request - same as retrieval test but filtered by doc_id + // Build search request - same as retrieval test but filtered by doc_id. + // KbIDs is pinned to the document's own dataset: searching every KB of the + // tenant also returned rows whose KB context is a different dataset (chunks + // carry img_id "-"), so a document of dataset A could list + // dataset B's chunks. Mirrors Python's [dataset_id] scope. searchReq := &types.SearchRequest{ IndexNames: []string{indexName}, MatchExprs: matchExprs, - KbIDs: kbIDs, + KbIDs: []string{doc.KbID}, Offset: (page - 1) * size, Limit: size, OrderBy: orderBy, diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index e9cfbd07c0..fdc2c20832 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -351,6 +351,241 @@ func TestListSortsChunksByDocumentPosition(t *testing.T) { } } +// TestListScopesSearchToDocumentDataset pins the dataset scope of the +// management chunk list: only the document's own KB may be searched. Passing +// every KB of the tenant made a document of dataset A list chunks that were +// written under dataset B's KB context (chunk rows carry img_id +// "-"), which the UI surfaced as "occasionally seeing another +// dataset's chunks". +func TestListScopesSearchToDocumentDataset(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + + userID := "user-1" + tenantID := "tenant-1" + datasetID := "kb-1" + otherDatasetID := "kb-2" + documentID := "doc-1" + insertChunkTestUserTenant(t, userID, tenantID) + insertChunkTestKB(t, datasetID, tenantID) + insertChunkTestKB(t, otherDatasetID, tenantID) + insertChunkTestDoc(t, documentID, datasetID) + + engine := &listChunksSearchEngine{} + svc := &ChunkService{ + docEngine: engine, + kbDAO: dao.NewKnowledgebaseDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + documentDAO: dao.NewDocumentDAO(), + } + + ctx := t.Context() + page := 1 + size := 30 + if _, err := svc.List(ctx, &service.ListChunksRequest{ + DatasetID: datasetID, + DocID: documentID, + Page: &page, + Size: &size, + }, userID); err != nil { + t.Fatalf("List() error = %v", err) + } + + if engine.searchReq == nil { + t.Fatal("expected Search to be called") + } + if !reflect.DeepEqual(engine.searchReq.KbIDs, []string{datasetID}) { + t.Fatalf("KbIDs = %#v, want the document's dataset only (%q)", engine.searchReq.KbIDs, datasetID) + } +} + +// TestGetScopesChunkLookupToDocumentDataset pins the detail endpoint's scope: +// the chunk is read from the dataset named in the route only, never from every +// KB of the tenant. +func TestGetScopesChunkLookupToDocumentDataset(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + + userID := "user-1" + tenantID := "tenant-1" + datasetID := "kb-1" + otherDatasetID := "kb-2" + documentID := "doc-1" + insertChunkTestUserTenant(t, userID, tenantID) + insertChunkTestKB(t, datasetID, tenantID) + insertChunkTestKB(t, otherDatasetID, tenantID) + insertChunkTestDoc(t, documentID, datasetID) + + engine := &getChunkTestEngine{chunk: map[string]interface{}{ + "id": "chunk-1", + "doc_id": documentID, + "kb_id": datasetID, + "docnm_kwd": "doc.txt", + "content_with_weight": "invoice body", + "content_ltks": []string{"invoice", "body"}, + }} + svc := &ChunkService{ + docEngine: engine, + kbDAO: dao.NewKnowledgebaseDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + } + + resp, err := svc.Get(t.Context(), &service.GetChunkRequest{ + DatasetID: datasetID, + DocumentID: documentID, + ChunkID: "chunk-1", + }, userID) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + if !reflect.DeepEqual(engine.datasetIDs, []string{datasetID}) { + t.Fatalf("datasetIDs = %#v, want the route's dataset only (%q)", engine.datasetIDs, datasetID) + } + if engine.chunkID != "chunk-1" { + t.Fatalf("chunkID = %q, want chunk-1", engine.chunkID) + } + if engine.indexName != "ragflow_"+tenantID { + t.Fatalf("indexName = %q, want ragflow_%s", engine.indexName, tenantID) + } + if got := resp.Chunk["content_with_weight"]; got != "invoice body" { + t.Fatalf("content_with_weight = %#v, want invoice body", got) + } + // Python returns the stored row (minus runtime fields), so the row identity + // fields survive the round trip. + if got := resp.Chunk["id"]; got != "chunk-1" { + t.Fatalf("id = %#v, want chunk-1", got) + } + if got := resp.Chunk["docnm_kwd"]; got != "doc.txt" { + t.Fatalf("docnm_kwd = %#v, want doc.txt", got) + } + if _, ok := resp.Chunk["content_ltks"]; ok { + t.Fatal("content_ltks should be stripped like Python's _strip_chunk_runtime_fields") + } +} + +// TestGetRejectsChunkOfAnotherDocument pins the doc_id guard: holding a chunk id +// is not authority to read it through a different document. +func TestGetRejectsChunkOfAnotherDocument(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + + userID := "user-1" + tenantID := "tenant-1" + datasetID := "kb-1" + documentID := "doc-1" + insertChunkTestUserTenant(t, userID, tenantID) + insertChunkTestKB(t, datasetID, tenantID) + insertChunkTestDoc(t, documentID, datasetID) + + engine := &getChunkTestEngine{chunk: map[string]interface{}{ + "doc_id": "doc-2", + "content_with_weight": "other document body", + }} + svc := &ChunkService{ + docEngine: engine, + kbDAO: dao.NewKnowledgebaseDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + } + + _, err := svc.Get(t.Context(), &service.GetChunkRequest{ + DatasetID: datasetID, + DocumentID: documentID, + ChunkID: "chunk-1", + }, userID) + if err == nil || !strings.Contains(err.Error(), "chunk not found") { + t.Fatalf("Get() error = %v, want chunk not found", err) + } +} + +// TestGetRequiresDocumentInRouteDataset pins the document guard: the chunk is +// only readable through a document that exists in the route's dataset, +// mirroring Python's DocumentService.query(id=document_id, kb_id=dataset_id). +func TestGetRequiresDocumentInRouteDataset(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + + userID := "user-1" + tenantID := "tenant-1" + datasetID := "kb-1" + otherDatasetID := "kb-2" + insertChunkTestUserTenant(t, userID, tenantID) + insertChunkTestKB(t, datasetID, tenantID) + insertChunkTestKB(t, otherDatasetID, tenantID) + insertChunkTestDoc(t, "doc-2", otherDatasetID) + + engine := &getChunkTestEngine{chunk: map[string]interface{}{ + "doc_id": "doc-2", + "content_with_weight": "other dataset document body", + }} + svc := &ChunkService{ + docEngine: engine, + kbDAO: dao.NewKnowledgebaseDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + } + + _, err := svc.Get(t.Context(), &service.GetChunkRequest{ + DatasetID: datasetID, + DocumentID: "doc-missing", + ChunkID: "chunk-1", + }, userID) + if err == nil || !strings.Contains(err.Error(), "document not found") { + t.Fatalf("Get() error = %v, want document not found", err) + } + + _, err = svc.Get(t.Context(), &service.GetChunkRequest{ + DatasetID: datasetID, + DocumentID: "doc-2", + ChunkID: "chunk-1", + }, userID) + if err == nil || !strings.Contains(err.Error(), "document does not belong to this dataset") { + t.Fatalf("Get() error = %v, want document does not belong to this dataset", err) + } +} + +// TestIsChunkRuntimeField pins the strip rule's Python-parity semantics +// (_vec$|_sm_|_tks|_ltks — only _vec is end-anchored): tokenized/vector runtime +// fields are dropped, stored fields survive. +func TestIsChunkRuntimeField(t *testing.T) { + cases := []struct { + name string + want bool + }{ + {"content_ltks", true}, + {"content_sm_ltks", true}, + {"tokenized_content_ltks", true}, + {"title_tks", true}, + {"title_sm_tks", true}, + {"authors_sm_tks", true}, + {"important_tks", true}, + {"question_tks", true}, + {"name_tks", true}, + {"tags_tks", true}, + {"q_1024_vec", true}, + {"q_2_vec", true}, + // Only _vec is anchored at the end, so an index-shaped name is kept. + {"q_2_vec_idx", false}, + {"content_with_weight", false}, + {"content", false}, + {"docnm_kwd", false}, + {"doc_id", false}, + {"kb_id", false}, + {"important_kwd", false}, + {"question_kwd", false}, + {"compile_kwd", false}, + {"tag_kwd", false}, + {"tag_feas", false}, + {"mom_id", false}, + {"position_int", false}, + {"img_id", false}, + {"available_int", false}, + } + for _, tc := range cases { + if got := isChunkRuntimeField(tc.name); got != tc.want { + t.Errorf("isChunkRuntimeField(%q) = %v, want %v", tc.name, got, tc.want) + } + } +} + func TestListBuildsMatchTextExprForKeywords(t *testing.T) { db := setupChunkTestDB(t) pushChunkTestDB(t, db) @@ -1275,6 +1510,28 @@ type listChunksSearchEngine struct { searchReq *types.SearchRequest } +type getChunkTestEngine struct { + parseTestDocEngine + chunk map[string]interface{} + err error + indexName string + chunkID string + datasetIDs []string +} + +func (e *getChunkTestEngine) GetChunk(_ context.Context, baseName, chunkID string, datasetIDs []string) (interface{}, error) { + e.indexName = baseName + e.chunkID = chunkID + e.datasetIDs = datasetIDs + if e.err != nil { + return nil, e.err + } + if e.chunk == nil { + return nil, nil + } + return e.chunk, nil +} + func (e *listChunksSearchEngine) Search(_ context.Context, req *types.SearchRequest) (*types.SearchResult, error) { e.searchReq = req return &types.SearchResult{ diff --git a/internal/service/chunk_types.go b/internal/service/chunk_types.go index 13935e2ec1..f9e64eaee2 100644 --- a/internal/service/chunk_types.go +++ b/internal/service/chunk_types.go @@ -53,7 +53,9 @@ type RetrievalTestResponse struct { // GetChunkRequest request for getting a chunk by ID type GetChunkRequest struct { - ChunkID string `json:"chunk_id"` + DatasetID string `json:"dataset_id"` + DocumentID string `json:"document_id"` + ChunkID string `json:"chunk_id"` } // GetChunkResponse response for getting a chunk