unstructured-ingest 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

test/integration/connectors/test_pinecone.py

@@ -0,0 +1,213 @@
+import json
+import os
+import re
+import time
+from pathlib import Path
+from typing import Generator
+from uuid import uuid4
+
+import pytest
+from pinecone import Pinecone, ServerlessSpec
+from pinecone.core.openapi.shared.exceptions import NotFoundException
+
+from test.integration.connectors.utils.constants import (
+    DESTINATION_TAG,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connectors.pinecone import (
+    CONNECTOR_TYPE,
+    PineconeAccessConfig,
+    PineconeConnectionConfig,
+    PineconeUploader,
+    PineconeUploaderConfig,
+    PineconeUploadStager,
+    PineconeUploadStagerConfig,
+)
+
+METADATA_BYTES_LIMIT = (
+    40960  # 40KB https://docs.pinecone.io/reference/quotas-and-limits#hard-limits
+)
+VECTOR_DIMENSION = 384
+SPEC = {"serverless": {"cloud": "aws", "region": "us-east-1"}}
+ALLOWED_METADATA_FIELD = "text"
+API_KEY = "PINECONE_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+def wait_for_delete(client: Pinecone, index_name: str, timeout=60, interval=1) -> None:
+    start = time.time()
+    while True and time.time() - start < timeout:
+        try:
+            description = client.describe_index(name=index_name)
+            logger.info(f"current index status: {description}")
+        except NotFoundException:
+            return
+        time.sleep(interval)
+
+    raise TimeoutError("time out waiting for index to delete")
+
+
+def wait_for_ready(client: Pinecone, index_name: str, timeout=60, interval=1) -> None:
+    def is_ready_status():
+        description = client.describe_index(name=index_name)
+        status = description["status"]
+        return status["ready"]
+
+    start = time.time()
+    is_ready = is_ready_status()
+    while not is_ready and time.time() - start < timeout:
+        time.sleep(interval)
+        is_ready = is_ready_status()
+    if not is_ready:
+        raise TimeoutError("time out waiting for index to be ready")
+
+
+@pytest.fixture
+def pinecone_index() -> Generator[str, None, None]:
+    pinecone = Pinecone(api_key=get_api_key())
+    random_id = str(uuid4()).split("-")[0]
+    index_name = f"ingest-test-{random_id}"
+    assert len(index_name) < 45
+    logger.info(f"Creating index: {index_name}")
+    try:
+        pinecone.create_index(
+            name=index_name,
+            dimension=384,
+            metric="cosine",
+            spec=ServerlessSpec(
+                cloud="aws",
+                region="us-east-1",
+            ),
+            deletion_protection="disabled",
+        )
+        wait_for_ready(client=pinecone, index_name=index_name)
+        yield index_name
+    except Exception as e:
+        logger.error(f"failed to create index {index_name}: {e}")
+    finally:
+        try:
+            logger.info(f"deleting index: {index_name}")
+            pinecone.delete_index(name=index_name)
+            wait_for_delete(client=pinecone, index_name=index_name)
+        except NotFoundException:
+            return
+
+
+def validate_pinecone_index(
+    index_name: str, expected_num_of_vectors: int, retries=30, interval=1
+) -> None:
+    # Because there's a delay for the index to catch up to the recent writes, add in a retry
+    pinecone = Pinecone(api_key=get_api_key())
+    index = pinecone.Index(name=index_name)
+    vector_count = -1
+    for i in range(retries):
+        index_stats = index.describe_index_stats()
+        vector_count = index_stats["total_vector_count"]
+        if vector_count == expected_num_of_vectors:
+            logger.info(f"expected {expected_num_of_vectors} == vector count {vector_count}")
+            break
+        logger.info(
+            f"retry attempt {i}: expected {expected_num_of_vectors} != vector count {vector_count}"
+        )
+        time.sleep(interval)
+    assert vector_count == expected_num_of_vectors
+
+
+@requires_env(API_KEY)
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+async def test_pinecone_destination(pinecone_index: str, upload_file: Path, temp_dir: Path):
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="pinecone_mock_id",
+    )
+    connection_config = PineconeConnectionConfig(
+        index_name=pinecone_index,
+        access_config=PineconeAccessConfig(api_key=get_api_key()),
+    )
+    stager_config = PineconeUploadStagerConfig()
+    stager = PineconeUploadStager(upload_stager_config=stager_config)
+    new_upload_file = stager.run(
+        elements_filepath=upload_file,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+        file_data=file_data,
+    )
+
+    upload_config = PineconeUploaderConfig()
+    uploader = PineconeUploader(connection_config=connection_config, upload_config=upload_config)
+    uploader.precheck()
+
+    if uploader.is_async():
+        await uploader.run_async(path=new_upload_file, file_data=file_data)
+    else:
+        uploader.run(path=new_upload_file, file_data=file_data)
+    with new_upload_file.open() as f:
+        staged_content = json.load(f)
+    expected_num_of_vectors = len(staged_content)
+    logger.info("validating first upload")
+    validate_pinecone_index(
+        index_name=pinecone_index, expected_num_of_vectors=expected_num_of_vectors
+    )
+
+    # Rerun uploader and make sure no duplicates exist
+    if uploader.is_async():
+        await uploader.run_async(path=new_upload_file, file_data=file_data)
+    else:
+        uploader.run(path=new_upload_file, file_data=file_data)
+    logger.info("validating second upload")
+    validate_pinecone_index(
+        index_name=pinecone_index, expected_num_of_vectors=expected_num_of_vectors
+    )
+
+
+@requires_env(API_KEY)
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+def test_large_metadata(pinecone_index: str, tmp_path: Path, upload_file: Path):
+    stager = PineconeUploadStager()
+    uploader = PineconeUploader(
+        connection_config=PineconeConnectionConfig(
+            access_config=PineconeAccessConfig(api_key=get_api_key()),
+            index_name=pinecone_index,
+        ),
+        upload_config=PineconeUploaderConfig(),
+    )
+    large_metadata_upload_file = tmp_path / "mock-upload-file.pdf.json"
+    large_metadata = {ALLOWED_METADATA_FIELD: "0" * 2 * METADATA_BYTES_LIMIT}
+
+    with open(upload_file) as file:
+        elements = json.load(file)
+
+    with open(large_metadata_upload_file, "w") as file:
+        mock_element = elements[0]
+        mock_element["metadata"] = large_metadata
+        json.dump([mock_element], file)
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(
+            fullpath=large_metadata_upload_file.name, filename=large_metadata_upload_file.name
+        ),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+    staged_file = stager.run(
+        file_data, large_metadata_upload_file, tmp_path, large_metadata_upload_file.name
+    )
+    try:
+        uploader.run(staged_file, file_data)
+    except DestinationConnectionError as e:
+        error_line = r"Metadata size is \d+ bytes, which exceeds the limit of \d+ bytes per vector"
+        if re.search(re.compile(error_line), str(e)) is None:
+            raise e
+        raise pytest.fail("Upload request failed due to metadata exceeding limits.")
+
+    validate_pinecone_index(pinecone_index, 1, interval=5)

