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

kiln_ai/utils/test_lock.py

@@ -0,0 +1,185 @@
+import asyncio
+
+from .lock import AsyncLockManager, shared_async_lock_manager
+
+
+async def test_same_key_returns_same_lock():
+    """Test that the same key returns the same lock object."""
+    locks = AsyncLockManager()
+
+    # Test that the same key gets the same lock entry
+    async with locks.acquire("test_key"):
+        # The lock should exist in the manager
+        snapshot = await locks.snapshot()
+        assert "test_key" in snapshot
+
+
+async def test_different_keys_return_different_locks():
+    """Test that different keys return different lock objects."""
+    locks = AsyncLockManager()
+
+    # Use different keys
+    async with locks.acquire("key1"):
+        async with locks.acquire("key2"):
+            snapshot = await locks.snapshot()
+            assert "key1" in snapshot
+            assert "key2" in snapshot
+
+
+async def test_lock_functionality():
+    """Test that the locks actually provide mutual exclusion."""
+    results = []
+    locks = AsyncLockManager()
+
+    async def worker(worker_id: int):
+        async with locks.acquire("shared_resource"):
+            # Record start
+            results.append(f"worker_{worker_id}_start")
+            await asyncio.sleep(0.1)  # Simulate work
+            # Record end
+            results.append(f"worker_{worker_id}_end")
+
+    # Run multiple workers concurrently
+    await asyncio.gather(*[worker(i) for i in range(3)])
+
+    # Verify that the work was done exclusively
+    # Each worker's start should be immediately followed by its end
+    i = 0
+    while i < len(results):
+        start_event = results[i]
+        end_event = results[i + 1]
+
+        # Extract worker ID from start event
+        worker_id = start_event.split("_")[1]
+        expected_end = f"worker_{worker_id}_end"
+
+        assert end_event == expected_end, (
+            f"Non-exclusive access detected. Expected {expected_end}, got {end_event}. Full results: {results}"
+        )
+        i += 2
+
+
+async def test_lock_cleanup():
+    """Test that locks are automatically cleaned up when no longer needed."""
+    locks = AsyncLockManager()
+
+    # Use a lock
+    async with locks.acquire("cleanup_test"):
+        pass
+
+    # Check that the lock was cleaned up
+    snapshot = await locks.snapshot()
+    assert "cleanup_test" not in snapshot
+
+
+async def test_multiple_holders_cleanup():
+    """Test that locks are cleaned up when multiple holders finish."""
+    locks = AsyncLockManager()
+
+    # Create multiple tasks that will hold the lock sequentially
+    async def holder(holder_id: int):
+        async with locks.acquire("multi_holder"):
+            await asyncio.sleep(0.05)
+            return f"holder_{holder_id}_done"
+
+    # Run multiple holders sequentially (not concurrently to avoid deadlock)
+    results = []
+    for i in range(3):
+        result = await holder(i)
+        results.append(result)
+
+    # Check that all holders completed
+    assert len(results) == 3
+    assert all(result.startswith("holder_") for result in results)
+
+    # Check that the lock was cleaned up
+    snapshot = await locks.snapshot()
+    assert "multi_holder" not in snapshot
+
+
+async def test_global_instance():
+    """Test that the global shared_async_lock_manager instance works correctly."""
+    results = []
+
+    async def worker(worker_id: int):
+        async with shared_async_lock_manager.acquire("global_test"):
+            results.append(f"worker_{worker_id}_start")
+            await asyncio.sleep(0.05)
+            results.append(f"worker_{worker_id}_end")
+
+    # Run multiple workers sequentially to avoid deadlock
+    for i in range(2):
+        await worker(i)
+
+    # Verify sequential access
+    assert len(results) == 4
+    assert results[0] == "worker_0_start"
+    assert results[1] == "worker_0_end"
+    assert results[2] == "worker_1_start"
+    assert results[3] == "worker_1_end"
+
+
+async def test_timeout():
+    """Test that timeout functionality works correctly."""
+    locks = AsyncLockManager()
+
+    # Hold the lock for a while
+    async def holder():
+        async with locks.acquire("timeout_test"):
+            await asyncio.sleep(0.3)
+
+    # Try to acquire with a short timeout
+    async def waiter():
+        try:
+            async with locks.acquire("timeout_test", timeout=0.1):
+                assert False, "Should have timed out"
+        except asyncio.TimeoutError:
+            return "timed_out"
+
+    # Start holder first
+    holder_task = asyncio.create_task(holder())
+    await asyncio.sleep(0.05)  # Let holder acquire the lock
+
+    # Then try waiter
+    result = await waiter()
+    assert result == "timed_out"
+
+    # Wait for holder to finish
+    await holder_task
+
+
+async def test_cancellation():
+    """Test that cancellation is handled correctly."""
+    locks = AsyncLockManager()
+
+    # Hold the lock
+    async def holder():
+        async with locks.acquire("cancel_test"):
+            await asyncio.sleep(0.3)
+
+    # Try to acquire but get cancelled
+    async def waiter():
+        try:
+            async with locks.acquire("cancel_test"):
+                assert False, "Should not acquire lock"
+        except asyncio.CancelledError:
+            return "cancelled"
+
+    # Start holder first
+    holder_task = asyncio.create_task(holder())
+    await asyncio.sleep(0.05)  # Let holder acquire the lock
+
+    # Start waiter and then cancel it
+    waiter_task = asyncio.create_task(waiter())
+    await asyncio.sleep(0.05)
+    waiter_task.cancel()
+
+    # Check result
+    try:
+        result = await waiter_task
+        assert result == "cancelled"
+    except asyncio.CancelledError:
+        pass  # Expected
+
+    # Wait for holder to finish
+    await holder_task

