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

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,
+    }