opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-19 10:11:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: migrate RAG index sessions to SQLite (#40080)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Asuka Minato
2026-08-10 08:43:18 +00:00
committed by GitHub
co-authored by autofix-ci[bot]
parent 648806bc9c
commit 96655436f4
5 changed files with 467 additions and 584 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py
+103 -257
View File
@@ -5,7 +5,7 @@ Tests cover all public methods and error paths of the DatasetDocumentStore class
which provides document storage and retrieval functionality for datasets in the RAG system.
"""
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock
import pytest
from sqlalchemy import func, select
@@ -19,6 +19,7 @@ TENANT_ID = "00000000-0000-0000-0000-000000000001"
DATASET_ID = "00000000-0000-0000-0000-000000000002"
DOCUMENT_ID = "00000000-0000-0000-0000-000000000003"
USER_ID = "00000000-0000-0000-0000-000000000004"
ATTACHMENT_ID = "00000000-0000-0000-0000-000000000005"
def _dataset() -> Dataset:
@@ -69,7 +70,15 @@ def _segment(
)
class TestDatasetDocumentStoreInit:
class _UsesSQLiteSession:
session: Session
@pytest.fixture(autouse=True)
def _inject_sqlite_session(self, sqlite_session: Session) -> None:
self.session = sqlite_session
class TestDatasetDocumentStoreInit(_UsesSQLiteSession):
"""Tests for DatasetDocumentStore initialization."""
def test_init_with_all_parameters(self):
@@ -107,7 +116,7 @@ class TestDatasetDocumentStoreInit:
assert store.dataset_id == "test-dataset-id"
class TestDatasetDocumentStoreSerialization:
class TestDatasetDocumentStoreSerialization(_UsesSQLiteSession):
"""Tests for to_dict and from_dict methods."""
def test_to_dict(self):
@@ -142,23 +151,20 @@ class TestDatasetDocumentStoreSerialization:
assert store._document_id == "test-doc"
class TestDatasetDocumentStoreDocs:
class TestDatasetDocumentStoreDocs(_UsesSQLiteSession):
"""Tests for the docs property."""
def test_docs_returns_document_dict(self):
"""Test that docs property returns a dictionary of documents."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_segment = _segment(index_node_id="node-1")
mock_session = MagicMock()
mock_session.scalars.return_value.all.return_value = [mock_segment]
mock_session = self.session
mock_session.add(mock_segment)
mock_session.flush()
store = DatasetDocumentStore(
dataset=mock_dataset,
dataset=_dataset(),
user_id="test-user-id",
)
@@ -174,8 +180,7 @@ class TestDatasetDocumentStoreDocs:
id="test-dataset-id",
)
mock_session = MagicMock()
mock_session.scalars.return_value.all.return_value = []
mock_session = self.session
store = DatasetDocumentStore(
dataset=mock_dataset,
@@ -187,12 +192,7 @@ class TestDatasetDocumentStoreDocs:
assert result == {}
@pytest.mark.parametrize(
"sqlite_session",
[(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
indirect=True,
)
class TestDatasetDocumentStoreAddDocuments:
class TestDatasetDocumentStoreAddDocuments(_UsesSQLiteSession):
"""Tests for add_documents method."""
def test_add_documents_new_document_with_token_count(self, sqlite_session: Session):
@@ -336,268 +336,147 @@ class TestDatasetDocumentStoreAddDocuments:
assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegment)) == 0
class TestDatasetDocumentStoreExists:
class TestDatasetDocumentStoreExists(_UsesSQLiteSession):
"""Tests for document_exists method."""
def test_document_exists_returns_true(self):
"""Test document_exists returns True when segment exists."""
mock_dataset = Dataset(
id="test-dataset-id",
)
_persist_segment(self.session)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
mock_segment = MagicMock()
mock_session = MagicMock()
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.document_exists("doc-1", session=mock_session)
assert result is True
assert store.document_exists("doc-1", session=self.session) is True
def test_document_exists_returns_false(self):
"""Test document_exists returns False when segment doesn't exist."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.document_exists("doc-1", session=mock_session)
assert result is False
assert store.document_exists("doc-1", session=self.session) is False
class TestDatasetDocumentStoreGetDocument:
class TestDatasetDocumentStoreGetDocument(_UsesSQLiteSession):
"""Tests for get_document method."""
def test_get_document_success(self):
"""Test getting a document successfully."""
mock_dataset = Dataset(
id="test-dataset-id",
)
_persist_segment(self.session, index_node_id="node-1")
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
mock_segment = _segment(index_node_id="node-1")
mock_session = MagicMock()
result = store.get_document("node-1", session=self.session, raise_error=False)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.get_document("node-1", session=mock_session, raise_error=False)
assert isinstance(result, Document)
assert result.page_content == "Test content"
assert isinstance(result, Document)
assert result.page_content == "Test content"
def test_get_document_returns_none_when_not_found(self):
"""Test get_document returns None when not found and raise_error=False."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.get_document("nonexistent", session=self.session, raise_error=False)
result = store.get_document("nonexistent", session=mock_session, raise_error=False)
assert result is None
assert result is None
def test_get_document_raises_when_not_found(self):
"""Test get_document raises ValueError when not found and raise_error=True."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
with pytest.raises(ValueError, match="not found"):
store.get_document("nonexistent", session=mock_session, raise_error=True)
with pytest.raises(ValueError, match="not found"):
store.get_document("nonexistent", session=self.session, raise_error=True)
class TestDatasetDocumentStoreDeleteDocument:
class TestDatasetDocumentStoreDeleteDocument(_UsesSQLiteSession):
"""Tests for delete_document method."""
def test_delete_document_success(self):
"""Test deleting a document successfully."""
mock_dataset = Dataset(
id="test-dataset-id",
)
segment = _persist_segment(self.session)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
mock_segment = MagicMock()
mock_session = MagicMock()
store.delete_document("doc-1", session=self.session)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
store.delete_document("doc-1", session=mock_session)
mock_session.delete.assert_called_with(mock_segment)
mock_session.flush.assert_called()
assert self.session.get(DocumentSegment, segment.id) is None
def test_delete_document_returns_none_when_not_found(self):
"""Test delete_document returns None when not found and raise_error=False."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.delete_document("nonexistent", session=self.session, raise_error=False)
result = store.delete_document("nonexistent", session=mock_session, raise_error=False)
assert result is None
assert result is None
def test_delete_document_raises_when_not_found(self):
"""Test delete_document raises ValueError when not found and raise_error=True."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
with pytest.raises(ValueError, match="not found"):
store.delete_document("nonexistent", session=mock_session, raise_error=True)
with pytest.raises(ValueError, match="not found"):
store.delete_document("nonexistent", session=self.session, raise_error=True)
class TestDatasetDocumentStoreHashOperations:
class TestDatasetDocumentStoreHashOperations(_UsesSQLiteSession):
"""Tests for set_document_hash and get_document_hash methods."""
def test_set_document_hash_success(self):
"""Test setting document hash successfully."""
mock_dataset = Dataset(
id="test-dataset-id",
)
segment = _persist_segment(self.session, index_node_hash="old-hash")
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
mock_segment = MagicMock()
mock_segment.index_node_hash = "old-hash"
mock_session = MagicMock()
store.set_document_hash("doc-1", "new-hash", session=self.session)
self.session.expire_all()
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
store.set_document_hash("doc-1", "new-hash", session=mock_session)
assert mock_segment.index_node_hash == "new-hash"
mock_session.flush.assert_called()
updated_segment = self.session.get(DocumentSegment, segment.id)
assert updated_segment is not None
assert updated_segment.index_node_hash == "new-hash"
def test_set_document_hash_returns_none_when_not_found(self):
"""Test set_document_hash returns None when segment not found."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.set_document_hash("nonexistent", "new-hash", session=self.session)
result = store.set_document_hash("nonexistent", "new-hash", session=mock_session)
assert result is None
assert result is None
def test_get_document_hash_success(self):
"""Test getting document hash successfully."""
mock_dataset = Dataset(
id="test-dataset-id",
)
_persist_segment(self.session, index_node_hash="test-hash")
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
mock_segment = MagicMock()
mock_segment.index_node_hash = "test-hash"
mock_session = MagicMock()
result = store.get_document_hash("doc-1", session=self.session)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.get_document_hash("doc-1", session=mock_session)
assert result == "test-hash"
assert result == "test-hash"
def test_get_document_hash_returns_none_when_not_found(self):
"""Test get_document_hash returns None when segment not found."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID)
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
)
result = store.get_document_hash("nonexistent", session=self.session)
result = store.get_document_hash("nonexistent", session=mock_session)
assert result is None
assert result is None
class TestDatasetDocumentStoreSegment:
class TestDatasetDocumentStoreSegment(_UsesSQLiteSession):
"""Tests for get_document_segment method."""
def test_get_document_segment_returns_segment(self):
"""Test getting a document segment."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_segment = _segment()
mock_session = MagicMock()
mock_session.scalar.return_value = mock_segment
mock_session = self.session
mock_session.add(mock_segment)
mock_session.flush()
store = DatasetDocumentStore(
dataset=mock_dataset,
dataset=_dataset(),
user_id="test-user-id",
)
@@ -608,15 +487,10 @@ class TestDatasetDocumentStoreSegment:
def test_get_document_segment_returns_none(self):
"""Test getting a non-existent document segment."""
mock_dataset = Dataset(
id="test-dataset-id",
)
mock_session = MagicMock()
mock_session.scalar.return_value = None
mock_session = self.session
store = DatasetDocumentStore(
dataset=mock_dataset,
dataset=_dataset(),
user_id="test-user-id",
)
@@ -625,102 +499,79 @@ class TestDatasetDocumentStoreSegment:
assert result is None
class TestDatasetDocumentStoreMultimodelBinding:
class TestDatasetDocumentStoreMultimodelBinding(_UsesSQLiteSession):
"""Tests for add_multimodel_documents_binding method."""
def test_add_multimodel_documents_binding_with_attachments(self):
"""Test adding multimodel document bindings."""
mock_dataset = Dataset(
id="test-dataset-id",
tenant_id="tenant-1",
)
attachment = AttachmentDocument(page_content="attachment", metadata={"doc_id": ATTACHMENT_ID})
mock_attachment = MagicMock(spec=AttachmentDocument)
mock_attachment.metadata = {"doc_id": "attachment-1"}
mock_session = MagicMock()
mock_session = self.session
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
dataset=_dataset(),
user_id=USER_ID,
document_id=DOCUMENT_ID,
)
store.add_multimodel_documents_binding("seg-1", [mock_attachment], session=mock_session)
store.add_multimodel_documents_binding("seg-1", [attachment], session=mock_session)
mock_session.flush()
mock_session.add.assert_called()
binding = mock_session.scalar(select(SegmentAttachmentBinding))
assert binding is not None
assert binding.segment_id == "seg-1"
assert binding.attachment_id == ATTACHMENT_ID
def test_add_multimodel_documents_binding_without_attachments(self):
"""Test adding bindings with None attachments."""
mock_dataset = Dataset(
id="test-dataset-id",
tenant_id="tenant-1",
)
mock_session = MagicMock()
mock_session = self.session
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
dataset=_dataset(),
user_id=USER_ID,
document_id=DOCUMENT_ID,
)
store.add_multimodel_documents_binding("seg-1", None, session=mock_session)
mock_session.add.assert_not_called()
assert mock_session.scalar(select(func.count()).select_from(SegmentAttachmentBinding)) == 0
def test_add_multimodel_documents_binding_with_empty_list(self):
"""Test adding bindings with empty list."""
mock_dataset = Dataset(
id="test-dataset-id",
tenant_id="tenant-1",
)
mock_session = MagicMock()
mock_session = self.session
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
dataset=_dataset(),
user_id=USER_ID,
document_id=DOCUMENT_ID,
)
store.add_multimodel_documents_binding("seg-1", [], session=mock_session)
mock_session.add.assert_not_called()
assert mock_session.scalar(select(func.count()).select_from(SegmentAttachmentBinding)) == 0
def test_add_multimodel_documents_binding_with_none_document_id(self):
"""Test that no bindings are added when document_id is None."""
mock_dataset = Dataset(
id="test-dataset-id",
tenant_id="tenant-1",
)
attachment = AttachmentDocument(page_content="attachment", metadata={"doc_id": ATTACHMENT_ID})
mock_attachment = MagicMock(spec=AttachmentDocument)
mock_attachment.metadata = {"doc_id": "attachment-1"}
mock_session = MagicMock()
mock_session = self.session
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
dataset=_dataset(),
user_id=USER_ID,
document_id=None,
)
store.add_multimodel_documents_binding("seg-1", [mock_attachment], session=mock_session)
store.add_multimodel_documents_binding("seg-1", [attachment], session=mock_session)
mock_session.add.assert_not_called()
assert mock_session.scalar(select(func.count()).select_from(SegmentAttachmentBinding)) == 0
@pytest.mark.parametrize(
"sqlite_session",
[(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
indirect=True,
)
class TestDatasetDocumentStoreAddDocumentsUpdateChild:
class TestDatasetDocumentStoreAddDocumentsUpdateChild(_UsesSQLiteSession):
"""Tests for add_documents when updating existing documents with children."""
def test_add_documents_update_existing_with_children(self, sqlite_session: Session):
@@ -768,12 +619,7 @@ class TestDatasetDocumentStoreAddDocumentsUpdateChild:
assert children[0].index_node_id == "child-1"
@pytest.mark.parametrize(
"sqlite_session",
[(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
indirect=True,
)
class TestDatasetDocumentStoreAddDocumentsUpdateAnswer:
class TestDatasetDocumentStoreAddDocumentsUpdateAnswer(_UsesSQLiteSession):
"""Tests for add_documents when updating existing documents with answer metadata."""
def test_add_documents_update_existing_with_answer(self, sqlite_session: Session):
api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py
+125 -108
View File
@@ -1,21 +1,36 @@
import logging
from contextlib import nullcontext
from datetime import datetime
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock, patch
import pytest
from sqlalchemy import event
from sqlalchemy.orm import Session, sessionmaker
from core.entities.knowledge_entities import PreviewDetail
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor
from core.rag.models.document import AttachmentDocument, Document
from extensions.storage.storage_type import StorageType
from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, ImagePromptMessageContent
from graphon.model_runtime.entities.model_entities import ModelFeature
from models.dataset import DocumentSegment, SegmentAttachmentBinding
from models.enums import CreatorUserRole
from models.model import UploadFile
class TestParagraphIndexProcessor:
session: Session
session_factory: sessionmaker[Session]
@pytest.fixture(autouse=True)
def _inject_sqlite_sessions(self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> None:
self.session = sqlite_session
self.session_factory = sqlite_session_factory
@pytest.fixture
def processor(self) -> ParagraphIndexProcessor:
return ParagraphIndexProcessor()
@@ -54,9 +69,34 @@ class TestParagraphIndexProcessor:
usage=LLMUsage.empty_usage(),
)
def _upload_file(
self,
*,
file_id: str,
tenant_id: str = "tenant-1",
name: str = "image.png",
extension: str = "png",
mime_type: str = "image/png",
) -> UploadFile:
upload_file = UploadFile(
tenant_id=tenant_id,
storage_type=StorageType.LOCAL,
key=f"key-{file_id}",
name=name,
size=1,
extension=extension,
mime_type=mime_type,
created_by_role=CreatorUserRole.ACCOUNT,
created_by="user-1",
created_at=datetime(2025, 1, 1),
used=True,
)
upload_file.id = file_id
return upload_file
def test_extract_forwards_automatic_flag(self, processor: ParagraphIndexProcessor) -> None:
extract_setting = Mock()
session = Mock()
session = self.session
expected_docs = [Document(page_content="chunk", metadata={})]
with patch(
@@ -69,7 +109,7 @@ class TestParagraphIndexProcessor:
mock_extract.assert_called_once_with(extract_setting=extract_setting, is_automatic=True, session=session)
def test_transform_validates_process_rule(self, processor: ParagraphIndexProcessor) -> None:
session = Mock()
session = self.session
with pytest.raises(ValueError, match="No process rule found"):
processor.transform([Document(page_content="text", metadata={})], process_rule=None, session=session)
@@ -82,7 +122,7 @@ class TestParagraphIndexProcessor:
self, processor: ParagraphIndexProcessor, process_rule: dict[str, Any]
) -> None:
rules_without_segmentation = SimpleNamespace(segmentation=None)
session = Mock()
session = self.session
with patch(
"core.rag.index_processor.processor.paragraph_index_processor.Rule.model_validate",
@@ -99,7 +139,7 @@ class TestParagraphIndexProcessor:
self, processor: ParagraphIndexProcessor, process_rule: dict[str, Any]
) -> None:
source_document = Document(page_content="source", metadata={"dataset_id": "dataset-1", "document_id": "doc-1"})
session = Mock()
session = self.session
splitter = Mock()
splitter.split_documents.return_value = [
Document(page_content=".first", metadata={}),
@@ -138,7 +178,7 @@ class TestParagraphIndexProcessor:
def test_transform_automatic_mode_uses_default_rules(self, processor: ParagraphIndexProcessor) -> None:
splitter = Mock()
splitter.split_documents.return_value = [Document(page_content="text", metadata={})]
session = Mock()
session = self.session
with (
patch(
@@ -173,7 +213,7 @@ class TestParagraphIndexProcessor:
) -> None:
docs = [Document(page_content="chunk", metadata={})]
multimodal_docs = [AttachmentDocument(page_content="image", metadata={})]
session = Mock()
session = self.session
with (
patch("core.rag.index_processor.processor.paragraph_index_processor.Vector") as mock_vector_cls,
@@ -191,7 +231,7 @@ class TestParagraphIndexProcessor:
) -> None:
dataset.indexing_technique = IndexTechniqueType.ECONOMY
docs = [Document(page_content="chunk", metadata={})]
session = Mock()
session = self.session
keywords_list = [["k1"], ["k2"]]
with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls:
@@ -204,7 +244,7 @@ class TestParagraphIndexProcessor:
) -> None:
dataset.indexing_technique = IndexTechniqueType.ECONOMY
docs = [Document(page_content="chunk", metadata={})]
session = Mock()
session = self.session
with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls:
processor.load(dataset, docs, session=session)
@@ -212,10 +252,20 @@ class TestParagraphIndexProcessor:
mock_keyword_cls.return_value.add_texts.assert_called_once_with(docs, session)
def test_clean_deletes_summaries_and_vector(self, processor: ParagraphIndexProcessor, dataset: Mock) -> None:
scalars_result = Mock()
scalars_result.all.return_value = [SimpleNamespace(id="seg-1")]
session = Mock()
session.scalars.return_value = scalars_result
session = self.session
segment = DocumentSegment(
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
document_id="doc-1",
position=1,
content="segment",
word_count=1,
tokens=1,
created_by="user-1",
index_node_id="node-1",
)
session.add(segment)
session.flush()
with (
patch(
@@ -226,14 +276,14 @@ class TestParagraphIndexProcessor:
vector = mock_vector_cls.return_value
processor.clean(dataset, ["node-1"], delete_summaries=True, session=session)
mock_summary.assert_called_once_with(dataset, ["seg-1"], session=session)
mock_summary.assert_called_once_with(dataset, [segment.id], session=session)
vector.delete_by_ids.assert_called_once_with(["node-1"])
def test_clean_economy_deletes_summaries_and_keywords(
self, processor: ParagraphIndexProcessor, dataset: Mock
) -> None:
dataset.indexing_technique = IndexTechniqueType.ECONOMY
session = Mock()
session = self.session
with (
patch(
@@ -248,7 +298,7 @@ class TestParagraphIndexProcessor:
def test_clean_deletes_keywords_by_ids(self, processor: ParagraphIndexProcessor, dataset: Mock) -> None:
dataset.indexing_technique = IndexTechniqueType.ECONOMY
session = Mock()
session = self.session
with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls:
processor.clean(dataset, ["node-2"], with_keywords=True, session=session)
@@ -257,9 +307,9 @@ class TestParagraphIndexProcessor:
def test_index_list_chunks_high_quality(
self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock
) -> None:
session = Mock()
session = self.session
phase_events: list[str] = []
session.commit.side_effect = lambda: phase_events.append("commit")
event.listen(session, "after_commit", lambda _session: phase_events.append("commit"))
with (
patch(
"core.rag.index_processor.processor.paragraph_index_processor.helper.generate_text_hash",
@@ -299,9 +349,9 @@ class TestParagraphIndexProcessor:
self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock
) -> None:
dataset.indexing_technique = IndexTechniqueType.ECONOMY
session = Mock()
session = self.session
phase_events: list[str] = []
session.commit.side_effect = lambda: phase_events.append("commit")
event.listen(session, "after_commit", lambda _session: phase_events.append("commit"))
with (
patch(
"core.rag.index_processor.processor.paragraph_index_processor.helper.generate_text_hash",
@@ -334,8 +384,8 @@ class TestParagraphIndexProcessor:
)
chunk_without_files = SimpleNamespace(content="content-2", files=None)
structure = SimpleNamespace(general_chunks=[chunk_with_files, chunk_without_files])
session = Mock()
account_session = Mock()
session = self.session
account_session = self.session_factory()
with (
patch(
@@ -374,7 +424,7 @@ class TestParagraphIndexProcessor:
self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock
) -> None:
structure = SimpleNamespace(general_chunks=[SimpleNamespace(content="content", files=None)])
session = Mock()
session = self.session
with (
patch(
@@ -403,8 +453,8 @@ class TestParagraphIndexProcessor:
def test_generate_summary_preview_success_and_failure(self, processor: ParagraphIndexProcessor) -> None:
preview_items = [PreviewDetail(content="chunk-1"), PreviewDetail(content="chunk-2")]
session = Mock()
worker_sessions = [Mock(), Mock()]
session = self.session
worker_sessions = [self.session_factory(), self.session_factory()]
with (
patch(
@@ -440,7 +490,9 @@ class TestParagraphIndexProcessor:
patch("flask.current_app", fake_current_app),
patch.object(processor, "generate_summary", return_value=("summary", LLMUsage.empty_usage())),
):
result = processor.generate_summary_preview("tenant-1", preview_items, {"enable": True}, session=Mock())
result = processor.generate_summary_preview(
"tenant-1", preview_items, {"enable": True}, session=self.session
)
assert result[0].summary == "summary"
@@ -456,16 +508,16 @@ class TestParagraphIndexProcessor:
patch("concurrent.futures.wait", side_effect=[(set(), {future}), (set(), set())]),
):
with pytest.raises(ValueError, match="timeout"):
processor.generate_summary_preview("tenant-1", preview_items, {"enable": True}, session=Mock())
processor.generate_summary_preview("tenant-1", preview_items, {"enable": True}, session=self.session)
future.cancel.assert_called_once()
def test_generate_summary_validates_input(self) -> None:
with pytest.raises(ValueError, match="must be enabled"):
ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": False}, session=Mock())
ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": False}, session=self.session)
with pytest.raises(ValueError, match="model_name and model_provider_name"):
ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": True}, session=Mock())
ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": True}, session=self.session)
def test_generate_summary_text_only_flow(self, caplog: pytest.LogCaptureFixture) -> None:
model_instance = Mock()
@@ -495,7 +547,7 @@ class TestParagraphIndexProcessor:
"text content",
{"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"},
document_language="English",
session=Mock(),
session=self.session,
)
assert summary == "text summary"
@@ -537,7 +589,7 @@ class TestParagraphIndexProcessor:
"text content",
{"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"},
segment_id="seg-1",
session=Mock(),
session=self.session,
)
assert summary == "vision summary"
@@ -580,7 +632,7 @@ class TestParagraphIndexProcessor:
"tenant-1",
"text content",
{"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"},
session=Mock(),
session=self.session,
)
assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) == 1
@@ -594,30 +646,19 @@ class TestParagraphIndexProcessor:
"![img2](/files/22222222-2222-2222-2222-222222222222/file-preview) "
"![tool](/files/tools/33333333-3333-3333-3333-333333333333.png)"
)
image_upload = SimpleNamespace(
id="11111111-1111-1111-1111-111111111111",
tenant_id="tenant-1",
name="image.png",
mime_type="image/png",
extension="png",
source_url="",
size=1,
key="key",
session = self.session
session.add_all(
[
self._upload_file(file_id="11111111-1111-1111-1111-111111111111"),
self._upload_file(
file_id="22222222-2222-2222-2222-222222222222",
name="file.txt",
extension="txt",
mime_type="text/plain",
),
]
)
non_image_upload = SimpleNamespace(
id="22222222-2222-2222-2222-222222222222",
tenant_id="tenant-1",
name="file.txt",
mime_type="text/plain",
extension="txt",
source_url="",
size=1,
key="key",
)
scalars_result = Mock()
scalars_result.all.return_value = [image_upload, non_image_upload]
session = Mock()
session.scalars.return_value = scalars_result
session.flush()
with (
patch(
@@ -633,28 +674,14 @@ class TestParagraphIndexProcessor:
assert not any(record.levelno == logging.WARNING for record in caplog.records)
def test_extract_images_from_text_returns_empty_when_no_matches(self) -> None:
scalars_result = Mock()
scalars_result.all.return_value = []
session = Mock()
session.scalars.return_value = scalars_result
session = self.session
assert ParagraphIndexProcessor._extract_images_from_text("tenant-1", "no images here", session) == []
def test_extract_images_from_text_logs_when_build_fails(self, caplog: pytest.LogCaptureFixture) -> None:
text = "![img](/files/11111111-1111-1111-1111-111111111111/image-preview)"
image_upload = SimpleNamespace(
id="11111111-1111-1111-1111-111111111111",
tenant_id="tenant-1",
name="image.png",
mime_type="image/png",
extension="png",
source_url="",
size=1,
key="key",
)
scalars_result = Mock()
scalars_result.all.return_value = [image_upload]
session = Mock()
session.scalars.return_value = scalars_result
session = self.session
session.add(self._upload_file(file_id="11111111-1111-1111-1111-111111111111"))
session.flush()
with (
patch(
@@ -669,49 +696,39 @@ class TestParagraphIndexProcessor:
assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) == 1
def test_extract_images_from_segment_attachments(self, caplog: pytest.LogCaptureFixture) -> None:
image_upload = SimpleNamespace(
id="file-1",
name="image",
extension="png",
mime_type="image/png",
source_url="",
size=1,
key="k1",
)
bad_upload = SimpleNamespace(
id="file-2",
name="broken",
extension=None,
mime_type="image/png",
source_url="",
size=1,
key="k2",
)
non_image_upload = SimpleNamespace(
id="file-3",
name="text",
extension="txt",
mime_type="text/plain",
source_url="",
size=1,
key="k3",
)
execute_result = Mock()
execute_result.all.return_value = [(None, image_upload), (None, bad_upload), (None, non_image_upload)]
session = Mock()
session.execute.return_value = execute_result
session = self.session
uploads = [
self._upload_file(file_id="file-1", name="image"),
self._upload_file(file_id="file-2", name="broken"),
self._upload_file(file_id="file-3", name="text", extension="txt", mime_type="text/plain"),
]
bindings = [
SegmentAttachmentBinding(
tenant_id="tenant-1",
dataset_id="dataset-1",
document_id="doc-1",
segment_id="seg-1",
attachment_id=upload.id,
)
for upload in uploads
]
session.add_all([*uploads, *bindings])
session.flush()
with caplog.at_level(logging.WARNING, logger="core.rag.index_processor.processor.paragraph_index_processor"):
with (
patch(
"core.rag.index_processor.processor.paragraph_index_processor.File",
side_effect=[SimpleNamespace(id="file-1"), RuntimeError("bad file")],
),
caplog.at_level(logging.WARNING, logger="core.rag.index_processor.processor.paragraph_index_processor"),
):
files = ParagraphIndexProcessor._extract_images_from_segment_attachments("tenant-1", "seg-1", session)
assert len(files) == 1
assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) == 1
def test_extract_images_from_segment_attachments_empty(self) -> None:
execute_result = Mock()
execute_result.all.return_value = []
session = Mock()
session.execute.return_value = execute_result
session = self.session
empty_files = ParagraphIndexProcessor._extract_images_from_segment_attachments("tenant-1", "seg-1", session)
api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py
+80 -48
View File
@@ -3,15 +3,26 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch
import pytest
from sqlalchemy import event
from sqlalchemy.orm import Session, sessionmaker
from core.entities.knowledge_entities import PreviewDetail
from core.rag.entities import ParentMode, Rule, Segmentation
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor
from core.rag.models.document import AttachmentDocument, ChildDocument, Document
from models.dataset import ChildChunk, DatasetProcessRule, DocumentSegment
class TestParentChildIndexProcessor:
session: Session
session_factory: sessionmaker[Session]
@pytest.fixture(autouse=True)
def _inject_sqlite_sessions(self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> None:
self.session = sqlite_session
self.session_factory = sqlite_session_factory
@pytest.fixture
def processor(self) -> ParentChildIndexProcessor:
return ParentChildIndexProcessor()
@@ -50,7 +61,7 @@ class TestParentChildIndexProcessor:
def test_extract_forwards_automatic_flag(self, processor: ParentChildIndexProcessor) -> None:
extract_setting = Mock()
session = Mock()
session = self.session
expected = [Document(page_content="chunk", metadata={})]
with patch(
@@ -63,7 +74,7 @@ class TestParentChildIndexProcessor:
mock_extract.assert_called_once_with(extract_setting=extract_setting, is_automatic=True, session=session)
def test_transform_validates_process_rule(self, processor: ParentChildIndexProcessor) -> None:
session = MagicMock()
session = self.session
with pytest.raises(ValueError, match="No process rule found"):
processor.transform([Document(page_content="text", metadata={})], process_rule=None, session=session)
@@ -82,7 +93,7 @@ class TestParentChildIndexProcessor:
processor.transform(
[Document(page_content="text", metadata={})],
process_rule={"mode": "custom", "rules": {"enabled": True}},
session=MagicMock(),
session=self.session,
)
def test_transform_paragraph_builds_parent_and_child_docs(self, processor: ParentChildIndexProcessor) -> None:
@@ -117,7 +128,7 @@ class TestParentChildIndexProcessor:
[parent_document],
process_rule={"mode": "custom", "rules": {"enabled": True}},
preview=False,
session=MagicMock(),
session=self.session,
)
assert len(result) == 1
@@ -154,7 +165,7 @@ class TestParentChildIndexProcessor:
documents,
process_rule={"mode": "custom", "rules": {"enabled": True}},
preview=True,
session=MagicMock(),
session=self.session,
)
assert len(result) == 10
@@ -188,7 +199,7 @@ class TestParentChildIndexProcessor:
docs,
process_rule={"mode": "hierarchical", "rules": {"enabled": True}},
preview=True,
session=MagicMock(),
session=self.session,
)
assert len(result) == 1
@@ -205,7 +216,7 @@ class TestParentChildIndexProcessor:
],
)
multimodal_docs = [AttachmentDocument(page_content="image", metadata={})]
session = MagicMock()
session = self.session
with patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls:
vector = mock_vector_cls.return_value
@@ -219,7 +230,7 @@ class TestParentChildIndexProcessor:
vector.create_multimodal.assert_called_once_with(multimodal_docs)
def test_clean_with_precomputed_child_ids(self, processor: ParentChildIndexProcessor, dataset: Mock) -> None:
session = MagicMock()
session = self.session
with (
patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls,
@@ -234,16 +245,42 @@ class TestParentChildIndexProcessor:
)
vector.delete_by_ids.assert_called_once_with(["child-1", "child-2"])
session.execute.assert_called()
session.flush.assert_called_once()
assert session.query(ChildChunk).count() == 0
def test_clean_queries_child_ids_when_not_precomputed(
self, processor: ParentChildIndexProcessor, dataset: Mock
) -> None:
execute_result = Mock()
execute_result.all.return_value = [("child-1",), (None,), ("child-2",)]
session = MagicMock()
session.execute.return_value = execute_result
session = self.session
parent = DocumentSegment(
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
document_id="doc-1",
position=1,
content="parent",
word_count=1,
tokens=1,
created_by="user-1",
index_node_id="node-1",
)
session.add(parent)
session.flush()
session.add_all(
[
ChildChunk(
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
document_id="doc-1",
segment_id=parent.id,
position=index,
content=f"child-{index}",
word_count=1,
created_by="user-1",
index_node_id=f"child-{index}",
)
for index in (1, 2)
]
)
session.flush()
with (
patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls,
@@ -254,7 +291,7 @@ class TestParentChildIndexProcessor:
vector.delete_by_ids.assert_called_once_with(["child-1", "child-2"])
def test_clean_dataset_wide_cleanup(self, processor: ParentChildIndexProcessor, dataset: Mock) -> None:
session = MagicMock()
session = self.session
with (
patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls,
@@ -263,17 +300,23 @@ class TestParentChildIndexProcessor:
processor.clean(dataset, None, delete_child_chunks=True, session=session)
vector.delete.assert_called_once()
session.execute.assert_called()
session.flush.assert_called_once()
assert session.query(ChildChunk).count() == 0
def test_clean_deletes_summaries_when_requested(self, processor: ParentChildIndexProcessor, dataset: Mock) -> None:
scalars_result = Mock()
scalars_result.all.return_value = [SimpleNamespace(id="seg-1")]
session = MagicMock()
session.scalars.return_value = scalars_result
session_ctx = MagicMock()
session_ctx.__enter__.return_value = session
session_ctx.__exit__.return_value = False
session = self.session
segment = DocumentSegment(
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
document_id="doc-1",
position=1,
content="parent",
word_count=1,
tokens=1,
created_by="user-1",
index_node_id="node-1",
)
session.add(segment)
session.flush()
with (
patch(
@@ -283,7 +326,7 @@ class TestParentChildIndexProcessor:
):
processor.clean(dataset, ["node-1"], delete_summaries=True, precomputed_child_node_ids=[], session=session)
mock_summary.assert_called_once_with(dataset, ["seg-1"], session=session)
mock_summary.assert_called_once_with(dataset, [segment.id], session=session)
def test_clean_deletes_all_summaries_when_node_ids_missing(
self, processor: ParentChildIndexProcessor, dataset: Mock
@@ -294,7 +337,7 @@ class TestParentChildIndexProcessor:
) as mock_summary,
patch("core.rag.index_processor.processor.parent_child_index_processor.Vector"),
):
session = MagicMock()
session = self.session
processor.clean(dataset, None, delete_summaries=True, session=session)
mock_summary.assert_called_once_with(dataset, None, session=session)
@@ -341,20 +384,15 @@ class TestParentChildIndexProcessor:
)
],
)
dataset_rule = SimpleNamespace(id="rule-1")
session = MagicMock()
session = self.session
phase_events: list[str] = []
session.commit.side_effect = lambda: phase_events.append("commit")
event.listen(session, "after_commit", lambda _session: phase_events.append("commit"))
with (
patch(
"core.rag.index_processor.processor.parent_child_index_processor.ParentChildStructureChunk.model_validate",
return_value=parent_childs,
),
patch(
"core.rag.index_processor.processor.parent_child_index_processor.DatasetProcessRule",
return_value=dataset_rule,
),
patch(
"core.rag.index_processor.processor.parent_child_index_processor.helper.generate_text_hash",
side_effect=lambda text: f"hash-{text}",
@@ -373,9 +411,8 @@ class TestParentChildIndexProcessor:
processor.index(dataset, dataset_document, {"parent_child_chunks": []}, session)
assert phase_events == ["count", "store", "commit", "vector"]
assert dataset_document.dataset_process_rule_id == "rule-1"
session.add.assert_called_once_with(dataset_rule)
session.flush.assert_called_once()
dataset_rule = session.get(DatasetProcessRule, dataset_document.dataset_process_rule_id)
assert dataset_rule is not None
documents = mock_token_counter.call_args.kwargs["documents"]
assert [document.page_content for document in documents] == ["parent text"]
mock_token_counter.assert_called_once_with(dataset=dataset, documents=documents)
@@ -396,19 +433,14 @@ class TestParentChildIndexProcessor:
parent_mode=ParentMode.PARAGRAPH,
parent_child_chunks=[SimpleNamespace(parent_content="parent", child_contents=["child"], files=None)],
)
dataset_rule = SimpleNamespace(id="rule-1")
session = MagicMock()
account_session = MagicMock()
session = self.session
account_session = self.session_factory()
with (
patch(
"core.rag.index_processor.processor.parent_child_index_processor.ParentChildStructureChunk.model_validate",
return_value=parent_childs,
),
patch(
"core.rag.index_processor.processor.parent_child_index_processor.DatasetProcessRule",
return_value=dataset_rule,
),
patch(
"core.rag.index_processor.processor.parent_child_index_processor.helper.generate_text_hash",
return_value="hash",
@@ -480,8 +512,8 @@ class TestParentChildIndexProcessor:
def test_generate_summary_preview_sets_summaries(self, processor: ParentChildIndexProcessor) -> None:
preview_texts = [PreviewDetail(content="chunk-1"), PreviewDetail(content="chunk-2")]
session = MagicMock()
worker_sessions = [MagicMock(), MagicMock()]
session = self.session
worker_sessions = [self.session_factory(), self.session_factory()]
with (
patch(
@@ -513,7 +545,7 @@ class TestParentChildIndexProcessor:
side_effect=RuntimeError("summary failed"),
):
with pytest.raises(ValueError, match="Failed to generate summaries"):
processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=MagicMock())
processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=self.session)
def test_generate_summary_preview_falls_back_without_flask_context(
self, processor: ParentChildIndexProcessor
@@ -529,7 +561,7 @@ class TestParentChildIndexProcessor:
),
):
result = processor.generate_summary_preview(
"tenant-1", preview_texts, {"enable": True}, session=MagicMock()
"tenant-1", preview_texts, {"enable": True}, session=self.session
)
assert result[0].summary == "summary"
@@ -546,6 +578,6 @@ class TestParentChildIndexProcessor:
patch("concurrent.futures.wait", side_effect=[(set(), {future}), (set(), set())]),
):
with pytest.raises(ValueError, match="timeout"):
processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=MagicMock())
processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=self.session)
future.cancel.assert_called_once()
api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py
+39 -27
View File
@@ -1,10 +1,11 @@
import logging
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock, Mock, patch
from unittest.mock import Mock, patch
import pandas as pd
import pytest
from sqlalchemy import event
from sqlalchemy.orm import Session
from werkzeug.datastructures import FileStorage
@@ -15,8 +16,6 @@ from core.rag.models.document import AttachmentDocument, Document
from models.dataset import Dataset, DocumentSegment
from models.dataset import Document as DatasetDocument
TABLES = (DocumentSegment,)
class _ImmediateThread:
def __init__(self, target, args=(), kwargs=None):
@@ -62,9 +61,9 @@ class TestQAIndexProcessor:
segmentation = SimpleNamespace(max_tokens=256, chunk_overlap=10, separator="\n")
return SimpleNamespace(segmentation=segmentation)
def test_extract_forwards_automatic_flag(self, processor: QAIndexProcessor) -> None:
def test_extract_forwards_automatic_flag(self, processor: QAIndexProcessor, sqlite_session: Session) -> None:
extract_setting = Mock()
session = Mock()
session = sqlite_session
expected_docs = [Document(page_content="chunk", metadata={})]
with patch("core.rag.index_processor.processor.qa_index_processor.ExtractProcessor.extract") as mock_extract:
@@ -75,22 +74,22 @@ class TestQAIndexProcessor:
assert docs == expected_docs
mock_extract.assert_called_once_with(extract_setting=extract_setting, is_automatic=True, session=session)
def test_transform_rejects_none_process_rule(self, processor: QAIndexProcessor) -> None:
session = MagicMock()
def test_transform_rejects_none_process_rule(self, processor: QAIndexProcessor, sqlite_session: Session) -> None:
session = sqlite_session
with pytest.raises(ValueError, match="No process rule found"):
processor.transform([Document(page_content="text", metadata={})], process_rule=None, session=session)
def test_transform_rejects_missing_rules_key(self, processor: QAIndexProcessor) -> None:
session = MagicMock()
def test_transform_rejects_missing_rules_key(self, processor: QAIndexProcessor, sqlite_session: Session) -> None:
session = sqlite_session
with pytest.raises(ValueError, match="No rules found in process rule"):
processor.transform(
[Document(page_content="text", metadata={})], process_rule={"mode": "custom"}, session=session
)
def test_transform_preview_calls_formatter_once(
self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app
self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app, sqlite_session: Session
) -> None:
session = MagicMock()
session = sqlite_session
document = Document(page_content="raw text", metadata={"dataset_id": "dataset-1", "document_id": "doc-1"})
split_node = Document(page_content=".question", metadata={})
splitter = Mock()
@@ -132,9 +131,9 @@ class TestQAIndexProcessor:
mock_format.assert_called_once()
def test_transform_non_preview_uses_thread_batches(
self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app
self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app, sqlite_session: Session
) -> None:
session = MagicMock()
session = sqlite_session
documents = [
Document(page_content="doc-1", metadata={"document_id": "doc-1", "dataset_id": "dataset-1"}),
Document(page_content="doc-2", metadata={"document_id": "doc-2", "dataset_id": "dataset-1"}),
@@ -216,8 +215,10 @@ class TestQAIndexProcessor:
with pytest.raises(ValueError, match="bad csv"):
processor.format_by_template(csv_file)
def test_load_creates_vectors_for_high_quality_dataset(self, processor: QAIndexProcessor, dataset: Dataset) -> None:
session = MagicMock()
def test_load_creates_vectors_for_high_quality_dataset(
self, processor: QAIndexProcessor, dataset: Dataset, sqlite_session: Session
) -> None:
session = sqlite_session
docs = [Document(page_content="Q1", metadata={"answer": "A1"})]
multimodal_docs = [AttachmentDocument(page_content="image", metadata={})]
@@ -229,8 +230,10 @@ class TestQAIndexProcessor:
vector.create.assert_called_once_with(docs)
vector.create_multimodal.assert_called_once_with(multimodal_docs)
def test_load_skips_vector_for_non_high_quality(self, processor: QAIndexProcessor, dataset: Dataset) -> None:
session = MagicMock()
def test_load_skips_vector_for_non_high_quality(
self, processor: QAIndexProcessor, dataset: Dataset, sqlite_session: Session
) -> None:
session = sqlite_session
dataset.indexing_technique = IndexTechniqueType.ECONOMY
docs = [Document(page_content="Q1", metadata={"answer": "A1"})]
@@ -239,7 +242,6 @@ class TestQAIndexProcessor:
mock_vector_cls.assert_not_called()
@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True)
def test_clean_handles_summary_deletion_and_vector_cleanup(
self,
processor: QAIndexProcessor,
@@ -294,8 +296,10 @@ class TestQAIndexProcessor:
mock_summary.assert_called_once_with(dataset, [matching_segment.id], session=sqlite_session)
vector.delete_by_ids.assert_called_once_with(["node-1"])
def test_clean_handles_dataset_wide_cleanup(self, processor: QAIndexProcessor, dataset: Dataset) -> None:
session = MagicMock()
def test_clean_handles_dataset_wide_cleanup(
self, processor: QAIndexProcessor, dataset: Dataset, sqlite_session: Session
) -> None:
session = sqlite_session
with (
patch(
"core.rag.index_processor.processor.qa_index_processor.SummaryIndexService.delete_summaries_for_segments"
@@ -309,11 +313,15 @@ class TestQAIndexProcessor:
vector.delete.assert_called_once()
def test_index_adds_documents_and_vectors_for_high_quality(
self, processor: QAIndexProcessor, dataset: Dataset, dataset_document: DatasetDocument
self,
processor: QAIndexProcessor,
dataset: Dataset,
dataset_document: DatasetDocument,
sqlite_session: Session,
) -> None:
session = MagicMock()
session = sqlite_session
phase_events: list[str] = []
session.commit.side_effect = lambda: phase_events.append("commit")
event.listen(session, "after_commit", lambda _session: phase_events.append("commit"))
qa_chunks = SimpleNamespace(
qa_chunks=[
SimpleNamespace(question="Q1", answer="A1"),
@@ -353,9 +361,13 @@ class TestQAIndexProcessor:
mock_vector_cls.return_value.create.assert_called_once()
def test_index_requires_high_quality(
self, processor: QAIndexProcessor, dataset: Dataset, dataset_document: DatasetDocument
self,
processor: QAIndexProcessor,
dataset: Dataset,
dataset_document: DatasetDocument,
sqlite_session: Session,
) -> None:
session = MagicMock()
session = sqlite_session
dataset.indexing_technique = IndexTechniqueType.ECONOMY
qa_chunks = SimpleNamespace(qa_chunks=[SimpleNamespace(question="Q1", answer="A1")])
@@ -389,10 +401,10 @@ class TestQAIndexProcessor:
assert preview["total_segments"] == 1
assert preview["qa_preview"] == [{"question": "Q1", "answer": "A1"}]
def test_generate_summary_preview_returns_input(self, processor: QAIndexProcessor) -> None:
def test_generate_summary_preview_returns_input(self, processor: QAIndexProcessor, sqlite_session: Session) -> None:
preview_items = [PreviewDetail(content="Q1")]
assert (
processor.generate_summary_preview("tenant-1", preview_items, {"enable": False}, session=MagicMock())
processor.generate_summary_preview("tenant-1", preview_items, {"enable": False}, session=sqlite_session)
is preview_items
)
api/tests/unit_tests/core/rag/indexing/test_index_processor.py
+120 -144
View File
@@ -1,12 +1,51 @@
import datetime
import uuid
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from sqlalchemy import event
from sqlalchemy.orm import Session, sessionmaker
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.index_processor.index_processor import IndexProcessor
from core.workflow.nodes.knowledge_index.protocols import Preview, PreviewItem
from models.dataset import Dataset, Document
from models.dataset import Dataset, Document, DocumentSegment
from models.enums import DataSourceType, DocumentCreatedFrom, SegmentStatus
def _persist_dataset_and_document(
session: Session,
*,
indexing_technique: IndexTechniqueType = IndexTechniqueType.HIGH_QUALITY,
summary_index_setting: dict | None = None,
) -> tuple[Dataset, Document]:
tenant_id = str(uuid.uuid4())
created_by = str(uuid.uuid4())
dataset = Dataset(
id=str(uuid.uuid4()),
tenant_id=tenant_id,
name="Dataset",
data_source_type=DataSourceType.UPLOAD_FILE,
indexing_technique=indexing_technique,
chunk_structure="text_model",
summary_index_setting=summary_index_setting,
created_by=created_by,
)
document = Document(
id=str(uuid.uuid4()),
tenant_id=tenant_id,
dataset_id=dataset.id,
position=1,
data_source_type=DataSourceType.UPLOAD_FILE,
batch="batch-1",
name="Document",
created_from=DocumentCreatedFrom.WEB,
created_by=created_by,
doc_language="English",
)
session.add_all([dataset, document])
session.flush()
return dataset, document
class TestIndexProcessor:
@@ -22,28 +61,10 @@ class TestIndexProcessor:
assert preview.qa_preview[0].question == "Q1"
assert preview.qa_preview[0].answer == "A1"
def test_index_and_clean_ends_transactions_around_index_io(self) -> None:
document = SimpleNamespace(
id="document-1",
name="Document",
created_at=datetime.datetime(2026, 1, 1),
indexing_latency=None,
indexing_status=None,
completed_at=None,
word_count=0,
need_summary=False,
)
dataset = SimpleNamespace(
id="dataset-1",
tenant_id="tenant-1",
name="Dataset",
chunk_structure="text_model",
summary_index_setting=None,
)
session = MagicMock()
session.scalar.side_effect = [dataset, document, 3]
def test_index_and_clean_ends_transactions_around_index_io(self, sqlite_session: Session) -> None:
dataset, document = _persist_dataset_and_document(sqlite_session)
phase_events: list[str] = []
session.commit.side_effect = lambda: phase_events.append("commit")
event.listen(sqlite_session, "after_commit", lambda _session: phase_events.append("commit"))
index_processor = MagicMock()
index_processor.index.side_effect = lambda *args: phase_events.append("index")
@@ -65,7 +86,7 @@ class TestIndexProcessor:
original_document_id="",
chunks=chunks,
batch="batch-1",
session=session,
session=sqlite_session,
)
assert phase_events == ["commit", "index", "commit"]
@@ -75,30 +96,13 @@ class TestIndexProcessor:
chunk_structure=dataset.chunk_structure,
chunks=chunks,
include_summaries=False,
session=session,
session=sqlite_session,
)
def test_index_and_clean_skips_admission_for_replacement_without_existing_vector_points(self) -> None:
document = SimpleNamespace(
id="document-1",
name="Document",
created_at=datetime.datetime(2026, 1, 1),
indexing_latency=None,
indexing_status=None,
completed_at=None,
word_count=0,
need_summary=False,
)
dataset = SimpleNamespace(
id="dataset-1",
tenant_id="tenant-1",
name="Dataset",
chunk_structure="text_model",
summary_index_setting=None,
)
session = MagicMock()
session.scalar.side_effect = [dataset, document, 3]
session.scalars.return_value.all.return_value = []
def test_index_and_clean_skips_admission_for_replacement_without_existing_vector_points(
self, sqlite_session: Session
) -> None:
dataset, document = _persist_dataset_and_document(sqlite_session)
index_processor = MagicMock()
processor = IndexProcessor()
chunks = {"general_chunks": ["content"]}
@@ -114,39 +118,50 @@ class TestIndexProcessor:
original_document_id=document.id,
chunks=chunks,
batch="batch-1",
session=session,
session=sqlite_session,
)
admission_service_class.assert_not_called()
def test_index_and_clean_scopes_replacement_queries_to_dataset_owner(self) -> None:
dataset = SimpleNamespace(
id="dataset-1",
tenant_id="tenant-1",
name="Dataset",
summary_index_setting=None,
chunk_structure="text_model",
def test_index_and_clean_scopes_replacement_queries_to_dataset_owner(self, sqlite_session: Session) -> None:
dataset, document = _persist_dataset_and_document(sqlite_session)
original_document = Document(
id=str(uuid.uuid4()),
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
position=2,
data_source_type=DataSourceType.UPLOAD_FILE,
batch="batch-1",
name="Original document",
created_from=DocumentCreatedFrom.WEB,
created_by=dataset.created_by,
)
document = SimpleNamespace(
id="doc-1",
tenant_id="tenant-1",
dataset_id="dataset-1",
name="Document",
created_at=datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC),
segment = DocumentSegment(
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
document_id=original_document.id,
position=1,
content="old content",
word_count=3,
tokens=3,
created_by=dataset.created_by,
index_node_id="node-1",
status=SegmentStatus.COMPLETED,
)
segment = SimpleNamespace(index_node_id="node-1")
session = MagicMock()
def resolve_owner(statement):
entity = statement.column_descriptions[0]["entity"]
if entity is Dataset:
return dataset
if entity is Document:
return document
return 3
session.scalar.side_effect = resolve_owner
session.scalars.return_value.all.return_value = [segment]
control_segment = DocumentSegment(
tenant_id=str(uuid.uuid4()),
dataset_id=str(uuid.uuid4()),
document_id=original_document.id,
position=1,
content="other tenant content",
word_count=4,
tokens=4,
created_by=str(uuid.uuid4()),
index_node_id="other-node",
status=SegmentStatus.COMPLETED,
)
sqlite_session.add_all([original_document, segment, control_segment])
sqlite_session.flush()
processor = IndexProcessor()
with (
@@ -155,91 +170,54 @@ class TestIndexProcessor:
):
index_backend = index_processor_factory.return_value.init_index_processor.return_value
processor.index_and_clean(
dataset_id="dataset-1",
document_id="doc-1",
original_document_id="original-doc",
dataset_id=dataset.id,
document_id=document.id,
original_document_id=original_document.id,
chunks={},
batch="batch-1",
session=session,
session=sqlite_session,
)
document_statement = next(
call.args[0]
for call in session.scalar.call_args_list
if call.args[0].column_descriptions[0]["entity"] is Document
)
segment_statement = session.scalars.call_args.args[0]
delete_statement = session.execute.call_args_list[0].args[0]
word_count_statement = session.scalar.call_args_list[-1].args[0]
segment_update_statement = session.execute.call_args_list[1].args[0]
document_owner = {"doc-1", "dataset-1", "tenant-1"}
original_document_owner = {"original-doc", "dataset-1", "tenant-1"}
assert document_owner <= set(document_statement.compile().params.values())
assert original_document_owner <= set(segment_statement.compile().params.values())
assert original_document_owner <= set(delete_statement.compile().params.values())
assert document_owner <= set(word_count_statement.compile().params.values())
assert document_owner <= set(segment_update_statement.compile().params.values())
assert sqlite_session.get(DocumentSegment, segment.id) is None
assert sqlite_session.get(DocumentSegment, control_segment.id) is control_segment
assert sqlite_session.get(Document, document.id).indexing_status == "completed"
index_backend.clean.assert_called_once_with(
dataset,
["node-1"],
with_keywords=True,
delete_child_chunks=True,
session=session,
session=sqlite_session,
)
index_backend.index.assert_called_once_with(dataset, document, {}, session)
index_backend.index.assert_called_once_with(dataset, document, {}, sqlite_session)
admission_service_class.assert_not_called()
def test_get_preview_output_scopes_document_to_dataset_owner(self) -> None:
dataset = SimpleNamespace(
id="dataset-1",
tenant_id="tenant-1",
indexing_technique=IndexTechniqueType.ECONOMY,
summary_index_setting=None,
)
document = SimpleNamespace(doc_language="English")
session = MagicMock()
def resolve_owner(statement):
entity = statement.column_descriptions[0]["entity"]
if entity is Dataset:
return dataset
if entity is Document:
return document
raise AssertionError(f"Unexpected entity: {entity}")
session.scalar.side_effect = resolve_owner
def test_get_preview_output_scopes_document_to_dataset_owner(self, sqlite_session: Session) -> None:
dataset, document = _persist_dataset_and_document(sqlite_session, indexing_technique=IndexTechniqueType.ECONOMY)
processor = IndexProcessor()
expected_preview = MagicMock()
with patch.object(processor, "format_preview", return_value=expected_preview):
result = processor.get_preview_output(
chunks={},
dataset_id="dataset-1",
document_id="doc-1",
dataset_id=dataset.id,
document_id=document.id,
chunk_structure="text_model",
summary_index_setting=None,
session=session,
session=sqlite_session,
)
document_statement = next(
call.args[0]
for call in session.scalar.call_args_list
if call.args[0].column_descriptions[0]["entity"] is Document
)
assert {"doc-1", "dataset-1", "tenant-1"} <= set(document_statement.compile().params.values())
assert result is expected_preview
def test_preview_summary_workers_use_independent_sessions(self) -> None:
caller_session = MagicMock()
phase_events: list[str] = []
caller_session.commit.side_effect = lambda: phase_events.append("commit")
caller_session.scalar.return_value = SimpleNamespace(
indexing_technique=IndexTechniqueType.HIGH_QUALITY,
def test_preview_summary_workers_use_independent_sessions(
self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
dataset, _ = _persist_dataset_and_document(
sqlite_session,
summary_index_setting={"enable": True},
tenant_id="tenant-1",
)
worker_sessions = [MagicMock(), MagicMock()]
phase_events: list[str] = []
event.listen(sqlite_session, "after_commit", lambda _session: phase_events.append("commit"))
worker_sessions: list[Session] = []
preview = Preview(
chunk_structure="text_model",
total_segments=2,
@@ -247,11 +225,11 @@ class TestIndexProcessor:
)
flask_app = SimpleNamespace(app_context=lambda: nullcontext())
processor = IndexProcessor()
worker_contexts = iter(nullcontext(worker_session) for worker_session in worker_sessions)
def create_worker_session():
def generate_summary(*_args, **kwargs):
phase_events.append("worker")
return next(worker_contexts)
worker_sessions.append(kwargs["session"])
return "summary", None
with (
patch.object(processor, "format_preview", return_value=preview),
@@ -260,27 +238,25 @@ class TestIndexProcessor:
SimpleNamespace(_get_current_object=lambda: flask_app),
),
patch(
"core.rag.index_processor.index_processor.session_factory.create_session",
side_effect=create_worker_session,
"core.rag.index_processor.index_processor.session_factory",
SimpleNamespace(create_session=sqlite_session_factory),
),
patch(
"core.rag.index_processor.index_processor.ParagraphIndexProcessor.generate_summary",
return_value=("summary", None),
side_effect=generate_summary,
) as generate_summary,
):
result = processor.get_preview_output(
chunks=[],
dataset_id="dataset-1",
dataset_id=dataset.id,
document_id="",
chunk_structure="text_model",
summary_index_setting={"enable": True},
session=caller_session,
session=sqlite_session,
)
assert all(item.summary == "summary" for item in result.preview)
assert phase_events == ["commit", "worker", "worker"]
call_sessions = [call.kwargs["session"] for call in generate_summary.call_args_list]
assert all(call_session is not caller_session for call_session in call_sessions)
assert all(
any(call_session is worker_session for worker_session in worker_sessions) for call_session in call_sessions
)
assert all(call_session is not sqlite_session for call_session in call_sessions)
assert call_sessions == worker_sessions