diff --git a/api/apps/restful_apis/file2document_api.py b/api/apps/restful_apis/file2document_api.py index baf8bdfeb3..f5009b648d 100644 --- a/api/apps/restful_apis/file2document_api.py +++ b/api/apps/restful_apis/file2document_api.py @@ -39,34 +39,33 @@ logger = logging.getLogger(__name__) def _convert_files(file_ids, kb_ids, user_id, mode): """Synchronous worker: add missing links or replace existing links.""" replace_existing = mode == "replace" + kb_ids = set(kb_ids) for id in file_ids: e, file = FileService.get_by_id(id) if not e: continue - existing_links = File2DocumentService.get_by_file_id(id) existing_kb_ids = set() - if replace_existing: - for inform in existing_links: - doc_id = inform.document_id - e, doc = DocumentService.get_by_id(doc_id) - if e and doc: - tenant_id = DocumentService.get_tenant_id(doc_id) + existing_links = File2DocumentService.get_by_file_id(id) + for inform in existing_links: + e, doc = DocumentService.get_by_id(inform.document_id) + if e and doc: + existing_kb_ids.add(doc.kb_id) + + # Delete existing link to KB if it is replaced + if replace_existing and doc.kb_id not in kb_ids: + logger.info("Unlink file_id=%s kb_id=%s", id, doc.kb_id) + tenant_id = DocumentService.get_tenant_id(doc.id) if not tenant_id: raise RuntimeError("Tenant not found!") if not DocumentService.remove_document(doc, tenant_id): raise RuntimeError("Database error (Document removal)!") - File2DocumentService.delete_by_document_id(doc_id) - if existing_links: - File2DocumentService.delete_by_file_id(id) - else: - for inform in existing_links: - e, doc = DocumentService.get_by_id(inform.document_id) - if e and doc: - existing_kb_ids.add(doc.kb_id) + File2DocumentService.delete_by_document_id(doc.id) for kb_id in kb_ids: + # Skip if the file is already linked to this KB if kb_id in existing_kb_ids: + logger.info("skip existing file_id=%s kb_id=%s", id, kb_id) continue e, kb = KnowledgebaseService.get_by_id(kb_id) if not e: diff --git a/internal/service/document/file2document.go b/internal/service/document/file2document.go index c25faf7b20..dc50877153 100644 --- a/internal/service/document/file2document.go +++ b/internal/service/document/file2document.go @@ -76,9 +76,6 @@ type LinkToDatasetsRequest struct { // LinkToDatasets validates inputs, expands folders, checks permissions, and // schedules convertFiles in a goroutine — mirroring Python convert(). -// mode is "add" (link on top of existing KBs) or "replace" (remove existing -// links first); any other value behaves as "replace", matching Python. -// Returns immediately (fire-and-forget for the heavy DB work). // // On validation failure it returns a sentinel error (see ErrLink* above) so the // handler can map it to a Python-compatible response without leaking internals. @@ -162,12 +159,15 @@ func (s *File2DocumentService) LinkToDatasets(ctx context.Context, userID string return nil } -// convertFiles mirrors Python _convert_files: for each file, depending on mode, -// either remove existing documents/mappings (replace) or keep them and skip -// already-linked KBs (add), then create a new document in each target KB and -// a fresh mapping. +// convertFiles mirrors Python _convert_files: keep existing links, remove links +// outside the requested KBs in replace mode, then add any missing links. func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs []string, userID, mode string) error { replaceExisting := mode != "add" + kbIDs = dedupeStrings(kbIDs) + requestedKBIDs := make(map[string]struct{}, len(kbIDs)) + for _, kbID := range kbIDs { + requestedKBIDs[kbID] = struct{}{} + } for _, fileID := range fileIDs { mappings, err := s.file2DocumentDAO.GetByFileID(ctx, dao.DB, fileID) if err != nil { @@ -175,39 +175,26 @@ func (s *File2DocumentService) convertFiles(ctx context.Context, fileIDs, kbIDs } existingKBIDs := make(map[string]struct{}) - if replaceExisting { - // Remove existing documents linked to this file. Routing through - // DocumentService.RemoveDocumentKeepFile ensures KB doc_num/chunk_num/ - // token_num counters are decremented (mirrors Python remove_document) - // while preserving the file record itself for re-linking. - for _, m := range mappings { - if m.DocumentID == nil { - continue - } - if err = s.documentSvc.RemoveDocumentKeepFile(ctx, *m.DocumentID); err != nil { - common.Warn("convertFiles: RemoveDocumentKeepFile failed", - zap.String("docID", *m.DocumentID), zap.Error(err)) - } + for _, m := range mappings { + if m.DocumentID == nil { + continue } - // Drop the file2document mappings for this file (mirrors Python - // File2DocumentService.delete_by_file_id, done once per file). - if err = s.file2DocumentDAO.DeleteByFileID(ctx, dao.DB, fileID); err != nil { - common.Warn("convertFiles: DeleteByFileID failed", zap.String("fileID", fileID), zap.Error(err)) + doc, getErr := s.documentDAO.GetByID(ctx, dao.DB, *m.DocumentID) + if getErr != nil || doc == nil { + continue } - } else { - // "add" mode: collect KB IDs already linked to this file so we - // skip them when creating new documents below. Existing links - // are preserved (mirrors Python _convert_files add path). - var doc *entity.Document - for _, m := range mappings { - if m.DocumentID == nil { + existingKBIDs[doc.KbID] = struct{}{} + if replaceExisting { + if _, exists := requestedKBIDs[doc.KbID]; exists { continue } - doc, err = s.documentDAO.GetByID(ctx, dao.DB, *m.DocumentID) - if err != nil || doc == nil { + if err = s.documentSvc.RemoveDocumentKeepFile(ctx, doc.ID); err != nil { + common.Warn("convertFiles: RemoveDocumentKeepFile failed", zap.String("docID", doc.ID), zap.Error(err)) continue } - existingKBIDs[doc.KbID] = struct{}{} + if err = s.file2DocumentDAO.DeleteByDocumentID(ctx, dao.DB, doc.ID); err != nil { + common.Warn("convertFiles: DeleteByDocumentID failed", zap.String("docID", doc.ID), zap.Error(err)) + } } } diff --git a/test/testcases/restful_api/test_file_routes_unit.py b/test/testcases/restful_api/test_file_routes_unit.py index f8dfb91ba7..5988d04591 100644 --- a/test/testcases/restful_api/test_file_routes_unit.py +++ b/test/testcases/restful_api/test_file_routes_unit.py @@ -686,7 +686,7 @@ def test_convert_files_mode_add_and_replace_unit(monkeypatch): assert len(inserted) == 2 assert removed == [("doc-f1", "tenant-1"), ("doc-f2", "tenant-1")] assert deleted_doc_links == ["doc-f1", "doc-f2"] - assert deleted_file_links == ["f1", "f2"] + assert deleted_file_links == [] @pytest.mark.p2