test: replace indexing_technique string literals with IndexTechnique (#34042)

This commit is contained in:
tmimmanuel 2026-03-25 04:39:58 +01:00 committed by GitHub
parent cb28885205
commit a946015ebf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
15 changed files with 120 additions and 114 deletions

View File

@ -4,6 +4,7 @@ from unittest.mock import Mock, patch
import pytest import pytest
from core.entities.knowledge_entities import PreviewDetail from core.entities.knowledge_entities import PreviewDetail
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor
from core.rag.models.document import AttachmentDocument, Document from core.rag.models.document import AttachmentDocument, Document
from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMUsage from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMUsage
@ -21,7 +22,7 @@ class TestParagraphIndexProcessor:
dataset = Mock() dataset = Mock()
dataset.id = "dataset-1" dataset.id = "dataset-1"
dataset.tenant_id = "tenant-1" dataset.tenant_id = "tenant-1"
dataset.indexing_technique = "high_quality" dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
dataset.is_multimodal = True dataset.is_multimodal = True
return dataset return dataset
@ -167,7 +168,7 @@ class TestParagraphIndexProcessor:
def test_load_uses_keyword_add_texts_with_keywords_when_economy( def test_load_uses_keyword_add_texts_with_keywords_when_economy(
self, processor: ParagraphIndexProcessor, dataset: Mock self, processor: ParagraphIndexProcessor, dataset: Mock
) -> None: ) -> None:
dataset.indexing_technique = "economy" dataset.indexing_technique = IndexTechniqueType.ECONOMY
docs = [Document(page_content="chunk", metadata={})] docs = [Document(page_content="chunk", metadata={})]
with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls: with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls:
@ -178,7 +179,7 @@ class TestParagraphIndexProcessor:
def test_load_uses_keyword_add_texts_without_keywords_when_economy( def test_load_uses_keyword_add_texts_without_keywords_when_economy(
self, processor: ParagraphIndexProcessor, dataset: Mock self, processor: ParagraphIndexProcessor, dataset: Mock
) -> None: ) -> None:
dataset.indexing_technique = "economy" dataset.indexing_technique = IndexTechniqueType.ECONOMY
docs = [Document(page_content="chunk", metadata={})] docs = [Document(page_content="chunk", metadata={})]
with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls: with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls:
@ -208,7 +209,7 @@ class TestParagraphIndexProcessor:
def test_clean_economy_deletes_summaries_and_keywords( def test_clean_economy_deletes_summaries_and_keywords(
self, processor: ParagraphIndexProcessor, dataset: Mock self, processor: ParagraphIndexProcessor, dataset: Mock
) -> None: ) -> None:
dataset.indexing_technique = "economy" dataset.indexing_technique = IndexTechniqueType.ECONOMY
with ( with (
patch( patch(
@ -222,7 +223,7 @@ class TestParagraphIndexProcessor:
mock_keyword_cls.return_value.delete.assert_called_once() mock_keyword_cls.return_value.delete.assert_called_once()
def test_clean_deletes_keywords_by_ids(self, processor: ParagraphIndexProcessor, dataset: Mock) -> None: def test_clean_deletes_keywords_by_ids(self, processor: ParagraphIndexProcessor, dataset: Mock) -> None:
dataset.indexing_technique = "economy" dataset.indexing_technique = IndexTechniqueType.ECONOMY
with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls: with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls:
processor.clean(dataset, ["node-2"], with_keywords=True) processor.clean(dataset, ["node-2"], with_keywords=True)
@ -267,7 +268,7 @@ class TestParagraphIndexProcessor:
def test_index_list_chunks_economy( def test_index_list_chunks_economy(
self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock
) -> None: ) -> None:
dataset.indexing_technique = "economy" dataset.indexing_technique = IndexTechniqueType.ECONOMY
with ( with (
patch( patch(
"core.rag.index_processor.processor.paragraph_index_processor.helper.generate_text_hash", "core.rag.index_processor.processor.paragraph_index_processor.helper.generate_text_hash",

View File

@ -4,6 +4,7 @@ from unittest.mock import MagicMock, Mock, patch
import pytest import pytest
from core.entities.knowledge_entities import PreviewDetail from core.entities.knowledge_entities import PreviewDetail
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor
from core.rag.models.document import AttachmentDocument, ChildDocument, Document from core.rag.models.document import AttachmentDocument, ChildDocument, Document
from services.entities.knowledge_entities.knowledge_entities import ParentMode from services.entities.knowledge_entities.knowledge_entities import ParentMode
@ -19,7 +20,7 @@ class TestParentChildIndexProcessor:
dataset = Mock() dataset = Mock()
dataset.id = "dataset-1" dataset.id = "dataset-1"
dataset.tenant_id = "tenant-1" dataset.tenant_id = "tenant-1"
dataset.indexing_technique = "high_quality" dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
dataset.is_multimodal = True dataset.is_multimodal = True
return dataset return dataset

View File

@ -6,6 +6,7 @@ import pytest
from werkzeug.datastructures import FileStorage from werkzeug.datastructures import FileStorage
from core.entities.knowledge_entities import PreviewDetail from core.entities.knowledge_entities import PreviewDetail
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.index_processor.processor.qa_index_processor import QAIndexProcessor from core.rag.index_processor.processor.qa_index_processor import QAIndexProcessor
from core.rag.models.document import AttachmentDocument, Document from core.rag.models.document import AttachmentDocument, Document
@ -33,7 +34,7 @@ class TestQAIndexProcessor:
dataset = Mock() dataset = Mock()
dataset.id = "dataset-1" dataset.id = "dataset-1"
dataset.tenant_id = "tenant-1" dataset.tenant_id = "tenant-1"
dataset.indexing_technique = "high_quality" dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
dataset.is_multimodal = True dataset.is_multimodal = True
return dataset return dataset
@ -207,7 +208,7 @@ class TestQAIndexProcessor:
vector.create_multimodal.assert_called_once_with(multimodal_docs) vector.create_multimodal.assert_called_once_with(multimodal_docs)
def test_load_skips_vector_for_non_high_quality(self, processor: QAIndexProcessor, dataset: Mock) -> None: def test_load_skips_vector_for_non_high_quality(self, processor: QAIndexProcessor, dataset: Mock) -> None:
dataset.indexing_technique = "economy" dataset.indexing_technique = IndexTechniqueType.ECONOMY
docs = [Document(page_content="Q1", metadata={"answer": "A1"})] docs = [Document(page_content="Q1", metadata={"answer": "A1"})]
with patch("core.rag.index_processor.processor.qa_index_processor.Vector") as mock_vector_cls: with patch("core.rag.index_processor.processor.qa_index_processor.Vector") as mock_vector_cls:
@ -298,7 +299,7 @@ class TestQAIndexProcessor:
def test_index_requires_high_quality( def test_index_requires_high_quality(
self, processor: QAIndexProcessor, dataset: Mock, dataset_document: Mock self, processor: QAIndexProcessor, dataset: Mock, dataset_document: Mock
) -> None: ) -> None:
dataset.indexing_technique = "economy" dataset.indexing_technique = IndexTechniqueType.ECONOMY
qa_chunks = SimpleNamespace(qa_chunks=[SimpleNamespace(question="Q1", answer="A1")]) qa_chunks = SimpleNamespace(qa_chunks=[SimpleNamespace(question="Q1", answer="A1")])
with ( with (

View File

@ -61,7 +61,7 @@ from core.indexing_runner import (
DocumentIsPausedError, DocumentIsPausedError,
IndexingRunner, IndexingRunner,
) )
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from core.rag.models.document import ChildDocument, Document from core.rag.models.document import ChildDocument, Document
from dify_graph.model_runtime.entities.model_entities import ModelType from dify_graph.model_runtime.entities.model_entities import ModelType
from libs.datetime_utils import naive_utc_now from libs.datetime_utils import naive_utc_now
@ -76,7 +76,7 @@ from models.dataset import Document as DatasetDocument
def create_mock_dataset( def create_mock_dataset(
dataset_id: str | None = None, dataset_id: str | None = None,
tenant_id: str | None = None, tenant_id: str | None = None,
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
embedding_provider: str = "openai", embedding_provider: str = "openai",
embedding_model: str = "text-embedding-ada-002", embedding_model: str = "text-embedding-ada-002",
) -> Mock: ) -> Mock:
@ -458,7 +458,7 @@ class TestIndexingRunnerTransform:
dataset = Mock(spec=Dataset) dataset = Mock(spec=Dataset)
dataset.id = str(uuid.uuid4()) dataset.id = str(uuid.uuid4())
dataset.tenant_id = str(uuid.uuid4()) dataset.tenant_id = str(uuid.uuid4())
dataset.indexing_technique = "high_quality" dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
dataset.embedding_model_provider = "openai" dataset.embedding_model_provider = "openai"
dataset.embedding_model = "text-embedding-ada-002" dataset.embedding_model = "text-embedding-ada-002"
return dataset return dataset
@ -521,7 +521,7 @@ class TestIndexingRunnerTransform:
"""Test transformation with economy indexing (no embeddings).""" """Test transformation with economy indexing (no embeddings)."""
# Arrange # Arrange
runner = IndexingRunner() runner = IndexingRunner()
sample_dataset.indexing_technique = "economy" sample_dataset.indexing_technique = IndexTechniqueType.ECONOMY
mock_processor = MagicMock() mock_processor = MagicMock()
transformed_docs = [ transformed_docs = [
@ -605,7 +605,7 @@ class TestIndexingRunnerLoad:
dataset = Mock(spec=Dataset) dataset = Mock(spec=Dataset)
dataset.id = str(uuid.uuid4()) dataset.id = str(uuid.uuid4())
dataset.tenant_id = str(uuid.uuid4()) dataset.tenant_id = str(uuid.uuid4())
dataset.indexing_technique = "high_quality" dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
dataset.embedding_model_provider = "openai" dataset.embedding_model_provider = "openai"
dataset.embedding_model = "text-embedding-ada-002" dataset.embedding_model = "text-embedding-ada-002"
return dataset return dataset
@ -674,7 +674,7 @@ class TestIndexingRunnerLoad:
"""Test loading with economy indexing (keyword only).""" """Test loading with economy indexing (keyword only)."""
# Arrange # Arrange
runner = IndexingRunner() runner = IndexingRunner()
sample_dataset.indexing_technique = "economy" sample_dataset.indexing_technique = IndexTechniqueType.ECONOMY
mock_processor = MagicMock() mock_processor = MagicMock()
@ -701,7 +701,7 @@ class TestIndexingRunnerLoad:
# Arrange # Arrange
runner = IndexingRunner() runner = IndexingRunner()
sample_dataset_document.doc_form = IndexStructureType.PARENT_CHILD_INDEX sample_dataset_document.doc_form = IndexStructureType.PARENT_CHILD_INDEX
sample_dataset.indexing_technique = "high_quality" sample_dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
# Add child documents # Add child documents
for doc in sample_documents: for doc in sample_documents:
@ -795,7 +795,7 @@ class TestIndexingRunnerRun:
mock_dataset = Mock(spec=Dataset) mock_dataset = Mock(spec=Dataset)
mock_dataset.id = doc.dataset_id mock_dataset.id = doc.dataset_id
mock_dataset.tenant_id = doc.tenant_id mock_dataset.tenant_id = doc.tenant_id
mock_dataset.indexing_technique = "economy" mock_dataset.indexing_technique = IndexTechniqueType.ECONOMY
mock_dependencies["db"].session.query.return_value.filter_by.return_value.first.return_value = mock_dataset mock_dependencies["db"].session.query.return_value.filter_by.return_value.first.return_value = mock_dataset
mock_process_rule = Mock(spec=DatasetProcessRule) mock_process_rule = Mock(spec=DatasetProcessRule)
@ -949,7 +949,7 @@ class TestIndexingRunnerRun:
mock_dependencies["db"].session.get.side_effect = get_side_effect mock_dependencies["db"].session.get.side_effect = get_side_effect
mock_dataset = Mock(spec=Dataset) mock_dataset = Mock(spec=Dataset)
mock_dataset.indexing_technique = "economy" mock_dataset.indexing_technique = IndexTechniqueType.ECONOMY
mock_dependencies["db"].session.query.return_value.filter_by.return_value.first.return_value = mock_dataset mock_dependencies["db"].session.query.return_value.filter_by.return_value.first.return_value = mock_dataset
mock_process_rule = Mock(spec=DatasetProcessRule) mock_process_rule = Mock(spec=DatasetProcessRule)

View File

@ -5,6 +5,7 @@ from unittest.mock import Mock
import pytest import pytest
from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.workflow.nodes.knowledge_index.entities import KnowledgeIndexNodeData from core.workflow.nodes.knowledge_index.entities import KnowledgeIndexNodeData
from core.workflow.nodes.knowledge_index.exc import KnowledgeIndexNodeError from core.workflow.nodes.knowledge_index.exc import KnowledgeIndexNodeError
from core.workflow.nodes.knowledge_index.knowledge_index_node import KnowledgeIndexNode from core.workflow.nodes.knowledge_index.knowledge_index_node import KnowledgeIndexNode
@ -78,7 +79,7 @@ def sample_node_data():
type="knowledge-index", type="knowledge-index",
chunk_structure="general_structure", chunk_structure="general_structure",
index_chunk_variable_selector=["start", "chunks"], index_chunk_variable_selector=["start", "chunks"],
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
summary_index_setting=None, summary_index_setting=None,
) )

View File

@ -15,6 +15,7 @@ from datetime import UTC, datetime
from unittest.mock import patch from unittest.mock import patch
from uuid import uuid4 from uuid import uuid4
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from models.dataset import ( from models.dataset import (
AppDatasetJoin, AppDatasetJoin,
ChildChunk, ChildChunk,
@ -67,14 +68,14 @@ class TestDatasetModelValidation:
data_source_type=DataSourceType.UPLOAD_FILE, data_source_type=DataSourceType.UPLOAD_FILE,
created_by=str(uuid4()), created_by=str(uuid4()),
description="Test description", description="Test description",
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
embedding_model="text-embedding-ada-002", embedding_model="text-embedding-ada-002",
embedding_model_provider="openai", embedding_model_provider="openai",
) )
# Assert # Assert
assert dataset.description == "Test description" assert dataset.description == "Test description"
assert dataset.indexing_technique == "high_quality" assert dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY
assert dataset.embedding_model == "text-embedding-ada-002" assert dataset.embedding_model == "text-embedding-ada-002"
assert dataset.embedding_model_provider == "openai" assert dataset.embedding_model_provider == "openai"
@ -86,21 +87,21 @@ class TestDatasetModelValidation:
name="High Quality Dataset", name="High Quality Dataset",
data_source_type=DataSourceType.UPLOAD_FILE, data_source_type=DataSourceType.UPLOAD_FILE,
created_by=str(uuid4()), created_by=str(uuid4()),
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
) )
dataset_economy = Dataset( dataset_economy = Dataset(
tenant_id=str(uuid4()), tenant_id=str(uuid4()),
name="Economy Dataset", name="Economy Dataset",
data_source_type=DataSourceType.UPLOAD_FILE, data_source_type=DataSourceType.UPLOAD_FILE,
created_by=str(uuid4()), created_by=str(uuid4()),
indexing_technique="economy", indexing_technique=IndexTechniqueType.ECONOMY,
) )
# Assert # Assert
assert dataset_high_quality.indexing_technique == "high_quality" assert dataset_high_quality.indexing_technique == IndexTechniqueType.HIGH_QUALITY
assert dataset_economy.indexing_technique == "economy" assert dataset_economy.indexing_technique == IndexTechniqueType.ECONOMY
assert "high_quality" in Dataset.INDEXING_TECHNIQUE_LIST assert IndexTechniqueType.HIGH_QUALITY in Dataset.INDEXING_TECHNIQUE_LIST
assert "economy" in Dataset.INDEXING_TECHNIQUE_LIST assert IndexTechniqueType.ECONOMY in Dataset.INDEXING_TECHNIQUE_LIST
def test_dataset_provider_validation(self): def test_dataset_provider_validation(self):
"""Test dataset provider values.""" """Test dataset provider values."""
@ -983,7 +984,7 @@ class TestModelIntegration:
name="Test Dataset", name="Test Dataset",
data_source_type=DataSourceType.UPLOAD_FILE, data_source_type=DataSourceType.UPLOAD_FILE,
created_by=created_by, created_by=created_by,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
) )
dataset.id = dataset_id dataset.id = dataset_id
@ -1019,7 +1020,7 @@ class TestModelIntegration:
assert document.dataset_id == dataset_id assert document.dataset_id == dataset_id
assert segment.dataset_id == dataset_id assert segment.dataset_id == dataset_id
assert segment.document_id == document_id assert segment.document_id == document_id
assert dataset.indexing_technique == "high_quality" assert dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY
assert document.word_count == 100 assert document.word_count == 100
assert segment.status == SegmentStatus.COMPLETED assert segment.status == SegmentStatus.COMPLETED

View File

@ -97,6 +97,7 @@ from unittest.mock import Mock, create_autospec, patch
import pytest import pytest
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from models import Account, TenantAccountRole from models import Account, TenantAccountRole
from models.dataset import ( from models.dataset import (
AppDatasetJoin, AppDatasetJoin,
@ -149,7 +150,7 @@ class DatasetUpdateDeleteTestDataFactory:
name: str = "Test Dataset", name: str = "Test Dataset",
description: str = "Test description", description: str = "Test description",
tenant_id: str = "tenant-123", tenant_id: str = "tenant-123",
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider: str | None = "openai", embedding_model_provider: str | None = "openai",
embedding_model: str | None = "text-embedding-ada-002", embedding_model: str | None = "text-embedding-ada-002",
collection_binding_id: str | None = "binding-123", collection_binding_id: str | None = "binding-123",
@ -237,7 +238,7 @@ class DatasetUpdateDeleteTestDataFactory:
@staticmethod @staticmethod
def create_knowledge_configuration_mock( def create_knowledge_configuration_mock(
chunk_structure: str = "tree", chunk_structure: str = "tree",
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider: str = "openai", embedding_model_provider: str = "openai",
embedding_model: str = "text-embedding-ada-002", embedding_model: str = "text-embedding-ada-002",
keyword_number: int = 10, keyword_number: int = 10,
@ -630,12 +631,12 @@ class TestDatasetServiceUpdateRagPipelineDatasetSettings:
dataset_id="dataset-123", dataset_id="dataset-123",
runtime_mode="rag_pipeline", runtime_mode="rag_pipeline",
chunk_structure="tree", chunk_structure="tree",
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
) )
knowledge_config = DatasetUpdateDeleteTestDataFactory.create_knowledge_configuration_mock( knowledge_config = DatasetUpdateDeleteTestDataFactory.create_knowledge_configuration_mock(
chunk_structure="list", chunk_structure="list",
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider="openai", embedding_model_provider="openai",
embedding_model="text-embedding-ada-002", embedding_model="text-embedding-ada-002",
) )
@ -671,7 +672,7 @@ class TestDatasetServiceUpdateRagPipelineDatasetSettings:
# Assert # Assert
assert dataset.chunk_structure == "list" assert dataset.chunk_structure == "list"
assert dataset.indexing_technique == "high_quality" assert dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY
assert dataset.embedding_model == "text-embedding-ada-002" assert dataset.embedding_model == "text-embedding-ada-002"
assert dataset.embedding_model_provider == "openai" assert dataset.embedding_model_provider == "openai"
assert dataset.collection_binding_id == "binding-123" assert dataset.collection_binding_id == "binding-123"
@ -698,12 +699,12 @@ class TestDatasetServiceUpdateRagPipelineDatasetSettings:
dataset_id="dataset-123", dataset_id="dataset-123",
runtime_mode="rag_pipeline", runtime_mode="rag_pipeline",
chunk_structure="tree", # Existing structure chunk_structure="tree", # Existing structure
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
) )
knowledge_config = DatasetUpdateDeleteTestDataFactory.create_knowledge_configuration_mock( knowledge_config = DatasetUpdateDeleteTestDataFactory.create_knowledge_configuration_mock(
chunk_structure="list", # Different structure chunk_structure="list", # Different structure
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
) )
mock_session.merge.return_value = dataset mock_session.merge.return_value = dataset
@ -735,11 +736,11 @@ class TestDatasetServiceUpdateRagPipelineDatasetSettings:
dataset = DatasetUpdateDeleteTestDataFactory.create_dataset_mock( dataset = DatasetUpdateDeleteTestDataFactory.create_dataset_mock(
dataset_id="dataset-123", dataset_id="dataset-123",
runtime_mode="rag_pipeline", runtime_mode="rag_pipeline",
indexing_technique="high_quality", # Current technique indexing_technique=IndexTechniqueType.HIGH_QUALITY, # Current technique
) )
knowledge_config = DatasetUpdateDeleteTestDataFactory.create_knowledge_configuration_mock( knowledge_config = DatasetUpdateDeleteTestDataFactory.create_knowledge_configuration_mock(
indexing_technique="economy", # Trying to change to economy indexing_technique=IndexTechniqueType.ECONOMY, # Trying to change to economy
) )
mock_session.merge.return_value = dataset mock_session.merge.return_value = dataset

View File

@ -111,7 +111,7 @@ from unittest.mock import Mock, patch
import pytest import pytest
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from dify_graph.model_runtime.entities.model_entities import ModelType from dify_graph.model_runtime.entities.model_entities import ModelType
from models.dataset import Dataset, DatasetProcessRule, Document from models.dataset import Dataset, DatasetProcessRule, Document
from services.dataset_service import DatasetService, DocumentService from services.dataset_service import DatasetService, DocumentService
@ -154,7 +154,7 @@ class DocumentValidationTestDataFactory:
dataset_id: str = "dataset-123", dataset_id: str = "dataset-123",
tenant_id: str = "tenant-123", tenant_id: str = "tenant-123",
doc_form: str | None = None, doc_form: str | None = None,
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider: str = "openai", embedding_model_provider: str = "openai",
embedding_model: str = "text-embedding-ada-002", embedding_model: str = "text-embedding-ada-002",
**kwargs, **kwargs,
@ -190,7 +190,7 @@ class DocumentValidationTestDataFactory:
data_source: DataSource | None = None, data_source: DataSource | None = None,
process_rule: ProcessRule | None = None, process_rule: ProcessRule | None = None,
doc_form: str = IndexStructureType.PARAGRAPH_INDEX, doc_form: str = IndexStructureType.PARAGRAPH_INDEX,
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
**kwargs, **kwargs,
) -> Mock: ) -> Mock:
""" """
@ -448,7 +448,7 @@ class TestDatasetServiceCheckDatasetModelSetting:
""" """
# Arrange # Arrange
dataset = DocumentValidationTestDataFactory.create_dataset_mock( dataset = DocumentValidationTestDataFactory.create_dataset_mock(
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider="openai", embedding_model_provider="openai",
embedding_model="text-embedding-ada-002", embedding_model="text-embedding-ada-002",
) )
@ -481,7 +481,7 @@ class TestDatasetServiceCheckDatasetModelSetting:
- No errors are raised - No errors are raised
""" """
# Arrange # Arrange
dataset = DocumentValidationTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = DocumentValidationTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
# Act (should not raise) # Act (should not raise)
DatasetService.check_dataset_model_setting(dataset) DatasetService.check_dataset_model_setting(dataset)
@ -503,7 +503,7 @@ class TestDatasetServiceCheckDatasetModelSetting:
""" """
# Arrange # Arrange
dataset = DocumentValidationTestDataFactory.create_dataset_mock( dataset = DocumentValidationTestDataFactory.create_dataset_mock(
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider="openai", embedding_model_provider="openai",
embedding_model="invalid-model", embedding_model="invalid-model",
) )
@ -533,7 +533,7 @@ class TestDatasetServiceCheckDatasetModelSetting:
""" """
# Arrange # Arrange
dataset = DocumentValidationTestDataFactory.create_dataset_mock( dataset = DocumentValidationTestDataFactory.create_dataset_mock(
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider="openai", embedding_model_provider="openai",
embedding_model="text-embedding-ada-002", embedding_model="text-embedding-ada-002",
) )

View File

@ -2,7 +2,7 @@ from unittest.mock import MagicMock, Mock, patch
import pytest import pytest
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from models.account import Account from models.account import Account
from models.dataset import ChildChunk, Dataset, Document, DocumentSegment from models.dataset import ChildChunk, Dataset, Document, DocumentSegment
from models.enums import SegmentType from models.enums import SegmentType
@ -111,7 +111,7 @@ class SegmentTestDataFactory:
def create_dataset_mock( def create_dataset_mock(
dataset_id: str = "dataset-123", dataset_id: str = "dataset-123",
tenant_id: str = "tenant-123", tenant_id: str = "tenant-123",
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
embedding_model: str = "text-embedding-ada-002", embedding_model: str = "text-embedding-ada-002",
embedding_model_provider: str = "openai", embedding_model_provider: str = "openai",
**kwargs, **kwargs,
@ -163,7 +163,7 @@ class TestSegmentServiceCreateSegment:
"""Test successful creation of a segment.""" """Test successful creation of a segment."""
# Arrange # Arrange
document = SegmentTestDataFactory.create_document_mock(word_count=100) document = SegmentTestDataFactory.create_document_mock(word_count=100)
dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
args = {"content": "New segment content", "keywords": ["test", "segment"]} args = {"content": "New segment content", "keywords": ["test", "segment"]}
mock_query = MagicMock() mock_query = MagicMock()
@ -212,7 +212,7 @@ class TestSegmentServiceCreateSegment:
"""Test creation of segment with QA model (requires answer).""" """Test creation of segment with QA model (requires answer)."""
# Arrange # Arrange
document = SegmentTestDataFactory.create_document_mock(doc_form=IndexStructureType.QA_INDEX, word_count=100) document = SegmentTestDataFactory.create_document_mock(doc_form=IndexStructureType.QA_INDEX, word_count=100)
dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
args = {"content": "What is AI?", "answer": "AI is Artificial Intelligence", "keywords": ["ai"]} args = {"content": "What is AI?", "answer": "AI is Artificial Intelligence", "keywords": ["ai"]}
mock_query = MagicMock() mock_query = MagicMock()
@ -247,7 +247,7 @@ class TestSegmentServiceCreateSegment:
"""Test creation of segment with high quality indexing technique.""" """Test creation of segment with high quality indexing technique."""
# Arrange # Arrange
document = SegmentTestDataFactory.create_document_mock(word_count=100) document = SegmentTestDataFactory.create_document_mock(word_count=100)
dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique="high_quality") dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
args = {"content": "New segment content", "keywords": ["test"]} args = {"content": "New segment content", "keywords": ["test"]}
mock_query = MagicMock() mock_query = MagicMock()
@ -289,7 +289,7 @@ class TestSegmentServiceCreateSegment:
"""Test segment creation when vector indexing fails.""" """Test segment creation when vector indexing fails."""
# Arrange # Arrange
document = SegmentTestDataFactory.create_document_mock(word_count=100) document = SegmentTestDataFactory.create_document_mock(word_count=100)
dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
args = {"content": "New segment content", "keywords": ["test"]} args = {"content": "New segment content", "keywords": ["test"]}
mock_query = MagicMock() mock_query = MagicMock()
@ -342,7 +342,7 @@ class TestSegmentServiceUpdateSegment:
# Arrange # Arrange
segment = SegmentTestDataFactory.create_segment_mock(enabled=True, word_count=10) segment = SegmentTestDataFactory.create_segment_mock(enabled=True, word_count=10)
document = SegmentTestDataFactory.create_document_mock(word_count=100) document = SegmentTestDataFactory.create_document_mock(word_count=100)
dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
args = SegmentUpdateArgs(content="Updated content", keywords=["updated"]) args = SegmentUpdateArgs(content="Updated content", keywords=["updated"])
mock_db_session.query.return_value.where.return_value.first.return_value = segment mock_db_session.query.return_value.where.return_value.first.return_value = segment
@ -431,7 +431,7 @@ class TestSegmentServiceUpdateSegment:
# Arrange # Arrange
segment = SegmentTestDataFactory.create_segment_mock(enabled=True, word_count=10) segment = SegmentTestDataFactory.create_segment_mock(enabled=True, word_count=10)
document = SegmentTestDataFactory.create_document_mock(doc_form=IndexStructureType.QA_INDEX, word_count=100) document = SegmentTestDataFactory.create_document_mock(doc_form=IndexStructureType.QA_INDEX, word_count=100)
dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = SegmentTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
args = SegmentUpdateArgs(content="Updated question", answer="Updated answer", keywords=["qa"]) args = SegmentUpdateArgs(content="Updated question", answer="Updated answer", keywords=["qa"])
mock_db_session.query.return_value.where.return_value.first.return_value = segment mock_db_session.query.return_value.where.return_value.first.return_value = segment

View File

@ -4,7 +4,7 @@ from unittest.mock import Mock, create_autospec
import pytest import pytest
from redis.exceptions import LockNotOwnedError from redis.exceptions import LockNotOwnedError
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from models.account import Account from models.account import Account
from models.dataset import Dataset, Document from models.dataset import Dataset, Document
from services.dataset_service import DocumentService, SegmentService from services.dataset_service import DocumentService, SegmentService
@ -71,7 +71,7 @@ def test_save_document_with_dataset_id_ignores_lock_not_owned(
dataset.id = "ds-1" dataset.id = "ds-1"
dataset.tenant_id = fake_current_user.current_tenant_id dataset.tenant_id = fake_current_user.current_tenant_id
dataset.data_source_type = "upload_file" dataset.data_source_type = "upload_file"
dataset.indexing_technique = "high_quality" # so we skip re-initialization branch dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY # so we skip re-initialization branch
# Minimal knowledge_config stub that satisfies pre-lock code # Minimal knowledge_config stub that satisfies pre-lock code
info_list = types.SimpleNamespace(data_source_type="upload_file") info_list = types.SimpleNamespace(data_source_type="upload_file")
@ -80,7 +80,7 @@ def test_save_document_with_dataset_id_ignores_lock_not_owned(
doc_form=IndexStructureType.QA_INDEX, doc_form=IndexStructureType.QA_INDEX,
original_document_id=None, # go into "new document" branch original_document_id=None, # go into "new document" branch
data_source=data_source, data_source=data_source,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
embedding_model=None, embedding_model=None,
embedding_model_provider=None, embedding_model_provider=None,
retrieval_model=None, retrieval_model=None,
@ -126,7 +126,7 @@ def test_add_segment_ignores_lock_not_owned(
dataset = create_autospec(Dataset, instance=True) dataset = create_autospec(Dataset, instance=True)
dataset.id = "ds-1" dataset.id = "ds-1"
dataset.tenant_id = fake_current_user.current_tenant_id dataset.tenant_id = fake_current_user.current_tenant_id
dataset.indexing_technique = "economy" # skip embedding/token calculation branch dataset.indexing_technique = IndexTechniqueType.ECONOMY # skip embedding/token calculation branch
document = create_autospec(Document, instance=True) document = create_autospec(Document, instance=True)
document.id = "doc-1" document.id = "doc-1"
@ -169,7 +169,7 @@ def test_multi_create_segment_ignores_lock_not_owned(
dataset = create_autospec(Dataset, instance=True) dataset = create_autospec(Dataset, instance=True)
dataset.id = "ds-1" dataset.id = "ds-1"
dataset.tenant_id = fake_current_user.current_tenant_id dataset.tenant_id = fake_current_user.current_tenant_id
dataset.indexing_technique = "economy" # again, skip high_quality path dataset.indexing_technique = IndexTechniqueType.ECONOMY # again, skip high_quality path
document = create_autospec(Document, instance=True) document = create_autospec(Document, instance=True)
document.id = "doc-1" document.id = "doc-1"

View File

@ -11,7 +11,7 @@ from unittest.mock import MagicMock
import pytest import pytest
import services.summary_index_service as summary_module import services.summary_index_service as summary_module
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from models.enums import SegmentStatus, SummaryStatus from models.enums import SegmentStatus, SummaryStatus
from services.summary_index_service import SummaryIndexService from services.summary_index_service import SummaryIndexService
@ -27,7 +27,7 @@ class _SessionContext:
return None return None
def _dataset(*, indexing_technique: str = "high_quality") -> MagicMock: def _dataset(*, indexing_technique: str = IndexTechniqueType.HIGH_QUALITY) -> MagicMock:
dataset = MagicMock(name="dataset") dataset = MagicMock(name="dataset")
dataset.id = "dataset-1" dataset.id = "dataset-1"
dataset.tenant_id = "tenant-1" dataset.tenant_id = "tenant-1"
@ -169,7 +169,8 @@ def test_create_summary_record_creates_new(monkeypatch: pytest.MonkeyPatch) -> N
def test_vectorize_summary_skips_non_high_quality(monkeypatch: pytest.MonkeyPatch) -> None: def test_vectorize_summary_skips_non_high_quality(monkeypatch: pytest.MonkeyPatch) -> None:
vector_cls = MagicMock() vector_cls = MagicMock()
monkeypatch.setattr(summary_module, "Vector", vector_cls) monkeypatch.setattr(summary_module, "Vector", vector_cls)
SummaryIndexService.vectorize_summary(_summary_record(), _segment(), _dataset(indexing_technique="economy")) dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY)
SummaryIndexService.vectorize_summary(_summary_record(), _segment(), dataset)
vector_cls.assert_not_called() vector_cls.assert_not_called()
@ -621,7 +622,7 @@ def test_generate_and_vectorize_summary_creates_missing_record_and_logs_usage(mo
def test_generate_summaries_for_document_skip_conditions(monkeypatch: pytest.MonkeyPatch) -> None: def test_generate_summaries_for_document_skip_conditions(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _dataset(indexing_technique="economy") dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY)
document = MagicMock(spec=summary_module.DatasetDocument) document = MagicMock(spec=summary_module.DatasetDocument)
document.id = "doc-1" document.id = "doc-1"
document.doc_form = IndexStructureType.PARAGRAPH_INDEX document.doc_form = IndexStructureType.PARAGRAPH_INDEX
@ -778,7 +779,7 @@ def test_disable_summaries_for_segments_no_summaries_noop(monkeypatch: pytest.Mo
def test_enable_summaries_for_segments_skips_non_high_quality() -> None: def test_enable_summaries_for_segments_skips_non_high_quality() -> None:
SummaryIndexService.enable_summaries_for_segments(_dataset(indexing_technique="economy")) SummaryIndexService.enable_summaries_for_segments(_dataset(indexing_technique=IndexTechniqueType.ECONOMY))
def test_enable_summaries_for_segments_revectorizes_and_enables(monkeypatch: pytest.MonkeyPatch) -> None: def test_enable_summaries_for_segments_revectorizes_and_enables(monkeypatch: pytest.MonkeyPatch) -> None:
@ -932,9 +933,8 @@ def test_delete_summaries_for_segments_no_summaries_noop(monkeypatch: pytest.Mon
def test_update_summary_for_segment_skip_conditions() -> None: def test_update_summary_for_segment_skip_conditions() -> None:
assert ( economy_dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY)
SummaryIndexService.update_summary_for_segment(_segment(), _dataset(indexing_technique="economy"), "x") is None assert SummaryIndexService.update_summary_for_segment(_segment(), economy_dataset, "x") is None
)
seg = _segment(has_document=True) seg = _segment(has_document=True)
seg.document.doc_form = IndexStructureType.QA_INDEX seg.document.doc_form = IndexStructureType.QA_INDEX
assert SummaryIndexService.update_summary_for_segment(seg, _dataset(), "x") is None assert SummaryIndexService.update_summary_for_segment(seg, _dataset(), "x") is None

View File

@ -9,7 +9,7 @@ from unittest.mock import MagicMock
import pytest import pytest
import services.vector_service as vector_service_module import services.vector_service as vector_service_module
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from services.vector_service import VectorService from services.vector_service import VectorService
@ -32,7 +32,7 @@ class _ParentDocStub:
def _make_dataset( def _make_dataset(
*, *,
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
doc_form: str = IndexStructureType.PARAGRAPH_INDEX, doc_form: str = IndexStructureType.PARAGRAPH_INDEX,
tenant_id: str = "tenant-1", tenant_id: str = "tenant-1",
dataset_id: str = "dataset-1", dataset_id: str = "dataset-1",
@ -192,7 +192,7 @@ def test_create_segments_vector_parent_child_calls_generate_child_chunks_with_ex
dataset = _make_dataset( dataset = _make_dataset(
doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX,
embedding_model_provider="openai", embedding_model_provider="openai",
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
) )
segment = _make_segment() segment = _make_segment()
@ -241,7 +241,7 @@ def test_create_segments_vector_parent_child_uses_default_embedding_model_when_p
dataset = _make_dataset( dataset = _make_dataset(
doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX,
embedding_model_provider=None, embedding_model_provider=None,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
) )
segment = _make_segment() segment = _make_segment()
@ -329,7 +329,7 @@ def test_create_segments_vector_parent_child_missing_processing_rule_raises(monk
def test_create_segments_vector_parent_child_non_high_quality_raises(monkeypatch: pytest.MonkeyPatch) -> None: def test_create_segments_vector_parent_child_non_high_quality_raises(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset( dataset = _make_dataset(
doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX,
indexing_technique="economy", indexing_technique=IndexTechniqueType.ECONOMY,
) )
segment = _make_segment() segment = _make_segment()
dataset_document = MagicMock() dataset_document = MagicMock()
@ -348,7 +348,7 @@ def test_create_segments_vector_parent_child_non_high_quality_raises(monkeypatch
def test_update_segment_vector_high_quality_uses_vector(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_segment_vector_high_quality_uses_vector(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="high_quality") dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
segment = _make_segment() segment = _make_segment()
vector_instance = MagicMock() vector_instance = MagicMock()
@ -364,7 +364,7 @@ def test_update_segment_vector_high_quality_uses_vector(monkeypatch: pytest.Monk
def test_update_segment_vector_economy_uses_keyword_with_keywords_list(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_segment_vector_economy_uses_keyword_with_keywords_list(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="economy") dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY)
segment = _make_segment() segment = _make_segment()
keyword_instance = MagicMock() keyword_instance = MagicMock()
@ -380,7 +380,7 @@ def test_update_segment_vector_economy_uses_keyword_with_keywords_list(monkeypat
def test_update_segment_vector_economy_uses_keyword_without_keywords_list(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_segment_vector_economy_uses_keyword_without_keywords_list(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="economy") dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY)
segment = _make_segment() segment = _make_segment()
keyword_instance = MagicMock() keyword_instance = MagicMock()
@ -473,7 +473,7 @@ def test_generate_child_chunks_commits_even_when_no_children(monkeypatch: pytest
def test_create_child_chunk_vector_high_quality_adds_texts(monkeypatch: pytest.MonkeyPatch) -> None: def test_create_child_chunk_vector_high_quality_adds_texts(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="high_quality") dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
child_chunk = MagicMock() child_chunk = MagicMock()
child_chunk.content = "child" child_chunk.content = "child"
child_chunk.index_node_id = "id" child_chunk.index_node_id = "id"
@ -489,7 +489,7 @@ def test_create_child_chunk_vector_high_quality_adds_texts(monkeypatch: pytest.M
def test_create_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch) -> None: def test_create_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="economy") dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY)
vector_cls = MagicMock() vector_cls = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", vector_cls) monkeypatch.setattr(vector_service_module, "Vector", vector_cls)
@ -505,7 +505,7 @@ def test_create_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch)
def test_update_child_chunk_vector_high_quality_updates_vector(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_child_chunk_vector_high_quality_updates_vector(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="high_quality") dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
new_chunk = MagicMock() new_chunk = MagicMock()
new_chunk.content = "n" new_chunk.content = "n"
@ -536,7 +536,7 @@ def test_update_child_chunk_vector_high_quality_updates_vector(monkeypatch: pyte
def test_update_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="economy") dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY)
vector_cls = MagicMock() vector_cls = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", vector_cls) monkeypatch.setattr(vector_service_module, "Vector", vector_cls)
VectorService.update_child_chunk_vector([], [], [], dataset) VectorService.update_child_chunk_vector([], [], [], dataset)
@ -561,7 +561,7 @@ def test_delete_child_chunk_vector_deletes_by_id(monkeypatch: pytest.MonkeyPatch
def test_update_multimodel_vector_returns_when_not_high_quality(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_multimodel_vector_returns_when_not_high_quality(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="economy", is_multimodal=True) dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY, is_multimodal=True)
segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}]) segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}])
vector_cls = MagicMock() vector_cls = MagicMock()
@ -575,7 +575,7 @@ def test_update_multimodel_vector_returns_when_not_high_quality(monkeypatch: pyt
def test_update_multimodel_vector_returns_when_no_actual_change(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_multimodel_vector_returns_when_no_actual_change(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="high_quality", is_multimodal=True) dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}, {"id": "b"}]) segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}, {"id": "b"}])
vector_cls = MagicMock() vector_cls = MagicMock()
@ -591,7 +591,7 @@ def test_update_multimodel_vector_returns_when_no_actual_change(monkeypatch: pyt
def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids( def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
dataset = _make_dataset(indexing_technique="high_quality", is_multimodal=True) dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}, {"id": "old-2"}]) segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}, {"id": "old-2"}])
vector_instance = MagicMock(name="vector_instance") vector_instance = MagicMock(name="vector_instance")
@ -612,7 +612,7 @@ def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids(
def test_update_multimodel_vector_commits_when_no_upload_files_found(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_multimodel_vector_commits_when_no_upload_files_found(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="high_quality", is_multimodal=True) dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}]) segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}])
vector_instance = MagicMock() vector_instance = MagicMock()
@ -630,7 +630,7 @@ def test_update_multimodel_vector_commits_when_no_upload_files_found(monkeypatch
def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_upload_files( def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_upload_files(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
dataset = _make_dataset(indexing_technique="high_quality", is_multimodal=True) dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}]) segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}])
vector_instance = MagicMock() vector_instance = MagicMock()
@ -663,7 +663,7 @@ def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_up
def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops( def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
dataset = _make_dataset(indexing_technique="high_quality", is_multimodal=False) dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=False)
segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}]) segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}])
vector_instance = MagicMock() vector_instance = MagicMock()
@ -683,7 +683,7 @@ def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops
def test_update_multimodel_vector_rolls_back_and_reraises_on_error(monkeypatch: pytest.MonkeyPatch) -> None: def test_update_multimodel_vector_rolls_back_and_reraises_on_error(monkeypatch: pytest.MonkeyPatch) -> None:
dataset = _make_dataset(indexing_technique="high_quality", is_multimodal=True) dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}]) segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}])
vector_instance = MagicMock() vector_instance = MagicMock()

