unstructured-ingest 0.0.18__py3-none-any.whl → 0.0.21__py3-none-any.whl

unstructured_ingest/v2/unstructured_api.py

@@ -0,0 +1,87 @@
+import asyncio
+from dataclasses import fields
+from functools import partial
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+from unstructured_ingest.v2.logger import logger
+
+if TYPE_CHECKING:
+    from unstructured_client.models.operations import PartitionRequest
+
+
+def create_partition_request(filename: Path, parameters_dict: dict) -> "PartitionRequest":
+    """Given a filename and a dict of API parameters, return a PartitionRequest for use
+    by unstructured-client. Remove any params that aren't recognized by the SDK.
+
+    Args:
+        filename: Path to the file being partitioned
+        parameters_dict: A mapping of all API params we want to send
+
+    Returns: A PartitionRequest containing the file and all valid params
+    """
+    from unstructured_client.models.operations import PartitionRequest
+    from unstructured_client.models.shared import Files, PartitionParameters
+
+    # NOTE(austin): PartitionParameters is a Pydantic model in v0.26.0
+    # Prior to this it was a dataclass which doesn't have .__fields
+    try:
+        possible_fields = PartitionParameters.__fields__
+    except AttributeError:
+        possible_fields = [f.name for f in fields(PartitionParameters)]
+
+    filtered_partition_request = {k: v for k, v in parameters_dict.items() if k in possible_fields}
+    if len(filtered_partition_request) != len(parameters_dict):
+        logger.debug(
+            "Following fields were omitted due to not being "
+            "supported by the currently used unstructured client: {}".format(
+                ", ".join([v for v in parameters_dict if v not in filtered_partition_request])
+            )
+        )
+
+    logger.debug(f"using hosted partitioner with kwargs: {parameters_dict}")
+
+    with open(filename, "rb") as f:
+        files = Files(
+            content=f.read(),
+            file_name=str(filename.resolve()),
+        )
+        filtered_partition_request["files"] = files
+
+    partition_params = PartitionParameters(**filtered_partition_request)
+
+    return PartitionRequest(partition_parameters=partition_params)
+
+
+async def call_api(
+    server_url: Optional[str], api_key: Optional[str], filename: Path, api_parameters: dict
+) -> list[dict]:
+    """Call the Unstructured API using unstructured-client.
+
+    Args:
+        server_url: The base URL where the API is hosted
+        api_key: The user's API key (can be empty if this is a self hosted API)
+        filename: Path to the file being partitioned
+        api_parameters: A dict containing the requested API parameters
+
+    Returns: A list of the file's elements, or an empty list if there was an error
+    """
+    from unstructured_client import UnstructuredClient
+
+    client = UnstructuredClient(
+        server_url=server_url,
+        api_key_auth=api_key,
+    )
+    partition_request = create_partition_request(filename=filename, parameters_dict=api_parameters)
+
+    # TODO when client supports async, run without using run_in_executor
+    # isolate the IO heavy call
+    loop = asyncio.get_event_loop()
+
+    # Note(austin) - The partition calls needs request to be a keyword arg
+    # We have to use partial to do this, we can't pass request=request into run_in_executor
+    partition_call = partial(client.general.partition, request=partition_request)
+
+    res = await loop.run_in_executor(None, partition_call)
+
+    return res.elements or []