kiln_ai/utils/test_mime_type.py

@@ -0,0 +1,66 @@
+from kiln_ai.utils.mime_type import guess_mime_type
+
+
+def test_mov_files():
+    assert guess_mime_type("video.mov") == "video/quicktime"
+    assert guess_mime_type("my_video.mov") == "video/quicktime"
+    assert guess_mime_type("path/to/video.mov") == "video/quicktime"
+
+
+def test_mp3_files():
+    assert guess_mime_type("song.mp3") == "audio/mpeg"
+    assert guess_mime_type("music_file.mp3") == "audio/mpeg"
+    assert guess_mime_type("audio/track.mp3") == "audio/mpeg"
+
+
+def test_wav_files():
+    assert guess_mime_type("sound.wav") == "audio/wav"
+    assert guess_mime_type("audio_file.wav") == "audio/wav"
+    assert guess_mime_type("sounds/effect.wav") == "audio/wav"
+
+
+def test_mp4_files():
+    assert guess_mime_type("movie.mp4") == "video/mp4"
+    assert guess_mime_type("video_file.mp4") == "video/mp4"
+    assert guess_mime_type("videos/clip.mp4") == "video/mp4"
+
+
+def test_case_insensitive_extensions():
+    assert guess_mime_type("video.MOV") == "video/quicktime"
+    assert guess_mime_type("song.MP3") == "audio/mpeg"
+    assert guess_mime_type("sound.WAV") == "audio/wav"
+    assert guess_mime_type("movie.MP4") == "video/mp4"
+
+
+def test_standard_mimetypes_fallback():
+    assert guess_mime_type("document.pdf") == "application/pdf"
+    assert guess_mime_type("image.jpg") == "image/jpeg"
+    assert guess_mime_type("image.png") == "image/png"
+    assert guess_mime_type("text.txt") == "text/plain"
+    assert guess_mime_type("data.json") == "application/json"
+
+
+def test_unknown_extensions():
+    assert guess_mime_type("file.invalidmime") is None
+    assert guess_mime_type("no_extension") is None
+
+
+def test_edge_cases():
+    # Files with multiple dots
+    assert guess_mime_type("video.backup.mov") == "video/quicktime"
+    assert guess_mime_type("song.remix.mp3") == "audio/mpeg"
+
+    # Files with dots in the middle
+    assert guess_mime_type("my.video.mov") == "video/quicktime"
+    assert guess_mime_type("track.1.mp3") == "audio/mpeg"
+
+    # Empty filename
+    assert guess_mime_type("") is None
+
+    # Just extension
+    assert guess_mime_type(".mov") == "video/quicktime"
+    assert guess_mime_type(".mp3") == "audio/mpeg"
+
+
+def test_priority_order():
+    assert guess_mime_type("file.mov.mp3") == "audio/mpeg"

