opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-25 13:38:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: avoid duplicate token counting during dataset indexing (#39466)

Browse Source
This commit is contained in:
Byron.wang 2026-07-23 18:56:03 +08:00 committed by GitHub
parent 2beafdc457
commit 43254c1ded
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 633 additions and 382 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

90
api/core/indexing_runner.py
View File

@ -21,6 +21,7 @@ from core.model_manager import ModelInstance, ModelManager
from core.rag.cleaner.clean_processor import CleanProcessor
from core.rag.datasource.keyword.keyword_factory import Keyword
from core.rag.docstore.dataset_docstore import DatasetDocumentStore
from core.rag.embedding.token_counter import calculate_segment_token_counts
from core.rag.extractor.entity.datasource_type import DatasourceType
from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
@ -113,17 +114,25 @@ class IndexingRunner:
current_user=current_user,
session=session,
)
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
total_tokens = sum(token_counts)
# save segment
self._load_segments(dataset, requeried_document, documents, session)
self._load_segments(
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
token_counts=token_counts,
)
session.commit()
# load
self._load(
index_processor=index_processor,
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
session=session,
total_tokens=total_tokens,
)
except DocumentIsPausedError:
raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
@ -190,17 +199,25 @@ class IndexingRunner:
current_user=current_user,
session=session,
)
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
total_tokens = sum(token_counts)
# save segment
self._load_segments(dataset, requeried_document, documents, session)
self._load_segments(
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
token_counts=token_counts,
)
session.commit()
# load
self._load(
index_processor=index_processor,
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
session=session,
total_tokens=total_tokens,
)
except DocumentIsPausedError:
raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
@ -225,7 +242,7 @@ class IndexingRunner:
if not dataset:
raise ValueError("no dataset found")
# get exist document_segment list and delete
# get existing document segments
document_segments = session.scalars(
select(DocumentSegment).where(
DocumentSegment.dataset_id == dataset.id,
@ -264,15 +281,15 @@ class IndexingRunner:
child_documents.append(child_document)
document.children = child_documents
documents.append(document)
# Preserve the full document total even when only incomplete segments are re-indexed.
total_tokens = sum(document_segment.tokens for document_segment in document_segments)
# build index
index_type = requeried_document.doc_form
index_processor = IndexProcessorFactory(index_type).init_index_processor()
self._load(
index_processor=index_processor,
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
session=session,
total_tokens=total_tokens,
)
except DocumentIsPausedError:
raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
@ -601,28 +618,16 @@ class IndexingRunner:
def _load(
self,
index_processor: BaseIndexProcessor,
session: Session,
dataset: Dataset,
dataset_document: DatasetDocument,
documents: list[Document],
session: Session,
):
"""
insert index and update document/segment status to completed
"""
total_tokens: int,
) -> None:
"""Build indexes and mark the document complete using the token total computed before hash sharding."""
embedding_model_instance = None
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
embedding_model_instance = self._get_model_manager(dataset.tenant_id).get_model_instance(
tenant_id=dataset.tenant_id,
provider=dataset.embedding_model_provider,
model_type=ModelType.TEXT_EMBEDDING,
model=dataset.embedding_model,
)
# chunk nodes by chunk size
# Build indexes using the existing hash-based worker groups.
indexing_start_at = time.perf_counter()
tokens = 0
create_keyword_thread = None
if (
dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
@ -659,12 +664,11 @@ class IndexingRunner:
chunk_documents,
dataset.id,
dataset_document.id,
embedding_model_instance,
)
)
for future in futures:
tokens += future.result()
future.result()
if (
dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
and dataset.indexing_technique == IndexTechniqueType.ECONOMY
@ -679,7 +683,7 @@ class IndexingRunner:
document_id=dataset_document.id,
after_indexing_status=IndexingStatus.COMPLETED,
extra_update_params={
DatasetDocument.tokens: tokens,
DatasetDocument.tokens: total_tokens,
DatasetDocument.completed_at: naive_utc_now(),
DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
DatasetDocument.error: None,
@ -720,8 +724,7 @@ class IndexingRunner:
chunk_documents: list[Document],
dataset_id: str,
dataset_document_id: str,
embedding_model_instance: ModelInstance | None,
):
) -> None:
with flask_app.app_context():
with session_factory.create_session() as session:
dataset = session.get(Dataset, dataset_id)
@ -735,11 +738,6 @@ class IndexingRunner:
# check document is paused
self._check_document_paused_status(dataset_document.id)
tokens = 0
if embedding_model_instance:
page_content_list = [document.page_content for document in chunk_documents]
tokens += sum(embedding_model_instance.get_text_embedding_num_tokens(page_content_list))
multimodal_documents = []
for document in chunk_documents:
if document.attachments and dataset.is_multimodal:
@ -773,8 +771,6 @@ class IndexingRunner:
session.commit()
return tokens
@staticmethod
def _check_document_paused_status(document_id: str):
indexing_cache_key = f"document_{document_id}_is_paused"
@ -864,8 +860,14 @@ class IndexingRunner:
return documents
def _load_segments(
self, dataset: Dataset, dataset_document: DatasetDocument, documents: list[Document], session: Session
):
self,
session: Session,
dataset: Dataset,
dataset_document: DatasetDocument,
documents: list[Document],
token_counts: list[int],
) -> None:
"""Persist transformed documents and their precomputed token counts before indexing starts."""
# save node to document segment
doc_store = DatasetDocumentStore(
dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id
@ -873,9 +875,10 @@ class IndexingRunner:
# add document segments
doc_store.add_documents(
session=session,
docs=documents,
save_child=dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX,
session=session,
token_counts=token_counts,
)
# update document status to indexing
@ -900,7 +903,6 @@ class IndexingRunner:
DocumentSegment.indexing_at: naive_utc_now(),
},
)
pass
class DocumentIsPausedError(Exception):

27
api/core/rag/docstore/dataset_docstore.py
View File

