unstructured-ingest 0.0.25__py3-none-any.whl → 0.1.1__py3-none-any.whl

test/integration/connectors/test_s3.py

@@ -0,0 +1,152 @@
+import os
+import tempfile
+import uuid
+from pathlib import Path
+
+import pytest
+
+from test.integration.connectors.utils.constants import (
+    DESTINATION_TAG,
+    SOURCE_TAG,
+    env_setup_path,
+)
+from test.integration.connectors.utils.docker_compose import docker_compose_context
+from test.integration.connectors.utils.validation import (
+    ValidationConfigs,
+    source_connector_validation,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.error import (
+    SourceConnectionError,
+)
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.fsspec.s3 import (
+    CONNECTOR_TYPE,
+    S3AccessConfig,
+    S3ConnectionConfig,
+    S3Downloader,
+    S3DownloaderConfig,
+    S3Indexer,
+    S3IndexerConfig,
+    S3Uploader,
+    S3UploaderConfig,
+)
+
+
+def validate_predownload_file_data(file_data: FileData):
+    assert file_data.connector_type == CONNECTOR_TYPE
+    assert file_data.local_download_path is None
+
+
+def validate_postdownload_file_data(file_data: FileData):
+    assert file_data.connector_type == CONNECTOR_TYPE
+    assert file_data.local_download_path is not None
+
+
+@pytest.fixture
+def anon_connection_config() -> S3ConnectionConfig:
+    return S3ConnectionConfig(access_config=S3AccessConfig(), anonymous=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+async def test_s3_source(anon_connection_config: S3ConnectionConfig):
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-dev-tech-fixtures/small-pdf-set/")
+    with tempfile.TemporaryDirectory() as tempdir:
+        tempdir_path = Path(tempdir)
+        download_config = S3DownloaderConfig(download_dir=tempdir_path)
+        indexer = S3Indexer(connection_config=anon_connection_config, index_config=indexer_config)
+        downloader = S3Downloader(
+            connection_config=anon_connection_config, download_config=download_config
+        )
+        await source_connector_validation(
+            indexer=indexer,
+            downloader=downloader,
+            configs=ValidationConfigs(
+                test_id="s3",
+                predownload_file_data_check=validate_predownload_file_data,
+                postdownload_file_data_check=validate_postdownload_file_data,
+                expected_num_files=4,
+            ),
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+async def test_s3_source_no_access(anon_connection_config: S3ConnectionConfig):
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-ingest-test-fixtures/destination/")
+    indexer = S3Indexer(connection_config=anon_connection_config, index_config=indexer_config)
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "minio")
+async def test_s3_minio_source(anon_connection_config: S3ConnectionConfig):
+    anon_connection_config.endpoint_url = "http://localhost:9000"
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-dev-tech-fixtures/")
+    with docker_compose_context(docker_compose_path=env_setup_path / "minio"):
+        with tempfile.TemporaryDirectory() as tempdir:
+            tempdir_path = Path(tempdir)
+            download_config = S3DownloaderConfig(download_dir=tempdir_path)
+            indexer = S3Indexer(
+                connection_config=anon_connection_config, index_config=indexer_config
+            )
+            downloader = S3Downloader(
+                connection_config=anon_connection_config, download_config=download_config
+            )
+            await source_connector_validation(
+                indexer=indexer,
+                downloader=downloader,
+                configs=ValidationConfigs(
+                    test_id="s3-minio",
+                    predownload_file_data_check=validate_predownload_file_data,
+                    postdownload_file_data_check=validate_postdownload_file_data,
+                    expected_num_files=1,
+                    exclude_fields_extend=[
+                        "metadata.date_modified",
+                        "metadata.date_created",
+                        "additional_metadata.LastModified",
+                    ],
+                ),
+            )
+
+
+def get_aws_credentials() -> dict:
+    access_key = os.getenv("S3_INGEST_TEST_ACCESS_KEY", None)
+    assert access_key
+    secret_key = os.getenv("S3_INGEST_TEST_SECRET_KEY", None)
+    assert secret_key
+    return {"aws_access_key_id": access_key, "aws_secret_access_key": secret_key}
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("S3_INGEST_TEST_ACCESS_KEY", "S3_INGEST_TEST_SECRET_KEY")
+async def test_s3_destination(upload_file: Path):
+    aws_credentials = get_aws_credentials()
+    s3_bucket = "s3://utic-ingest-test-fixtures"
+    destination_path = f"{s3_bucket}/destination/{uuid.uuid4()}"
+    connection_config = S3ConnectionConfig(
+        access_config=S3AccessConfig(
+            key=aws_credentials["aws_access_key_id"],
+            secret=aws_credentials["aws_secret_access_key"],
+        ),
+    )
+    upload_config = S3UploaderConfig(remote_url=destination_path)
+    uploader = S3Uploader(connection_config=connection_config, upload_config=upload_config)
+    s3fs = uploader.fs
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    try:
+        if uploader.is_async():
+            await uploader.run_async(path=upload_file, file_data=file_data)
+        else:
+            uploader.run(path=upload_file, file_data=file_data)
+        uploaded_files = s3fs.ls(path=destination_path)
+        assert len(uploaded_files) == 1
+    finally:
+        s3fs.rm(path=destination_path, recursive=True)

test/integration/connectors/utils/constants.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+
+SOURCE_TAG = "source"
+DESTINATION_TAG = "destination"
+
+env_setup_path = Path(__file__).parents[1] / "env_setup"
+expected_results_path = Path(__file__).parents[1] / "expected_results"

test/integration/connectors/utils/docker_compose.py

@@ -0,0 +1,44 @@
+import subprocess
+from contextlib import contextmanager
+from pathlib import Path
+
+
+@contextmanager
+def docker_compose_context(docker_compose_path: Path):
+    # Dynamically run a specific docker compose file and make sure it gets cleanup by
+    # by leveraging a context manager. Uses subprocess to map docker compose commands
+    # to the underlying shell.
+    assert docker_compose_path.exists()
+    if docker_compose_path.is_dir():
+        if (docker_compose_path / "docker-compose.yml").exists():
+            docker_compose_path = docker_compose_path / "docker-compose.yml"
+        elif (docker_compose_path / "docker-compose.yaml").exists():
+            docker_compose_path = docker_compose_path / "docker-compose.yaml"
+    assert docker_compose_path.is_file()
+    resp = None
+    try:
+        cmd = f"docker compose -f {docker_compose_path.resolve()} up -d --wait"
+        print(f"Running command: {cmd}")
+        resp = subprocess.run(
+            cmd,
+            shell=True,
+            capture_output=True,
+        )
+        # Return code from docker compose using --wait can be 1 even if no error
+        yield
+    except Exception as e:
+        if resp:
+            print("STDOUT: {}".format(resp.stdout.decode("utf-8")))
+            print("STDERR: {}".format(resp.stderr.decode("utf-8")))
+        raise e
+    finally:
+        cmd = f"docker compose -f {docker_compose_path.resolve()} down --remove-orphans -v"
+        print(f"Running command: {cmd}")
+        final_resp = subprocess.run(
+            cmd,
+            shell=True,
+            capture_output=True,
+        )
+        if final_resp.returncode != 0:
+            print("STDOUT: {}".format(final_resp.stdout.decode("utf-8")))
+            print("STDERR: {}".format(final_resp.stderr.decode("utf-8")))

test/integration/connectors/utils/validation.py

@@ -0,0 +1,203 @@
+import json
+import os
+import shutil
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import Callable, Optional
+
+from deepdiff import DeepDiff
+
+from test.integration.connectors.utils.constants import expected_results_path
+from unstructured_ingest.v2.interfaces import Downloader, FileData, Indexer
+
+
+@dataclass
+class ValidationConfigs:
+    test_id: str
+    expected_num_files: Optional[int] = None
+    predownload_file_data_check: Optional[Callable[[FileData], None]] = None
+    postdownload_file_data_check: Optional[Callable[[FileData], None]] = None
+    exclude_fields: list[str] = field(
+        default_factory=lambda: ["local_download_path", "metadata.date_processed"]
+    )
+    exclude_fields_extend: list[str] = field(default_factory=list)
+
+    def get_exclude_fields(self) -> list[str]:
+        exclude_fields = self.exclude_fields
+        exclude_fields.extend(self.exclude_fields_extend)
+        return exclude_fields
+
+    def run_file_data_validation(
+        self, predownload_file_data: FileData, postdownload_file_data: FileData
+    ):
+        if predownload_file_data_check := self.predownload_file_data_check:
+            predownload_file_data_check(predownload_file_data)
+        if postdownload_file_data_check := self.postdownload_file_data_check:
+            postdownload_file_data_check(postdownload_file_data)
+
+    def run_download_dir_validation(self, download_dir: Path):
+        if expected_num_files := self.expected_num_files:
+            downloaded_files = [p for p in download_dir.rglob("*") if p.is_file()]
+            assert len(downloaded_files) == expected_num_files
+
+    def test_output_dir(self) -> Path:
+        return expected_results_path / self.test_id
+
+    def omit_ignored_fields(self, data: dict) -> dict:
+        exclude_fields = self.get_exclude_fields()
+        # Ignore fields that dynamically change every time the tests run
+        copied_data = data.copy()
+        for exclude_field in exclude_fields:
+            exclude_field_vals = exclude_field.split(".")
+            if len(exclude_field_vals) == 1:
+                current_val = copied_data
+                drop_field = exclude_field_vals[0]
+                copied_data.pop(exclude_field_vals[0], None)
+            else:
+                current_val = copied_data
+                for val in exclude_field_vals[:-1]:
+                    current_val = current_val.get(val, {})
+                drop_field = exclude_field_vals[-1]
+            if drop_field == "*":
+                current_val.clear()
+            else:
+                current_val.pop(drop_field, None)
+        return copied_data
+
+
+def get_files(dir_path: Path) -> list[str]:
+    return [
+        str(f).replace(str(dir_path), "").lstrip("/") for f in dir_path.iterdir() if f.is_file()
+    ]
+
+
+def check_files(expected_output_dir: Path, all_file_data: list[FileData]):
+    expected_files = get_files(dir_path=expected_output_dir)
+    current_files = [f"{file_data.identifier}.json" for file_data in all_file_data]
+    diff = set(expected_files) ^ set(current_files)
+    assert not diff, "diff in files that exist: {}".format(", ".join(diff))
+
+
+def check_contents(
+    expected_output_dir: Path, all_file_data: list[FileData], configs: ValidationConfigs
+):
+    found_diff = False
+    for file_data in all_file_data:
+        file_data_path = expected_output_dir / f"{file_data.identifier}.json"
+        with file_data_path.open("r") as file:
+            expected_file_data_contents = json.load(file)
+        current_file_data_contents = file_data.to_dict()
+        expected_file_data_contents = configs.omit_ignored_fields(expected_file_data_contents)
+        current_file_data_contents = configs.omit_ignored_fields(current_file_data_contents)
+        diff = DeepDiff(expected_file_data_contents, current_file_data_contents)
+        if diff:
+            found_diff = True
+            print(diff.to_json(indent=2))
+    assert not found_diff, f"Diffs found between files: {found_diff}"
+
+
+def run_expected_results_validation(
+    expected_output_dir: Path, all_file_data: list[FileData], configs: ValidationConfigs
+):
+    check_files(expected_output_dir=expected_output_dir, all_file_data=all_file_data)
+    check_contents(
+        expected_output_dir=expected_output_dir, all_file_data=all_file_data, configs=configs
+    )
+
+
+def run_directory_structure_validation(expected_output_dir: Path, download_files: list[str]):
+    directory_record = expected_output_dir / "directory_structure.json"
+    with directory_record.open("r") as directory_file:
+        directory_file_contents = json.load(directory_file)
+    directory_structure = directory_file_contents["directory_structure"]
+    assert directory_structure == download_files
+
+
+def update_fixtures(output_dir: Path, download_dir: Path, all_file_data: list[FileData]):
+    # Delete current files
+    shutil.rmtree(path=output_dir, ignore_errors=True)
+    output_dir.mkdir(parents=True)
+    # Rewrite the current file data
+    file_data_output_path = output_dir / "file_data"
+    file_data_output_path.mkdir(parents=True)
+    for file_data in all_file_data:
+        file_data_path = file_data_output_path / f"{file_data.identifier}.json"
+        with file_data_path.open(mode="w") as f:
+            json.dump(file_data.to_dict(), f, indent=2)
+
+    # Record file structure of download directory
+    download_files = get_files(dir_path=download_dir)
+    download_files.sort()
+    download_dir_record = output_dir / "directory_structure.json"
+    with download_dir_record.open(mode="w") as f:
+        json.dump({"directory_structure": download_files}, f, indent=2)
+
+
+def run_all_validations(
+    configs: ValidationConfigs,
+    predownload_file_data: list[FileData],
+    postdownload_file_data: list[FileData],
+    download_dir: Path,
+    test_output_dir: Path,
+):
+    for pre_data, post_data in zip(predownload_file_data, postdownload_file_data):
+        configs.run_file_data_validation(
+            predownload_file_data=pre_data, postdownload_file_data=post_data
+        )
+    configs.run_download_dir_validation(download_dir=download_dir)
+    run_expected_results_validation(
+        expected_output_dir=test_output_dir / "file_data",
+        all_file_data=postdownload_file_data,
+        configs=configs,
+    )
+    download_files = get_files(dir_path=download_dir)
+    download_files.sort()
+    run_directory_structure_validation(
+        expected_output_dir=configs.test_output_dir(), download_files=download_files
+    )
+
+
+async def source_connector_validation(
+    indexer: Indexer,
+    downloader: Downloader,
+    configs: ValidationConfigs,
+    overwrite_fixtures: bool = os.getenv("OVERWRITE_FIXTURES", "False").lower() == "true",
+) -> None:
+    # Run common validations on the process of running a source connector, supporting dynamic
+    # validators that get passed in along with comparisons on the saved expected values.
+    # If overwrite_fixtures is st to True, will ignore all validators but instead overwrite the
+    # expected values with what gets generated by this test.
+    all_predownload_file_data = []
+    all_postdownload_file_data = []
+    indexer.precheck()
+    download_dir = downloader.download_config.download_dir
+    test_output_dir = configs.test_output_dir()
+    for file_data in indexer.run():
+        assert file_data
+        predownload_file_data = replace(file_data)
+        all_predownload_file_data.append(predownload_file_data)
+        if downloader.is_async():
+            resp = await downloader.run_async(file_data=file_data)
+        else:
+            resp = downloader.run(file_data=file_data)
+        if isinstance(resp, list):
+            for r in resp:
+                postdownload_file_data = replace(r["file_data"])
+                all_postdownload_file_data.append(postdownload_file_data)
+        else:
+            postdownload_file_data = replace(resp["file_data"])
+            all_postdownload_file_data.append(postdownload_file_data)
+    if not overwrite_fixtures:
+        run_all_validations(
+            configs=configs,
+            predownload_file_data=all_predownload_file_data,
+            postdownload_file_data=all_postdownload_file_data,
+            download_dir=download_dir,
+            test_output_dir=test_output_dir,
+        )
+    else:
+        update_fixtures(
+            output_dir=test_output_dir,
+            download_dir=download_dir,
+            all_file_data=all_postdownload_file_data,
+        )

test/integration/embedders/conftest.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+
+import pytest
+
+
+@pytest.fixture
+def embedder_file() -> Path:
+    int_test_dir = Path(__file__).parent
+    assets_dir = int_test_dir / "assets"
+    embedder_file = assets_dir / "DA-1p-with-duplicate-pages.pdf.json"
+    assert embedder_file.exists()
+    assert embedder_file.is_file()
+    return embedder_file

test/integration/embedders/test_bedrock.py

@@ -0,0 +1,49 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.bedrock import BedrockEmbeddingConfig, BedrockEmbeddingEncoder
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+
+def get_aws_credentials() -> dict:
+    access_key = os.getenv("AWS_ACCESS_KEY_ID", None)
+    assert access_key
+    secret_key = os.getenv("AWS_SECRET_ACCESS_KEY", None)
+    assert secret_key
+    return {"aws_access_key_id": access_key, "aws_secret_access_key": secret_key}
+
+
+@requires_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY")
+def test_bedrock_embedder(embedder_file: Path):
+    aws_credentials = get_aws_credentials()
+    embedder_config = EmbedderConfig(
+        embedding_provider="aws-bedrock",
+        embedding_aws_access_key_id=aws_credentials["aws_access_key_id"],
+        embedding_aws_secret_access_key=aws_credentials["aws_secret_access_key"],
+    )
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY")
+def test_raw_bedrock_embedder(embedder_file: Path):
+    aws_credentials = get_aws_credentials()
+    embedder = BedrockEmbeddingEncoder(
+        config=BedrockEmbeddingConfig(
+            aws_access_key_id=aws_credentials["aws_access_key_id"],
+            aws_secret_access_key=aws_credentials["aws_secret_access_key"],
+        )
+    )
+    validate_raw_embedder(
+        embedder=embedder,
+        embedder_file=embedder_file,
+        expected_dimensions=(1536,),
+        expected_is_unit_vector=False,
+    )

test/integration/embedders/test_huggingface.py

@@ -0,0 +1,26 @@
+import json
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from unstructured_ingest.embed.huggingface import (
+    HuggingFaceEmbeddingConfig,
+    HuggingFaceEmbeddingEncoder,
+)
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+
+def test_huggingface_embedder(embedder_file: Path):
+    embedder_config = EmbedderConfig(embedding_provider="huggingface")
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+def test_raw_hugginface_embedder(embedder_file: Path):
+    embedder = HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig())
+    validate_raw_embedder(
+        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(384,)
+    )

test/integration/embedders/test_mixedbread.py

@@ -0,0 +1,47 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.mixedbreadai import (
+    MixedbreadAIEmbeddingConfig,
+    MixedbreadAIEmbeddingEncoder,
+)
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "MXBAI_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+@requires_env(API_KEY)
+def test_mixedbread_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder_config = EmbedderConfig(embedding_provider="mixedbread-ai", embedding_api_key=api_key)
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY)
+def test_raw_mixedbread_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder = MixedbreadAIEmbeddingEncoder(
+        config=MixedbreadAIEmbeddingConfig(
+            api_key=api_key,
+        )
+    )
+    validate_raw_embedder(
+        embedder=embedder,
+        embedder_file=embedder_file,
+        expected_dimensions=(1024,),
+        expected_is_unit_vector=False,
+    )