kiln_ai/utils/test_open_ai_types.py

@@ -8,9 +8,13 @@ from openai.types.chat import (
 from openai.types.chat import (
     ChatCompletionMessageParam as OpenAIChatCompletionMessageParam,
 )
+from openai.types.chat import (
+    ChatCompletionToolMessageParam as OpenAIChatCompletionToolMessageParam,
+)
 
 from kiln_ai.utils.open_ai_types import (
     ChatCompletionAssistantMessageParamWrapper,
+    ChatCompletionToolMessageParamWrapper,
 )
 from kiln_ai.utils.open_ai_types import (
     ChatCompletionMessageParam as KilnChatCompletionMessageParam,
@@ -45,10 +49,40 @@ def test_assistant_message_param_properties_match():
     )
 
 
+def test_tool_message_param_properties_match():
+    """
+    Test that ChatCompletionToolMessageParamWrapper has all the same properties
+    as OpenAI's ChatCompletionToolMessageParam, plus the kiln_task_tool_data property.
+
+    This will catch any changes to the OpenAI types that we haven't updated our wrapper for.
+    """
+    # Get annotations for both types
+    openai_annotations = OpenAIChatCompletionToolMessageParam.__annotations__
+    kiln_annotations = ChatCompletionToolMessageParamWrapper.__annotations__
+
+    # Check that both have the same property names
+    openai_properties = set(openai_annotations.keys())
+    kiln_properties = set(kiln_annotations.keys())
+
+    # Kiln task tool data is an added property. Confirm it's there and remove it from the comparison.
+    assert "kiln_task_tool_data" in kiln_properties, (
+        "Kiln should have kiln_task_tool_data"
+    )
+    kiln_properties.remove("kiln_task_tool_data")
+
+    assert openai_properties == kiln_properties, (
+        f"Property names don't match. "
+        f"OpenAI has: {openai_properties}, "
+        f"Kiln has: {kiln_properties}, "
+        f"Missing from Kiln: {openai_properties - kiln_properties}, "
+        f"Extra in Kiln: {kiln_properties - openai_properties}"
+    )
+
+
 def test_chat_completion_message_param_union_compatibility():
     """
     Test that our ChatCompletionMessageParam union contains the same types as OpenAI's,
-    except with our wrapper instead of the original assistant message param.
+    except with our wrappers instead of the original assistant and tool message params.
     """
     # Get the union members for both types
     openai_union_args = get_args(OpenAIChatCompletionMessageParam)
@@ -70,10 +104,16 @@ def test_chat_completion_message_param_union_compatibility():
     openai_type_names = {arg.__name__ for arg in openai_union_args}
     kiln_type_names = {arg.__name__ for arg in kiln_union_args}
 
-    # Expected difference: OpenAI has ChatCompletionAssistantMessageParam,
-    # Kiln has ChatCompletionAssistantMessageParamWrapper
-    expected_openai_only = {"ChatCompletionAssistantMessageParam"}
-    expected_kiln_only = {"ChatCompletionAssistantMessageParamWrapper"}
+    # Expected differences: OpenAI has ChatCompletionAssistantMessageParam and ChatCompletionToolMessageParam,
+    # Kiln has ChatCompletionAssistantMessageParamWrapper and ChatCompletionToolMessageParamWrapper
+    expected_openai_only = {
+        "ChatCompletionAssistantMessageParam",
+        "ChatCompletionToolMessageParam",
+    }
+    expected_kiln_only = {
+        "ChatCompletionAssistantMessageParamWrapper",
+        "ChatCompletionToolMessageParamWrapper",
+    }
 
     openai_only = openai_type_names - kiln_type_names
     kiln_only = kiln_type_names - openai_type_names
@@ -91,7 +131,6 @@ def test_chat_completion_message_param_union_compatibility():
         "ChatCompletionDeveloperMessageParam",
         "ChatCompletionSystemMessageParam",
         "ChatCompletionUserMessageParam",
-        "ChatCompletionToolMessageParam",
         "ChatCompletionFunctionMessageParam",
     }
 
