unstructured-ingest 0.3.13__py3-none-any.whl

test/integration/connectors/test_pinecone.py

@@ -0,0 +1,288 @@
+import json
+import math
+import os
+import re
+import time
+from pathlib import Path
+from typing import Generator
+from uuid import uuid4
+
+import pytest
+from _pytest.fixtures import TopRequest
+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.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+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,
+    MAX_QUERY_RESULTS,
+    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, (
+        f"vector count from index ({vector_count}) doesn't "
+        f"match expected number: {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()
+
+    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
+    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.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@pytest.mark.skip(reason="TODO: get this to work")
+async def test_pinecone_destination_large_index(
+    pinecone_index: str, upload_file: Path, temp_dir: Path
+):
+    new_file = temp_dir / "large_file.json"
+    with upload_file.open() as f:
+        upload_content = json.load(f)
+
+    min_entries = math.ceil((MAX_QUERY_RESULTS * 2) / len(upload_content))
+    new_content = (upload_content * min_entries)[: (2 * MAX_QUERY_RESULTS)]
+    print(f"Creating large index content with {len(new_content)} records")
+    with new_file.open("w") as f:
+        json.dump(new_content, f)
+
+    expected_num_of_vectors = len(new_content)
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=new_file.name, filename=new_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=new_file,
+        output_dir=temp_dir,
+        output_filename=new_file.name,
+        file_data=file_data,
+    )
+
+    upload_config = PineconeUploaderConfig()
+    uploader = PineconeUploader(connection_config=connection_config, upload_config=upload_config)
+    uploader.precheck()
+
+    uploader.run(path=new_upload_file, file_data=file_data)
+    validate_pinecone_index(
+        index_name=pinecone_index, expected_num_of_vectors=expected_num_of_vectors
+    )
+    # Rerun uploader and make sure no duplicates exist
+    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(
+        elements_filepath=large_metadata_upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=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)
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_pinecone_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = PineconeUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/test_qdrant.py

