unstructured-ingest 0.0.24__py3-none-any.whl → 0.1.0__py3-none-any.whl

test/integration/utils.py

@@ -0,0 +1,15 @@
+import os
+
+import pytest
+
+
+def requires_env(*envs):
+    if len(envs) == 1:
+        env = envs[0]
+        return pytest.mark.skipif(
+            env not in os.environ, reason=f"Environment variable not set: {env}"
+        )
+    return pytest.mark.skipif(
+        not all(env in os.environ for env in envs),
+        reason="All required environment variables not set: {}".format(", ".join(envs)),
+    )

test/unit/embed/test_mixedbreadai.py

@@ -0,0 +1,41 @@
+from unstructured_ingest.embed.mixedbreadai import (
+    MixedbreadAIEmbeddingConfig,
+    MixedbreadAIEmbeddingEncoder,
+)
+
+
+def test_embed_documents_does_not_break_element_to_dict(mocker):
+    mock_client = mocker.MagicMock()
+
+    def mock_embeddings(
+        model,
+        normalized,
+        encoding_format,
+        truncation_strategy,
+        request_options,
+        input,
+    ):
+        mock_response = mocker.MagicMock()
+        mock_response.data = [mocker.MagicMock(embedding=[i, i + 1]) for i in range(len(input))]
+        return mock_response
+
+    mock_client.embeddings.side_effect = mock_embeddings
+
+    # Mock get_client to return our mock_client
+    mocker.patch.object(MixedbreadAIEmbeddingConfig, "get_client", return_value=mock_client)
+
+    encoder = MixedbreadAIEmbeddingEncoder(
+        config=MixedbreadAIEmbeddingConfig(
+            api_key="api_key", model_name="mixedbread-ai/mxbai-embed-large-v1"
+        )
+    )
+
+    raw_elements = [{"text": f"This is sentence {i+1}"} for i in range(2)]
+    elements = encoder.embed_documents(
+        elements=raw_elements,
+    )
+    assert len(elements) == 2
+    assert elements[0]["text"] == "This is sentence 1"
+    assert elements[1]["text"] == "This is sentence 2"
+    assert elements[0]["embeddings"] is not None
+    assert elements[1]["embeddings"] is not None

test/unit/embed/test_octoai.py

@@ -0,0 +1,20 @@
+from unstructured_ingest.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
+
+
+def test_embed_documents_does_not_break_element_to_dict(mocker):
+    # Mocked client with the desired behavior for embed_documents
+    mock_client = mocker.MagicMock()
+    mock_client.embed_documents.return_value = [1, 2]
+
+    # Mock get_client to return our mock_client
+    mocker.patch.object(OctoAiEmbeddingConfig, "get_client", return_value=mock_client)
+
+    encoder = OctoAIEmbeddingEncoder(config=OctoAiEmbeddingConfig(api_key="api_key"))
+    raw_elements = [{"text": f"This is sentence {i+1}"} for i in range(2)]
+
+    elements = encoder.embed_documents(
+        elements=raw_elements,
+    )
+    assert len(elements) == 2
+    assert elements[0]["text"] == "This is sentence 1"
+    assert elements[1]["text"] == "This is sentence 2"

test/unit/embed/test_openai.py

@@ -0,0 +1,20 @@
+from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+
+
+def test_embed_documents_does_not_break_element_to_dict(mocker):
+    # Mocked client with the desired behavior for embed_documents
+    mock_client = mocker.MagicMock()
+    mock_client.embed_documents.return_value = [1, 2]
+
+    # Mock get_client to return our mock_client
+    mocker.patch.object(OpenAIEmbeddingConfig, "get_client", return_value=mock_client)
+
+    encoder = OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(api_key="api_key"))
+    raw_elements = [{"text": f"This is sentence {i+1}"} for i in range(2)]
+
+    elements = encoder.embed_documents(
+        elements=raw_elements,
+    )
+    assert len(elements) == 2
+    assert elements[0]["text"] == "This is sentence 1"
+    assert elements[1]["text"] == "This is sentence 2"