test/integration/embedders/test_octoai.py

@@ -0,0 +1,41 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "OCTOAI_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+@requires_env(API_KEY)
+def test_octoai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder_config = EmbedderConfig(embedding_provider="octoai", embedding_api_key=api_key)
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY)
+def test_raw_octoai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder = OctoAIEmbeddingEncoder(
+        config=OctoAiEmbeddingConfig(
+            api_key=api_key,
+        )
+    )
+    validate_raw_embedder(
+        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1024,)
+    )

test/integration/embedders/test_openai.py

@@ -0,0 +1,41 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "OPENAI_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+@requires_env(API_KEY)
+def test_openai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder_config = EmbedderConfig(embedding_provider="openai", embedding_api_key=api_key)
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY)
+def test_raw_openai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder = OpenAIEmbeddingEncoder(
+        config=OpenAIEmbeddingConfig(
+            api_key=api_key,
+        )
+    )
+    validate_raw_embedder(
+        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1536,)
+    )

test/integration/embedders/test_vertexai.py

@@ -0,0 +1,41 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.vertexai import VertexAIEmbeddingConfig, VertexAIEmbeddingEncoder
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "VERTEXAI_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+@requires_env(API_KEY)
+def test_vertexai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder_config = EmbedderConfig(embedding_provider="vertexai", embedding_api_key=api_key)
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY)
+def test_raw_vertexai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder = VertexAIEmbeddingEncoder(
+        config=VertexAIEmbeddingConfig(
+            api_key=api_key,
+        )
+    )
+    validate_raw_embedder(
+        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(768,)
+    )