View File

@ -121,7 +121,7 @@ import pytest
from core.rag.datasource.vdb.vector_base import BaseVector from core.rag.datasource.vdb.vector_base import BaseVector
from core.rag.datasource.vdb.vector_factory import Vector from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.datasource.vdb.vector_type import VectorType from core.rag.datasource.vdb.vector_type import VectorType
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from core.rag.models.document import Document from core.rag.models.document import Document
from models.dataset import ChildChunk, Dataset, DatasetDocument, DatasetProcessRule, DocumentSegment from models.dataset import ChildChunk, Dataset, DatasetDocument, DatasetProcessRule, DocumentSegment
from services.vector_service import VectorService from services.vector_service import VectorService
@ -153,7 +153,7 @@ class VectorServiceTestDataFactory:
dataset_id: str = "dataset-123", dataset_id: str = "dataset-123",
tenant_id: str = "tenant-123", tenant_id: str = "tenant-123",
doc_form: str = IndexStructureType.PARAGRAPH_INDEX, doc_form: str = IndexStructureType.PARAGRAPH_INDEX,
indexing_technique: str = "high_quality", indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
embedding_model_provider: str = "openai", embedding_model_provider: str = "openai",
embedding_model: str = "text-embedding-ada-002", embedding_model: str = "text-embedding-ada-002",
index_struct_dict: dict | None = None, index_struct_dict: dict | None = None,
@ -494,7 +494,7 @@ class TestVectorService:
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock( dataset = VectorServiceTestDataFactory.create_dataset_mock(
doc_form=IndexStructureType.PARAGRAPH_INDEX, indexing_technique="high_quality" doc_form=IndexStructureType.PARAGRAPH_INDEX, indexing_technique=IndexTechniqueType.HIGH_QUALITY
) )
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -535,7 +535,7 @@ class TestVectorService:
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock( dataset = VectorServiceTestDataFactory.create_dataset_mock(
doc_form="parent_child_model", indexing_technique="high_quality" doc_form="parent_child_model", indexing_technique=IndexTechniqueType.HIGH_QUALITY
) )
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -568,7 +568,7 @@ class TestVectorService:
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock( dataset = VectorServiceTestDataFactory.create_dataset_mock(
doc_form="parent_child_model", indexing_technique="high_quality" doc_form="parent_child_model", indexing_technique=IndexTechniqueType.HIGH_QUALITY
) )
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -591,7 +591,7 @@ class TestVectorService:
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock( dataset = VectorServiceTestDataFactory.create_dataset_mock(
doc_form="parent_child_model", indexing_technique="high_quality" doc_form="parent_child_model", indexing_technique=IndexTechniqueType.HIGH_QUALITY
) )
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -616,7 +616,7 @@ class TestVectorService:
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock( dataset = VectorServiceTestDataFactory.create_dataset_mock(
doc_form="parent_child_model", indexing_technique="economy" doc_form="parent_child_model", indexing_technique=IndexTechniqueType.ECONOMY
) )
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -669,7 +669,7 @@ class TestVectorService:
store when using high_quality indexing. store when using high_quality indexing.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="high_quality") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -695,7 +695,7 @@ class TestVectorService:
index when using economy indexing with keywords. index when using economy indexing with keywords.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -731,7 +731,7 @@ class TestVectorService:
index when using economy indexing without keywords. index when using economy indexing without keywords.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
segment = VectorServiceTestDataFactory.create_document_segment_mock() segment = VectorServiceTestDataFactory.create_document_segment_mock()
@ -895,7 +895,7 @@ class TestVectorService:
when using high_quality indexing. when using high_quality indexing.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="high_quality") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock() child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock()
@ -923,7 +923,7 @@ class TestVectorService:
using economy indexing. using economy indexing.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock() child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock()
@ -951,7 +951,7 @@ class TestVectorService:
when there are new chunks, updated chunks, and deleted chunks. when there are new chunks, updated chunks, and deleted chunks.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="high_quality") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
new_chunk = VectorServiceTestDataFactory.create_child_chunk_mock(chunk_id="new-chunk-1") new_chunk = VectorServiceTestDataFactory.create_child_chunk_mock(chunk_id="new-chunk-1")
@ -993,7 +993,7 @@ class TestVectorService:
add_texts is called, not delete_by_ids. add_texts is called, not delete_by_ids.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="high_quality") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
new_chunk = VectorServiceTestDataFactory.create_child_chunk_mock() new_chunk = VectorServiceTestDataFactory.create_child_chunk_mock()
@ -1019,7 +1019,7 @@ class TestVectorService:
delete_by_ids is called, not add_texts. delete_by_ids is called, not add_texts.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="high_quality") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
delete_chunk = VectorServiceTestDataFactory.create_child_chunk_mock() delete_chunk = VectorServiceTestDataFactory.create_child_chunk_mock()
@ -1045,7 +1045,7 @@ class TestVectorService:
using economy indexing. using economy indexing.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
new_chunk = VectorServiceTestDataFactory.create_child_chunk_mock() new_chunk = VectorServiceTestDataFactory.create_child_chunk_mock()
@ -1075,7 +1075,7 @@ class TestVectorService:
when using high_quality indexing. when using high_quality indexing.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="high_quality") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock() child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock()
@ -1099,7 +1099,7 @@ class TestVectorService:
using economy indexing. using economy indexing.
""" """
# Arrange # Arrange
dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique="economy") dataset = VectorServiceTestDataFactory.create_dataset_mock(indexing_technique=IndexTechniqueType.ECONOMY)
child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock() child_chunk = VectorServiceTestDataFactory.create_child_chunk_mock()

