diff --git a/api/models/enums.py b/api/models/enums.py index cdec7b2f122..93d2ac02b0e 100644 --- a/api/models/enums.py +++ b/api/models/enums.py @@ -113,6 +113,7 @@ class DataSourceType(StrEnum): WEBSITE_CRAWL = "website_crawl" LOCAL_FILE = "local_file" ONLINE_DOCUMENT = "online_document" + ONLINE_DRIVE = "online_drive" class ProcessRuleMode(StrEnum): diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py index 06face41fe7..ca09324bf42 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py @@ -8,6 +8,7 @@ import core.app.apps.pipeline.pipeline_generator as module from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom from core.datasource.entities.datasource_entities import DatasourceProviderType +from models.enums import DataSourceType class FakeRagPipelineGenerateEntity(SimpleNamespace): @@ -558,6 +559,24 @@ def test_build_document_sets_metadata_for_builtin_fields(generator, mocker): assert document.doc_metadata +def test_build_document_supports_online_drive_datasource_type(generator): + document = generator._build_document( + tenant_id="tenant", + dataset_id="ds", + built_in_field_enabled=True, + datasource_type=DatasourceProviderType.ONLINE_DRIVE, + datasource_info={"id": "file-1", "bucket": "bucket-1", "name": "drive.pdf", "type": "file"}, + created_from="rag-pipeline", + position=1, + account=_build_user(), + batch="batch", + document_form="text", + ) + + assert DataSourceType(document.data_source_type) == DataSourceType.ONLINE_DRIVE + assert document.name == "drive.pdf" + + def test_build_document_invalid_datasource_type(generator): with pytest.raises(ValueError): generator._build_document(