unstructured-ingest 0.3.13__py3-none-any.whl

test/integration/partitioners/test_partitioner.py

@@ -0,0 +1,75 @@
+import json
+import os
+from pathlib import Path
+
+import pytest
+
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.errors import UserError
+from unstructured_ingest.v2.processes.partitioner import Partitioner, PartitionerConfig
+
+int_test_dir = Path(__file__).parent
+assets_dir = int_test_dir / "assets"
+
+all_partition_files = [path for path in assets_dir.iterdir() if path.is_file()]
+non_image_partition_files = [
+    path for path in all_partition_files if path.suffix not in [".jpg", ".png", ".tif"]
+]
+image_partition_files = [
+    path for path in all_partition_files if path not in non_image_partition_files
+]
+
+
+@pytest.mark.parametrize(
+    "partition_file", all_partition_files, ids=[path.name for path in all_partition_files]
+)
+@requires_env("UNSTRUCTURED_API_KEY", "UNSTRUCTURED_API_URL")
+@pytest.mark.asyncio
+async def test_partitioner_api_hi_res(partition_file: Path):
+    api_key = os.getenv("UNSTRUCTURED_API_KEY")
+    api_url = os.getenv("UNSTRUCTURED_API_URL")
+    partitioner_config = PartitionerConfig(
+        strategy="hi_res", partition_by_api=True, api_key=api_key, partition_endpoint=api_url
+    )
+    partitioner = Partitioner(config=partitioner_config)
+    results = await partitioner.run_async(filename=partition_file)
+    results_dir = int_test_dir / "results"
+    results_dir.mkdir(exist_ok=True)
+    results_path = results_dir / f"{partition_file.name}.json"
+    with results_path.open("w") as f:
+        json.dump(results, f, indent=2)
+    assert results
+
+
+@pytest.mark.parametrize(
+    "partition_file",
+    non_image_partition_files,
+    ids=[path.name for path in non_image_partition_files],
+)
+@requires_env("UNSTRUCTURED_API_KEY", "UNSTRUCTURED_API_URL")
+@pytest.mark.asyncio
+async def test_partitioner_api_fast(partition_file: Path):
+    api_key = os.getenv("UNSTRUCTURED_API_KEY")
+    api_url = os.getenv("UNSTRUCTURED_API_URL")
+    partitioner_config = PartitionerConfig(
+        strategy="fast", partition_by_api=True, api_key=api_key, partition_endpoint=api_url
+    )
+    partitioner = Partitioner(config=partitioner_config)
+    results = await partitioner.run_async(filename=partition_file)
+    assert results
+
+
+@pytest.mark.parametrize(
+    "partition_file", image_partition_files, ids=[path.name for path in image_partition_files]
+)
+@requires_env("UNSTRUCTURED_API_KEY", "UNSTRUCTURED_API_URL")
+@pytest.mark.asyncio
+async def test_partitioner_api_fast_error(partition_file: Path):
+    api_key = os.getenv("UNSTRUCTURED_API_KEY")
+    api_url = os.getenv("UNSTRUCTURED_API_URL")
+    partitioner_config = PartitionerConfig(
+        strategy="fast", partition_by_api=True, api_key=api_key, partition_endpoint=api_url
+    )
+    partitioner = Partitioner(config=partitioner_config)
+    with pytest.raises(UserError):
+        await partitioner.run_async(filename=partition_file)

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,42 @@
+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)
+    mocker.patch.object(MixedbreadAIEmbeddingEncoder, "get_request_options", return_value={})
+
+    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,27 @@
+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_data = []
+    for i in range(2):
+        data = mocker.MagicMock()
+        data.embedding = [1, 2]
+        mock_data.append(data)
+    mock_response = mocker.MagicMock()
+    mock_response.data = mock_data
+    mock_client.embeddings.create.return_value = mock_response
+
+    # 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_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_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

test/unit/test_utils.py

