unstructured-ingest 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

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

test/unit/v2/embedders/test_huggingface.py

@@ -0,0 +1,48 @@
+import random
+from typing import Any
+
+import faker
+import pytest
+
+from test.unit.v2.utils.data_generator import generate_random_dictionary
+from unstructured_ingest.embed.huggingface import (
+    HuggingFaceEmbeddingConfig,
+    HuggingFaceEmbeddingEncoder,
+)
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {}
+    if random.random() < 0.5:
+        params["embed_model_name"] = fake.word() if random.random() < 0.5 else None
+        params["embedder_model_kwargs"] = (
+            generate_random_dictionary(key_type=str, value_type=Any)
+            if random.random() < 0.5
+            else None
+        )
+        params["encode_kwargs"] = (
+            generate_random_dictionary(key_type=str, value_type=Any)
+            if random.random() < 0.5
+            else None
+        )
+        params["cache_folder"] = fake.file_path() if random.random() < 0.5 else None
+    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 = HuggingFaceEmbeddingConfig.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 = HuggingFaceEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = HuggingFaceEmbeddingEncoder(config=embedder_config)
+    assert embedder

test/unit/v2/embedders/test_mixedbread.py

@@ -0,0 +1,37 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.embed.mixedbreadai import (
+    MixedbreadAIEmbeddingConfig,
+    MixedbreadAIEmbeddingEncoder,
+)
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {
+        "api_key": fake.password(),
+    }
+    if random.random() < 0.5:
+        params["embedder_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 = MixedbreadAIEmbeddingConfig.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 = MixedbreadAIEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = MixedbreadAIEmbeddingEncoder(config=embedder_config)
+    assert embedder

test/unit/v2/embedders/test_octoai.py

@@ -0,0 +1,35 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {
+        "api_key": fake.password(),
+    }
+    if random.random() < 0.5:
+        params["embedder_model_name"] = fake.word()
+        params["base_url"] = fake.url()
+    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 = OctoAiEmbeddingConfig.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 = OctoAiEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = OctoAIEmbeddingEncoder(config=embedder_config)
+    assert embedder

test/unit/v2/embedders/test_openai.py

@@ -0,0 +1,35 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {
+        "api_key": fake.password(),
+    }
+    if random.random() < 0.5:
+        params["embedder_model_name"] = fake.word()
+        params["base_url"] = fake.url()
+    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 = OpenAIEmbeddingConfig.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 = OpenAIEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = OpenAIEmbeddingEncoder(config=embedder_config)
+    assert embedder

test/unit/v2/embedders/test_togetherai.py

@@ -0,0 +1,37 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.embed.togetherai import (
+    TogetherAIEmbeddingConfig,
+    TogetherAIEmbeddingEncoder,
+)
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {
+        "api_key": fake.password(),
+    }
+    if random.random() < 0.5:
+        params["embedder_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 = TogetherAIEmbeddingConfig.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 = TogetherAIEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = TogetherAIEmbeddingEncoder(config=embedder_config)
+    assert embedder

test/unit/v2/embedders/test_vertexai.py

@@ -0,0 +1,37 @@
+import json
+import random
+from typing import Any
+
+import faker
+import pytest
+
+from test.unit.v2.utils.data_generator import generate_random_dictionary
+from unstructured_ingest.embed.vertexai import VertexAIEmbeddingConfig, VertexAIEmbeddingEncoder
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {
+        "api_key": json.dumps(generate_random_dictionary(key_type=str, value_type=Any)),
+    }
+    if random.random() < 0.5:
+        params["embedder_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 = VertexAIEmbeddingConfig.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 = VertexAIEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = VertexAIEmbeddingEncoder(config=embedder_config)
+    assert embedder

test/unit/v2/embedders/test_voyageai.py

@@ -0,0 +1,38 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.embed.voyageai import VoyageAIEmbeddingConfig, VoyageAIEmbeddingEncoder
+
+fake = faker.Faker()
+
+
+def generate_embedder_config_params() -> dict:
+    params = {
+        "api_key": fake.password(),
+    }
+    if random.random() < 0.5:
+        params["embedder_model_name"] = fake.word()
+        params["batch_size"] = fake.random_int()
+        params["truncation"] = fake.boolean()
+        params["max_retries"] = fake.random_int()
+        params["timeout_in_seconds"] = fake.random_int()
+    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 = VoyageAIEmbeddingConfig.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 = VoyageAIEmbeddingConfig.model_validate(embedder_config_params)
+    embedder = VoyageAIEmbeddingEncoder(config=embedder_config)
+    assert embedder

test/unit/v2/partitioners/test_partitioner.py

@@ -0,0 +1,63 @@
+import random
+from typing import Any
+
+import faker
+import pytest
+
+from test.unit.v2.utils.data_generator import generate_random_dictionary
+from unstructured_ingest.v2.processes.partitioner import Partitioner, PartitionerConfig
+
+fake = faker.Faker()
+
+
+def generate_partitioner_config_params() -> dict:
+    params = {
+        "strategy": random.choice(["fast", "hi_res", "auto"]),
+        "ocr_languages": fake.words() if random.random() < 0.5 else None,
+        "encoding": fake.word() if random.random() < 0.5 else None,
+        "additional_partition_args": (
+            generate_random_dictionary(key_type=str, value_type=Any)
+            if random.random() < 0.5
+            else None
+        ),
+        "skip_infer_table_types": fake.words() if random.random() < 0.5 else None,
+        "flatten_metadata": fake.boolean(),
+        "hi_res_model_name": fake.word() if random.random() < 0.5 else None,
+    }
+    random_val = random.random()
+    # Randomly set the fields_include to a random list[str]
+    if random_val < 0.5:
+        params["fields_include"] = fake.words()
+
+    # Randomly set the metadata_exclude or metadata_include to a valid
+    # list[str] or don't set it at all
+    if random.random() < (1 / 3):
+        params["metadata_exclude"] = fake.words()
+    elif random_val < (2 / 3):
+        params["metadata_include"] = fake.words()
+
+    # Randomly set the values associated with calling the api, or not at all
+    if random.random() < 0.5:
+        params["partition_by_api"]: True
+        params["partition_endpoint"] = fake.url()
+        params["api_key"] = fake.password()
+    else:
+        params["partition_by_api"]: False
+    return params
+
+
+@pytest.mark.parametrize(
+    "partition_config_params", [generate_partitioner_config_params() for i in range(10)]
+)
+def test_partition_config(partition_config_params: dict):
+    partition_config = PartitionerConfig.model_validate(partition_config_params)
+    assert partition_config
+
+
+@pytest.mark.parametrize(
+    "partition_config_params", [generate_partitioner_config_params() for i in range(10)]
+)
+def test_partitioner(partition_config_params: dict):
+    partition_config = PartitionerConfig.model_validate(partition_config_params)
+    partitioner = Partitioner(config=partition_config)
+    assert partitioner

test/unit/v2/utils/data_generator.py

@@ -0,0 +1,32 @@
+import random
+from typing import Any, Type
+
+from faker import Faker
+
+fake = Faker()
+
+type_to_random_value_map = {
+    str: fake.sentence,
+    int: fake.random_int,
+    float: fake.random_digit,
+    bool: fake.boolean,
+}
+type_to_random_value_map_key = type_to_random_value_map.copy()
+type_to_random_value_map_key[str] = fake.word
+
+
+def generate_random_dictionary(key_type: Type = str, value_type: Type = str) -> dict:
+    d = {}
+    num_keys = random.randint(1, 3)
+    for i in range(num_keys):
+        key = type_to_random_value_map_key[key_type]()
+        current_value_type = value_type
+        if current_value_type == Any:
+            current_value_type = random.choice(list(type_to_random_value_map.keys()) + [dict])
+        value = (
+            generate_random_dictionary(key_type=key_type, value_type=value_type)
+            if current_value_type is dict
+            else type_to_random_value_map[current_value_type]()
+        )
+        d[key] = value
+    return d

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.1"  # pragma: no cover
+__version__ = "0.3.0"  # pragma: no cover

unstructured_ingest/cli/cmds/__init__.py

@@ -9,7 +9,7 @@ from unstructured_ingest.cli.cmds.fsspec.sftp import get_base_src_cmd as sftp_ba
 from .airtable import get_base_src_cmd as airtable_base_src_cmd
 from .astradb import get_base_dest_cmd as astradb_base_dest_cmd
 from .astradb import get_base_src_cmd as astradb_base_src_cmd
-from .azure_cognitive_search import get_base_dest_cmd as azure_cognitive_search_base_dest_cmd
+from .azure_ai_search import get_base_dest_cmd as azure_ai_search_base_dest_cmd
 from .biomed import get_base_src_cmd as biomed_base_src_cmd
 from .chroma import get_base_dest_cmd as chroma_base_dest_cmd
 from .clarifai import get_base_dest_cmd as clarifai_base_dest_cmd
@@ -118,7 +118,7 @@ base_dest_cmd_fns: t.List[t.Callable[[], "BaseDestCmd"]] = [
     gcs_base_dest_cmd,
     kafka_base_dest_cmd,
     s3_base_dest_cmd,
-    azure_cognitive_search_base_dest_cmd,
+    azure_ai_search_base_dest_cmd,
     delta_table_dest_cmd,
     sql_base_dest_cmd,
     weaviate_dest_cmd,

unstructured_ingest/cli/cmds/{azure_cognitive_search.py → azure_ai_search.py}

@@ -6,14 +6,14 @@ import click
 from unstructured_ingest.cli.interfaces import (
     CliConfig,
 )
-from unstructured_ingest.connector.azure_cognitive_search import (
-    AzureCognitiveSearchWriteConfig,
-    SimpleAzureCognitiveSearchStorageConfig,
+from unstructured_ingest.connector.azure_ai_search import (
+    AzureAISearchWriteConfig,
+    SimpleAzureAISearchStorageConfig,
 )
 
 
 @dataclass
-class AzureCognitiveSearchCliConfig(SimpleAzureCognitiveSearchStorageConfig, CliConfig):
+class AzureAISearchCliConfig(SimpleAzureAISearchStorageConfig, CliConfig):
     @staticmethod
     def get_cli_options() -> t.List[click.Option]:
         options = [
@@ -39,7 +39,7 @@ class AzureCognitiveSearchCliConfig(SimpleAzureCognitiveSearchStorageConfig, Cli
 
 
 @dataclass
-class AzureCognitiveSearchCliWriteConfig(AzureCognitiveSearchWriteConfig, CliConfig):
+class AzureAISearchCliWriteConfig(AzureAISearchWriteConfig, CliConfig):
     @staticmethod
     def get_cli_options() -> t.List[click.Option]:
         options = [
@@ -57,9 +57,9 @@ def get_base_dest_cmd():
     from unstructured_ingest.cli.base.dest import BaseDestCmd
 
     cmd_cls = BaseDestCmd(
-        cmd_name="azure-cognitive-search",
-        cli_config=AzureCognitiveSearchCliConfig,
-        additional_cli_options=[AzureCognitiveSearchCliWriteConfig],
-        write_config=AzureCognitiveSearchCliWriteConfig,
+        cmd_name="azure-ai-search",
+        cli_config=AzureAISearchCliConfig,
+        additional_cli_options=[AzureAISearchCliWriteConfig],
+        write_config=AzureAISearchCliWriteConfig,
     )
     return cmd_cls

unstructured_ingest/connector/{azure_cognitive_search.py → azure_ai_search.py}

@@ -19,28 +19,28 @@ if t.TYPE_CHECKING:
 
 
 @dataclass
-class AzureCognitiveSearchAccessConfig(AccessConfig):
+class AzureAiSearchAccessConfig(AccessConfig):
     key: str = enhanced_field(sensitive=True)
 
 
 @dataclass
-class SimpleAzureCognitiveSearchStorageConfig(BaseConnectorConfig):
+class SimpleAzureAISearchStorageConfig(BaseConnectorConfig):
     endpoint: str
-    access_config: AzureCognitiveSearchAccessConfig
+    access_config: AzureAiSearchAccessConfig
 
 
 @dataclass
-class AzureCognitiveSearchWriteConfig(WriteConfig):
+class AzureAISearchWriteConfig(WriteConfig):
     index: str
 
 
 @dataclass
-class AzureCognitiveSearchDestinationConnector(BaseDestinationConnector):
-    write_config: AzureCognitiveSearchWriteConfig
-    connector_config: SimpleAzureCognitiveSearchStorageConfig
+class AzureAISearchDestinationConnector(BaseDestinationConnector):
+    write_config: AzureAISearchWriteConfig
+    connector_config: SimpleAzureAISearchStorageConfig
     _client: t.Optional["SearchClient"] = field(init=False, default=None)
 
-    @requires_dependencies(["azure.search"], extras="azure-cognitive-search")
+    @requires_dependencies(["azure.search"], extras="azure-ai-search")
     def generate_client(self) -> "SearchClient":
         from azure.core.credentials import AzureKeyCredential
         from azure.search.documents import SearchClient
@@ -112,7 +112,7 @@ class AzureCognitiveSearchDestinationConnector(BaseDestinationConnector):
         if page_number := data.get("metadata", {}).get("page_number"):
             data["metadata"]["page_number"] = str(page_number)
 
-    @requires_dependencies(["azure"], extras="azure-cognitive-search")
+    @requires_dependencies(["azure"], extras="azure-ai-search")
     def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
         import azure.core.exceptions
 

unstructured_ingest/connector/kafka.py

@@ -181,7 +181,6 @@ class KafkaSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
                     logger.debug(f"found {len(collected)} messages, stopping")
                     consumer.commit(asynchronous=False)
                     break
-
         return [
             KafkaIngestDoc(
                 connector_config=self.connector_config,

unstructured_ingest/interfaces.py

@@ -21,6 +21,7 @@ from unstructured_ingest.enhanced_dataclass.core import _asdict
 from unstructured_ingest.error import PartitionError, SourceConnectionError
 from unstructured_ingest.logger import logger
 from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.v2.unstructured_api import call_api
 
 if TYPE_CHECKING:
     from unstructured.documents.elements import Element
@@ -565,6 +566,7 @@ class BaseSingleIngestDoc(BaseIngestDoc, IngestDocJsonMixin, ABC):
     ) -> list["Element"]:
         from unstructured.documents.elements import DataSourceMetadata
         from unstructured.partition.auto import partition
+        from unstructured.staging.base import elements_from_dicts
 
         if not partition_config.partition_by_api:
             logger.debug("Using local partition")
@@ -582,18 +584,16 @@ class BaseSingleIngestDoc(BaseIngestDoc, IngestDocJsonMixin, ABC):
                 **partition_kwargs,
             )
         else:
-            from unstructured.partition.api import partition_via_api
-
             endpoint = partition_config.partition_endpoint
 
             logger.debug(f"using remote partition ({endpoint})")
-
-            elements = partition_via_api(
-                filename=str(self.filename),
+            elements_dicts = call_api(
+                server_url=endpoint,
                 api_key=partition_config.api_key,
-                api_url=endpoint,
-                **partition_kwargs,
+                filename=Path(self.filename),
+                api_parameters=partition_kwargs,
             )
+            elements = elements_from_dicts(elements_dicts)
             # TODO: add m_data_source_metadata to unstructured-api pipeline_api and then
             # pass the stringified json here
         return elements

unstructured_ingest/runner/writers/__init__.py

@@ -1,7 +1,7 @@
 import typing as t
 
 from .astradb import AstraDBWriter
-from .azure_cognitive_search import AzureCognitiveSearchWriter
+from .azure_ai_search import AzureAiSearchWriter
 from .base_writer import Writer
 from .chroma import ChromaWriter
 from .clarifai import ClarifaiWriter
@@ -25,7 +25,7 @@ from .weaviate import WeaviateWriter
 writer_map: t.Dict[str, t.Type[Writer]] = {
     "astradb": AstraDBWriter,
     "azure": AzureWriter,
-    "azure_cognitive_search": AzureCognitiveSearchWriter,
+    "azure_ai_search": AzureAiSearchWriter,
     "box": BoxWriter,
     "chroma": ChromaWriter,
     "clarifai": ClarifaiWriter,

unstructured_ingest/runner/writers/azure_ai_search.py

@@ -0,0 +1,24 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.azure_ai_search import (
+        AzureAISearchWriteConfig,
+        SimpleAzureAISearchStorageConfig,
+    )
+
+
+@dataclass
+class AzureAiSearchWriter(Writer):
+    connector_config: "SimpleAzureAISearchStorageConfig"
+    write_config: "AzureAISearchWriteConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.azure_ai_search import (
+            AzureAISearchDestinationConnector,
+        )
+
+        return AzureAISearchDestinationConnector

unstructured_ingest/v2/constants.py

@@ -0,0 +1,2 @@
+# Used to append to metadata for uploaders that store element-level data
+RECORD_ID_LABEL = "record_id"

unstructured_ingest/v2/processes/chunker.py

@@ -9,7 +9,7 @@ from unstructured_ingest.utils.chunking import assign_and_map_hash_ids
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces.process import BaseProcess
 from unstructured_ingest.v2.logger import logger
-from unstructured_ingest.v2.unstructured_api import call_api
+from unstructured_ingest.v2.unstructured_api import call_api_async
 
 CHUNK_MAX_CHARS_DEFAULT: int = 500
 CHUNK_MULTI_PAGE_DEFAULT: bool = True
@@ -112,7 +112,7 @@ class Chunker(BaseProcess, ABC):
 
     @requires_dependencies(dependencies=["unstructured_client"], extras="remote")
     async def run_async(self, elements_filepath: Path, **kwargs: Any) -> list[dict]:
-        elements = await call_api(
+        elements = await call_api_async(
             server_url=self.config.chunking_endpoint,
             api_key=self.config.chunk_api_key.get_secret_value(),
             filename=elements_filepath,