test/unit/embed/test_vertexai.py

@@ -0,0 +1,25 @@
+from unstructured_ingest.embed.vertexai import VertexAIEmbeddingConfig, VertexAIEmbeddingEncoder
+
+
+def test_embed_documents_does_not_break_element_to_dict(mocker):
+    # Mocked client with the desired behavior for embed_documents
+    mock_responses = []
+    for i in [1, 2]:
+        mock_response = mocker.Mock()
+        mocker.patch.object(mock_response, "values", i)
+        mock_responses.append(mock_response)
+
+    mock_client = mocker.MagicMock()
+    mock_client.get_embeddings.return_value = mock_responses
+
+    # Mock create_client to return our mock_client
+    mocker.patch.object(VertexAIEmbeddingConfig, "get_client", return_value=mock_client)
+    encoder = VertexAIEmbeddingEncoder(config=VertexAIEmbeddingConfig(api_key={"api_key": "value"}))
+    raw_elements = [{"text": f"This is sentence {i+1}"} for i in range(2)]
+
+    elements = encoder.embed_documents(
+        elements=raw_elements,
+    )
+    assert len(elements) == 2
+    assert elements[0]["text"] == "This is sentence 1"
+    assert elements[1]["text"] == "This is sentence 2"

test/unit/embed/test_voyageai.py

@@ -0,0 +1,24 @@
+from unstructured_ingest.embed.voyageai import VoyageAIEmbeddingConfig, VoyageAIEmbeddingEncoder
+
+
+def test_embed_documents_does_not_break_element_to_dict(mocker):
+    # Mocked client with the desired behavior for embed_documents
+    mock_response = mocker.MagicMock()
+    mocker.patch.object(mock_response, "embeddings", [1, 2])
+    mock_client = mocker.MagicMock()
+    mock_client.embed.return_value = mock_response
+
+    # Mock get_client to return our mock_client
+    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
+
+    encoder = VoyageAIEmbeddingEncoder(
+        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-law-2")
+    )
+    raw_elements = [{"text": f"This is sentence {i+1}"} for i in range(2)]
+
+    elements = encoder.embed_documents(
+        elements=raw_elements,
+    )
+    assert len(elements) == 2
+    assert elements[0]["text"] == "This is sentence 1"
+    assert elements[1]["text"] == "This is sentence 2"

test/unit/test_chunking_utils.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+import pytest
+from unstructured.chunking import dispatch
+from unstructured.documents.elements import assign_and_map_hash_ids
+from unstructured.partition.auto import partition
+
+from unstructured_ingest.utils.chunking import (
+    assign_and_map_hash_ids as new_assign_and_map_hash_ids,
+)
+
+test_file_path = Path(__file__).resolve()
+project_root = test_file_path.parents[2]
+docs_path = project_root / "example-docs"
+
+
+@pytest.mark.parametrize(
+    "chunking_strategy",
+    ["basic", "by_title"],
+)
+def test_assign_and_map_hash_ids(chunking_strategy):
+    # Make sure the new logic working on dict content matches the
+    # results if using the unstructured version
+    file_path = docs_path / "book-war-and-peace-1p.txt"
+    elements = partition(filename=str(file_path.resolve()), strategy="fast")
+    chunked_elements = dispatch.chunk(elements=elements, chunking_strategy=chunking_strategy)
+    chunked_elements_copy = chunked_elements.copy()
+
+    hashed_chunked_elements = assign_and_map_hash_ids(chunked_elements)
+    og_chunked_elements_dicts = [e.to_dict() for e in hashed_chunked_elements]
+
+    new_chunked_elements_dicts = [e.to_dict() for e in chunked_elements_copy]
+    new_chunked_elements_dicts = new_assign_and_map_hash_ids(new_chunked_elements_dicts)
+
+    for e1, e2 in zip(og_chunked_elements_dicts, new_chunked_elements_dicts):
+        assert e1 == e2