@@ -0,0 +1,215 @@
+import json
+import os
+import uuid
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import AsyncGenerator
+
+import pytest
+from _pytest.fixtures import TopRequest
+from qdrant_client import AsyncQdrantClient
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from test.integration.connectors.utils.docker import container_context
+from test.integration.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.qdrant.cloud import (
+    CloudQdrantAccessConfig,
+    CloudQdrantConnectionConfig,
+    CloudQdrantUploader,
+    CloudQdrantUploaderConfig,
+    CloudQdrantUploadStager,
+    CloudQdrantUploadStagerConfig,
+)
+from unstructured_ingest.v2.processes.connectors.qdrant.local import (
+    CONNECTOR_TYPE as LOCAL_CONNECTOR_TYPE,
+)
+from unstructured_ingest.v2.processes.connectors.qdrant.local import (
+    LocalQdrantConnectionConfig,
+    LocalQdrantUploader,
+    LocalQdrantUploaderConfig,
+    LocalQdrantUploadStager,
+    LocalQdrantUploadStagerConfig,
+)
+from unstructured_ingest.v2.processes.connectors.qdrant.server import (
+    CONNECTOR_TYPE as SERVER_CONNECTOR_TYPE,
+)
+from unstructured_ingest.v2.processes.connectors.qdrant.server import (
+    ServerQdrantConnectionConfig,
+    ServerQdrantUploader,
+    ServerQdrantUploaderConfig,
+    ServerQdrantUploadStager,
+    ServerQdrantUploadStagerConfig,
+)
+
+COLLECTION_NAME = f"test-coll-{uuid.uuid4().hex[:12]}"
+VECTORS_CONFIG = {"size": 384, "distance": "Cosine"}
+
+
+@asynccontextmanager
+async def qdrant_client(client_params: dict) -> AsyncGenerator[AsyncQdrantClient, None]:
+    client = AsyncQdrantClient(**client_params)
+    try:
+        yield client
+    finally:
+        await client.close()
+
+
+async def validate_upload(client: AsyncQdrantClient, upload_file: Path):
+    with upload_file.open() as upload_fp:
+        elements = json.load(upload_fp)
+    expected_point_count = len(elements)
+    first_element = elements[0]
+    expected_text = first_element["text"]
+    embeddings = first_element["embeddings"]
+    collection = await client.get_collection(COLLECTION_NAME)
+    assert collection.points_count == expected_point_count
+
+    response = await client.query_points(COLLECTION_NAME, query=embeddings, limit=1)
+    assert response.points[0].payload is not None
+    assert response.points[0].payload["text"] == expected_text
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(LOCAL_CONNECTOR_TYPE, DESTINATION_TAG, "qdrant")
+async def test_qdrant_destination_local(upload_file: Path, tmp_path: Path):
+    connection_kwargs = {"path": str(tmp_path / "qdrant")}
+    async with qdrant_client(connection_kwargs) as client:
+        await client.create_collection(COLLECTION_NAME, vectors_config=VECTORS_CONFIG)
+    AsyncQdrantClient(**connection_kwargs)
+    stager = LocalQdrantUploadStager(
+        upload_stager_config=LocalQdrantUploadStagerConfig(),
+    )
+    uploader = LocalQdrantUploader(
+        connection_config=LocalQdrantConnectionConfig(**connection_kwargs),
+        upload_config=LocalQdrantUploaderConfig(collection_name=COLLECTION_NAME),
+    )
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=LOCAL_CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+
+    staged_upload_file = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+
+    if uploader.is_async():
+        await uploader.run_async(path=staged_upload_file, file_data=file_data)
+    else:
+        uploader.run(path=upload_file, file_data=file_data)
+    async with qdrant_client(connection_kwargs) as client:
+        await validate_upload(client=client, upload_file=upload_file)
+
+
+@pytest.fixture
+def docker_context():
+    with container_context(image="qdrant/qdrant:latest", ports={"6333": "6333"}) as container:
+        yield container
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(SERVER_CONNECTOR_TYPE, DESTINATION_TAG, "qdrant")
+async def test_qdrant_destination_server(upload_file: Path, tmp_path: Path, docker_context):
+    connection_kwargs = {"location": "http://localhost:6333"}
+    async with qdrant_client(connection_kwargs) as client:
+        await client.create_collection(COLLECTION_NAME, vectors_config=VECTORS_CONFIG)
+    AsyncQdrantClient(**connection_kwargs)
+    stager = ServerQdrantUploadStager(
+        upload_stager_config=ServerQdrantUploadStagerConfig(),
+    )
+    uploader = ServerQdrantUploader(
+        connection_config=ServerQdrantConnectionConfig(**connection_kwargs),
+        upload_config=ServerQdrantUploaderConfig(collection_name=COLLECTION_NAME),
+    )
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=SERVER_CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+
+    staged_upload_file = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+    uploader.precheck()
+    if uploader.is_async():
+        await uploader.run_async(path=staged_upload_file, file_data=file_data)
+    else:
+        uploader.run(path=upload_file, file_data=file_data)
+    async with qdrant_client(connection_kwargs) as client:
+        await validate_upload(client=client, upload_file=upload_file)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(SERVER_CONNECTOR_TYPE, DESTINATION_TAG, "qdrant")
+@requires_env("QDRANT_API_KEY", "QDRANT_SERVER_URL")
+async def test_qdrant_destination_cloud(upload_file: Path, tmp_path: Path):
+    server_url = os.environ["QDRANT_SERVER_URL"]
+    api_key = os.environ["QDRANT_API_KEY"]
+    connection_kwargs = {"location": server_url, "api_key": api_key}
+    async with qdrant_client(connection_kwargs) as client:
+        await client.create_collection(COLLECTION_NAME, vectors_config=VECTORS_CONFIG)
+    AsyncQdrantClient(**connection_kwargs)
+
+    stager = CloudQdrantUploadStager(
+        upload_stager_config=CloudQdrantUploadStagerConfig(),
+    )
+    uploader = CloudQdrantUploader(
+        connection_config=CloudQdrantConnectionConfig(
+            url=server_url,
+            access_config=CloudQdrantAccessConfig(
+                api_key=api_key,
+            ),
+        ),
+        upload_config=CloudQdrantUploaderConfig(collection_name=COLLECTION_NAME),
+    )
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=SERVER_CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+
+    staged_upload_file = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+    uploader.precheck()
+    if uploader.is_async():
+        await uploader.run_async(path=staged_upload_file, file_data=file_data)
+    else:
+        uploader.run(path=staged_upload_file, file_data=file_data)
+    async with qdrant_client(connection_kwargs) as client:
+        await validate_upload(client=client, upload_file=upload_file)
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_qdrant_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = LocalQdrantUploadStager(
+        upload_stager_config=LocalQdrantUploadStagerConfig(),
+    )
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=LOCAL_CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/test_redis.py

