kiln-ai 0.20.1__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/adapters/extractors/test_base_extractor.py

@@ -0,0 +1,244 @@
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai.adapters.extractors.base_extractor import (
+    BaseExtractor,
+    ExtractionInput,
+    ExtractionOutput,
+)
+from kiln_ai.datamodel.extraction import ExtractorConfig, ExtractorType, OutputFormat
+
+
+class MockBaseExtractor(BaseExtractor):
+    async def _extract(self, input: ExtractionInput) -> ExtractionOutput:
+        return ExtractionOutput(
+            is_passthrough=False,
+            content="mock concrete extractor output",
+            content_format=OutputFormat.MARKDOWN,
+        )
+
+
+@pytest.fixture
+def mock_litellm_properties():
+    return {
+        "prompt_document": "mock prompt for document",
+        "prompt_image": "mock prompt for image",
+        "prompt_video": "mock prompt for video",
+        "prompt_audio": "mock prompt for audio",
+    }
+
+
+@pytest.fixture
+def mock_extractor(mock_litellm_properties):
+    return MockBaseExtractor(
+        ExtractorConfig(
+            name="mock",
+            model_provider_name="gemini_api",
+            model_name="gemini-2.0-flash",
+            extractor_type=ExtractorType.LITELLM,
+            output_format=OutputFormat.MARKDOWN,
+            properties=mock_litellm_properties,
+        )
+    )
+
+
+def mock_extractor_with_passthroughs(
+    properties: dict[str, Any],
+    mimetypes: list[OutputFormat],
+    output_format: OutputFormat,
+):
+    return MockBaseExtractor(
+        ExtractorConfig(
+            name="mock",
+            model_provider_name="gemini_api",
+            model_name="gemini-2.0-flash",
+            extractor_type=ExtractorType.LITELLM,
+            passthrough_mimetypes=mimetypes,
+            output_format=output_format,
+            properties=properties,
+        )
+    )
+
+
+def test_should_passthrough(mock_litellm_properties):
+    extractor = mock_extractor_with_passthroughs(
+        mock_litellm_properties,
+        [OutputFormat.TEXT, OutputFormat.MARKDOWN],
+        OutputFormat.TEXT,
+    )
+
+    # should passthrough
+    assert extractor._should_passthrough("text/plain")
+    assert extractor._should_passthrough("text/markdown")
+
+    # should not passthrough
+    assert not extractor._should_passthrough("image/png")
+    assert not extractor._should_passthrough("application/pdf")
+    assert not extractor._should_passthrough("text/html")
+    assert not extractor._should_passthrough("image/jpeg")
+
+
+async def test_extract_passthrough(mock_litellm_properties):
+    """
+    Tests that when a file's MIME type is configured for passthrough, the extractor skips
+    the concrete extraction method and returns the file's contents directly with the
+    correct passthrough output format.
+    """
+    extractor = mock_extractor_with_passthroughs(
+        mock_litellm_properties,
+        [OutputFormat.TEXT, OutputFormat.MARKDOWN],
+        OutputFormat.TEXT,
+    )
+    with (
+        patch.object(
+            extractor,
+            "_extract",
+            return_value=ExtractionOutput(
+                is_passthrough=False,
+                content="mock concrete extractor output",
+                content_format=OutputFormat.TEXT,
+            ),
+        ) as mock_extract,
+        patch(
+            "pathlib.Path.read_text",
+            return_value=b"test content",
+        ),
+    ):
+        result = await extractor.extract(
+            ExtractionInput(
+                path="test.txt",
+                mime_type="text/plain",
+            )
+        )
+
+        # Verify _extract was not called
+        mock_extract.assert_not_called()
+
+        # Verify correct passthrough result
+        assert result.is_passthrough
+        assert result.content == "test content"
+        assert result.content_format == OutputFormat.TEXT
+
+
+@pytest.mark.parametrize(
+    "output_format",
+    [
+        "text/plain",
+        "text/markdown",
+    ],
+)
+async def test_extract_passthrough_output_format(
+    mock_litellm_properties, output_format
+):
+    extractor = mock_extractor_with_passthroughs(
+        mock_litellm_properties,
+        [OutputFormat.TEXT, OutputFormat.MARKDOWN],
+        output_format,
+    )
+    with (
+        patch.object(
+            extractor,
+            "_extract",
+            return_value=ExtractionOutput(
+                is_passthrough=False,
+                content="mock concrete extractor output",
+                content_format=output_format,
+            ),
+        ) as mock_extract,
+        patch(
+            "pathlib.Path.read_text",
+            return_value="test content",
+        ),
+    ):
+        result = await extractor.extract(
+            ExtractionInput(
+                path="test.txt",
+                mime_type="text/plain",
+            )
+        )
+
+        # Verify _extract was not called
+        mock_extract.assert_not_called()
+
+        # Verify correct passthrough result
+        assert result.is_passthrough
+        assert result.content == "test content"
+        assert result.content_format == output_format
+
+
+@pytest.mark.parametrize(
+    "path, mime_type, output_format",
+    [
+        ("test.mp3", "audio/mpeg", OutputFormat.TEXT),
+        ("test.png", "image/png", OutputFormat.TEXT),
+        ("test.pdf", "application/pdf", OutputFormat.TEXT),
+        ("test.txt", "text/plain", OutputFormat.MARKDOWN),
+        ("test.txt", "text/markdown", OutputFormat.MARKDOWN),
+        ("test.html", "text/html", OutputFormat.MARKDOWN),
+    ],
+)
+async def test_extract_non_passthrough(
+    mock_extractor, path: str, mime_type: str, output_format: OutputFormat
+):
+    with (
+        patch.object(
+            mock_extractor,
+            "_extract",
+            return_value=ExtractionOutput(
+                is_passthrough=False,
+                content="mock concrete extractor output",
+                content_format=output_format,
+            ),
+        ) as mock_extract,
+    ):
+        # first we call the base class extract method
+        result = await mock_extractor.extract(
+            ExtractionInput(
+                path=path,
+                mime_type=mime_type,
+            )
+        )
+
+        # then we call the subclass _extract method and add validated mime_type
+        mock_extract.assert_called_once_with(
+            ExtractionInput(
+                path=path,
+                mime_type=mime_type,
+            )
+        )
+
+        assert not result.is_passthrough
+        assert result.content == "mock concrete extractor output"
+        assert result.content_format == output_format
+
+
+async def test_default_output_format(mock_litellm_properties):
+    config = ExtractorConfig(
+        name="mock",
+        model_provider_name="gemini_api",
+        model_name="gemini-2.0-flash",
+        extractor_type=ExtractorType.LITELLM,
+        properties=mock_litellm_properties,
+    )
+    assert config.output_format == OutputFormat.MARKDOWN
+
+
+async def test_extract_failure_from_concrete_extractor(mock_extractor):
+    with patch.object(
+        mock_extractor,
+        "_extract",
+        side_effect=Exception("error from concrete extractor"),
+    ):
+        with pytest.raises(ValueError, match="error from concrete extractor"):
+            await mock_extractor.extract(
+                ExtractionInput(
+                    path="test.txt",
+                    mime_type="text/plain",
+                )
+            )
+
+
+async def test_output_format(mock_extractor):
+    assert mock_extractor.output_format() == OutputFormat.MARKDOWN