@ -6,10 +6,7 @@ from typing import Any
from sqlalchemy import delete, func, select
from sqlalchemy.orm import Session
from core.model_manager import ModelManager
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.models.document import AttachmentDocument, Document
from graphon.model_runtime.entities.model_entities import ModelType
from models.dataset import ChildChunk, Dataset, DocumentSegment, SegmentAttachmentBinding
from models.enums import SegmentType
@ -69,34 +66,22 @@ class DatasetDocumentStore:
def add_documents(
self,
docs: Sequence[Document],
session: Session,
docs: Sequence[Document],
token_counts: list[int],
allow_update: bool = True,
save_child: bool = False,
):
) -> None:
document_token_pairs = list(zip(docs, token_counts, strict=True))
max_position = session.scalar(
select(func.max(DocumentSegment.position)).where(DocumentSegment.document_id == self._document_id)
)
if max_position is None:
max_position = 0
embedding_model = None
if self._dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
model_manager = ModelManager.for_tenant(tenant_id=self._dataset.tenant_id)
embedding_model = model_manager.get_model_instance(
tenant_id=self._dataset.tenant_id,
provider=self._dataset.embedding_model_provider,
model_type=ModelType.TEXT_EMBEDDING,
model=self._dataset.embedding_model,
)
if embedding_model:
page_content_list = [doc.page_content for doc in docs]
tokens_list = embedding_model.get_text_embedding_num_tokens(page_content_list)
else:
tokens_list = [0] * len(docs)
for doc, tokens in zip(docs, tokens_list):
for doc, tokens in document_token_pairs:
if not isinstance(doc, Document):
raise ValueError("doc must be a Document")

25
api/core/rag/embedding/token_counter.py Normal file
View File

@ -0,0 +1,25 @@
"""Token counting for document segments."""
from core.model_manager import ModelManager
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.models.document import Document
from graphon.model_runtime.entities.model_entities import ModelType
from models.dataset import Dataset
def calculate_segment_token_counts(dataset: Dataset, documents: list[Document]) -> list[int]:
"""Return one token count per document, invoking the embedding model only for high-quality indexes."""
if not documents:
return []
if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY:
return [0] * len(documents)
model_manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id)
embedding_model = model_manager.get_model_instance(
tenant_id=dataset.tenant_id,
provider=dataset.embedding_model_provider,
model_type=ModelType.TEXT_EMBEDDING,
model=dataset.embedding_model,
)
return embedding_model.get_text_embedding_num_tokens([document.page_content for document in documents])

9
api/core/rag/index_processor/processor/paragraph_index_processor.py
View File

@ -19,6 +19,7 @@ from core.rag.cleaner.clean_processor import CleanProcessor
from core.rag.datasource.keyword.keyword_factory import Keyword
from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.docstore.dataset_docstore import DatasetDocumentStore
from core.rag.embedding.token_counter import calculate_segment_token_counts
from core.rag.entities import Rule
from core.rag.extractor.entity.extract_setting import ExtractSetting
from core.rag.extractor.extract_processor import ExtractProcessor
@ -244,10 +245,16 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
all_multimodal_documents.extend(doc.attachments)
documents.append(doc)
if documents:
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
# save node to document segment
doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
# add document segments
doc_store.add_documents(docs=documents, save_child=False, session=session)
doc_store.add_documents(
session=session,
docs=documents,
token_counts=token_counts,
save_child=False,
)
session.commit()
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
vector = Vector(dataset, session=session)

9
api/core/rag/index_processor/processor/parent_child_index_processor.py
View File