View File

@ -16,7 +16,7 @@ from unittest.mock import MagicMock, patch
import pytest import pytest
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from models.enums import DataSourceType from models.enums import DataSourceType
from tasks.clean_dataset_task import clean_dataset_task from tasks.clean_dataset_task import clean_dataset_task
@ -184,7 +184,7 @@ class TestErrorHandling:
clean_dataset_task( clean_dataset_task(
dataset_id=dataset_id, dataset_id=dataset_id,
tenant_id=tenant_id, tenant_id=tenant_id,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph"}', index_struct='{"type": "paragraph"}',
collection_binding_id=collection_binding_id, collection_binding_id=collection_binding_id,
doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_form=IndexStructureType.PARAGRAPH_INDEX,
@ -229,7 +229,7 @@ class TestPipelineAndWorkflowDeletion:
clean_dataset_task( clean_dataset_task(
dataset_id=dataset_id, dataset_id=dataset_id,
tenant_id=tenant_id, tenant_id=tenant_id,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph"}', index_struct='{"type": "paragraph"}',
collection_binding_id=collection_binding_id, collection_binding_id=collection_binding_id,
doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_form=IndexStructureType.PARAGRAPH_INDEX,
@ -265,7 +265,7 @@ class TestPipelineAndWorkflowDeletion:
clean_dataset_task( clean_dataset_task(
dataset_id=dataset_id, dataset_id=dataset_id,
tenant_id=tenant_id, tenant_id=tenant_id,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph"}', index_struct='{"type": "paragraph"}',
collection_binding_id=collection_binding_id, collection_binding_id=collection_binding_id,
doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_form=IndexStructureType.PARAGRAPH_INDEX,
@ -321,7 +321,7 @@ class TestSegmentAttachmentCleanup:
clean_dataset_task( clean_dataset_task(
dataset_id=dataset_id, dataset_id=dataset_id,
tenant_id=tenant_id, tenant_id=tenant_id,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph"}', index_struct='{"type": "paragraph"}',
collection_binding_id=collection_binding_id, collection_binding_id=collection_binding_id,
doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_form=IndexStructureType.PARAGRAPH_INDEX,
@ -366,7 +366,7 @@ class TestSegmentAttachmentCleanup:
clean_dataset_task( clean_dataset_task(
dataset_id=dataset_id, dataset_id=dataset_id,
tenant_id=tenant_id, tenant_id=tenant_id,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph"}', index_struct='{"type": "paragraph"}',
collection_binding_id=collection_binding_id, collection_binding_id=collection_binding_id,
doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_form=IndexStructureType.PARAGRAPH_INDEX,
@ -408,7 +408,7 @@ class TestEdgeCases:
clean_dataset_task( clean_dataset_task(
dataset_id=dataset_id, dataset_id=dataset_id,
tenant_id=tenant_id, tenant_id=tenant_id,
indexing_technique="high_quality", indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph"}', index_struct='{"type": "paragraph"}',
collection_binding_id=collection_binding_id, collection_binding_id=collection_binding_id,
doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_form=IndexStructureType.PARAGRAPH_INDEX,
@ -445,7 +445,7 @@ class TestIndexProcessorParameters:
- Dataset object with correct attributes is passed - Dataset object with correct attributes is passed
""" """
# Arrange # Arrange
indexing_technique = "high_quality" indexing_technique = IndexTechniqueType.HIGH_QUALITY
index_struct = '{"type": "paragraph"}' index_struct = '{"type": "paragraph"}'
# Act # Act

View File

@ -15,7 +15,7 @@ from unittest.mock import MagicMock, Mock, patch
import pytest import pytest
from core.indexing_runner import DocumentIsPausedError from core.indexing_runner import DocumentIsPausedError
from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from core.rag.pipeline.queue import TenantIsolatedTaskQueue from core.rag.pipeline.queue import TenantIsolatedTaskQueue
from enums.cloud_plan import CloudPlan from enums.cloud_plan import CloudPlan
from extensions.ext_redis import redis_client from extensions.ext_redis import redis_client
@ -209,7 +209,7 @@ def mock_dataset(dataset_id, tenant_id):
dataset = Mock(spec=Dataset) dataset = Mock(spec=Dataset)
dataset.id = dataset_id dataset.id = dataset_id
dataset.tenant_id = tenant_id dataset.tenant_id = tenant_id
dataset.indexing_technique = "high_quality" dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
dataset.embedding_model_provider = "openai" dataset.embedding_model_provider = "openai"
dataset.embedding_model = "text-embedding-ada-002" dataset.embedding_model = "text-embedding-ada-002"
return dataset return dataset