@@ -0,0 +1,119 @@
+import asyncio
+import json
+import os
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pytest
+from redis import exceptions as redis_exceptions
+from redis.asyncio import Redis, from_url
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.redisdb import (
+    CONNECTOR_TYPE as REDIS_CONNECTOR_TYPE,
+)
+from unstructured_ingest.v2.processes.connectors.redisdb import (
+    RedisAccessConfig,
+    RedisConnectionConfig,
+    RedisUploader,
+    RedisUploaderConfig,
+)
+
+
+async def delete_record(client: Redis, element_id: str) -> None:
+    await client.delete(element_id)
+
+
+async def validate_upload(client: Redis, first_element: dict):
+    element_id = first_element["element_id"]
+    expected_text = first_element["text"]
+    expected_embeddings = first_element["embeddings"]
+    async with client.pipeline(transaction=True) as pipe:
+        try:
+            response = await pipe.json().get(element_id, "$").execute()
+            response = response[0][0]
+        except redis_exceptions.ResponseError:
+            response = await pipe.get(element_id).execute()
+            response = json.loads(response[0])
+
+    embedding_similarity = np.linalg.norm(
+        np.array(response["embeddings"]) - np.array(expected_embeddings)
+    )
+
+    assert response is not None
+    assert response["element_id"] == element_id
+    assert response["text"] == expected_text
+    assert embedding_similarity < 1e-10
+
+
+async def redis_destination_test(
+    upload_file: Path,
+    tmp_path: Path,
+    connection_kwargs: dict,
+    uri: Optional[str] = None,
+    password: Optional[str] = None,
+):
+    uploader = RedisUploader(
+        connection_config=RedisConnectionConfig(
+            **connection_kwargs, access_config=RedisAccessConfig(uri=uri, password=password)
+        ),
+        upload_config=RedisUploaderConfig(batch_size=10),
+    )
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=REDIS_CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+    with upload_file.open() as upload_fp:
+        elements = json.load(upload_fp)
+    first_element = elements[0]
+
+    try:
+        if uploader.is_async():
+            await uploader.run_data_async(data=elements, file_data=file_data)
+
+        if uri:
+            async with from_url(uri) as client:
+                await validate_upload(client=client, first_element=first_element)
+        else:
+            async with Redis(**connection_kwargs, password=password) as client:
+                await validate_upload(client=client, first_element=first_element)
+    except Exception as e:
+        raise e
+    finally:
+        if uri:
+            async with from_url(uri) as client:
+                tasks = [delete_record(client, element["element_id"]) for element in elements]
+                await asyncio.gather(*tasks)
+        else:
+            async with Redis(**connection_kwargs, password=password) as client:
+                tasks = [delete_record(client, element["element_id"]) for element in elements]
+                await asyncio.gather(*tasks)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(REDIS_CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("AZURE_REDIS_INGEST_TEST_PASSWORD")
+async def test_redis_destination_azure_with_password(upload_file: Path, tmp_path: Path):
+    connection_kwargs = {
+        "host": "utic-dashboard-dev.redis.cache.windows.net",
+        "port": 6380,
+        "db": 0,
+        "ssl": True,
+    }
+    redis_pw = os.environ["AZURE_REDIS_INGEST_TEST_PASSWORD"]
+    await redis_destination_test(upload_file, tmp_path, connection_kwargs, password=redis_pw)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(REDIS_CONNECTOR_TYPE, DESTINATION_TAG, "redis")
+@requires_env("AZURE_REDIS_INGEST_TEST_PASSWORD")
+async def test_redis_destination_azure_with_uri(upload_file: Path, tmp_path: Path):
+    connection_kwargs = {}
+    redis_pw = os.environ["AZURE_REDIS_INGEST_TEST_PASSWORD"]
+    uri = f"rediss://:{redis_pw}@utic-dashboard-dev.redis.cache.windows.net:6380/0"
+    await redis_destination_test(upload_file, tmp_path, connection_kwargs, uri=uri)