test/unit/test_error.py

@@ -0,0 +1,27 @@
+import pytest
+
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    PartitionError,
+    SourceConnectionError,
+)
+
+
+@pytest.mark.parametrize(
+    ("error_class", "exception_type", "error_message"),
+    [
+        (SourceConnectionError, ValueError, "Simulated connection error"),
+        (DestinationConnectionError, RuntimeError, "Simulated connection error"),
+        (PartitionError, FileNotFoundError, "Simulated partition error"),
+    ],
+)
+def test_custom_error_decorator(error_class, exception_type, error_message):
+    @error_class.wrap
+    def simulate_error():
+        raise exception_type(error_message)
+
+    with pytest.raises(error_class) as context:
+        simulate_error()
+
+    expected_error_string = error_class.error_string.format(error_message)
+    assert str(context.value) == expected_error_string

test/unit/test_interfaces.py

@@ -0,0 +1,280 @@
+from __future__ import annotations
+
+import pathlib
+from dataclasses import dataclass
+from typing import Any, Dict
+
+import pytest
+from unstructured.documents.elements import DataSourceMetadata
+from unstructured.partition.auto import partition
+from unstructured.staging.base import elements_to_dicts
+
+from unstructured_ingest.interfaces import (
+    BaseConnectorConfig,
+    BaseSingleIngestDoc,
+    ChunkingConfig,
+    PartitionConfig,
+    ProcessorConfig,
+    ReadConfig,
+)
+
+DIRECTORY = pathlib.Path(__file__).parents[2].resolve()
+EXAMPLE_DOCS_DIRECTORY = DIRECTORY / "example-docs"
+TEST_DOWNLOAD_DIR = "/tmp"
+TEST_OUTPUT_DIR = "/tmp"
+TEST_ID = "test"
+TEST_FILE_PATH = str(EXAMPLE_DOCS_DIRECTORY / "book-war-and-peace-1p.txt")
+
+
+@dataclass
+class ExampleConfig(BaseConnectorConfig):
+    id: str
+    path: str
+
+
+TEST_CONFIG = ExampleConfig(id=TEST_ID, path=TEST_FILE_PATH)
+TEST_SOURCE_URL = "test-source-url"
+TEST_VERSION = "1.1.1"
+TEST_RECORD_LOCATOR = {"id": "data-source-id"}
+TEST_DATE_CREATED = "2021-01-01T00:00:00"
+TEST_DATE_MODIFIED = "2021-01-02T00:00:00"
+TEST_DATE_PROCESSED = "2022-12-13T15:44:08"
+
+
+@dataclass
+class ExampleIngestDoc(BaseSingleIngestDoc):
+    connector_config: ExampleConfig
+
+    @property
+    def filename(self):
+        return TEST_FILE_PATH
+
+    @property
+    def _output_filename(self):
+        return TEST_FILE_PATH + ".json"
+
+    @property
+    def source_url(self) -> str:
+        return TEST_SOURCE_URL
+
+    @property
+    def version(self) -> str:
+        return TEST_VERSION
+
+    @property
+    def record_locator(self) -> Dict[str, Any]:
+        return TEST_RECORD_LOCATOR
+
+    @property
+    def date_created(self) -> str:
+        return TEST_DATE_CREATED
+
+    @property
+    def date_modified(self) -> str:
+        return TEST_DATE_MODIFIED
+
+    @property
+    def exists(self) -> bool:
+        return True
+
+    def cleanup_file(self):
+        pass
+
+    def get_file(self):
+        pass
+
+    def has_output(self):
+        return True
+
+    def write_result(self, result):
+        pass
+
+
+@pytest.fixture
+def partition_test_results():
+    # Reusable partition test results, calculated only once
+    result = partition(
+        filename=str(TEST_FILE_PATH),
+        data_source_metadata=DataSourceMetadata(
+            url=TEST_SOURCE_URL,
+            version=TEST_VERSION,
+            record_locator=TEST_RECORD_LOCATOR,
+            date_created=TEST_DATE_CREATED,
+            date_modified=TEST_DATE_MODIFIED,
+            date_processed=TEST_DATE_PROCESSED,
+        ),
+    )
+    return result
+
+
+@pytest.fixture
+def partition_file_test_results(partition_test_results):
+    # Reusable partition_file test results, calculated only once
+    return elements_to_dicts(partition_test_results)
+
+
+def test_partition_file():
+    """Validate partition_file returns a list of dictionaries with the expected keys,
+    metadatakeys, and data source metadata values."""
+    test_ingest_doc = ExampleIngestDoc(
+        connector_config=TEST_CONFIG,
+        read_config=ReadConfig(download_dir=TEST_DOWNLOAD_DIR),
+        processor_config=ProcessorConfig(output_dir=TEST_OUTPUT_DIR),
+    )
+    test_ingest_doc._date_processed = TEST_DATE_PROCESSED
+    elements = test_ingest_doc.partition_file(partition_config=PartitionConfig())
+    element_dicts = elements_to_dicts(elements)
+    assert len(element_dicts)
+    expected_keys = {
+        "element_id",
+        "text",
+        "type",
+        "metadata",
+    }
+    # The document in TEST_FILE_PATH does not have elements with coordinates so
+    # partition is not expected to return coordinates metadata.
+    expected_metadata_keys = {
+        "data_source",
+        "filename",
+        "file_directory",
+        "filetype",
+        "languages",
+        "last_modified",
+    }
+    for elem in element_dicts:
+        # Parent IDs are non-deterministic - remove them from the test
+        elem["metadata"].pop("parent_id", None)
+
+        assert expected_keys == set(elem.keys())
+        assert expected_metadata_keys == set(elem["metadata"].keys())
+        data_source_metadata = elem["metadata"]["data_source"]
+        assert data_source_metadata["url"] == TEST_SOURCE_URL
+        assert data_source_metadata["version"] == TEST_VERSION
+        assert data_source_metadata["record_locator"] == TEST_RECORD_LOCATOR
+        assert data_source_metadata["date_created"] == TEST_DATE_CREATED
+        assert data_source_metadata["date_modified"] == TEST_DATE_MODIFIED
+        assert data_source_metadata["date_processed"] == TEST_DATE_PROCESSED
+
+
+def test_process_file_fields_include_default(mocker, partition_test_results):
+    """Validate when metadata_include and metadata_exclude are not set, all fields:
+    ("element_id", "text", "type", "metadata") are included"""
+    mock_partition = mocker.patch(
+        "unstructured.partition.auto.partition",
+        return_value=partition_test_results,
+    )
+    test_ingest_doc = ExampleIngestDoc(
+        connector_config=TEST_CONFIG,
+        read_config=ReadConfig(download_dir=TEST_DOWNLOAD_DIR),
+        processor_config=ProcessorConfig(output_dir=TEST_OUTPUT_DIR),
+    )
+    elements = test_ingest_doc.partition_file(partition_config=PartitionConfig())
+    element_dicts = elements_to_dicts(elements)
+    assert len(element_dicts)
+    assert mock_partition.call_count == 1
+    for elem in element_dicts:
+        # Parent IDs are non-deterministic - remove them from the test
+        elem["metadata"].pop("parent_id", None)
+
+        assert {"element_id", "text", "type", "metadata"} == set(elem.keys())
+        data_source_metadata = elem["metadata"]["data_source"]
+        assert data_source_metadata["url"] == TEST_SOURCE_URL
+        assert data_source_metadata["version"] == TEST_VERSION
+        assert data_source_metadata["record_locator"] == TEST_RECORD_LOCATOR
+        assert data_source_metadata["date_created"] == TEST_DATE_CREATED
+        assert data_source_metadata["date_modified"] == TEST_DATE_MODIFIED
+        assert data_source_metadata["date_processed"] == TEST_DATE_PROCESSED
+
+
+def test_process_file_metadata_includes_filename_and_filetype(
+    mocker,
+    partition_test_results,
+):
+    """Validate when metadata_include is set to "filename,filetype",
+    only filename is included in metadata"""
+    mocker.patch(
+        "unstructured.partition.auto.partition",
+        return_value=partition_test_results,
+    )
+    partition_config = PartitionConfig(
+        metadata_include=["filename", "filetype"],
+    )
+    test_ingest_doc = ExampleIngestDoc(
+        connector_config=TEST_CONFIG,
+        read_config=ReadConfig(download_dir=TEST_DOWNLOAD_DIR),
+        processor_config=ProcessorConfig(output_dir=TEST_OUTPUT_DIR),
+    )
+    isd_elems = test_ingest_doc.process_file(partition_config=partition_config)
+    assert len(isd_elems)
+    for elem in isd_elems:
+        # Parent IDs are non-deterministic - remove them from the test
+        elem["metadata"].pop("parent_id", None)
+
+        assert set(elem["metadata"].keys()) == {"filename", "filetype"}
+
+
+def test_process_file_metadata_exclude_filename_pagenum(mocker, partition_test_results):
+    """Validate when metadata_exclude is set to "filename,page_number",
+    neither filename nor page_number are included in metadata"""
+    mocker.patch(
+        "unstructured.partition.auto.partition",
+        return_value=partition_test_results,
+    )
+    partition_config = PartitionConfig(
+        metadata_exclude=["filename", "page_number"],
+    )
+    test_ingest_doc = ExampleIngestDoc(
+        connector_config=TEST_CONFIG,
+        read_config=ReadConfig(download_dir=TEST_DOWNLOAD_DIR),
+        processor_config=ProcessorConfig(
+            output_dir=TEST_OUTPUT_DIR,
+        ),
+    )
+    isd_elems = test_ingest_doc.process_file(partition_config=partition_config)
+    assert len(isd_elems)
+    for elem in isd_elems:
+        assert "filename" not in elem["metadata"]
+        assert "page_number" not in elem["metadata"]
+
+
+def test_process_file_flatten_metadata(mocker, partition_test_results):
+    mocker.patch(
+        "unstructured.partition.auto.partition",
+        return_value=partition_test_results,
+    )
+    partition_config = PartitionConfig(
+        metadata_include=["filename", "file_directory", "filetype"],
+        flatten_metadata=True,
+    )
+    test_ingest_doc = ExampleIngestDoc(
+        connector_config=TEST_CONFIG,
+        read_config=ReadConfig(download_dir=TEST_DOWNLOAD_DIR),
+        processor_config=ProcessorConfig(
+            output_dir=TEST_OUTPUT_DIR,
+        ),
+    )
+    isd_elems = test_ingest_doc.process_file(partition_config=partition_config)
+    expected_keys = {"element_id", "text", "type", "filename", "file_directory", "filetype"}
+    for elem in isd_elems:
+        assert expected_keys == set(elem.keys())
+
+
+class DescribeChunkingConfig:
+    """Unit tests for unstructured.ingest.interfaces.ChunkingConfig"""
+
+    def it_accepts_chunking_strategy_by_itself(self):
+        config = ChunkingConfig(chunking_strategy="basic")
+        assert config.chunking_strategy == "basic"
+
+    def it_defaults_to_chunk_by_title_if_only_chunk_elements_is_True(self):
+        config = ChunkingConfig(chunk_elements=True)
+        assert config.chunking_strategy == "by_title"
+
+    def but_it_defaults_to_chunking_strategy_over_chunk_elements(self):
+        config = ChunkingConfig(chunk_elements=True, chunking_strategy="basic")
+        assert config.chunking_strategy == "basic"
+
+    def it_silently_accepts_unrecognized_chunker(self, caplog: pytest.LogCaptureFixture):
+        config = ChunkingConfig(chunking_strategy="foobar")
+        assert config.chunking_strategy == "foobar"
+        assert caplog.text == ""