@@ -100,17 +139,17 @@ def test_chat_completion_message_param_union_compatibility():
     )
 
 
-def test_wrapper_can_be_instantiated():
-    """Test that our wrapper can be instantiated with the same data as the original."""
-    # Create a sample message that should work with both types
-    sample_message: ChatCompletionAssistantMessageParamWrapper = {
+def test_assistant_message_wrapper_can_be_instantiated():
+    """Test that our assistant message wrapper can be instantiated with the same data as the original."""
+    # Test basic assistant message
+    sample_assistant_message: ChatCompletionAssistantMessageParamWrapper = {
         "role": "assistant",
         "content": "Hello, world!",
     }
 
     # This should work without type errors (runtime test)
-    assert sample_message["role"] == "assistant"
-    assert sample_message.get("content") == "Hello, world!"
+    assert sample_assistant_message["role"] == "assistant"
+    assert sample_assistant_message.get("content") == "Hello, world!"
 
     # Test with tool calls using List instead of Iterable
     sample_with_tools: ChatCompletionAssistantMessageParamWrapper = {
@@ -129,3 +168,40 @@ def test_wrapper_can_be_instantiated():
     tool_calls = sample_with_tools.get("tool_calls", [])
     if tool_calls:
         assert tool_calls[0]["id"] == "call_123"
+
+
+def test_tool_message_wrapper_can_be_instantiated():
+    """Test that our tool message wrapper can be instantiated with the same data as the original."""
+    # Test basic tool message
+    sample_tool_message: ChatCompletionToolMessageParamWrapper = {
+        "role": "tool",
+        "content": "Tool response",
+        "tool_call_id": "call_123",
+    }
+
+    assert sample_tool_message["role"] == "tool"
+    assert sample_tool_message.get("content") == "Tool response"
+    assert sample_tool_message.get("tool_call_id") == "call_123"
+
+    # Test with kiln_task_tool_data
+    sample_with_kiln_data: ChatCompletionToolMessageParamWrapper = {
+        "role": "tool",
+        "content": "Tool response",
+        "tool_call_id": "call_123",
+        "kiln_task_tool_data": "project_123:::tool_456:::task_789:::run_101",
+    }
+
+    assert (
+        sample_with_kiln_data.get("kiln_task_tool_data")
+        == "project_123:::tool_456:::task_789:::run_101"
+    )
+
+    # Test with kiln_task_tool_data as None
+    sample_with_none_kiln_data: ChatCompletionToolMessageParamWrapper = {
+        "role": "tool",
+        "content": "Tool response",
+        "tool_call_id": "call_123",
+        "kiln_task_tool_data": None,
+    }
+
+    assert sample_with_none_kiln_data.get("kiln_task_tool_data") is None

kiln_ai/utils/test_pdf_utils.py

@@ -0,0 +1,86 @@
+import tempfile
+from pathlib import Path
+
+import pytest
+from pypdf import PdfReader
+
+from conftest import MockFileFactoryMimeType
+from kiln_ai.utils.pdf_utils import convert_pdf_to_images, split_pdf_into_pages
+
+
+async def test_split_pdf_into_pages_success(mock_file_factory):
+    """Test that split_pdf_into_pages successfully splits a PDF into individual pages."""
+    test_file = mock_file_factory(MockFileFactoryMimeType.PDF)
+
+    async with split_pdf_into_pages(test_file) as page_paths:
+        # Verify we get the expected number of pages (test PDF has 2 pages)
+        assert len(page_paths) == 2
+
+        # Verify all page files exist
+        for page_path in page_paths:
+            assert page_path.exists()
+            assert page_path.suffix == ".pdf"
+
+        # Verify page files are named correctly
+        assert page_paths[0].name == "page_1.pdf"
+        assert page_paths[1].name == "page_2.pdf"
+
+        # Verify each page file is a valid PDF with exactly 1 page
+        for page_path in page_paths:
+            with open(page_path, "rb") as file:
+                reader = PdfReader(file)
+                assert len(reader.pages) == 1
+
+    # Verify cleanup: all page files should be removed after context exit
+    for page_path in page_paths:
+        assert not page_path.exists()
+
+
+async def test_split_pdf_into_pages_cleanup_on_exception(mock_file_factory):
+    """Test that temporary files are cleaned up even when an exception occurs during normal usage."""
+    test_file = mock_file_factory(MockFileFactoryMimeType.PDF)
+    captured_page_paths = []
+
+    # Test that cleanup happens even when an exception occurs during the with block
+    with pytest.raises(RuntimeError, match="Simulated error during usage"):
+        async with split_pdf_into_pages(test_file) as page_paths:
+            # Capture the page paths before the exception
+            captured_page_paths.extend(page_paths)
+            # Simulate an exception during normal usage of the context manager
+            raise RuntimeError("Simulated error during usage")
+
+    # Verify cleanup happened: the specific page files we created should be gone
+    for page_path in captured_page_paths:
+        assert not page_path.exists()
+
+    # Also verify the temporary directory itself is gone
+    if captured_page_paths:
+        temp_dir = captured_page_paths[0].parent
+        assert not temp_dir.exists()
+
+
+async def test_split_pdf_into_pages_temporary_directory_creation(mock_file_factory):
+    """Test that temporary directories are created with the correct prefix."""
+    test_file = mock_file_factory(MockFileFactoryMimeType.PDF)
+    captured_temp_dirs = []
+
+    async with split_pdf_into_pages(test_file) as page_paths:
+        # Check that page paths are in a directory with the expected prefix
+        temp_dir = page_paths[0].parent
+        captured_temp_dirs.append(temp_dir)
+        assert "kiln_pdf_pages_" in temp_dir.name
+        assert temp_dir.exists()
+
+    # Verify the temporary directory is cleaned up
+    for temp_dir in captured_temp_dirs:
+        assert not temp_dir.exists()
+
+
+async def test_convert_pdf_to_images(mock_file_factory):
+    """Test that convert_pdf_to_images successfully converts a PDF into individual images."""
+    test_file = mock_file_factory(MockFileFactoryMimeType.PDF)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        images = await convert_pdf_to_images(test_file, Path(temp_dir))
+        assert len(images) == 2
+        assert all(image.exists() for image in images)
+        assert all(image.suffix == ".png" for image in images)

kiln_ai/utils/test_uuid.py

@@ -0,0 +1,111 @@
+import uuid
+
+import pytest
+
+from kiln_ai.utils.uuid import string_to_uuid
+
+
+class TestStringToUuid:
+    """Test the string_to_uuid function for consistency and correctness."""
+
+    def test_same_string_produces_same_uuid(self):
+        """Test that the same string consistently produces the same UUID."""
+        test_string = "hello world"
+
+        uuid1 = string_to_uuid(test_string)
+        uuid2 = string_to_uuid(test_string)
+        uuid3 = string_to_uuid(test_string)
+
+        assert uuid1 == uuid2 == uuid3
+        assert isinstance(uuid1, uuid.UUID)
+
+    def test_different_strings_produce_different_uuids(self):
+        """Test that different strings produce different UUIDs."""
+        uuid1 = string_to_uuid("hello")
+        uuid2 = string_to_uuid("world")
+        uuid3 = string_to_uuid("hello world")
+
+        assert uuid1 != uuid2
+        assert uuid1 != uuid3
+        assert uuid2 != uuid3
+
+    def test_case_sensitivity(self):
+        """Test that string case affects the generated UUID."""
+        uuid_lower = string_to_uuid("hello")
+        uuid_upper = string_to_uuid("HELLO")
+        uuid_mixed = string_to_uuid("Hello")
+
+        assert uuid_lower != uuid_upper
+        assert uuid_lower != uuid_mixed
+        assert uuid_upper != uuid_mixed
+
+    @pytest.mark.parametrize(
+        "test_string",
+        [
+            "",
+            "a",
+            "test string with spaces",
+            "string_with_underscores",
+            "string-with-dashes",
+            "string.with.dots",
+            "string/with/slashes",
+            "string@with#special$characters!",
+            "1234567890",
+            "string with 数字 and unicode 🚀",
+            "\n\t\r",  # whitespace characters
+            "a" * 1000,  # very long string
+        ],
+    )
+    def test_various_string_inputs(self, test_string):
+        """Test that various string inputs produce consistent UUIDs."""
+        uuid1 = string_to_uuid(test_string)
+        uuid2 = string_to_uuid(test_string)
+
+        assert uuid1 == uuid2
+        assert isinstance(uuid1, uuid.UUID)
+
+    def test_uuid_format_is_valid(self):
+        """Test that the generated UUID is a valid UUID5."""
+        test_string = "test"
+        result_uuid = string_to_uuid(test_string)
+
+        # UUID5 should have version 5
+        assert result_uuid.version == 5
+
+        # Should be a valid UUID string format
+        uuid_str = str(result_uuid)
+        assert len(uuid_str) == 36
+        assert uuid_str.count("-") == 4
+
+        # Should be able to recreate UUID from string
+        recreated_uuid = uuid.UUID(uuid_str)
+        assert recreated_uuid == result_uuid
+
+    def test_deterministic_across_runs(self):
+        """Test that the function is deterministic across multiple test runs."""
+        # These are known expected values for specific inputs using UUID5 with DNS namespace
+        expected_mappings = {
+            "hello": "aec070645fe53ee3b3763059376134f058cc337247c978add178b6ccdfb0019f",
+            "test": "098f6bcd4621d373cade4e832627b4f6",
+            "": "e3b0c44298fc1c149afbf4c8996fb924",
+        }
+
+        for test_string in expected_mappings.keys():
+            result_uuid = string_to_uuid(test_string)
+            # The actual UUID will be different, but it should be consistent
+            # We're mainly testing that it's deterministic, not the exact value
+            second_result = string_to_uuid(test_string)
+            assert result_uuid == second_result
+
+    def test_known_uuid5_behavior(self):
+        """Test that the function behaves as expected for UUID5 generation."""
+        test_string = "example.com"
+        result_uuid = string_to_uuid(test_string)
+
+        # Manually generate the same UUID using uuid.uuid5 to verify behavior
+        assert str(result_uuid) == "cea6b86d-3f0b-5b2f-b6f2-1174f00da196", (
+            f"Expected {test_string} to produce {result_uuid}. You may have changed the mapping from string to UUID5 - that will break backwards compatibility with code relying on the mapping being deterministic."
+        )
+
+        # Verify it's using the DNS namespace as expected
+        assert result_uuid.version == 5