kiln_ai/adapters/extractors/test_encoding.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+
+import pytest
+
+from conftest import MockFileFactoryMimeType
+from kiln_ai.adapters.extractors.encoding import from_base64_url, to_base64_url
+
+
+async def test_to_base64_url(mock_file_factory):
+    mock_file = mock_file_factory(MockFileFactoryMimeType.JPEG)
+
+    byte_data = Path(mock_file).read_bytes()
+
+    # encode the byte data
+    base64_url = to_base64_url("image/jpeg", byte_data)
+    assert base64_url.startswith("data:image/jpeg;base64,")
+
+    # decode the base64 url
+    assert from_base64_url(base64_url) == byte_data
+
+
+def test_from_base64_url_invalid_format_no_data_prefix():
+    """Test that from_base64_url raises ValueError when input doesn't start with 'data:'"""
+    with pytest.raises(ValueError, match="Invalid base64 URL format"):
+        from_base64_url("not-a-data-url")
+
+
+def test_from_base64_url_invalid_format_no_comma():
+    """Test that from_base64_url raises ValueError when input doesn't contain a comma"""
+    with pytest.raises(ValueError, match="Invalid base64 URL format"):
+        from_base64_url("data:image/jpeg;base64")
+
+
+def test_from_base64_url_invalid_parts():
+    """Test that from_base64_url raises ValueError when splitting by comma doesn't result in exactly 2 parts"""
+    with pytest.raises(ValueError, match="Invalid base64 URL format"):
+        from_base64_url("data:image/jpeg;base64,part1,part2")
+
+
+def test_from_base64_url_base64_decode_failure():
+    """Test that from_base64_url raises ValueError when base64 decoding fails"""
+    with pytest.raises(ValueError, match="Failed to decode base64 data"):
+        from_base64_url("data:image/jpeg;base64,invalid-base64-data!")
+
+
+def test_from_base64_url_valid_format():
+    """Test that from_base64_url works with valid base64 URL format"""
+    # Create a simple valid base64 URL
+    test_data = b"Hello, World!"
+    base64_encoded = "SGVsbG8sIFdvcmxkIQ=="
+    base64_url = f"data:text/plain;base64,{base64_encoded}"
+
+    result = from_base64_url(base64_url)
+    assert result == test_data

