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

test/integration/connectors/test_milvus.py

@@ -1,208 +0,0 @@
-import json
-import time
-from pathlib import Path
-
-import docker
-import pytest
-from _pytest.fixtures import TopRequest
-from pymilvus import (
-    CollectionSchema,
-    DataType,
-    FieldSchema,
-    MilvusClient,
-)
-from pymilvus.milvus_client import IndexParams
-
-from test.integration.connectors.utils.constants import (
-    DESTINATION_TAG,
-    VECTOR_DB_TAG,
-    env_setup_path,
-)
-from test.integration.connectors.utils.docker import healthcheck_wait
-from test.integration.connectors.utils.docker_compose import docker_compose_context
-from test.integration.connectors.utils.validation.destination import (
-    StagerValidationConfigs,
-    stager_validation,
-)
-from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.error import DestinationConnectionError
-from unstructured_ingest.processes.connectors.milvus import (
-    CONNECTOR_TYPE,
-    MilvusConnectionConfig,
-    MilvusUploader,
-    MilvusUploaderConfig,
-    MilvusUploadStager,
-)
-
-DB_NAME = "test_database"
-EXISTENT_COLLECTION_NAME = "test_collection"
-NONEXISTENT_COLLECTION_NAME = "nonexistent_collection"
-DB_URI = "http://localhost:19530"
-
-
-def get_schema() -> CollectionSchema:
-    id_field = FieldSchema(
-        name="id", dtype=DataType.INT64, description="primary field", is_primary=True, auto_id=True
-    )
-    embeddings_field = FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=384)
-    record_id_field = FieldSchema(name="record_id", dtype=DataType.VARCHAR, max_length=64)
-
-    schema = CollectionSchema(
-        enable_dynamic_field=True,
-        fields=[
-            id_field,
-            record_id_field,
-            embeddings_field,
-        ],
-    )
-
-    return schema
-
-
-def get_index_params() -> IndexParams:
-    index_params = IndexParams()
-    index_params.add_index(field_name="embeddings", index_type="AUTOINDEX", metric_type="COSINE")
-    index_params.add_index(field_name="record_id", index_type="Trie")
-    return index_params
-
-
-# NOTE: Precheck tests are read-only so they don't interfere with destination test,
-# using scope="module" we can limit number of times the docker-compose has to be run
-@pytest.fixture(scope="module")
-def collection():
-    docker_client = docker.from_env()
-    with docker_compose_context(docker_compose_path=env_setup_path / "milvus"):
-        milvus_container = docker_client.containers.get("milvus-standalone")
-        healthcheck_wait(container=milvus_container)
-        milvus_client = MilvusClient(uri=DB_URI)
-        try:
-            # Create the database
-            database_resp = milvus_client._get_connection().create_database(db_name=DB_NAME)
-            milvus_client.using_database(db_name=DB_NAME)
-
-            print(f"Created database {DB_NAME}: {database_resp}")
-
-            # Create the collection
-            schema = get_schema()
-            index_params = get_index_params()
-            collection_resp = milvus_client.create_collection(
-                collection_name=EXISTENT_COLLECTION_NAME, schema=schema, index_params=index_params
-            )
-            print(f"Created collection {EXISTENT_COLLECTION_NAME}: {collection_resp}")
-            yield EXISTENT_COLLECTION_NAME
-        finally:
-            milvus_client.close()
-
-
-def get_count(client: MilvusClient) -> int:
-    count_field = "count(*)"
-    resp = client.query(collection_name="test_collection", output_fields=[count_field])
-    return resp[0][count_field]
-
-
-def validate_count(
-    client: MilvusClient, expected_count: int, retries: int = 10, interval: int = 1
-) -> None:
-    current_count = get_count(client=client)
-    retry_count = 0
-    while current_count != expected_count and retry_count < retries:
-        time.sleep(interval)
-        current_count = get_count(client=client)
-        retry_count += 1
-    assert current_count == expected_count, (
-        f"Expected count ({expected_count}) doesn't match how "
-        f"much came back from collection: {current_count}"
-    )
-
-
-@pytest.mark.asyncio
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-async def test_milvus_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 = MilvusUploadStager()
-    uploader = MilvusUploader(
-        connection_config=MilvusConnectionConfig(uri=DB_URI),
-        upload_config=MilvusUploaderConfig(collection_name=collection, db_name=DB_NAME),
-    )
-    staged_filepath = stager.run(
-        elements_filepath=upload_file,
-        file_data=file_data,
-        output_dir=tmp_path,
-        output_filename=upload_file.name,
-    )
-    uploader.precheck()
-    uploader.run(path=staged_filepath, file_data=file_data)
-
-    # Run validation
-    with staged_filepath.open() as f:
-        staged_elements = json.load(f)
-    expected_count = len(staged_elements)
-    with uploader.get_client() as client:
-        validate_count(client=client, expected_count=expected_count)
-
-    # Rerun and make sure the same documents get updated
-    uploader.run(path=staged_filepath, file_data=file_data)
-    with uploader.get_client() as client:
-        validate_count(client=client, expected_count=expected_count)
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-def test_precheck_succeeds(collection: str):
-    uploader = MilvusUploader(
-        connection_config=MilvusConnectionConfig(uri=DB_URI),
-        upload_config=MilvusUploaderConfig(db_name=DB_NAME, collection_name=collection),
-    )
-    uploader.precheck()
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-def test_precheck_fails_on_nonexistent_collection(collection: str):
-    uploader = MilvusUploader(
-        connection_config=MilvusConnectionConfig(uri=DB_URI),
-        upload_config=MilvusUploaderConfig(
-            db_name=DB_NAME, collection_name=NONEXISTENT_COLLECTION_NAME
-        ),
-    )
-    with pytest.raises(
-        DestinationConnectionError,
-        match=f"Collection '{NONEXISTENT_COLLECTION_NAME}' does not exist",
-    ):
-        uploader.precheck()
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-def test_precheck_fails_on_nonexisting_db(collection: str):
-    uploader = MilvusUploader(
-        connection_config=MilvusConnectionConfig(uri=DB_URI),
-        upload_config=MilvusUploaderConfig(db_name="nonexisting_db", collection_name=collection),
-    )
-    with pytest.raises(
-        DestinationConnectionError,
-        match="database not found",
-    ):
-        uploader.precheck()
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
-@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
-def test_milvus_stager(
-    request: TopRequest,
-    upload_file_str: str,
-    tmp_path: Path,
-):
-    upload_file: Path = request.getfixturevalue(upload_file_str)
-    stager = MilvusUploadStager()
-    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_mongodb.py