@@ -0,0 +1,184 @@
+import json
+import typing as t
+from dataclasses import dataclass, field
+from datetime import datetime
+
+import pytest
+import pytz
+
+from unstructured_ingest.cli.utils import extract_config
+from unstructured_ingest.interfaces import BaseConfig
+from unstructured_ingest.utils.string_and_date_utils import (
+    ensure_isoformat_datetime,
+    json_to_dict,
+    truncate_string_bytes,
+)
+
+
+@dataclass
+class A(BaseConfig):
+    a: str
+
+
+@dataclass
+class B(BaseConfig):
+    a: A
+    b: int
+
+
+flat_data = {"a": "test", "b": 4, "c": True}
+
+
+def test_extract_config_concrete():
+    @dataclass
+    class C(BaseConfig):
+        b: B
+        c: bool
+
+    c = extract_config(flat_data=flat_data, config=C)
+    expected_result = {"b": {"a": {"a": "test"}, "b": 4}, "c": True}
+    assert c.to_json(sort_keys=True) == json.dumps(expected_result, sort_keys=True)
+
+
+def test_extract_config_optional():
+    @dataclass
+    class C(BaseConfig):
+        c: bool
+        b: t.Optional[B] = None
+
+    c = extract_config(flat_data=flat_data, config=C)
+    expected_result = {"b": {"a": {"a": "test"}, "b": 4}, "c": True}
+    assert c.to_json(sort_keys=True) == json.dumps(expected_result, sort_keys=True)
+
+
+def test_extract_config_union():
+    @dataclass
+    class C(BaseConfig):
+        c: bool
+        b: t.Optional[t.Union[B, int]] = None
+
+    c = extract_config(flat_data=flat_data, config=C)
+    expected_result = {"b": 4, "c": True}
+    assert c.to_json(sort_keys=True) == json.dumps(expected_result, sort_keys=True)
+
+
+def test_extract_config_list():
+    @dataclass
+    class C(BaseConfig):
+        c: t.List[int]
+        b: B
+
+    flat_data = {"a": "test", "b": 4, "c": [1, 2, 3]}
+    c = extract_config(flat_data=flat_data, config=C)
+    expected_result = {"b": {"a": {"a": "test"}, "b": 4}, "c": [1, 2, 3]}
+    assert c.to_json(sort_keys=True) == json.dumps(expected_result, sort_keys=True)
+
+
+def test_extract_config_optional_list():
+    @dataclass
+    class C(BaseConfig):
+        b: B
+        c: t.Optional[t.List[int]] = None
+
+    flat_data = {"a": "test", "b": 4, "c": [1, 2, 3]}
+    c = extract_config(flat_data=flat_data, config=C)
+    expected_result = {"b": {"a": {"a": "test"}, "b": 4}, "c": [1, 2, 3]}
+    assert c.to_json(sort_keys=True) == json.dumps(expected_result, sort_keys=True)
+
+
+def test_extract_config_dataclass_list():
+    @dataclass
+    class C(BaseConfig):
+        c: bool
+        b: t.List[B] = field(default_factory=list)
+
+    flat_data = {"a": "test", "c": True}
+    c = extract_config(flat_data=flat_data, config=C)
+    expected_result = {"b": [], "c": True}
+    assert c.to_json(sort_keys=True) == json.dumps(expected_result, sort_keys=True)
+
+
+def test_extract_config_dict():
+    @dataclass
+    class C(BaseConfig):
+        c: bool
+        b: t.Dict[str, B] = field(default_factory=dict)
+
+    flat_data = {"c": True}
+    c = extract_config(flat_data=flat_data, config=C)
+    expected_result = {"c": True, "b": {}}
+    assert c.to_json(sort_keys=True) == json.dumps(expected_result, sort_keys=True)
+
+
+def test_json_to_dict_valid_json():
+    json_string = '{"key": "value"}'
+    expected_result = {"key": "value"}
+    assert json_to_dict(json_string) == expected_result
+    assert isinstance(json_to_dict(json_string), dict)
+
+
+def test_json_to_dict_malformed_json():
+    json_string = '{"key": "value"'
+    expected_result = '{"key": "value"'
+    assert json_to_dict(json_string) == expected_result
+    assert isinstance(json_to_dict(json_string), str)
+
+
+def test_json_to_dict_single_quotes():
+    json_string = "{'key': 'value'}"
+    expected_result = {"key": "value"}
+    assert json_to_dict(json_string) == expected_result
+    assert isinstance(json_to_dict(json_string), dict)
+
+
+def test_json_to_dict_path():
+    json_string = "/path/to/file.json"
+    expected_result = "/path/to/file.json"
+    assert json_to_dict(json_string) == expected_result
+    assert isinstance(json_to_dict(json_string), str)
+
+
+def test_ensure_isoformat_datetime_for_datetime():
+    dt = ensure_isoformat_datetime(datetime(2021, 1, 1, 12, 0, 0))
+    assert dt == "2021-01-01T12:00:00"
+
+
+def test_ensure_isoformat_datetime_for_datetime_with_tz():
+    dt = ensure_isoformat_datetime(datetime(2021, 1, 1, 12, 0, 0, tzinfo=pytz.UTC))
+    assert dt == "2021-01-01T12:00:00+00:00"
+
+
+def test_ensure_isoformat_datetime_for_string():
+    dt = ensure_isoformat_datetime("2021-01-01T12:00:00")
+    assert dt == "2021-01-01T12:00:00"
+
+
+def test_ensure_isoformat_datetime_for_string2():
+    dt = ensure_isoformat_datetime("2021-01-01T12:00:00+00:00")
+    assert dt == "2021-01-01T12:00:00+00:00"
+
+
+def test_ensure_isoformat_datetime_fails_on_string():
+    with pytest.raises(ValueError):
+        ensure_isoformat_datetime("bad timestamp")
+
+
+def test_ensure_isoformat_datetime_fails_on_int():
+    with pytest.raises(TypeError):
+        ensure_isoformat_datetime(1111)
+
+
+def test_truncate_string_bytes_return_truncated_string():
+    test_string = "abcdef안녕하세요ghijklmn방갑습니opqrstu 더 길어지면 안되는 문자열vwxyz"
+    max_bytes = 11
+    result = truncate_string_bytes(test_string, max_bytes)
+    assert result == "abcdef안"
+    assert len(result.encode("utf-8")) <= max_bytes
+
+
+def test_truncate_string_bytes_return_untouched_string():
+    test_string = "abcdef"
+    max_bytes = 11
+    result = truncate_string_bytes(test_string, max_bytes)
+    assert result == "abcdef"
+    assert len(result.encode("utf-8")) <= max_bytes