kiln_ai/adapters/extractors/test_extractor_registry.py

@@ -0,0 +1,181 @@
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai.adapters.extractors.extractor_registry import extractor_adapter_from_type
+from kiln_ai.adapters.extractors.litellm_extractor import LitellmExtractor
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.provider_tools import LiteLlmCoreConfig
+from kiln_ai.datamodel.extraction import ExtractorConfig, ExtractorType
+
+
+@pytest.fixture
+def mock_provider_configs():
+    with patch("kiln_ai.utils.config.Config.shared") as mock_config:
+        mock_config.return_value.open_ai_api_key = "test-openai-key"
+        mock_config.return_value.gemini_api_key = "test-gemini-key"
+        mock_config.return_value.anthropic_api_key = "test-anthropic-key"
+        mock_config.return_value.bedrock_access_key = "test-amazon-bedrock-key"
+        mock_config.return_value.bedrock_secret_key = "test-amazon-bedrock-secret-key"
+        mock_config.return_value.fireworks_api_key = "test-fireworks-key"
+        mock_config.return_value.groq_api_key = "test-groq-key"
+        mock_config.return_value.huggingface_api_key = "test-huggingface-key"
+        yield mock_config
+
+
+def test_extractor_adapter_from_type(mock_provider_configs):
+    extractor = extractor_adapter_from_type(
+        ExtractorType.LITELLM,
+        ExtractorConfig(
+            name="test-extractor",
+            extractor_type=ExtractorType.LITELLM,
+            model_provider_name="gemini_api",
+            model_name="gemini-2.0-flash",
+            properties={
+                "prompt_document": "Extract the text from the document",
+                "prompt_image": "Extract the text from the image",
+                "prompt_video": "Extract the text from the video",
+                "prompt_audio": "Extract the text from the audio",
+            },
+        ),
+    )
+    assert isinstance(extractor, LitellmExtractor)
+    assert extractor.extractor_config.model_name == "gemini-2.0-flash"
+    assert extractor.extractor_config.model_provider_name == "gemini_api"
+
+
+@patch(
+    "kiln_ai.adapters.extractors.extractor_registry.lite_llm_core_config_for_provider"
+)
+def test_extractor_adapter_from_type_uses_litellm_core_config(
+    mock_get_litellm_core_config,
+):
+    """Test that extractor receives auth details from provider_tools."""
+    mock_litellm_core_config = LiteLlmCoreConfig(
+        base_url="https://test.com",
+        additional_body_options={"api_key": "test-key"},
+        default_headers={},
+    )
+    mock_get_litellm_core_config.return_value = mock_litellm_core_config
+
+    extractor = extractor_adapter_from_type(
+        ExtractorType.LITELLM,
+        ExtractorConfig(
+            name="test-extractor",
+            extractor_type=ExtractorType.LITELLM,
+            model_provider_name="openai",
+            model_name="gpt-4",
+            properties={
+                "prompt_document": "Extract the text from the document",
+                "prompt_image": "Extract the text from the image",
+                "prompt_video": "Extract the text from the video",
+                "prompt_audio": "Extract the text from the audio",
+            },
+        ),
+    )
+
+    assert isinstance(extractor, LitellmExtractor)
+    assert extractor.litellm_core_config == mock_litellm_core_config
+    mock_get_litellm_core_config.assert_called_once_with(ModelProviderName.openai)
+
+
+def test_extractor_adapter_from_type_invalid_provider():
+    """Test that invalid model provider names raise a clear error."""
+    with pytest.raises(
+        ValueError, match="Unsupported model provider name: invalid_provider"
+    ):
+        extractor_adapter_from_type(
+            ExtractorType.LITELLM,
+            ExtractorConfig(
+                name="test-extractor",
+                extractor_type=ExtractorType.LITELLM,
+                model_provider_name="invalid_provider",
+                model_name="some-model",
+                properties={
+                    "prompt_document": "Extract the text from the document",
+                    "prompt_image": "Extract the text from the image",
+                    "prompt_video": "Extract the text from the video",
+                    "prompt_audio": "Extract the text from the audio",
+                },
+            ),
+        )
+
+
+def test_extractor_adapter_from_type_invalid():
+    with pytest.raises(ValueError, match="Unhandled enum value: fake_type"):
+        extractor_adapter_from_type(
+            "fake_type",
+            ExtractorConfig(
+                name="test-extractor",
+                extractor_type=ExtractorType.LITELLM,
+                model_provider_name="invalid_provider",
+                model_name="some-model",
+                properties={
+                    "prompt_document": "Extract the text from the document",
+                    "prompt_image": "Extract the text from the image",
+                    "prompt_video": "Extract the text from the video",
+                    "prompt_audio": "Extract the text from the audio",
+                },
+            ),
+        )
+
+
+@pytest.mark.parametrize(
+    "provider_name",
+    [
+        "openai",
+        "anthropic",
+        "gemini_api",
+        "amazon_bedrock",
+        "fireworks_ai",
+        "groq",
+        "huggingface",
+    ],
+)
+def test_extractor_adapter_from_type_different_providers(
+    provider_name, mock_provider_configs
+):
+    """Test that different providers work correctly."""
+    extractor = extractor_adapter_from_type(
+        ExtractorType.LITELLM,
+        ExtractorConfig(
+            name="test-extractor",
+            extractor_type=ExtractorType.LITELLM,
+            model_provider_name=provider_name,
+            model_name="test-model",
+            properties={
+                "prompt_document": "Extract the text from the document",
+                "prompt_image": "Extract the text from the image",
+                "prompt_video": "Extract the text from the video",
+                "prompt_audio": "Extract the text from the audio",
+            },
+        ),
+    )
+
+    assert isinstance(extractor, LitellmExtractor)
+    assert extractor.extractor_config.model_provider_name == provider_name
+
+
+def test_extractor_adapter_from_type_no_config_found(mock_provider_configs):
+    with patch(
+        "kiln_ai.adapters.extractors.extractor_registry.lite_llm_core_config_for_provider"
+    ) as mock_lite_llm_core_config_for_provider:
+        mock_lite_llm_core_config_for_provider.return_value = None
+        with pytest.raises(
+            ValueError, match="No configuration found for core provider: openai"
+        ):
+            extractor_adapter_from_type(
+                ExtractorType.LITELLM,
+                ExtractorConfig(
+                    name="test-extractor",
+                    extractor_type=ExtractorType.LITELLM,
+                    model_provider_name="openai",
+                    model_name="gpt-4",
+                    properties={
+                        "prompt_document": "Extract the text from the document",
+                        "prompt_image": "Extract the text from the image",
+                        "prompt_video": "Extract the text from the video",
+                        "prompt_audio": "Extract the text from the audio",
+                    },
+                ),
+            )