test/unit/test_interfaces_v2.py

@@ -0,0 +1,26 @@
+import pytest
+from pydantic import Secret, ValidationError
+
+from unstructured_ingest.v2.interfaces import AccessConfig, ConnectionConfig
+
+
+def test_failing_connection_config():
+    class MyAccessConfig(AccessConfig):
+        sensitive_value: str
+
+    class MyConnectionConfig(ConnectionConfig):
+        access_config: MyAccessConfig
+
+    with pytest.raises(ValidationError):
+        MyConnectionConfig(access_config=MyAccessConfig(sensitive_value="this"))
+
+
+def test_happy_path_connection_config():
+    class MyAccessConfig(AccessConfig):
+        sensitive_value: str
+
+    class MyConnectionConfig(ConnectionConfig):
+        access_config: Secret[MyAccessConfig]
+
+    connection_config = MyConnectionConfig(access_config=MyAccessConfig(sensitive_value="this"))
+    assert connection_config

test/unit/test_logger.py

@@ -0,0 +1,78 @@
+import json
+
+import pytest
+
+from unstructured_ingest.logger import (
+    default_is_data_sensitive,
+    hide_sensitive_fields,
+    redact_jsons,
+)
+
+
+@pytest.mark.parametrize(
+    ("key", "value", "is_sensitive"),
+    [
+        ("username", "john_smith", False),
+        ("password", "13?H%", True),
+        ("token", "123", True),
+        ("AWS_CREDENTIAL", "aws_credential", True),
+        ("AWS_KEY", None, False),
+    ],
+)
+def test_default_is_sensitive(key, value, is_sensitive):
+    assert default_is_data_sensitive(key, value) == is_sensitive
+
+
+def test_hide_sensitive_fields():
+    d = {
+        "username": "john_smith",
+        "password": "13?H%",
+        "inner": {
+            "token": "123",
+            "AWS_KEY": None,
+            "inner_j_string": json.dumps(
+                {"account_name": "secret name", "client_id": 123, "timestamp": 123}
+            ),
+        },
+    }
+    redacted_d = hide_sensitive_fields(d)
+    expected_d = {
+        "password": "*******",
+        "username": "john_smith",
+        "inner": {
+            "token": "*******",
+            "AWS_KEY": None,
+            "inner_j_string": json.dumps(
+                {"account_name": "*******", "client_id": "*******", "timestamp": 123}
+            ),
+        },
+    }
+    assert redacted_d == expected_d
+
+
+def test_redact_jsons():
+    d1 = {
+        "username": "john_smith",
+        "password": "13?H%",
+        "inner": {
+            "token": "123",
+            "AWS_KEY": None,
+            "inner_j_string": json.dumps(
+                {"account_name": "secret name", "client_id": 123, "timestamp": 123}
+            ),
+        },
+    }
+
+    d2 = {"username": "tim67", "update_time": 456}
+    d3 = {"account_name": "top secret", "host": "http://localhost:8888"}
+
+    sensitive_string = f"Some topic secret info ({json.dumps(d1)} regarding {d2} and {d3})"
+    expected_string = (
+        'Some topic secret info ({"username": "john_smith", "password": "*******", '
+        '"inner": {"token": "*******", "AWS_KEY": null, "inner_j_string": '
+        '"{\\"account_name\\": \\"*******\\", \\"client_id\\": \\"*******\\", '
+        '\\"timestamp\\": 123}"}} regarding {"username": "tim67", "update_time": 456} '
+        'and {"account_name": "*******", "host": "http://localhost:8888"})'
+    )
+    redacted_string = redact_jsons(sensitive_string)
+    assert redacted_string == expected_string