test/unit/v2/chunkers/test_chunkers.py

@@ -0,0 +1,49 @@
+import random
+
+import faker
+import pytest
+
+from unstructured_ingest.v2.processes.chunker import Chunker, ChunkerConfig
+
+fake = faker.Faker()
+
+
+def generate_chunker_config_params() -> dict:
+    params = {}
+    random_val = random.random()
+    if random_val < 0.5:
+        params["chunking_strategy"] = fake.word() if random.random() < 0.5 else None
+        params["chunk_combine_text_under_n_chars"] = (
+            fake.random_int() if random.random() < 0.5 else None
+        )
+        params["chunk_include_orig_elements"] = fake.boolean() if random.random() < 0.5 else None
+        params["chunk_max_characters"] = fake.random_int()
+        params["chunk_multipage_sections"] = fake.boolean()
+        params["chunk_new_after_n_chars"] = fake.random_int() if random.random() < 0.5 else None
+        params["chunk_overlap"] = fake.random_int() if random.random() < 0.5 else None
+        params["chunk_overlap_all"] = fake.boolean() if random.random() < 0.5 else None
+    if random_val < 0.5:
+        params["chunk_by_api"] = True
+        params["chunking_endpoint"] = fake.url()
+        params["chunk_api_key"] = fake.password()
+    else:
+        params["chunk_by_api"] = False
+
+    return params
+
+
+@pytest.mark.parametrize(
+    "partition_config_params", [generate_chunker_config_params() for i in range(10)]
+)
+def test_chunker_config(partition_config_params: dict):
+    chunker_config = ChunkerConfig.model_validate(partition_config_params)
+    assert chunker_config
+
+
+@pytest.mark.parametrize(
+    "partition_config_params", [generate_chunker_config_params() for i in range(10)]
+)
+def test_chunker(partition_config_params: dict):
+    chunker_config = ChunkerConfig.model_validate(partition_config_params)
+    chunker = Chunker(config=chunker_config)
+    assert chunker

test/unit/v2/connectors/test_confluence.py

@@ -0,0 +1,39 @@
+import pytest
+from pydantic import ValidationError
+
+from unstructured_ingest.v2.processes.connectors.confluence import (
+    ConfluenceAccessConfig,
+    ConfluenceConnectionConfig,
+)
+
+
+def test_connection_config_multiple_auth():
+    with pytest.raises(ValidationError):
+        ConfluenceConnectionConfig(
+            access_config=ConfluenceAccessConfig(
+                api_token="api_token",
+                access_token="access_token",
+            ),
+            user_email="user_email",
+            url="url",
+        )
+
+
+def test_connection_config_no_auth():
+    with pytest.raises(ValidationError):
+        ConfluenceConnectionConfig(access_config=ConfluenceAccessConfig(), url="url")
+
+
+def test_connection_config_basic_auth():
+    ConfluenceConnectionConfig(
+        access_config=ConfluenceAccessConfig(api_token="api_token"),
+        url="url",
+        user_email="user_email",
+    )
+
+
+def test_connection_config_pat_auth():
+    ConfluenceConnectionConfig(
+        access_config=ConfluenceAccessConfig(access_token="access_token"),
+        url="url",
+    )

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