@ -15,6 +15,7 @@ from core.model_manager import ModelInstance
from core.rag.cleaner.clean_processor import CleanProcessor
from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.docstore.dataset_docstore import DatasetDocumentStore
from core.rag.embedding.token_counter import calculate_segment_token_counts
from core.rag.entities import ParentMode, Rule
from core.rag.extractor.entity.extract_setting import ExtractSetting
from core.rag.extractor.extract_processor import ExtractProcessor
@ -304,6 +305,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
doc.attachments = self._get_content_files(doc, current_user=account, session=session)
documents.append(doc)
if documents:
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
# update document parent mode
dataset_process_rule = DatasetProcessRule(
dataset_id=dataset.id,
@ -321,7 +323,12 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
# save node to document segment
doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
# add document segments
doc_store.add_documents(docs=documents, save_child=True, session=session)
doc_store.add_documents(
session=session,
docs=documents,
token_counts=token_counts,
save_child=True,
)
session.commit()
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
all_child_documents = []

9
api/core/rag/index_processor/processor/qa_index_processor.py
View File

@ -17,6 +17,7 @@ from core.llm_generator.llm_generator import LLMGenerator
from core.rag.cleaner.clean_processor import CleanProcessor
from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.docstore.dataset_docstore import DatasetDocumentStore
from core.rag.embedding.token_counter import calculate_segment_token_counts
from core.rag.entities import Rule
from core.rag.extractor.entity.extract_setting import ExtractSetting
from core.rag.extractor.extract_processor import ExtractProcessor
@ -205,9 +206,15 @@ class QAIndexProcessor(BaseIndexProcessor):
doc = Document(page_content=qa_chunk.question, metadata=metadata)
documents.append(doc)
if documents:
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
# save node to document segment
doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
doc_store.add_documents(docs=documents, save_child=False, session=session)
doc_store.add_documents(
session=session,
docs=documents,
token_counts=token_counts,
save_child=False,
)
session.commit()
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
vector = Vector(dataset, session=session)

469
api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py
View File

@ -8,10 +8,49 @@ which provides document storage and retrieval functionality for datasets in the
from unittest.mock import MagicMock, patch
import pytest
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from core.rag.docstore.dataset_docstore import DatasetDocumentStore, DocumentSegment
from core.rag.models.document import AttachmentDocument, Document
from models.dataset import Dataset
from core.rag.models.document import AttachmentDocument, ChildDocument, Document
from models.dataset import ChildChunk, Dataset, SegmentAttachmentBinding
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"
def _dataset() -> Dataset:
dataset = MagicMock(spec=Dataset)
dataset.id = DATASET_ID
dataset.tenant_id = TENANT_ID
return dataset
def _persist_segment(
session: Session,
*,
index_node_id: str = "doc-1",
index_node_hash: str = "hash-1",
content: str = "Test content",
tokens: int = 5,
) -> DocumentSegment:
segment = DocumentSegment(
tenant_id=TENANT_ID,
dataset_id=DATASET_ID,
document_id=DOCUMENT_ID,
position=1,
content=content,
word_count=len(content),
tokens=tokens,
created_by=USER_ID,
index_node_id=index_node_id,
index_node_hash=index_node_hash,
)
session.add(segment)
session.flush()
return segment
class TestDatasetDocumentStoreInit:
@ -132,228 +171,153 @@ class TestDatasetDocumentStoreDocs:
assert result == {}
@pytest.mark.parametrize(
"sqlite_session",
[(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
indirect=True,
)
class TestDatasetDocumentStoreAddDocuments:
"""Tests for add_documents method."""
def test_add_documents_new_document_with_embedding(self):
"""Test adding new documents with embedding model."""
def test_add_documents_new_document_with_token_count(self, sqlite_session: Session):
"""Test adding a new document with a precomputed token count."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_dataset.indexing_technique = "high_quality"
mock_dataset.embedding_model_provider = "provider"
mock_dataset.embedding_model = "model"
document = Document(
page_content="Test content",
metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Test content"
mock_doc.metadata = {
"doc_id": "doc-1",
"doc_hash": "hash-1",
}
mock_doc.attachments = None
mock_doc.children = None
store.add_documents(session=sqlite_session, docs=[document], token_counts=[10])
sqlite_session.expire_all()
mock_model_instance = MagicMock()
mock_model_instance.get_text_embedding_num_tokens.return_value = [10]
segment = sqlite_session.scalar(
select(DocumentSegment).where(
DocumentSegment.dataset_id == DATASET_ID,
DocumentSegment.index_node_id == "doc-1",
)
)
assert segment is not None
assert segment.content == "Test content"
assert segment.tokens == 10
assert segment.position == 1
with (
patch("core.rag.docstore.dataset_docstore.ModelManager.for_tenant") as mock_manager_class,
):
mock_session = MagicMock()
mock_session.scalar.return_value = None
mock_manager = MagicMock()
mock_manager.get_model_instance.return_value = mock_model_instance
mock_manager_class.return_value = mock_manager
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
)
store.add_documents([mock_doc], session=mock_session)
mock_session.add.assert_called()
mock_session.flush.assert_called()
def test_add_documents_update_existing_document(self):
def test_add_documents_update_existing_document(self, sqlite_session: Session):
"""Test updating existing document with allow_update=True."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_dataset.indexing_technique = "economy"
mock_dataset.embedding_model_provider = None
mock_dataset.embedding_model = None
existing_segment = _persist_segment(sqlite_session)
document = Document(
page_content="Updated content",
metadata={"doc_id": "doc-1", "doc_hash": "new-hash"},
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Updated content"
mock_doc.metadata = {
"doc_id": "doc-1",
"doc_hash": "new-hash",
}
mock_doc.attachments = None
mock_doc.children = None
store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
sqlite_session.expire_all()
mock_existing_segment = MagicMock()
mock_existing_segment.id = "seg-1"
updated_segment = sqlite_session.get(DocumentSegment, existing_segment.id)
assert updated_segment is not None
assert updated_segment.content == "Updated content"
assert updated_segment.index_node_hash == "new-hash"
assert updated_segment.tokens == 0
mock_session = MagicMock()
mock_session.scalar.return_value = 5
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
)
store.add_documents([mock_doc], session=mock_session)
mock_session.flush.assert_called()
def test_add_documents_raises_when_not_allowed(self):
def test_add_documents_raises_when_not_allowed(self, sqlite_session: Session):
"""Test that adding existing doc without allow_update raises ValueError."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_dataset.indexing_technique = "economy"
_persist_segment(sqlite_session)
document = Document(
page_content="Test content",
metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Test content"
mock_doc.metadata = {
"doc_id": "doc-1",
"doc_hash": "hash-1",
}
mock_doc.attachments = None
mock_doc.children = None
mock_existing_segment = MagicMock()
mock_session = MagicMock()
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
with pytest.raises(ValueError, match="already exists"):
store.add_documents(
session=sqlite_session,
docs=[document],
token_counts=[0],
allow_update=False,
)
with pytest.raises(ValueError, match="already exists"):
store.add_documents([mock_doc], session=mock_session, allow_update=False)
assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegment)) == 1
def test_add_documents_with_answer_metadata(self):
def test_add_documents_with_answer_metadata(self, sqlite_session: Session):
"""Test adding document with answer in metadata."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_dataset.indexing_technique = "economy"
document = Document(
page_content="Test content",
metadata={
"doc_id": "doc-1",
"doc_hash": "hash-1",
"answer": "Test answer",
},
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Test content"
mock_doc.metadata = {
"doc_id": "doc-1",
"doc_hash": "hash-1",
"answer": "Test answer",
}
mock_doc.attachments = None
mock_doc.children = None
store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
sqlite_session.expire_all()
mock_session = MagicMock()
mock_session.scalar.return_value = None
segment = sqlite_session.scalar(select(DocumentSegment))
assert segment is not None
assert segment.answer == "Test answer"
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
)
store.add_documents([mock_doc], session=mock_session)
mock_session.add.assert_called()
def test_add_documents_with_invalid_document_type(self):
def test_add_documents_with_invalid_document_type(self, sqlite_session: Session):
"""Test that non-Document raises ValueError."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_session = MagicMock()
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
with pytest.raises(ValueError, match="must be a Document"):
store.add_documents(["not a document"], session=mock_session)
store.add_documents(session=sqlite_session, docs=["not a document"], token_counts=[0]) # type: ignore[list-item]
def test_add_documents_with_none_metadata(self):
def test_add_documents_with_none_metadata(self, sqlite_session: Session):
"""Test that document with None metadata raises ValueError."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Test content"
mock_doc.metadata = None
mock_session = MagicMock()
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
)
document = MagicMock(spec=Document)
document.metadata = None
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
with pytest.raises(ValueError, match="metadata must be a dict"):
store.add_documents([mock_doc], session=mock_session)
store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
def test_add_documents_with_save_child(self):
def test_add_documents_with_save_child(self, sqlite_session: Session):
"""Test adding documents with save_child=True."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_dataset.indexing_technique = "economy"
mock_child = MagicMock(spec=Document)
mock_child.page_content = "Child content"
mock_child.metadata = {
"doc_id": "child-1",
"doc_hash": "child-hash",
}
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Test content"
mock_doc.metadata = {
"doc_id": "doc-1",
"doc_hash": "hash-1",
}
mock_doc.attachments = None
mock_doc.children = [mock_child]
mock_session = MagicMock()
mock_session.scalar.return_value = None
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
document = Document(
page_content="Test content",
metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
children=[
ChildDocument(
page_content="Child content",
metadata={"doc_id": "child-1", "doc_hash": "child-hash"},
)
],
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
store.add_documents([mock_doc], session=mock_session, save_child=True)
store.add_documents(
session=sqlite_session,
docs=[document],
token_counts=[0],
save_child=True,
)
sqlite_session.expire_all()
mock_session.add.assert_called()
child = sqlite_session.scalar(select(ChildChunk))
assert child is not None
assert child.content == "Child content"
assert child.index_node_id == "child-1"
def test_add_documents_rejects_mismatched_token_counts(self, sqlite_session: Session):
document = Document(
page_content="Test content",
metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
with pytest.raises(ValueError):
store.add_documents(session=sqlite_session, docs=[document], token_counts=[])
assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegment)) == 0
class TestDatasetDocumentStoreExists:
@ -722,88 +686,85 @@ class TestDatasetDocumentStoreMultimodelBinding:
mock_session.add.assert_not_called()
@pytest.mark.parametrize(
"sqlite_session",
[(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
indirect=True,
)
class TestDatasetDocumentStoreAddDocumentsUpdateChild:
"""Tests for add_documents when updating existing documents with children."""
def test_add_documents_update_existing_with_children(self):
def test_add_documents_update_existing_with_children(self, sqlite_session: Session):
"""Test updating existing document with save_child=True and children."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_dataset.indexing_technique = "economy"
mock_child = MagicMock(spec=Document)
mock_child.page_content = "Updated child content"
mock_child.metadata = {
"doc_id": "child-1",
"doc_hash": "new-child-hash",
}
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Updated content"
mock_doc.metadata = {
"doc_id": "doc-1",
"doc_hash": "new-hash",
}
mock_doc.attachments = None
mock_doc.children = [mock_child]
mock_existing_segment = MagicMock()
mock_existing_segment.id = "seg-1"
mock_session = MagicMock()
mock_session.scalar.return_value = 5
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
segment = _persist_segment(sqlite_session)
sqlite_session.add(
ChildChunk(
tenant_id=TENANT_ID,
dataset_id=DATASET_ID,
document_id=DOCUMENT_ID,
segment_id=segment.id,
position=1,
index_node_id="old-child",
index_node_hash="old-child-hash",
content="Old child content",
word_count=len("Old child content"),
created_by=USER_ID,
)
)
sqlite_session.flush()
document = Document(
page_content="Updated content",
metadata={"doc_id": "doc-1", "doc_hash": "new-hash"},
children=[
ChildDocument(
page_content="Updated child content",
metadata={"doc_id": "child-1", "doc_hash": "new-child-hash"},
)
],
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
store.add_documents([mock_doc], session=mock_session, save_child=True)
store.add_documents(
session=sqlite_session,
docs=[document],
token_counts=[0],
save_child=True,
)
sqlite_session.expire_all()
mock_session.execute.assert_called()
mock_session.flush.assert_called()
children = sqlite_session.scalars(select(ChildChunk).order_by(ChildChunk.position)).all()
assert len(children) == 1
assert children[0].content == "Updated child content"
assert children[0].index_node_id == "child-1"
@pytest.mark.parametrize(
"sqlite_session",
[(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
indirect=True,
)
class TestDatasetDocumentStoreAddDocumentsUpdateAnswer:
"""Tests for add_documents when updating existing documents with answer metadata."""
def test_add_documents_update_existing_with_answer(self):
def test_add_documents_update_existing_with_answer(self, sqlite_session: Session):
"""Test updating existing document with answer in metadata."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_dataset.indexing_technique = "economy"
existing_segment = _persist_segment(sqlite_session)
document = Document(
page_content="Updated content",
metadata={
"doc_id": "doc-1",
"doc_hash": "new-hash",
"answer": "Updated answer",
},
)
store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
mock_doc = MagicMock(spec=Document)
mock_doc.page_content = "Updated content"
mock_doc.metadata = {
"doc_id": "doc-1",
"doc_hash": "new-hash",
"answer": "Updated answer",
}
mock_doc.attachments = None
mock_doc.children = None
store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
sqlite_session.expire_all()
mock_existing_segment = MagicMock()
mock_existing_segment.id = "seg-1"
mock_session = MagicMock()
mock_session.scalar.return_value = 5
with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id="test-doc-id",
)
store.add_documents([mock_doc], session=mock_session)
mock_session.flush.assert_called()
updated_segment = sqlite_session.get(DocumentSegment, existing_segment.id)
assert updated_segment is not None
assert updated_segment.answer == "Updated answer"
assert updated_segment.tokens == 0

56
api/tests/unit_tests/core/rag/embedding/test_token_counter.py Normal file
View File

@ -0,0 +1,56 @@
from unittest.mock import Mock, patch
from core.rag.embedding.token_counter import calculate_segment_token_counts
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.models.document import Document
from models.dataset import Dataset
def test_high_quality_counts_each_document_once() -> None:
dataset = Mock(spec=Dataset)
dataset.tenant_id = "tenant-1"
dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
dataset.embedding_model_provider = "provider"
dataset.embedding_model = "model"
documents = [
Document(page_content="first", metadata={}),
Document(page_content="second", metadata={}),
Document(page_content="third", metadata={}),
]
with patch("core.rag.embedding.token_counter.ModelManager.for_tenant") as model_manager_factory:
embedding_model = model_manager_factory.return_value.get_model_instance.return_value
embedding_model.get_text_embedding_num_tokens.return_value = [11, 22, 33]
result = calculate_segment_token_counts(dataset=dataset, documents=documents)
assert result == [11, 22, 33]
model_manager_factory.assert_called_once_with(tenant_id=dataset.tenant_id)
model_manager_factory.return_value.get_model_instance.assert_called_once()
embedding_model.get_text_embedding_num_tokens.assert_called_once_with(["first", "second", "third"])
def test_economy_returns_zero_without_loading_model() -> None:
dataset = Mock(spec=Dataset)
dataset.indexing_technique = IndexTechniqueType.ECONOMY
documents = [
Document(page_content="first", metadata={}),
Document(page_content="second", metadata={}),
]
with patch("core.rag.embedding.token_counter.ModelManager.for_tenant") as model_manager_factory:
result = calculate_segment_token_counts(dataset=dataset, documents=documents)
assert result == [0, 0]
model_manager_factory.assert_not_called()
def test_empty_documents_return_without_loading_model() -> None:
dataset = Mock(spec=Dataset)
dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
with patch("core.rag.embedding.token_counter.ModelManager.for_tenant") as model_manager_factory:
result = calculate_segment_token_counts(dataset=dataset, documents=[])
assert result == []
model_manager_factory.assert_not_called()

27
api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py
View File

@ -271,14 +271,26 @@ class TestParagraphIndexProcessor:
patch(
"core.rag.index_processor.processor.paragraph_index_processor.DatasetDocumentStore"
) as mock_store_cls,
patch(
"core.rag.index_processor.processor.paragraph_index_processor.calculate_segment_token_counts"
) as mock_token_counter,
patch("core.rag.index_processor.processor.paragraph_index_processor.Vector") as mock_vector_cls,
):
mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [11, 22]
mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
mock_vector_cls.return_value.create.side_effect = lambda _documents: phase_events.append("vector")
processor.index(dataset, dataset_document, ["chunk-1", "chunk-2"], session)
assert phase_events == ["store", "commit", "vector"]
mock_store_cls.return_value.add_documents.assert_called_once()
assert phase_events == ["count", "store", "commit", "vector"]
documents = mock_token_counter.call_args.kwargs["documents"]
assert [document.page_content for document in documents] == ["chunk-1", "chunk-2"]
mock_token_counter.assert_called_once_with(dataset=dataset, documents=documents)
mock_store_cls.return_value.add_documents.assert_called_once_with(
session=session,
docs=documents,
token_counts=[11, 22],
save_child=False,
)
mock_vector_cls.assert_called_once_with(dataset, session=session)
mock_vector_cls.return_value.create.assert_called_once()
mock_vector_cls.return_value.create_multimodal.assert_called_once()
@ -299,13 +311,18 @@ class TestParagraphIndexProcessor:
patch(
"core.rag.index_processor.processor.paragraph_index_processor.DatasetDocumentStore"
) as mock_store_cls,
patch(
"core.rag.index_processor.processor.paragraph_index_processor.calculate_segment_token_counts"
) as mock_token_counter,
patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls,
):
mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [0]
mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
mock_keyword_cls.return_value.add_texts.side_effect = lambda *_args: phase_events.append("keyword")
processor.index(dataset, dataset_document, ["chunk-3"], session)
assert phase_events == ["store", "commit", "keyword"]
assert phase_events == ["count", "store", "commit", "keyword"]
mock_token_counter.assert_called_once()
mock_keyword_cls.return_value.add_texts.assert_called_once()
def test_index_multimodal_structure_handles_files_and_account_lookup(
@ -341,6 +358,10 @@ class TestParagraphIndexProcessor:
processor, "_get_content_files", return_value=[AttachmentDocument(page_content="img", metadata={})]
) as mock_files,
patch("core.rag.index_processor.processor.paragraph_index_processor.DatasetDocumentStore"),
patch(
"core.rag.index_processor.processor.paragraph_index_processor.calculate_segment_token_counts",
return_value=[11, 22],
),
patch("core.rag.index_processor.processor.paragraph_index_processor.Vector"),
):
processor.index(dataset, dataset_document, {"general_chunks": []}, session)

20
api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py
View File

@ -362,17 +362,29 @@ class TestParentChildIndexProcessor:
patch(
"core.rag.index_processor.processor.parent_child_index_processor.DatasetDocumentStore"
) as mock_store_cls,
patch(
"core.rag.index_processor.processor.parent_child_index_processor.calculate_segment_token_counts"
) as mock_token_counter,
patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls,
):
mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [11]
mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
mock_vector_cls.return_value.create.side_effect = lambda _documents: phase_events.append("vector")
processor.index(dataset, dataset_document, {"parent_child_chunks": []}, session)
assert phase_events == ["store", "commit", "vector"]
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()
mock_store_cls.return_value.add_documents.assert_called_once()
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)
mock_store_cls.return_value.add_documents.assert_called_once_with(
session=session,
docs=documents,
token_counts=[11],
save_child=True,
)
mock_vector_cls.assert_called_once_with(dataset, session=session)
assert mock_vector_cls.return_value.create.call_count == 1
mock_vector_cls.return_value.create_multimodal.assert_called_once()
@ -413,6 +425,10 @@ class TestParentChildIndexProcessor:
processor, "_get_content_files", return_value=[AttachmentDocument(page_content="image", metadata={})]
) as mock_files,
patch("core.rag.index_processor.processor.parent_child_index_processor.DatasetDocumentStore"),
patch(
"core.rag.index_processor.processor.parent_child_index_processor.calculate_segment_token_counts",
return_value=[11],
),
patch("core.rag.index_processor.processor.parent_child_index_processor.Vector"),
):
processor.index(dataset, dataset_document, {"parent_child_chunks": []}, session)

20
api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py
View File

@ -292,14 +292,26 @@ class TestQAIndexProcessor:
"core.rag.index_processor.processor.qa_index_processor.helper.generate_text_hash", return_value="hash"
),
patch("core.rag.index_processor.processor.qa_index_processor.DatasetDocumentStore") as mock_store_cls,
patch(
"core.rag.index_processor.processor.qa_index_processor.calculate_segment_token_counts"
) as mock_token_counter,
patch("core.rag.index_processor.processor.qa_index_processor.Vector") as mock_vector_cls,
):
mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [11, 22]
mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
mock_vector_cls.return_value.create.side_effect = lambda _documents: phase_events.append("vector")
processor.index(dataset, dataset_document, {"qa_chunks": []}, session)
assert phase_events == ["store", "commit", "vector"]
mock_store_cls.return_value.add_documents.assert_called_once()
assert phase_events == ["count", "store", "commit", "vector"]
documents = mock_token_counter.call_args.kwargs["documents"]
assert [document.page_content for document in documents] == ["Q1", "Q2"]
mock_token_counter.assert_called_once_with(dataset=dataset, documents=documents)
mock_store_cls.return_value.add_documents.assert_called_once_with(
session=session,
docs=documents,
token_counts=[11, 22],
save_child=False,
)
mock_vector_cls.return_value.create.assert_called_once()
def test_index_requires_high_quality(
@ -318,6 +330,10 @@ class TestQAIndexProcessor:
"core.rag.index_processor.processor.qa_index_processor.helper.generate_text_hash", return_value="hash"
),
patch("core.rag.index_processor.processor.qa_index_processor.DatasetDocumentStore"),
patch(
"core.rag.index_processor.processor.qa_index_processor.calculate_segment_token_counts",
return_value=[0],
),
):
with pytest.raises(ValueError, match="must be high quality"):
processor.index(dataset, dataset_document, {"qa_chunks": []}, session)

254
api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py
View File

@ -69,6 +69,7 @@ from graphon.model_runtime.entities.model_entities import ModelType
from libs.datetime_utils import naive_utc_now
from models.dataset import Dataset, DatasetProcessRule, DocumentSegment
from models.dataset import Document as DatasetDocument
from models.enums import SegmentStatus
from models.model import Account
# ============================================================================
@ -611,7 +612,7 @@ class TestIndexingRunnerLoad:
- Keyword index creation
- Multi-threaded processing
- Document segment status updates
- Token counting
- Precomputed token totals
- Error handling during loading
"""
@ -677,16 +678,10 @@ class TestIndexingRunnerLoad:
"""Test loading with high quality indexing (vector embeddings)."""
# Arrange
runner = IndexingRunner()
mock_embedding_instance = MagicMock()
mock_embedding_instance.get_text_embedding_num_tokens.return_value = 100
model_manager = mock_dependencies["model_manager"].return_value
model_manager.get_model_instance.return_value = mock_embedding_instance
mock_processor = MagicMock()
# Mock ThreadPoolExecutor
mock_future = MagicMock()
mock_future.result.return_value = 300 # Total tokens
mock_future.result.return_value = None
mock_executor_instance = MagicMock()
mock_executor_instance.__enter__.return_value = mock_executor_instance
mock_executor_instance.__exit__.return_value = None
@ -694,20 +689,51 @@ class TestIndexingRunnerLoad:
mock_dependencies["executor"].return_value = mock_executor_instance
# Mock update_document_index_status to avoid database calls
with patch.object(runner, "_update_document_index_status"):
with patch.object(runner, "_update_document_index_status") as mock_update_status:
# Act
runner._load(
mock_processor,
sample_dataset,
sample_dataset_document,
sample_documents,
mock_dependencies["session"],
session=mock_dependencies["session"],
dataset=sample_dataset,
dataset_document=sample_dataset_document,
documents=sample_documents,
total_tokens=300,
)
# Assert
model_manager.get_model_instance.assert_called_once()
mock_dependencies["model_manager"].assert_not_called()
# Verify executor was used for parallel processing
assert mock_executor_instance.submit.called
for submit_call in mock_executor_instance.submit.call_args_list:
assert submit_call.args[0] == runner._process_chunk
assert len(submit_call.args) == 6
mock_future.result.assert_called()
assert mock_update_status.call_args.kwargs["extra_update_params"][DatasetDocument.tokens] == 300
def test_load_propagates_worker_errors(
self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents
):
runner = IndexingRunner()
mock_future = MagicMock()
mock_future.result.side_effect = RuntimeError("index failed")
mock_executor_instance = MagicMock()
mock_executor_instance.__enter__.return_value = mock_executor_instance
mock_executor_instance.__exit__.return_value = None
mock_executor_instance.submit.return_value = mock_future
mock_dependencies["executor"].return_value = mock_executor_instance
with (
patch.object(runner, "_update_document_index_status") as mock_update_status,
pytest.raises(RuntimeError, match="index failed"),
):
runner._load(
session=mock_dependencies["session"],
dataset=sample_dataset,
dataset_document=sample_dataset_document,
documents=sample_documents,
total_tokens=300,
)
mock_update_status.assert_not_called()
def test_load_with_economy_indexing(
self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents
@ -717,8 +743,6 @@ class TestIndexingRunnerLoad:
runner = IndexingRunner()
sample_dataset.indexing_technique = IndexTechniqueType.ECONOMY
mock_processor = MagicMock()
# Mock thread for keyword indexing
mock_thread_instance = MagicMock()
mock_thread_instance.join = MagicMock()
@ -728,11 +752,11 @@ class TestIndexingRunnerLoad:
with patch.object(runner, "_update_document_index_status"):
# Act
runner._load(
mock_processor,
sample_dataset,
sample_dataset_document,
sample_documents,
mock_dependencies["session"],
session=mock_dependencies["session"],
dataset=sample_dataset,
dataset_document=sample_dataset_document,
documents=sample_documents,
total_tokens=0,
)
# Assert
@ -759,16 +783,9 @@ class TestIndexingRunnerLoad:
)
]
mock_embedding_instance = MagicMock()
mock_embedding_instance.get_text_embedding_num_tokens.return_value = 50
model_manager = mock_dependencies["model_manager"].return_value
model_manager.get_model_instance.return_value = mock_embedding_instance
mock_processor = MagicMock()
# Mock ThreadPoolExecutor
mock_future = MagicMock()
mock_future.result.return_value = 150
mock_future.result.return_value = None
mock_executor_instance = MagicMock()
mock_executor_instance.__enter__.return_value = mock_executor_instance
mock_executor_instance.__exit__.return_value = None
@ -779,14 +796,15 @@ class TestIndexingRunnerLoad:
with patch.object(runner, "_update_document_index_status"):
# Act
runner._load(
mock_processor,
sample_dataset,
sample_dataset_document,
sample_documents,
mock_dependencies["session"],
session=mock_dependencies["session"],
dataset=sample_dataset,
dataset_document=sample_dataset_document,
documents=sample_documents,
total_tokens=150,
)
# Assert
mock_dependencies["model_manager"].assert_not_called()
# Verify no keyword thread for parent-child index
mock_dependencies["thread"].assert_not_called()
@ -850,6 +868,7 @@ class TestIndexingRunnerRun:
segment.index_node_hash = "parent-hash"
segment.document_id = dataset_document.id
segment.dataset_id = dataset_document.dataset_id
segment.tokens = 12
segment.get_child_chunks.return_value = [
SimpleNamespace(content="child", index_node_id="child-node", index_node_hash="child-hash")
]
@ -862,6 +881,32 @@ class TestIndexingRunnerRun:
segment.get_child_chunks.assert_called_once_with(session=session)
assert load.call_args.kwargs["documents"][0].children[0].page_content == "child"
assert load.call_args.kwargs["total_tokens"] == 12
def test_run_in_indexing_status_uses_tokens_from_all_segments(self, mock_dependencies, sample_dataset_documents):
runner = IndexingRunner()
dataset_document = sample_dataset_documents[0]
dataset = Mock(spec=Dataset)
completed_segment = Mock(spec=DocumentSegment)
completed_segment.status = SegmentStatus.COMPLETED
completed_segment.tokens = 10
incomplete_segment = Mock(spec=DocumentSegment)
incomplete_segment.status = SegmentStatus.WAITING
incomplete_segment.tokens = 20
incomplete_segment.content = "pending"
incomplete_segment.index_node_id = "pending-node"
incomplete_segment.index_node_hash = "pending-hash"
incomplete_segment.document_id = dataset_document.id
incomplete_segment.dataset_id = dataset_document.dataset_id
session = mock_dependencies["session"]
session.get.side_effect = lambda model, _: dataset_document if model is DatasetDocument else dataset
session.scalars.return_value.all.return_value = [completed_segment, incomplete_segment]
with patch.object(runner, "_load") as load:
runner.run_in_indexing_status(dataset_document, session)
assert load.call_args.kwargs["documents"][0].page_content == "pending"
assert load.call_args.kwargs["total_tokens"] == 30
def test_run_success_single_document(self, mock_dependencies, sample_dataset_documents):
"""Test successful run with single document."""
@ -953,6 +998,98 @@ class TestIndexingRunnerRun:
with pytest.raises(DocumentIsPausedError):
runner.run([doc], mock_dependencies["session"])
def test_run_counts_each_transformed_document_once(self, mock_dependencies, sample_dataset_documents):
runner = IndexingRunner()
dataset_document = sample_dataset_documents[0]
dataset = Mock(spec=Dataset)
dataset.id = dataset_document.dataset_id
dataset.tenant_id = dataset_document.tenant_id
dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
current_user = Mock(spec=Account)
transformed_documents = [
Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}),
Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}),
]
model_dispatch = {
DatasetDocument: dataset_document,
Dataset: dataset,
Account: current_user,
}
mock_dependencies["session"].get.side_effect = lambda model, _: model_dispatch.get(model)
process_rule = Mock(spec=DatasetProcessRule)
process_rule.to_dict.return_value = {"mode": "automatic", "rules": {}}
mock_dependencies["session"].scalar.return_value = process_rule
with (
patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]),
patch.object(runner, "_transform", return_value=transformed_documents),
patch.object(runner, "_load_segments") as load_segments,
patch.object(runner, "_load") as load,
patch(
"core.indexing_runner.calculate_segment_token_counts",
return_value=[11, 22],
) as calculate_token_counts,
):
runner.run([dataset_document], mock_dependencies["session"])
calculate_token_counts.assert_called_once_with(dataset=dataset, documents=transformed_documents)
load_segments.assert_called_once_with(
session=mock_dependencies["session"],
dataset=dataset,
dataset_document=dataset_document,
documents=transformed_documents,
token_counts=[11, 22],
)
assert load.call_args.kwargs["total_tokens"] == 33
def test_run_in_splitting_status_counts_each_transformed_document_once(
self, mock_dependencies, sample_dataset_documents
):
runner = IndexingRunner()
dataset_document = sample_dataset_documents[0]
dataset_document.created_by = "user-1"
dataset = Mock(spec=Dataset)
dataset.id = dataset_document.dataset_id
dataset.tenant_id = dataset_document.tenant_id
dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
current_user = Mock(spec=Account)
transformed_documents = [
Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}),
Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}),
]
model_dispatch = {
DatasetDocument: dataset_document,
Dataset: dataset,
Account: current_user,
}
mock_dependencies["session"].get.side_effect = lambda model, _: model_dispatch.get(model)
mock_dependencies["session"].scalars.return_value.all.return_value = []
process_rule = Mock(spec=DatasetProcessRule)
process_rule.to_dict.return_value = {"mode": "automatic", "rules": {}}
mock_dependencies["session"].scalar.return_value = process_rule
with (
patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]),
patch.object(runner, "_transform", return_value=transformed_documents),
patch.object(runner, "_load_segments") as load_segments,
patch.object(runner, "_load") as load,
patch(
"core.indexing_runner.calculate_segment_token_counts",
return_value=[11, 22],
) as calculate_token_counts,
):
runner.run_in_splitting_status(dataset_document, mock_dependencies["session"])
calculate_token_counts.assert_called_once_with(dataset=dataset, documents=transformed_documents)
load_segments.assert_called_once_with(
session=mock_dependencies["session"],
dataset=dataset,
dataset_document=dataset_document,
documents=transformed_documents,
token_counts=[11, 22],
)
assert load.call_args.kwargs["total_tokens"] == 33
def test_run_handles_provider_token_error(self, mock_dependencies, sample_dataset_documents):
"""Test run handles ProviderTokenNotInitError and updates document status."""
# Arrange
@ -1395,7 +1532,11 @@ class TestIndexingRunnerLoadSegments:
):
# Act
runner._load_segments(
sample_dataset, sample_dataset_document, sample_documents, mock_dependencies["session"]
session=mock_dependencies["session"],
dataset=sample_dataset,
dataset_document=sample_dataset_document,
documents=sample_documents,
token_counts=[10, 20],
)
# Assert
@ -1405,7 +1546,10 @@ class TestIndexingRunnerLoadSegments:
document_id=sample_dataset_document.id,
)
mock_docstore_instance.add_documents.assert_called_once_with(
docs=sample_documents, save_child=False, session=mock_dependencies["session"]
session=mock_dependencies["session"],
docs=sample_documents,
save_child=False,
token_counts=[10, 20],
)
def test_load_segments_parent_child_index(
@ -1435,12 +1579,19 @@ class TestIndexingRunnerLoadSegments:
):
# Act
runner._load_segments(
sample_dataset, sample_dataset_document, sample_documents, mock_dependencies["session"]
session=mock_dependencies["session"],
dataset=sample_dataset,
dataset_document=sample_dataset_document,
documents=sample_documents,
token_counts=[10, 20],
)
# Assert
mock_docstore_instance.add_documents.assert_called_once_with(
docs=sample_documents, save_child=True, session=mock_dependencies["session"]
session=mock_dependencies["session"],
docs=sample_documents,
save_child=True,
token_counts=[10, 20],
)
def test_load_segments_updates_word_count(
@ -1462,7 +1613,11 @@ class TestIndexingRunnerLoadSegments:
):
# Act
runner._load_segments(
sample_dataset, sample_dataset_document, sample_documents, mock_dependencies["session"]
session=mock_dependencies["session"],
dataset=sample_dataset,
dataset_document=sample_dataset_document,
documents=sample_documents,
token_counts=[10, 20],
)
# Assert
@ -1565,7 +1720,6 @@ class TestIndexingRunnerProcessChunk:
"""Unit tests for chunk processing in parallel.
Tests cover:
- Token counting
- Vector index creation
- Segment status updates
- Pause detection during processing
@ -1590,16 +1744,12 @@ class TestIndexingRunnerProcessChunk:
app.app_context.return_value.__exit__ = MagicMock()
return app
def test_process_chunk_counts_tokens(self, mock_dependencies, mock_flask_app):
"""Test process chunk correctly counts tokens."""
def test_process_chunk_loads_index_and_completes_segments(self, mock_dependencies, mock_flask_app):
"""Test process chunk loads the index and completes segments without counting tokens."""
# Arrange
from core.indexing_runner import IndexingRunner
runner = IndexingRunner()
mock_embedding_instance = MagicMock()
# Mock to return an iterable that sums to 150 tokens
mock_embedding_instance.get_text_embedding_num_tokens.return_value = [75, 75]
mock_processor = MagicMock()
chunk_documents = [
Document(page_content="Chunk 1", metadata={"doc_id": "c1"}),
@ -1638,18 +1788,19 @@ class TestIndexingRunnerProcessChunk:
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Act - the method creates its own app_context and session
tokens = runner._process_chunk(
result = runner._process_chunk(
mock_flask_app,
IndexStructureType.PARAGRAPH_INDEX,
chunk_documents,
mock_dataset.id,
mock_dataset_document.id,
mock_embedding_instance,
)
# Assert
assert tokens == 150
assert result is None
mock_processor.load.assert_called_once()
mock_dependencies["session"].execute.assert_called_once()
mock_dependencies["session"].commit.assert_called_once()
def test_process_chunk_detects_pause(self, mock_dependencies, mock_flask_app):
"""Test process chunk detects document pause."""
@ -1657,8 +1808,6 @@ class TestIndexingRunnerProcessChunk:
from core.indexing_runner import IndexingRunner
runner = IndexingRunner()
mock_embedding_instance = MagicMock()
mock_processor = MagicMock()
chunk_documents = [Document(page_content="Chunk", metadata={"doc_id": "c1"})]
mock_dataset = Mock(spec=Dataset)
@ -1691,5 +1840,4 @@ class TestIndexingRunnerProcessChunk:
chunk_documents,
mock_dataset.id,
mock_dataset_document.id,
mock_embedding_instance,
)