unstructured-ingest 0.7.1__py3-none-any.whl → 1.0.1__py3-none-any.whl

test/integration/connectors/test_pinecone.py

@@ -1,387 +0,0 @@
-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, VECTOR_DB_TAG
-from test.integration.connectors.utils.validation.destination import (
-    StagerValidationConfigs,
-    stager_validation,
-)
-from test.integration.utils import requires_env
-from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.error import DestinationConnectionError
-from unstructured_ingest.logger import logger
-from unstructured_ingest.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,
-    namespace: str = "default",
-) -> 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, namespace=namespace)
-    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, VECTOR_DB_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, VECTOR_DB_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.asyncio
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-async def test_pinecone_destination_namespace(
-    pinecone_index: str, upload_file: Path, temp_dir: Path
-):
-    """
-    tests namespace functionality of destination connector.
-    """
-
-    # creates a file data structure.
-    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,
-    )
-
-    # here add namespace defintion
-    upload_config = PineconeUploaderConfig()
-    namespace_test_name = "user-1"
-    upload_config.namespace = namespace_test_name
-    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,
-        namespace=namespace_test_name,
-    )
-
-    # 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,
-        namespace=namespace_test_name,
-    )
-
-
-@requires_env(API_KEY)
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_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.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-@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,
-    )
-
-
-@requires_env(API_KEY)
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-def test_pinecone_create_destination(pinecone_index):
-    uploader = PineconeUploader(
-        connection_config=PineconeConnectionConfig(
-            access_config=PineconeAccessConfig(api_key=get_api_key())
-        ),
-        upload_config=PineconeUploaderConfig(),
-    )
-
-    random_id = str(uuid4()).split("-")[0]
-
-    index_name = f"test-create-destination-{random_id}"
-
-    assert not uploader.index_exists(index_name=index_name)
-
-    try:
-        uploader.create_destination(destination_name=index_name, vector_length=1536)
-    except Exception as e:
-        error_body = getattr(e, "body", None)
-        raise pytest.fail(f"failed to create destination: {e} {error_body}")
-
-    assert uploader.index_exists(index_name=index_name), "destination was not created successfully"
-
-    try:
-        pc = uploader.connection_config.get_client()
-        logger.info(f"deleting index for test create destination: {index_name}")
-        pc.delete_index(name=index_name)
-    except Exception as e:
-        raise pytest.fail(f"failed to cleanup / delete the destination: {e}")
-
-    assert not uploader.index_exists(index_name=index_name), "cleanup failed"

test/integration/connectors/test_qdrant.py

@@ -1,216 +0,0 @@
-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, VECTOR_DB_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.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.processes.connectors.qdrant.cloud import (
-    CloudQdrantAccessConfig,
-    CloudQdrantConnectionConfig,
-    CloudQdrantUploader,
-    CloudQdrantUploaderConfig,
-    CloudQdrantUploadStager,
-    CloudQdrantUploadStagerConfig,
-)
-from unstructured_ingest.processes.connectors.qdrant.local import (
-    CONNECTOR_TYPE as LOCAL_CONNECTOR_TYPE,
-)
-from unstructured_ingest.processes.connectors.qdrant.local import (
-    LocalQdrantConnectionConfig,
-    LocalQdrantUploader,
-    LocalQdrantUploaderConfig,
-    LocalQdrantUploadStager,
-    LocalQdrantUploadStagerConfig,
-)
-from unstructured_ingest.processes.connectors.qdrant.server import (
-    CONNECTOR_TYPE as SERVER_CONNECTOR_TYPE,
-)
-from unstructured_ingest.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", VECTOR_DB_TAG)
-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", VECTOR_DB_TAG)
-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", VECTOR_DB_TAG)
-@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.tags(SERVER_CONNECTOR_TYPE, DESTINATION_TAG, "qdrant", VECTOR_DB_TAG)
-@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,
-    )