test/integration/connectors/test_s3.py

@@ -71,6 +71,29 @@ async def test_s3_source(anon_connection_config: S3ConnectionConfig):
         )
 
 
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+async def test_s3_source_special_char(anon_connection_config: S3ConnectionConfig):
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-dev-tech-fixtures/special-characters/")
+    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-specialchar",
+                predownload_file_data_check=validate_predownload_file_data,
+                postdownload_file_data_check=validate_postdownload_file_data,
+                expected_num_files=1,
+            ),
+        )
+
+
 @pytest.mark.asyncio
 @pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
 async def test_s3_source_no_access(anon_connection_config: S3ConnectionConfig):

test/integration/connectors/utils/docker.py

@@ -1,9 +1,43 @@
 import time
 from contextlib import contextmanager
-from typing import Optional
+from typing import Optional, Union
 
 import docker
 from docker.models.containers import Container
+from pydantic import BaseModel, Field, field_serializer
+
+
+class HealthCheck(BaseModel):
+    test: Union[str, list[str]]
+    interval: int = Field(
+        gt=0, default=30, description="The time to wait between checks in seconds."
+    )
+    timeout: int = Field(
+        gt=0, default=30, description="The time to wait before considering the check to have hung."
+    )
+    retries: int = Field(
+        gt=0,
+        default=3,
+        description="The number of consecutive failures needed "
+        "to consider a container as unhealthy.",
+    )
+    start_period: int = Field(
+        gt=0,
+        default=0,
+        description="Start period for the container to initialize before starting health-retries countdown in seconds.",  # noqa: E501
+    )
+
+    @field_serializer("interval")
+    def serialize_interval(self, interval: int) -> int:
+        return int(interval * 10e8)
+
+    @field_serializer("timeout")
+    def serialize_timeout(self, timeout: int) -> int:
+        return int(timeout * 10e8)
+
+    @field_serializer("start_period")
+    def serialize_start_period(self, start_period: int) -> int:
+        return int(start_period * 10e8)
 
 
 def get_container(
@@ -12,7 +46,7 @@ def get_container(
     ports: dict,
     environment: Optional[dict] = None,
     volumes: Optional[dict] = None,
-    healthcheck: Optional[dict] = None,
+    healthcheck: Optional[HealthCheck] = None,
 ) -> Container:
     run_kwargs = {
         "image": image,
@@ -24,25 +58,49 @@ def get_container(
     if volumes:
         run_kwargs["volumes"] = volumes
     if healthcheck:
-        run_kwargs["healthcheck"] = healthcheck
+        run_kwargs["healthcheck"] = healthcheck.model_dump()
     container: Container = docker_client.containers.run(**run_kwargs)
     return container
 
 
-def has_healthcheck(container: Container) -> bool:
-    return container.attrs.get("Config", {}).get("Healthcheck", None) is not None
+def get_healthcheck(container: Container) -> Optional[HealthCheck]:
+    healthcheck_config = container.attrs.get("Config", {}).get("Healthcheck", None)
+    if not healthcheck_config:
+        return None
+    healthcheck_data = {
+        "test": healthcheck_config["Test"],
+    }
+    if interval := healthcheck_config.get("Interval"):
+        healthcheck_data["interval"] = interval / 10e8
+    if start_period := healthcheck_config.get("StartPeriod"):
+        healthcheck_data["start_period"] = start_period / 10e8
+    if retries := healthcheck_config.get("Retries"):
+        healthcheck_data["retries"] = retries
+    return HealthCheck.model_validate(healthcheck_data)
 
 
-def healthcheck_wait(container: Container, timeout: int = 10) -> None:
+def healthcheck_wait(
+    container: Container, retries: int = 30, interval: int = 1, start_period: Optional[int] = None
+) -> None:
+    if start_period:
+        time.sleep(start_period)
     health = container.health
-    start = time.time()
-    while health != "healthy" and time.time() - start < timeout:
-        time.sleep(1)
+    tries = 0
+    while health != "healthy" and tries < retries:
+        tries += 1
+        logs = container.attrs.get("State", {}).get("Health", {}).get("Log")
+        latest_log = logs[-1] if logs else None
+        print(
+            f"attempt {tries} - waiting for docker container "
+            f"to be healthy: {health} latest log: {latest_log}"
+        )
+        time.sleep(interval)
         container.reload()
         health = container.health
     if health != "healthy":
-        health_dict = container.attrs.get("State", {}).get("Health", {})
-        raise TimeoutError(f"Docker container never came up healthy: {health_dict}")
+        logs = container.attrs.get("State", {}).get("Health", {}).get("Log")
+        latest_log = logs[-1] if logs else None
+        raise TimeoutError(f"Docker container never came up healthy: {latest_log}")
 
 
 @contextmanager
@@ -51,11 +109,13 @@ def container_context(
     ports: dict,
     environment: Optional[dict] = None,
     volumes: Optional[dict] = None,
-    healthcheck: Optional[dict] = None,
-    healthcheck_timeout: int = 10,
+    healthcheck: Optional[HealthCheck] = None,
+    healthcheck_retries: int = 30,
     docker_client: Optional[docker.DockerClient] = None,
 ):
     docker_client = docker_client or docker.from_env()
+    print(f"pulling image {image}")
+    docker_client.images.pull(image)
     container: Optional[Container] = None
     try:
         container = get_container(
@@ -66,8 +126,14 @@ def container_context(
             volumes=volumes,
             healthcheck=healthcheck,
         )
-        if has_healthcheck(container):
-            healthcheck_wait(container=container, timeout=healthcheck_timeout)
+        if healthcheck_data := get_healthcheck(container):
+            # Mirror whatever healthcheck config set on container
+            healthcheck_wait(
+                container=container,
+                retries=healthcheck_retries,
+                start_period=healthcheck_data.start_period,
+                interval=healthcheck_data.interval,
+            )
         yield container
     except AssertionError as e:
         if container:

test/integration/connectors/utils/validation.py

@@ -240,6 +240,10 @@ def update_fixtures(
     # Rewrite the current file data
     if save_filedata:
         file_data_output_path = output_dir / "file_data"
+        print(
+            f"Writing {len(all_file_data)} file data to "
+            f"saved fixture location {file_data_output_path}"
+        )
         file_data_output_path.mkdir(parents=True, exist_ok=True)
         for file_data in all_file_data:
             file_data_path = file_data_output_path / f"{file_data.identifier}.json"
@@ -256,6 +260,10 @@ def update_fixtures(
     # If applicable, save raw downloads
     if save_downloads:
         raw_download_output_path = output_dir / "downloads"
+        print(
+            f"Writing {len(download_files)} downloaded files to "
+            f"saved fixture location {raw_download_output_path}"
+        )
         shutil.copytree(download_dir, raw_download_output_path)
 
 
@@ -328,6 +336,7 @@ async def source_connector_validation(
             postdownload_file_data = replace(resp["file_data"])
             all_postdownload_file_data.append(postdownload_file_data)
     if not overwrite_fixtures:
+        print("Running validation")
         run_all_validations(
             configs=configs,
             predownload_file_data=all_predownload_file_data,
@@ -336,6 +345,7 @@ async def source_connector_validation(
             test_output_dir=test_output_dir,
         )
     else:
+        print("Running fixtures update")
         update_fixtures(
             output_dir=test_output_dir,
             download_dir=download_dir,

test/integration/connectors/weaviate/conftest.py

@@ -0,0 +1,15 @@
+import json
+from pathlib import Path
+
+import pytest
+
+
+@pytest.fixture
+def collections_schema_config() -> dict:
+    int_test_dir = Path(__file__).parent
+    assets_dir = int_test_dir / "assets"
+    config_file = assets_dir / "elements.json"
+    assert config_file.exists()
+    assert config_file.is_file()
+    with config_file.open() as config_data:
+        return json.load(config_data)

test/integration/connectors/weaviate/test_local.py

@@ -0,0 +1,131 @@
+import json
+import time
+from pathlib import Path
+
+import pytest
+import requests
+import weaviate
+from weaviate.client import WeaviateClient
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from test.integration.connectors.utils.docker import container_context
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.weaviate.local import (
+    CONNECTOR_TYPE,
+    LocalWeaviateConnectionConfig,
+    LocalWeaviateUploader,
+    LocalWeaviateUploaderConfig,
+    LocalWeaviateUploadStager,
+)
+
+COLLECTION_NAME = "elements"
+
+
+def wait_for_container(timeout: int = 10, interval: int = 1) -> None:
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            requests.get("http://localhost:8080/v1/.well-known/read")
+            return
+        except Exception as e:
+            print(f"Failed to validate container healthy, sleeping for {interval} seconds: {e}")
+            time.sleep(interval)
+    raise TimeoutError("Docker container never came up healthy")
+
+
+@pytest.fixture
+def collection(collections_schema_config: dict) -> str:
+    with container_context(
+        image="semitechnologies/weaviate:1.27.3",
+        ports={8080: 8080, 50051: 50051},
+    ):
+        wait_for_container()
+        with weaviate.connect_to_local() as weaviate_client:
+            weaviate_client.collections.create_from_dict(config=collections_schema_config)
+        yield COLLECTION_NAME
+
+
+def get_count(client: WeaviateClient) -> int:
+    collection = client.collections.get(COLLECTION_NAME)
+    resp = collection.aggregate.over_all(total_count=True)
+    return resp.total_count
+
+
+def validate_count(expected_count: int, retries: int = 10, interval: int = 1) -> None:
+    with weaviate.connect_to_local() as weaviate_client:
+        current_count = get_count(client=weaviate_client)
+        retry_count = 0
+        while current_count != expected_count and retry_count < retries:
+            retry_count += 1
+            time.sleep(interval)
+            current_count = get_count(client=weaviate_client)
+        assert current_count == expected_count, (
+            f"Expected count ({expected_count}) doesn't match how "
+            f"much came back from collection: {current_count}"
+        )
+
+
+def run_uploader_and_validate(
+    uploader: LocalWeaviateUploader, path: Path, file_data: FileData, expected_count: int
+):
+    uploader.precheck()
+    uploader.run(path=path, file_data=file_data)
+    validate_count(expected_count=expected_count)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+def test_weaviate_local_destination(upload_file: Path, collection: str, tmp_path: Path):
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    stager = LocalWeaviateUploadStager()
+
+    staged_filepath = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+    dynamic_uploader = LocalWeaviateUploader(
+        upload_config=LocalWeaviateUploaderConfig(
+            collection=COLLECTION_NAME,
+        ),
+        connection_config=LocalWeaviateConnectionConfig(),
+    )
+    fixed_size_uploader = LocalWeaviateUploader(
+        upload_config=LocalWeaviateUploaderConfig(
+            collection=COLLECTION_NAME, batch_size=10, dynamic_batch=False
+        ),
+        connection_config=LocalWeaviateConnectionConfig(),
+    )
+    rate_limited_uploader = LocalWeaviateUploader(
+        upload_config=LocalWeaviateUploaderConfig(
+            collection=COLLECTION_NAME, requests_per_minute=50, dynamic_batch=False
+        ),
+        connection_config=LocalWeaviateConnectionConfig(),
+    )
+    with staged_filepath.open() as f:
+        staged_elements = json.load(f)
+    expected_count = len(staged_elements)
+
+    run_uploader_and_validate(
+        uploader=dynamic_uploader,
+        path=staged_filepath,
+        file_data=file_data,
+        expected_count=expected_count,
+    )
+    run_uploader_and_validate(
+        uploader=fixed_size_uploader,
+        path=staged_filepath,
+        file_data=file_data,
+        expected_count=expected_count,
+    )
+    run_uploader_and_validate(
+        uploader=rate_limited_uploader,
+        path=staged_filepath,
+        file_data=file_data,
+        expected_count=expected_count,
+    )

test/unit/v2/chunkers/test_chunkers.py

@@ -0,0 +1,49 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.v2.processes.chunker import Chunker, ChunkerConfig
+
+fake = faker.Faker()
+
+
+def generate_chunker_config_params() -> dict:
+    params = {}
+    random_val = random.random()
+    if random_val < 0.5:
+        params["chunking_strategy"] = fake.word() if random.random() < 0.5 else None
+        params["chunk_combine_text_under_n_chars"] = (
+            fake.random_int() if random.random() < 0.5 else None
+        )
+        params["chunk_include_orig_elements"] = fake.boolean() if random.random() < 0.5 else None
+        params["chunk_max_characters"] = fake.random_int()
+        params["chunk_multipage_sections"] = fake.boolean()
+        params["chunk_new_after_n_chars"] = fake.random_int() if random.random() < 0.5 else None
+        params["chunk_overlap"] = fake.random_int() if random.random() < 0.5 else None
+        params["chunk_overlap_all"] = fake.boolean() if random.random() < 0.5 else None
+    if random_val < 0.5:
+        params["chunk_by_api"] = True
+        params["chunking_endpoint"] = fake.url()
+        params["chunk_api_key"] = fake.password()
+    else:
+        params["chunk_by_api"] = False
+
+    return params
+
+
+@pytest.mark.parametrize(
+    "partition_config_params", [generate_chunker_config_params() for i in range(10)]
+)
+def test_chunker_config(partition_config_params: dict):
+    chunker_config = ChunkerConfig.model_validate(partition_config_params)
+    assert chunker_config
+
+
+@pytest.mark.parametrize(
+    "partition_config_params", [generate_chunker_config_params() for i in range(10)]
+)
+def test_chunker(partition_config_params: dict):
+    chunker_config = ChunkerConfig.model_validate(partition_config_params)
+    chunker = Chunker(config=chunker_config)
+    assert chunker

test/unit/v2/embedders/test_bedrock.py

@@ -0,0 +1,36 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.embed.bedrock import BedrockEmbeddingConfig, BedrockEmbeddingEncoder
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {
+        "aws_access_key_id": fake.password(),
+        "aws_secret_access_key": fake.password(),
+        "region_name": fake.city(),
+    }
+    if random.random() < 0.5:
+        params["embed_model_name"] = fake.word()
+    return params
+
+
+@pytest.mark.parametrize(
+    "embedder_config_params", [generate_embedder_config_params() for i in range(10)]
+)
+def test_embedder_config(embedder_config_params: dict):
+    embedder_config = BedrockEmbeddingConfig.model_validate(embedder_config_params)
+    assert embedder_config
+
+
+@pytest.mark.parametrize(
+    "embedder_config_params", [generate_embedder_config_params() for i in range(10)]
+)
+def test_embedder(embedder_config_params: dict):
+    embedder_config = BedrockEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = BedrockEmbeddingEncoder(config=embedder_config)
+    assert embedder