mirror of
https://github.com/langgenius/dify.git
synced 2026-04-21 11:17:38 +08:00
refactor: migrate SegmentAttachmentBinding to TypeBase (#34810)
This commit is contained in:
@ -244,7 +244,7 @@ class DatasetDocumentStore:
|
||||
return document_segment
|
||||
|
||||
def add_multimodel_documents_binding(self, segment_id: str, multimodel_documents: list[AttachmentDocument] | None):
|
||||
if multimodel_documents:
|
||||
if multimodel_documents and self._document_id is not None:
|
||||
for multimodel_document in multimodel_documents:
|
||||
binding = SegmentAttachmentBinding(
|
||||
tenant_id=self._dataset.tenant_id,
|
||||
|
||||
@ -1688,7 +1688,7 @@ class PipelineRecommendedPlugin(TypeBase):
|
||||
)
|
||||
|
||||
|
||||
class SegmentAttachmentBinding(Base):
|
||||
class SegmentAttachmentBinding(TypeBase):
|
||||
__tablename__ = "segment_attachment_bindings"
|
||||
__table_args__ = (
|
||||
sa.PrimaryKeyConstraint("id", name="segment_attachment_binding_pkey"),
|
||||
@ -1701,13 +1701,17 @@ class SegmentAttachmentBinding(Base):
|
||||
),
|
||||
sa.Index("segment_attachment_binding_attachment_idx", "attachment_id"),
|
||||
)
|
||||
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7()))
|
||||
id: Mapped[str] = mapped_column(
|
||||
StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
|
||||
)
|
||||
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
segment_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
attachment_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
created_at: Mapped[datetime] = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
|
||||
)
|
||||
|
||||
|
||||
class DocumentSegmentSummary(Base):
|
||||
|
||||
@ -721,6 +721,30 @@ class TestDatasetDocumentStoreMultimodelBinding:
|
||||
|
||||
mock_db.session.add.assert_not_called()
|
||||
|
||||
def test_add_multimodel_documents_binding_with_none_document_id(self):
|
||||
"""Test that no bindings are added when document_id is None."""
|
||||
|
||||
mock_dataset = MagicMock(spec=Dataset)
|
||||
mock_dataset.id = "test-dataset-id"
|
||||
mock_dataset.tenant_id = "tenant-1"
|
||||
|
||||
mock_attachment = MagicMock(spec=AttachmentDocument)
|
||||
mock_attachment.metadata = {"doc_id": "attachment-1"}
|
||||
|
||||
with patch("core.rag.docstore.dataset_docstore.db") as mock_db:
|
||||
mock_session = MagicMock()
|
||||
mock_db.session = mock_session
|
||||
|
||||
store = DatasetDocumentStore(
|
||||
dataset=mock_dataset,
|
||||
user_id="test-user-id",
|
||||
document_id=None,
|
||||
)
|
||||
|
||||
store.add_multimodel_documents_binding("seg-1", [mock_attachment])
|
||||
|
||||
mock_db.session.add.assert_not_called()
|
||||
|
||||
|
||||
class TestDatasetDocumentStoreAddDocumentsUpdateChild:
|
||||
"""Tests for add_documents when updating existing documents with children."""
|
||||
|
||||
Reference in New Issue
Block a user