From 43254c1ded92d6bd465896356ea2840804e53261 Mon Sep 17 00:00:00 2001 From: "Byron.wang" Date: Thu, 23 Jul 2026 18:56:03 +0800 Subject: [PATCH] chore: avoid duplicate token counting during dataset indexing (#39466) --- api/core/indexing_runner.py | 90 ++-- api/core/rag/docstore/dataset_docstore.py | 27 +- api/core/rag/embedding/token_counter.py | 25 + .../processor/paragraph_index_processor.py | 9 +- .../processor/parent_child_index_processor.py | 9 +- .../processor/qa_index_processor.py | 9 +- .../rag/docstore/test_dataset_docstore.py | 469 ++++++++---------- .../core/rag/embedding/test_token_counter.py | 56 +++ .../test_paragraph_index_processor.py | 27 +- .../test_parent_child_index_processor.py | 20 +- .../processor/test_qa_index_processor.py | 20 +- .../core/rag/indexing/test_indexing_runner.py | 254 ++++++++-- 12 files changed, 633 insertions(+), 382 deletions(-) create mode 100644 api/core/rag/embedding/token_counter.py create mode 100644 api/tests/unit_tests/core/rag/embedding/test_token_counter.py diff --git a/api/core/indexing_runner.py b/api/core/indexing_runner.py index 65f0b9a9523..92246b6614c 100644 --- a/api/core/indexing_runner.py +++ b/api/core/indexing_runner.py @@ -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): diff --git a/api/core/rag/docstore/dataset_docstore.py b/api/core/rag/docstore/dataset_docstore.py index 9b269e782e7..456df05daa4 100644 --- a/api/core/rag/docstore/dataset_docstore.py +++ b/api/core/rag/docstore/dataset_docstore.py @@ -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") diff --git a/api/core/rag/embedding/token_counter.py b/api/core/rag/embedding/token_counter.py new file mode 100644 index 00000000000..e5e25c0fd2a --- /dev/null +++ b/api/core/rag/embedding/token_counter.py @@ -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]) diff --git a/api/core/rag/index_processor/processor/paragraph_index_processor.py b/api/core/rag/index_processor/processor/paragraph_index_processor.py index 9f2ab8e8d73..6e4a1fef757 100644 --- a/api/core/rag/index_processor/processor/paragraph_index_processor.py +++ b/api/core/rag/index_processor/processor/paragraph_index_processor.py @@ -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) diff --git a/api/core/rag/index_processor/processor/parent_child_index_processor.py b/api/core/rag/index_processor/processor/parent_child_index_processor.py index fa2826f7a57..b31875c2f20 100644 --- a/api/core/rag/index_processor/processor/parent_child_index_processor.py +++ b/api/core/rag/index_processor/processor/parent_child_index_processor.py @@ -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 = [] diff --git a/api/core/rag/index_processor/processor/qa_index_processor.py b/api/core/rag/index_processor/processor/qa_index_processor.py index 8a70ec63773..4b55d0d75a9 100644 --- a/api/core/rag/index_processor/processor/qa_index_processor.py +++ b/api/core/rag/index_processor/processor/qa_index_processor.py @@ -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) diff --git a/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py b/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py index 740e4391cac..b7c51c77427 100644 --- a/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py +++ b/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py @@ -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 diff --git a/api/tests/unit_tests/core/rag/embedding/test_token_counter.py b/api/tests/unit_tests/core/rag/embedding/test_token_counter.py new file mode 100644 index 00000000000..0c134a94c90 --- /dev/null +++ b/api/tests/unit_tests/core/rag/embedding/test_token_counter.py @@ -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() diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py index 8ca844d0da0..6f9b0d0cc74 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py @@ -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) diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py index d2c9670b795..a2a5c67972d 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py @@ -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) diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py index ad7b843f89e..ec390fcc578 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py @@ -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) diff --git a/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py b/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py index bceadf3ee84..5307da6d343 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py +++ b/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py @@ -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, )