@@ -1,335 +0,0 @@
-import json
-import os
-import time
-import uuid
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Generator
-
-import pytest
-from pydantic import BaseModel, SecretStr
-from pymongo.collection import Collection
-from pymongo.database import Database
-from pymongo.mongo_client import MongoClient
-from pymongo.operations import SearchIndexModel
-
-from test.integration.connectors.utils.constants import DESTINATION_TAG, NOSQL_TAG, SOURCE_TAG
-from test.integration.connectors.utils.validation.source import (
-    SourceValidationConfigs,
-    source_connector_validation,
-)
-from test.integration.utils import requires_env
-from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
-from unstructured_ingest.processes.connectors.mongodb import (
-    CONNECTOR_TYPE,
-    MongoDBAccessConfig,
-    MongoDBConnectionConfig,
-    MongoDBDownloader,
-    MongoDBDownloaderConfig,
-    MongoDBIndexer,
-    MongoDBIndexerConfig,
-    MongoDBUploader,
-    MongoDBUploaderConfig,
-)
-
-SOURCE_COLLECTION = "sample-mongodb-data"
-
-
-class EnvData(BaseModel):
-    uri: SecretStr
-    database: str
-
-
-def get_env_data() -> EnvData:
-    uri = os.getenv("MONGODB_URI")
-    assert uri
-    database = os.getenv("MONGODB_DATABASE")
-    assert database
-    return EnvData(uri=uri, database=database)
-
-
-@contextmanager
-def get_client() -> Generator[MongoClient, None, None]:
-    uri = get_env_data().uri.get_secret_value()
-    with MongoClient(uri) as client:
-        assert client.admin.command("ping")
-        yield client
-
-
-def wait_for_collection(
-    database: Database, collection_name: str, retries: int = 10, interval: int = 1
-):
-    collections = database.list_collection_names()
-    attempts = 0
-    while collection_name not in collections and attempts < retries:
-        attempts += 1
-        print(
-            "Waiting for collection {} to be recognized: {}".format(
-                collection_name, ", ".join(collections)
-            )
-        )
-        time.sleep(interval)
-        collections = database.list_collection_names()
-    if collection_name not in collection_name:
-        raise TimeoutError(f"Collection {collection_name} was not recognized")
-
-
-def get_search_index_status(collection: Collection, index_name: str) -> str:
-    search_indexes = collection.list_search_indexes(name=index_name)
-    search_index = list(search_indexes)[0]
-    return search_index["status"]
-
-
-def wait_for_search_index(
-    collection: Collection, index_name: str, retries: int = 60, interval: int = 1
-):
-    current_status = get_search_index_status(collection, index_name)
-    attempts = 0
-    while current_status != "READY" and attempts < retries:
-        attempts += 1
-        print(f"attempt {attempts}: waiting for search index to be READY: {current_status}")
-        time.sleep(interval)
-        current_status = get_search_index_status(collection, index_name)
-
-    if current_status != "READY":
-        raise TimeoutError("search index never detected as READY")
-
-
-@pytest.fixture
-def destination_collection() -> Collection:
-    env_data = get_env_data()
-    collection_name = f"utic-test-output-{uuid.uuid4()}"
-    with get_client() as client:
-        database = client[env_data.database]
-        print(f"creating collection in database {database}: {collection_name}")
-        collection = database.create_collection(name=collection_name)
-        search_index_name = "embeddings"
-        collection.create_search_index(
-            model=SearchIndexModel(
-                name=search_index_name,
-                definition={
-                    "mappings": {
-                        "dynamic": True,
-                        "fields": {
-                            "embeddings": [
-                                {"type": "knnVector", "dimensions": 384, "similarity": "euclidean"}
-                            ]
-                        },
-                    }
-                },
-            )
-        )
-        collection.create_index("record_id")
-        wait_for_collection(database=database, collection_name=collection_name)
-        wait_for_search_index(collection=collection, index_name=search_index_name)
-        try:
-            yield collection
-        finally:
-            print(f"deleting collection: {collection_name}")
-            collection.drop()
-
-
-def validate_collection_count(
-    collection: Collection, expected_records: int, retries: int = 10, interval: int = 1
-) -> None:
-    count = collection.count_documents(filter={})
-    attempt = 0
-    while count != expected_records and attempt < retries:
-        attempt += 1
-        print(f"attempt {attempt} to get count of collection {count} to match {expected_records}")
-        time.sleep(interval)
-        count = collection.count_documents(filter={})
-    assert (
-        count == expected_records
-    ), f"expected count ({expected_records}) does not match how many records were found: {count}"
-
-
-def validate_collection_vector(
-    collection: Collection, embedding: list[float], text: str, retries: int = 30, interval: int = 1
-) -> None:
-    pipeline = [
-        {
-            "$vectorSearch": {
-                "index": "embeddings",
-                "path": "embeddings",
-                "queryVector": embedding,
-                "numCandidates": 150,
-                "limit": 10,
-            },
-        },
-        {"$project": {"_id": 0, "text": 1, "score": {"$meta": "vectorSearchScore"}}},
-    ]
-    attempts = 0
-    results = list(collection.aggregate(pipeline=pipeline))
-    while not results and attempts < retries:
-        attempts += 1
-        print(f"attempt {attempts}, waiting for valid results: {results}")
-        time.sleep(interval)
-        results = list(collection.aggregate(pipeline=pipeline))
-    if not results:
-        raise TimeoutError("Timed out waiting for valid results")
-    print(f"found results on attempt {attempts}")
-    top_result = results[0]
-    assert top_result["score"] == 1.0, "score detected should be 1: {}".format(top_result["score"])
-    assert top_result["text"] == text, "text detected should be {}, found: {}".format(
-        text, top_result["text"]
-    )
-    for r in results[1:]:
-        assert r["score"] < 1.0, "score detected should be less than 1: {}".format(r["score"])
-
-
-@pytest.mark.asyncio
-@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, NOSQL_TAG)
-@requires_env("MONGODB_URI", "MONGODB_DATABASE")
-async def test_mongodb_source(temp_dir: Path):
-    env_data = get_env_data()
-    indexer_config = MongoDBIndexerConfig(database=env_data.database, collection=SOURCE_COLLECTION)
-    download_config = MongoDBDownloaderConfig(download_dir=temp_dir)
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
-    )
-    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
-    downloader = MongoDBDownloader(
-        connection_config=connection_config, download_config=download_config
-    )
-    await source_connector_validation(
-        indexer=indexer,
-        downloader=downloader,
-        configs=SourceValidationConfigs(
-            test_id=CONNECTOR_TYPE,
-            expected_num_files=4,
-            validate_downloaded_files=True,
-            expected_number_indexed_file_data=1,
-        ),
-    )
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, NOSQL_TAG)
-def test_mongodb_indexer_precheck_fail_no_host():
-    indexer_config = MongoDBIndexerConfig(
-        database="non-existent-database", collection="non-existent-database"
-    )
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri="mongodb+srv://ingest-test.hgaig.mongodb"),
-    )
-    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
-    with pytest.raises(SourceConnectionError):
-        indexer.precheck()
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, NOSQL_TAG)
-@requires_env("MONGODB_URI", "MONGODB_DATABASE")
-def test_mongodb_indexer_precheck_fail_no_database():
-    env_data = get_env_data()
-    indexer_config = MongoDBIndexerConfig(
-        database="non-existent-database", collection=SOURCE_COLLECTION
-    )
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
-    )
-    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
-    with pytest.raises(SourceConnectionError):
-        indexer.precheck()
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, NOSQL_TAG)
-@requires_env("MONGODB_URI", "MONGODB_DATABASE")
-def test_mongodb_indexer_precheck_fail_no_collection():
-    env_data = get_env_data()
-    indexer_config = MongoDBIndexerConfig(
-        database=env_data.database, collection="non-existent-collection"
-    )
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
-    )
-    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
-    with pytest.raises(SourceConnectionError):
-        indexer.precheck()
-
-
-@pytest.mark.asyncio
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, NOSQL_TAG)
-@requires_env("MONGODB_URI", "MONGODB_DATABASE")
-async def test_mongodb_destination(
-    upload_file: Path,
-    destination_collection: Collection,
-    tmp_path: Path,
-):
-    env_data = get_env_data()
-    file_data = FileData(
-        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
-        connector_type=CONNECTOR_TYPE,
-        identifier="mongodb_mock_id",
-    )
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
-    )
-
-    upload_config = MongoDBUploaderConfig(
-        database=env_data.database,
-        collection=destination_collection.name,
-    )
-    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
-    uploader.precheck()
-    uploader.run(path=upload_file, file_data=file_data)
-
-    with upload_file.open() as f:
-        staged_elements = json.load(f)
-    expected_records = len(staged_elements)
-    validate_collection_count(collection=destination_collection, expected_records=expected_records)
-    first_element = staged_elements[0]
-    validate_collection_vector(
-        collection=destination_collection,
-        embedding=first_element["embeddings"],
-        text=first_element["text"],
-    )
-
-    uploader.run(path=upload_file, file_data=file_data)
-    validate_collection_count(collection=destination_collection, expected_records=expected_records)
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, NOSQL_TAG)
-def test_mongodb_uploader_precheck_fail_no_host():
-    upload_config = MongoDBUploaderConfig(
-        database="database",
-        collection="collection",
-    )
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri="mongodb+srv://ingest-test.hgaig.mongodb"),
-    )
-    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
-    with pytest.raises(DestinationConnectionError):
-        uploader.precheck()
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, NOSQL_TAG)
-@requires_env("MONGODB_URI", "MONGODB_DATABASE")
-def test_mongodb_uploader_precheck_fail_no_database():
-    env_data = get_env_data()
-    upload_config = MongoDBUploaderConfig(
-        database="database",
-        collection="collection",
-    )
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
-    )
-    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
-    with pytest.raises(DestinationConnectionError):
-        uploader.precheck()
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, NOSQL_TAG)
-@requires_env("MONGODB_URI", "MONGODB_DATABASE")
-def test_mongodb_uploader_precheck_fail_no_collection():
-    env_data = get_env_data()
-    upload_config = MongoDBUploaderConfig(
-        database=env_data.database,
-        collection="collection",
-    )
-    connection_config = MongoDBConnectionConfig(
-        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
-    )
-    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
-    with pytest.raises(DestinationConnectionError):
-        uploader.precheck()