unstructured-ingest 0.0.25__py3-none-any.whl → 0.1.1__py3-none-any.whl

test/integration/embedders/test_voyageai.py

@@ -0,0 +1,41 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.voyageai import VoyageAIEmbeddingConfig, VoyageAIEmbeddingEncoder
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "VOYAGEAI_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+@requires_env(API_KEY)
+def test_voyageai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder_config = EmbedderConfig(embedding_provider="voyageai", embedding_api_key=api_key)
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY)
+def test_raw_voyageai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder = VoyageAIEmbeddingEncoder(
+        config=VoyageAIEmbeddingConfig(
+            api_key=api_key,
+        )
+    )
+    validate_raw_embedder(
+        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1024,)
+    )

test/integration/embedders/togetherai.py

@@ -0,0 +1,43 @@
+import json
+import os
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.togetherai import (
+    TogetherAIEmbeddingConfig,
+    TogetherAIEmbeddingEncoder,
+)
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "TOGETHERAI_API_KEY"
+
+
+def get_api_key() -> str:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    return api_key
+
+
+@requires_env(API_KEY)
+def test_togetherai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder_config = EmbedderConfig(embedding_provider="togetherai", embedding_api_key=api_key)
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY)
+def test_raw_togetherai_embedder(embedder_file: Path):
+    api_key = get_api_key()
+    embedder = TogetherAIEmbeddingEncoder(config=TogetherAIEmbeddingConfig(api_key=api_key))
+    validate_raw_embedder(
+        embedder=embedder,
+        embedder_file=embedder_file,
+        expected_dimensions=(768,),
+        expected_is_unit_vector=False,
+    )

test/integration/embedders/utils.py

@@ -0,0 +1,44 @@
+import json
+from pathlib import Path
+from typing import Optional
+
+from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder
+
+
+def validate_embedding_output(original_elements: list[dict], output_elements: list[dict]):
+    """
+    Make sure the following characteristics are met:
+    * The same number of elements are returned
+    * For each element that had text, an embeddings entry was added in the output
+    * Other than the embedding, nothing about the element was changed
+    """
+    assert len(original_elements) == len(output_elements)
+    for original_element, output_element in zip(original_elements, output_elements):
+        if original_element.get("text"):
+            assert output_element.get("embeddings", None)
+        output_element.pop("embeddings", None)
+        assert original_element == output_element
+
+
+def validate_raw_embedder(
+    embedder: BaseEmbeddingEncoder,
+    embedder_file: Path,
+    expected_dimensions: Optional[tuple[int, ...]] = None,
+    expected_is_unit_vector: bool = True,
+):
+    with open(embedder_file) as f:
+        elements = json.load(f)
+    all_text = [element["text"] for element in elements]
+    single_text = all_text[0]
+    num_of_dimensions = embedder.num_of_dimensions
+    if expected_dimensions:
+        assert (
+            num_of_dimensions == expected_dimensions
+        ), f"number of dimensions {num_of_dimensions} didn't match expected: {expected_dimensions}"
+    is_unit_vector = embedder.is_unit_vector
+    assert is_unit_vector == expected_is_unit_vector
+    single_embedding = embedder.embed_query(query=single_text)
+    expected_length = num_of_dimensions[0]
+    assert len(single_embedding) == expected_length
+    embedded_elements = embedder.embed_documents(elements=elements)
+    validate_embedding_output(original_elements=elements, output_elements=embedded_elements)

test/integration/partitioners/test_partitioner.py

@@ -0,0 +1,75 @@
+import json
+import os
+from pathlib import Path
+
+import pytest
+from unstructured_client.models.errors.sdkerror import SDKError
+
+from test.integration.utils import requires_env
+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(SDKError):
+        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,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