diff --git a/api/repositories/knowledge/segment_repository.py b/api/repositories/knowledge/segment_repository.py index 0b478b331d833b..c9e3af3647443c 100644 --- a/api/repositories/knowledge/segment_repository.py +++ b/api/repositories/knowledge/segment_repository.py @@ -506,6 +506,32 @@ def get_indexing_snapshots( for segment_id in segment_ids ) + def get_exclusive_attachment_ids(self, segment_ref: SegmentRef, attachment_ids: Sequence[str]) -> tuple[str, ...]: + """Return candidates with no other segment binding in the same dataset. + + Accepts a trusted segment reference and candidate upload IDs. This read + does not mutate bindings or uploads. Disabled segments still own their + attachments; references in other datasets do not own this dataset's + vector collection. The result reflects bindings at read time. + """ + if not attachment_ids: + return () + with self._session_factory() as session: + shared_ids = set( + session.scalars( + select(SegmentAttachmentBinding.attachment_id).where( + SegmentAttachmentBinding.tenant_id == segment_ref.document.dataset.tenant_id, + SegmentAttachmentBinding.dataset_id == segment_ref.document.dataset.dataset_id, + SegmentAttachmentBinding.attachment_id.in_(attachment_ids), + or_( + SegmentAttachmentBinding.document_id != segment_ref.document.document_id, + SegmentAttachmentBinding.segment_id != segment_ref.segment_id, + ), + ) + ) + ) + return tuple(attachment_id for attachment_id in attachment_ids if attachment_id not in shared_ids) + def replace_attachments(self, segment_ref: SegmentRef, attachment_ids: Sequence[str]) -> None: with self._session_factory.begin() as session: if _get_segment(session, segment_ref) is None: diff --git a/api/services/knowledge/segments/indexing.py b/api/services/knowledge/segments/indexing.py index 462e19f5ecec4c..7768f8c77034af 100644 --- a/api/services/knowledge/segments/indexing.py +++ b/api/services/knowledge/segments/indexing.py @@ -249,16 +249,28 @@ def _attachment_documents(snapshot: SegmentIndexingSnapshot, files) -> list[Docu @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) def update_attachments(self, segment_ref: SegmentRef, attachment_ids: Sequence[str]) -> None: + """Replace an existing segment's attachments using tenant-owned uploads. + + Keep retained vectors and vectors referenced by another segment in this + dataset. Only removed, exclusively owned vectors can be deleted. External + index I/O completes before replacing bindings, so an index failure leaves + the previous bindings available for diagnosis and retry. + """ snapshot = self._segments.get_indexing_snapshot(segment_ref) if snapshot.dataset.indexing_technique != "high_quality" or set(snapshot.attachment_ids) == set(attachment_ids): return files = self._uploads.get_files(workspace_id=snapshot.dataset.tenant_id, file_ids=attachment_ids) if snapshot.dataset.is_multimodal: vector = self._vector(snapshot) - if snapshot.attachment_ids: - vector.delete_by_ids(list(snapshot.attachment_ids)) - if files: - vector.create_multimodal(self._attachment_documents(snapshot, files.values()), upload_files=files) + removed_ids = tuple(file_id for file_id in snapshot.attachment_ids if file_id not in files) + exclusive_ids = self._segments.get_exclusive_attachment_ids(segment_ref, removed_ids) + if exclusive_ids: + vector.delete_by_ids(list(exclusive_ids)) + added_files = {file_id: file for file_id, file in files.items() if file_id not in snapshot.attachment_ids} + if added_files: + vector.create_multimodal( + self._attachment_documents(snapshot, added_files.values()), upload_files=added_files + ) self._segments.replace_attachments(segment_ref, tuple(files)) def update_summary(self, segment_ref: SegmentRef, *, summary: str | None, content_changed: bool) -> None: diff --git a/api/tests/unit_tests/services/knowledge/segments/test_adapters.py b/api/tests/unit_tests/services/knowledge/segments/test_adapters.py index 67e457c65a2b49..91b8a9486e55de 100644 --- a/api/tests/unit_tests/services/knowledge/segments/test_adapters.py +++ b/api/tests/unit_tests/services/knowledge/segments/test_adapters.py @@ -821,6 +821,119 @@ def test_segment_update_preserves_or_replaces_attachments( assert ("image-1" in probe.deleted_ids) is removed_original +@pytest.mark.parametrize("other_document_id", ["document-1", "document-2"]) +@pytest.mark.parametrize("other_enabled", [True, False]) +def test_attachment_update_preserves_shared_vector_until_last_binding( + segment_update: SegmentUpdateEntry, + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + other_document_id: str, + other_enabled: bool, +) -> None: + app, repository, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + with sqlite_session_factory.begin() as session: + session.add(_upload("image-1", "workspace-1", "shared.png")) + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.is_multimodal = True + if other_document_id != "document-1": + session.add(_document(other_document_id, "workspace-1", "dataset-1")) + first = app.create_segment( + context, + dataset_id="dataset-1", + document_id="document-1", + values={"content": "first", "attachment_ids": ["image-1"]}, + ) + second = app.create_segment( + context, + dataset_id="dataset-1", + document_id=other_document_id, + values={"content": "second", "attachment_ids": ["image-1"]}, + ) + second_ref = DatasetRef("workspace-1", "dataset-1").document(other_document_id).segment(second.data.id) + repository.save_segment(second_ref, {"enabled": other_enabled}) + segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=first.data.id, + values={"content": "first", "attachment_ids": []}, + ) + assert "image-1" not in probe.deleted_ids + assert repository.get_indexing_snapshot(second_ref).attachment_ids == ("image-1",) + repository.save_segment(second_ref, {"enabled": True}) + segment_update( + context, + dataset_id="dataset-1", + document_id=other_document_id, + segment_id=second.data.id, + values={"content": "second", "attachment_ids": []}, + ) + assert probe.deleted_ids.count("image-1") == 1 + with sqlite_session_factory() as session: + assert session.get(UploadFile, "image-1") is not None + + +def test_attachment_update_does_not_rewrite_retained_vector( + segment_update: SegmentUpdateEntry, + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], +) -> None: + app, _, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + with sqlite_session_factory.begin() as session: + session.add_all([_upload("image-1", "workspace-1", "one.png"), _upload("image-2", "workspace-1", "two.png")]) + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.is_multimodal = True + created = app.create_segment( + context, + dataset_id="dataset-1", + document_id="document-1", + values={"content": "first", "attachment_ids": ["image-1", "image-2"]}, + ) + writes_before = probe.writes + segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "first", "attachment_ids": ["image-1"]}, + ) + assert probe.deleted_ids == ["image-2"] + assert probe.writes == writes_before + + +def test_attachment_reference_query_respects_vector_collection_scope( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + for workspace, dataset, document, segment, image in [ + ("workspace-1", "dataset-1", "document-1", "segment-1", "own"), + ("workspace-1", "dataset-2", "document-2", "segment-2", "other-dataset"), + ("workspace-2", "dataset-1", "document-2", "segment-2", "other-tenant"), + ("workspace-1", "dataset-1", "document-2", "segment-2", "shared"), + ]: + session.add( + SegmentAttachmentBinding( + tenant_id=workspace, + dataset_id=dataset, + document_id=document, + segment_id=segment, + attachment_id=image, + ) + ) + repository = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + ref = DatasetRef("workspace-1", "dataset-1").document("document-1").segment("segment-1") + assert repository.get_exclusive_attachment_ids(ref, ["own", "other-dataset", "other-tenant", "shared"]) == ( + "own", + "other-dataset", + "other-tenant", + ) + assert repository.get_exclusive_attachment_ids(ref, []) == () + + def test_segment_indexing_and_summary_use_closed_read_sessions( segment_update: SegmentUpdateEntry, indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session] ) -> None: