chore: import upstream snapshot with attribution

test_unstructured_ingest/python/test-azure-output.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+import click
+from azure.storage.blob import ContainerClient
+
+
+@click.group(name="azure-ingest")
+def cli():
+    pass
+
+
+@cli.command()
+@click.option("--connection-string", type=str, required=True)
+@click.option("--container", type=str, required=True)
+@click.option("--blob-path", type=str, required=True)
+def down(connection_string: str, container: str, blob_path: str):
+    container_client = ContainerClient.from_connection_string(
+        conn_str=connection_string, container_name=container
+    )
+    blob_list = [b.name for b in list(container_client.list_blobs(name_starts_with=blob_path))]
+    print(f"deleting all content from {container}/{blob_path}")
+    # Delete all content in folder first
+    container_client.delete_blobs(*[b for b in blob_list if b != blob_path])
+
+    # Delete folder itself
+    container_client.delete_blob(blob_path)
+
+
+@cli.command()
+@click.option("--connection-string", type=str, required=True)
+@click.option("--container", type=str, required=True)
+@click.option("--blob-path", type=str, required=True)
+@click.option("--expected-files", type=int, required=True)
+def check(connection_string: str, container: str, blob_path: str, expected_files: int):
+    container_client = ContainerClient.from_connection_string(
+        conn_str=connection_string, container_name=container
+    )
+    blob_json_list = [
+        b.name
+        for b in list(container_client.list_blobs(name_starts_with=blob_path))
+        if b.name.endswith("json")
+    ]
+    found = len(blob_json_list)
+    print(
+        f"Checking that the number of files found ({found}) "
+        f"matches what's expected: {expected_files}"
+    )
+    assert found == expected_files, (
+        f"number of files found ({found}) doesn't match what's expected: {expected_files}"
+    )
+    print("successfully checked the number of files!")
+
+
+if __name__ == "__main__":
+    cli()

test_unstructured_ingest/python/test-databricks-volumes.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+import json
+
+import click
+from databricks.sdk import WorkspaceClient
+
+
+@click.group()
+def cli():
+    pass
+
+
+def _get_volume_path(catalog: str, volume: str, volume_path: str):
+    return f"/Volumes/{catalog}/default/{volume}/{volume_path}"
+
+
+@cli.command()
+@click.option("--host", type=str, required=True)
+@click.option("--username", type=str, required=True)
+@click.option("--password", type=str, required=True)
+@click.option("--catalog", type=str, required=True)
+@click.option("--volume", type=str, required=True)
+@click.option("--volume-path", type=str, required=True)
+def test(
+    host: str,
+    username: str,
+    password: str,
+    catalog: str,
+    volume: str,
+    volume_path: str,
+):
+    client = WorkspaceClient(host=host, username=username, password=password)
+    files = list(
+        client.files.list_directory_contents(_get_volume_path(catalog, volume, volume_path))
+    )
+
+    assert len(files) == 1
+
+    resp = client.files.download(files[0].path)
+    data = json.loads(resp.contents.read())
+
+    assert len(data) == 5
+    assert [v["type"] for v in data] == [
+        "UncategorizedText",
+        "Title",
+        "NarrativeText",
+        "UncategorizedText",
+        "Title",
+    ]
+
+    print("Databricks test passed!")
+
+
+@cli.command()
+@click.option("--host", type=str, required=True)
+@click.option("--username", type=str, required=True)
+@click.option("--password", type=str, required=True)
+@click.option("--catalog", type=str, required=True)
+@click.option("--volume", type=str, required=True)
+@click.option("--volume-path", type=str, required=True)
+def cleanup(
+    host: str,
+    username: str,
+    password: str,
+    catalog: str,
+    volume: str,
+    volume_path: str,
+):
+    client = WorkspaceClient(host=host, username=username, password=password)
+
+    for file in client.files.list_directory_contents(
+        _get_volume_path(catalog, volume, volume_path)
+    ):
+        client.files.delete(file.path)
+    client.files.delete_directory(_get_volume_path(catalog, volume, volume_path))
+
+
+if __name__ == "__main__":
+    cli()

test_unstructured_ingest/python/test-gcs-output.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+import click
+from google.cloud import storage
+from google.oauth2 import service_account
+
+
+@click.group(name="gcs-ingest")
+def cli():
+    pass
+
+
+@cli.command()
+@click.option("--service-account-file", type=click.Path(), required=True)
+@click.option("--bucket", type=str, required=True)
+@click.option("--blob-path", type=str, required=True)
+def down(service_account_file: str, bucket: str, blob_path: str):
+    credentials = service_account.Credentials.from_service_account_file(
+        filename=service_account_file
+    )
+
+    storage_client = storage.Client(credentials=credentials)
+    for blob in storage_client.list_blobs(bucket_or_name=bucket, prefix=blob_path):
+        print(f"deleting {blob.name}")
+        blob.delete()
+
+
+@cli.command()
+@click.option("--service-account-file", type=click.Path(), required=True)
+@click.option("--bucket", type=str, required=True)
+@click.option("--blob-path", type=str, required=True)
+@click.option("--expected-files", type=int, required=True)
+def check(service_account_file: str, bucket: str, blob_path: str, expected_files: int):
+    credentials = service_account.Credentials.from_service_account_file(
+        filename=service_account_file
+    )
+
+    storage_client = storage.Client(credentials=credentials)
+    blob_json_list = [
+        f.name
+        for f in storage_client.list_blobs(bucket_or_name=bucket, prefix=blob_path)
+        if f.name.endswith("json")
+    ]
+    found = len(blob_json_list)
+    print(
+        f"Checking that the number of files found ({found}) "
+        f"matches what's expected: {expected_files}"
+    )
+    assert found == expected_files, (
+        f"number of files found ({found}) doesn't match what's expected: {expected_files}"
+    )
+    print("successfully checked the number of files!")
+
+
+if __name__ == "__main__":
+    cli()

test_unstructured_ingest/python/test-ingest-astradb-output.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+import click
+from astrapy.db import AstraDB
+
+
+def get_client(token, api_endpoint, collection_name) -> AstraDB:
+    # Initialize our vector db
+    astra_db = AstraDB(token=token, api_endpoint=api_endpoint)
+    astra_db_collection = astra_db.collection(collection_name)
+    return astra_db, astra_db_collection
+
+
+@click.group(name="astradb-ingest")
+@click.option("--token", type=str)
+@click.option("--api-endpoint", type=str)
+@click.option("--collection-name", type=str, default="collection_test")
+@click.option("--embedding-dimension", type=int, default=384)
+@click.pass_context
+def cli(ctx, token: str, api_endpoint: str, collection_name: str, embedding_dimension: int):
+    # ensure that ctx.obj exists and is a dict (in case `cli()` is called
+    ctx.ensure_object(dict)
+
+    ctx.obj["db"], ctx.obj["collection"] = get_client(token, api_endpoint, collection_name)
+
+
+@cli.command()
+@click.pass_context
+def check(ctx):
+    collection_name = ctx.parent.params["collection_name"]
+    print(f"Checking contents of Astra DB collection: {collection_name}")
+
+    astra_db_collection = ctx.obj["collection"]
+
+    # Tally up the embeddings
+    docs_count = astra_db_collection.count_documents()
+    number_of_embeddings = docs_count["status"]["count"]
+
+    # Print the results
+    expected_embeddings = 3
+    print(
+        f"# of embeddings in collection vs expected: {number_of_embeddings}/{expected_embeddings}"
+    )
+
+    # Check that the assertion is true
+    assert number_of_embeddings == expected_embeddings, (
+        f"Number of rows in generated table ({number_of_embeddings})"
+        f"doesn't match expected value: {expected_embeddings}"
+    )
+
+    # Grab an embedding from the collection and search against itself
+    # Should get the same document back as the most similar
+    find_one = astra_db_collection.find_one(projection={"*": 1})
+    random_vector = find_one["data"]["document"]["$vector"]
+    random_text = find_one["data"]["document"]["content"]
+
+    # Perform a similarity search
+    find_result = astra_db_collection.vector_find(
+        random_vector,
+        limit=1,
+        fields=["*"],
+    )
+
+    # Check that we retrieved the coded cleats copy data
+    assert find_result[0]["content"] == random_text
+    print("Vector search complete.")
+
+
+@cli.command()
+@click.pass_context
+def down(ctx):
+    astra_db = ctx.obj["db"]
+    collection_name = ctx.parent.params["collection_name"]
+    print(f"deleting collection: {collection_name}")
+    astra_db.delete_collection(collection_name)
+    print(f"successfully deleted collection: {collection_name}")
+
+
+if __name__ == "__main__":
+    cli()

test_unstructured_ingest/python/test-ingest-chroma-output.py

@@ -0,0 +1,29 @@
+import chromadb
+import click
+
+
+@click.command()
+@click.option("--collection-name", type=str)
+def run_check(collection_name):
+    print(f"Checking contents of Chroma collection: {collection_name}")
+
+    chroma_client = chromadb.HttpClient(host="localhost", port=8000)
+
+    collection = chroma_client.get_or_create_collection(name=collection_name)
+
+    number_of_embeddings = collection.count()
+    expected_embeddings = 3
+    print(
+        f"# of embeddings in collection vs expected: {number_of_embeddings}/{expected_embeddings}"
+    )
+
+    assert number_of_embeddings == expected_embeddings, (
+        f"Number of rows in generated table ({number_of_embeddings}) "
+        f"doesn't match expected value: {expected_embeddings}"
+    )
+
+    print("Table check complete")
+
+
+if __name__ == "__main__":
+    run_check()

test_unstructured_ingest/python/test-ingest-delta-table-output.py

@@ -0,0 +1,37 @@
+import click
+from deltalake import DeltaTable
+
+
+@click.command()
+@click.option("--table-uri", type=str)
+def run_check(table_uri):
+    print(f"Checking contents of table at {table_uri}")
+    delta_table = DeltaTable(
+        table_uri=table_uri,
+    )
+
+    df = delta_table.to_pandas()
+    EXPECTED_ROWS = 5
+    EXPECTED_COLUMNS = 19
+    print(f"Number of rows in table vs expected: {len(df)}/{EXPECTED_ROWS}")
+    print(f"Number of columns in table vs expected: {len(df.columns)}/{EXPECTED_COLUMNS}")
+    number_of_rows = len(df)
+    assert number_of_rows == EXPECTED_ROWS, (
+        f"number of rows in generated table ({number_of_rows}) "
+        f"doesn't match expected value: {EXPECTED_ROWS}"
+    )
+
+    """
+    The number of columns is associated with the flattened JSON structure of the partition output.
+    If this changes, it's most likely due to the metadata changing in the output.
+    """
+    number_of_columns = len(df.columns)
+    assert number_of_columns == EXPECTED_COLUMNS, (
+        f"number of columns in generated table ({number_of_columns}) doesn't "
+        f"match expected value: {EXPECTED_COLUMNS}"
+    )
+    print("table check complete")
+
+
+if __name__ == "__main__":
+    run_check()

test_unstructured_ingest/python/test-ingest-mongodb.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python
+import json
+import time
+
+import click
+from pymongo.mongo_client import MongoClient
+from pymongo.operations import SearchIndexModel
+
+
+def get_client(uri: str) -> MongoClient:
+    client = MongoClient(uri)
+    client.admin.command("ping")
+    print("Successfully connected to MongoDB")
+    return client
+
+
+@click.group(name="mongo-ingest")
+@click.option("--uri", type=str, required=True)
+@click.option("--database", type=str, required=True)
+@click.option("--collection", type=str, required=True)
+@click.pass_context
+def cli(ctx, uri: str, database: str, collection: str):
+    # ensure that ctx.obj exists and is a dict (in case `cli()` is called
+    # by means other than the `if` block below)
+    ctx.ensure_object(dict)
+
+    ctx.obj["client"] = get_client(uri)
+
+
+@cli.command()
+@click.pass_context
+def up(ctx):
+    client = ctx.obj["client"]
+    collection_name = ctx.parent.params["collection"]
+    db = client[ctx.parent.params["database"]]
+    print(f"creating collection {collection_name}")
+    collection = db.create_collection(name=collection_name)
+    print(f"successfully created collection: {collection_name}")
+    if "embeddings" in [c["name"] for c in collection.list_search_indexes()]:
+        print("search index already exists, skipping creation")
+        return
+
+    search_index_name = collection.create_search_index(
+        model=SearchIndexModel(
+            name="embeddings",
+            definition={
+                "mappings": {
+                    "dynamic": True,
+                    "fields": {
+                        "embeddings": [
+                            {"type": "knnVector", "dimensions": 384, "similarity": "euclidean"}
+                        ]
+                    },
+                }
+            },
+        )
+    )
+    print(f"Added search index: {search_index_name}")
+
+
+@cli.command()
+@click.pass_context
+def down(ctx):
+    collection_name = ctx.parent.params["collection"]
+    client = ctx.obj["client"]
+    db = client[ctx.parent.params["database"]]
+    if collection_name not in db.list_collection_names():
+        print(
+            "collection name {} does not exist amongst those in database: {}, "
+            "skipping deletion".format(collection_name, ", ".join(db.list_collection_names()))
+        )
+        return
+    print(f"deleting collection: {collection_name}")
+    collection = db[collection_name]
+    collection.drop()
+    print(f"successfully deleted collection: {collection}")
+
+
+@cli.command()
+@click.option("--expected-records", type=int, required=True)
+@click.pass_context
+def check(ctx, expected_records: int):
+    client = ctx.obj["client"]
+    db = client[ctx.parent.params["database"]]
+    collection = db[ctx.parent.params["collection"]]
+    count = collection.count_documents(filter={})
+    print(f"checking the count in the db ({count}) matches what's expected: {expected_records}")
+    assert count == expected_records, (
+        f"expected count ({expected_records}) does not match how many records were found: {count}"
+    )
+    print("successfully checked that the expected number of records was found in the db!")
+
+
+@cli.command()
+@click.option("--output-json", type=click.File())
+@click.pass_context
+def check_vector(ctx, output_json):
+    """
+    Checks the functionality of the vector search index by getting a score based on the
+    exact result of one of the embeddings. Makes sure that the search index itself has finished
+    indexing before running a query, then validated that the first item in the returned sorted
+    list has a score of 1.0 given that the exact embedding is used as a match, and all others
+    have a score less than 1.0.
+    """
+    # Get the first embedding from the output file
+    json_content = json.load(output_json)
+    exact_embedding = json_content[0]["embeddings"]
+    client = ctx.obj["client"]
+    db = client[ctx.parent.params["database"]]
+    collection = db[ctx.parent.params["collection"]]
+    vector_index_name = "embeddings"
+    status = [ind for ind in collection.list_search_indexes() if ind["name"] == vector_index_name][
+        0
+    ].get("status")
+    max_attempts = 30
+    attempts = 0
+    wait_seconds = 5
+    while status != "READY" and attempts < max_attempts:
+        print(
+            f"status of search index: {status}, waiting another {wait_seconds} "
+            f"seconds for it to be ready"
+        )
+        attempts += 1
+        time.sleep(wait_seconds)
+        status = [
+            ind for ind in collection.list_search_indexes() if ind["name"] == vector_index_name
+        ][0].get("status")
+    print(f"search index is ready to go ({status}), checking vector content")
+    pipeline = [
+        {
+            "$vectorSearch": {
+                "index": "embeddings",
+                "path": "embeddings",
+                "queryVector": exact_embedding,
+                "numCandidates": 150,
+                "limit": 10,
+            },
+        },
+        {"$project": {"_id": 0, "text": 1, "score": {"$meta": "vectorSearchScore"}}},
+    ]
+    result = list(collection.aggregate(pipeline=pipeline))
+    print(f"vector query result: {result}")
+    assert result[0]["score"] == 1.0, "score detected should be 1: {}".format(result[0]["score"])
+    for r in result[1:]:
+        assert r["score"] < 1.0, "score detected should be less than 1: {}".format(r["score"])
+    print("successfully validated vector content!")
+
+
+if __name__ == "__main__":
+    cli()

test_unstructured_ingest/python/test-ingest-sql-output.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+
+import sys
+
+N_ELEMENTS = 5
+
+
+def create_connection(db_type, database=None, port=None):
+    if db_type == "pgvector":
+        from psycopg2 import connect
+
+        return connect(
+            user="unstructured",
+            password="test",
+            dbname="elements",
+            host="localhost",
+            port=port,
+        )
+    elif db_type == "sqlite":
+        from sqlite3 import connect
+
+        return connect(database=database)
+    raise ValueError(f"Unsupported database {db_type} connection.")
+
+
+if __name__ == "__main__":
+    database_name = sys.argv[1]
+    db_url = None
+    port = None
+    if database_name == "sqlite":
+        db_url = sys.argv[2]
+    else:
+        port = sys.argv[2]
+
+    print(f"Running SQL output test for: {database_name}")
+    conn = create_connection(database_name, db_url, port)
+    query = "select count(*) from elements;"
+    cursor = conn.cursor()
+    cursor.execute(query)
+    count = cursor.fetchone()[0]
+
+    if database_name == "pgvector":
+        """Get embedding from database and then use it to
+        search for the closest vector (which should be itself)"""
+        cursor = conn.cursor()
+        cursor.execute("SELECT embeddings FROM elements order by text limit 1")
+        test_embedding = cursor.fetchone()[0]
+        similarity_query = (
+            f"SELECT text FROM elements ORDER BY embeddings <-> '{test_embedding}' LIMIT 1;"
+        )
+        cursor.execute(similarity_query)
+        res = cursor.fetchone()
+        assert res[0] == "Best Regards,"
+        print("Result of vector search against pgvector with embeddings successful")
+
+    try:
+        assert count == N_ELEMENTS
+    except AssertionError:
+        print(f"{database_name} dest check failed: got {count}, expected {N_ELEMENTS}")
+        raise
+    finally:
+        cursor.close()
+        conn.close()
+
+    print(f"SUCCESS: {database_name} dest check")

test_unstructured_ingest/python/test-ingest-weaviate-output.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+
+import weaviate
+
+weaviate_host_url = os.getenv("WEAVIATE_HOST_URL", "http://localhost:8080")
+class_name = os.getenv("WEAVIATE_CLASS_NAME", "Elements")
+N_ELEMENTS = 5
+
+if __name__ == "__main__":
+    print(f"Checking contents of class collection {class_name} at {weaviate_host_url}")
+
+    client = weaviate.Client(
+        url=weaviate_host_url,
+    )
+
+    response = client.query.aggregate(class_name).with_meta_count().do()
+    count = response["data"]["Aggregate"][class_name][0]["meta"]["count"]
+    try:
+        assert count == N_ELEMENTS
+    except AssertionError:
+        sys.exit(f"FAIL: weaviate dest check failed: got {count}, expected {N_ELEMENTS}")
+    print("SUCCESS: weaviate dest check")

test_unstructured_ingest/python/test-kafka-output.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python
+import socket
+from concurrent.futures import ThreadPoolExecutor
+
+import click
+from confluent_kafka import Consumer, TopicPartition
+
+
+@click.group(name="kafka-ingest")
+def cli():
+    pass
+
+
+def get_partition_size(consumer: Consumer, topic_name: str, partition_key: int):
+    topic_partition = TopicPartition(topic_name, partition_key)
+    low_offset, high_offset = consumer.get_watermark_offsets(topic_partition)
+    partition_size = high_offset - low_offset
+    return partition_size
+
+
+def get_topic_size(consumer: Consumer, topic_name: str):
+    print(f"Getting the number of messages in the topic {topic_name}")
+    topic = consumer.list_topics(topic=topic_name)
+    print(f"topic {topic}")
+    partitions = topic.topics[topic_name].partitions
+    workers, max_workers = [], len(partitions) or 1
+
+    with ThreadPoolExecutor(max_workers=max_workers) as e:
+        for partition_key in list(topic.topics[topic_name].partitions.keys()):
+            job = e.submit(get_partition_size, consumer, topic_name, partition_key)
+            workers.append(job)
+
+    topic_size = sum([w.result() for w in workers])
+    return topic_size
+
+
+@cli.command()
+@click.option("--bootstrap-server", type=str, required=True)
+@click.option("--topic", type=str, required=True)
+@click.option("--api-key", type=str, required=False)
+@click.option("--secret", type=str, required=False)
+@click.option("--confluent", type=bool, required=True, default=True)
+@click.option("--port", type=int, required=False, default=9092)
+def check(bootstrap_server: str, topic: str, api_key: str, secret: str, confluent: bool, port: int):
+    conf = {
+        "bootstrap.servers": f"{bootstrap_server}:{port}",
+        "client.id": socket.gethostname(),
+        "group.id": "your_group_id",
+        "enable.auto.commit": "true",
+        "auto.offset.reset": "earliest",
+    }
+
+    if confluent:
+        conf["security.protocol"] = "SASL_SSL"
+        conf["sasl.mechanism"] = "PLAIN"
+        conf["sasl.username"] = api_key
+        conf["sasl.password"] = secret
+
+    consumer = Consumer(conf)
+    print("Checking the number of messages in the topic")
+    topic_size = get_topic_size(consumer, topic)
+    expected = 16
+    print(
+        f"Checking that the number of messages found ({topic_size}) "
+        f"matches what's expected: {expected}"
+    )
+    assert topic_size == expected, (
+        f"number of messages found ({topic_size}) doesn't match what's expected: {expected}"
+    )
+    print("successfully checked the number of messages!")
+
+
+if __name__ == "__main__":
+    cli()

test_unstructured_ingest/python/test-produce-kafka-message.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python
+import base64
+import json
+import socket
+
+import click
+from confluent_kafka import Producer
+
+
+@click.group(name="kafka-ingest")
+def cli():
+    pass
+
+
+@cli.command()
+@click.option("--input-file", type=str, required=True)
+@click.option("--bootstrap-server", type=str, required=True)
+@click.option("--topic", type=str, required=True)
+@click.option("--api-key", type=str, required=False)
+@click.option("--secret", type=str, required=False)
+@click.option("--confluent", type=bool, required=False, default=True)
+@click.option("--port", type=int, required=False, default=9092)
+def up(
+    input_file: str,
+    bootstrap_server: str,
+    topic: str,
+    api_key: str,
+    secret: str,
+    confluent: bool,
+    port: int,
+):
+    conf = {
+        "bootstrap.servers": f"{bootstrap_server}:{port}",
+        "client.id": socket.gethostname(),
+        "message.max.bytes": 10485760,
+    }
+
+    print(f"Confluent setting: {confluent}")
+    if confluent:
+        conf["security.protocol"] = "SASL_SSL"
+        conf["sasl.mechanism"] = "PLAIN"
+        conf["sasl.username"] = api_key
+        conf["sasl.password"] = secret
+
+    producer = Producer(conf)
+
+    # Read the file in binary mode and encode content in base64
+    with open(input_file, "rb") as file:
+        file_content = base64.b64encode(file.read()).decode("utf-8")
+
+    print(f"Message is {len(file_content)} bytes long")
+    # Construct the message with filename and file content
+    message = json.dumps({"filename": input_file.split("/")[-1], "content": file_content}).encode(
+        "utf-8"
+    )
+
+    # Send the message to Kafka
+    producer.produce(topic, message)
+    producer.flush()
+    print(f"File and filename sent to Kafka topic {topic} successfully.")
+
+
+if __name__ == "__main__":
+    cli()