kiln_ai/adapters/extractors/test_extractor_runner.py

@@ -0,0 +1,181 @@
+from unittest.mock import AsyncMock
+
+import pytest
+
+from conftest import MockFileFactoryMimeType
+from kiln_ai.adapters.extractors.extractor_runner import ExtractorRunner
+from kiln_ai.datamodel.basemodel import KilnAttachmentModel
+from kiln_ai.datamodel.extraction import (
+    Document,
+    Extraction,
+    ExtractionSource,
+    ExtractorConfig,
+    ExtractorType,
+    FileInfo,
+    Kind,
+    OutputFormat,
+)
+from kiln_ai.datamodel.project import Project
+
+
+@pytest.fixture
+def mock_project(tmp_path):
+    project = Project(
+        name="test",
+        description="test",
+        path=tmp_path / "project.kiln",
+    )
+    project.save_to_file()
+    return project
+
+
+@pytest.fixture
+def mock_extractor_config(mock_project):
+    extractor_config = ExtractorConfig(
+        name="test",
+        description="test",
+        output_format=OutputFormat.TEXT,
+        passthrough_mimetypes=[],
+        extractor_type=ExtractorType.LITELLM,
+        model_provider_name="gemini_api",
+        model_name="gemini-2.0-flash",
+        parent=mock_project,
+        properties={
+            "prompt_document": "Extract the text from the document",
+            "prompt_image": "Extract the text from the image",
+            "prompt_video": "Extract the text from the video",
+            "prompt_audio": "Extract the text from the audio",
+        },
+    )
+    extractor_config.save_to_file()
+    return extractor_config
+
+
+@pytest.fixture
+def mock_document(mock_project, mock_file_factory) -> Document:
+    test_pdf_file = mock_file_factory(MockFileFactoryMimeType.PDF)
+    document = Document(
+        name="test",
+        description="test",
+        kind=Kind.DOCUMENT,
+        original_file=FileInfo(
+            filename="test.pdf",
+            size=100,
+            mime_type="application/pdf",
+            attachment=KilnAttachmentModel.from_file(test_pdf_file),
+        ),
+        parent=mock_project,
+    )
+    document.save_to_file()
+    return document
+
+
+@pytest.fixture
+def mock_extractor_runner(mock_extractor_config, mock_document):
+    return ExtractorRunner(
+        extractor_configs=[mock_extractor_config],
+        documents=[mock_document],
+    )
+
+
+# Test with and without concurrency
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_extractor_runner_status_updates(
+    mock_extractor_runner, concurrency
+):
+    # Real async testing!
+
+    job_count = 50
+    # Job objects are not the right type, but since we're mocking run_job, it doesn't matter
+    jobs = [{} for _ in range(job_count)]
+
+    # Mock collect_tasks to return our fake jobs
+    mock_extractor_runner.collect_jobs = lambda: jobs
+
+    # Mock run_job to return True immediately
+    mock_extractor_runner.run_job = AsyncMock(return_value=True)
+
+    # Expect the status updates in order, and 1 for each job
+    expected_completed_count = 0
+    async for progress in mock_extractor_runner.run(concurrency=concurrency):
+        assert progress.complete == expected_completed_count
+        expected_completed_count += 1
+        assert progress.errors == 0
+        assert progress.total == job_count
+
+    # Verify last status update was complete
+    assert expected_completed_count == job_count + 1
+
+    # Verify run_job was called for each job
+    assert mock_extractor_runner.run_job.call_count == job_count
+
+
+def test_collect_jobs_excludes_already_run_extraction(
+    mock_extractor_runner, mock_document, mock_extractor_config
+):
+    """Test that already run documents are excluded"""
+    Extraction(
+        parent=mock_document,
+        source=ExtractionSource.PROCESSED,
+        extractor_config_id="other-extractor-config-id",
+        output=KilnAttachmentModel.from_data("test extraction output", "text/plain"),
+    ).save_to_file()
+
+    # should get the one job, since the document was not already extracted with this extractor config
+    jobs = mock_extractor_runner.collect_jobs()
+    assert len(jobs) == 1
+    assert jobs[0].doc.id == mock_document.id
+    assert jobs[0].extractor_config.id == mock_extractor_config.id
+
+    # Create an extraction for this document
+    Extraction(
+        parent=mock_document,
+        source=ExtractionSource.PROCESSED,
+        extractor_config_id=mock_extractor_config.id,
+        output=KilnAttachmentModel.from_data("test extraction output", "text/plain"),
+    ).save_to_file()
+
+    jobs = mock_extractor_runner.collect_jobs()
+
+    # should now get no jobs since the document was already extracted with this extractor config
+    assert len(jobs) == 0
+
+
+def test_collect_jobs_multiple_extractor_configs(
+    mock_extractor_runner,
+    mock_document,
+    mock_extractor_config,
+    mock_project,
+):
+    """Test handling multiple extractor configs"""
+    second_config = ExtractorConfig(
+        name="test2",
+        description="test2",
+        output_format=OutputFormat.TEXT,
+        passthrough_mimetypes=[],
+        extractor_type=ExtractorType.LITELLM,
+        parent=mock_project,
+        model_provider_name="gemini_api",
+        model_name="gemini-2.0-flash",
+        properties={
+            "prompt_document": "Extract the text from the document",
+            "prompt_image": "Extract the text from the image",
+            "prompt_video": "Extract the text from the video",
+            "prompt_audio": "Extract the text from the audio",
+        },
+    )
+    second_config.save_to_file()
+
+    runner = ExtractorRunner(
+        extractor_configs=[mock_extractor_config, second_config],
+        documents=[mock_document],
+    )
+    jobs = runner.collect_jobs()
+
+    # Should get 2 jobs, one for each config
+    assert len(jobs) == 2
+    assert {job.extractor_config.id for job in jobs} == {
+        second_config.id,
+        mock_extractor_config.id,
+    }