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

kiln_ai/utils/test_filesystem_cache.py

@@ -0,0 +1,316 @@
+import tempfile
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from kiln_ai.utils.filesystem_cache import FilesystemCache, TemporaryFilesystemCache
+
+
+class TestFilesystemCache:
+    @pytest.fixture
+    def temp_dir(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            yield Path(tmp_dir)
+
+    @pytest.fixture
+    def cache(self, temp_dir):
+        return FilesystemCache(temp_dir)
+
+    def test_init(self, temp_dir):
+        cache = FilesystemCache(temp_dir)
+        assert cache.cache_dir_path == temp_dir
+        assert cache.validate_key is not None
+
+    def test_get_path_valid_key(self, cache):
+        key = "test_file"
+        expected_path = cache.cache_dir_path / key
+        assert cache.get_path(key) == expected_path
+
+    def test_get_path_invalid_key_empty(self, cache):
+        with pytest.raises(ValueError):
+            cache.get_path("")
+
+    def test_get_path_invalid_key_too_long(self, cache):
+        long_key = "x" * 121  # exceeds max_length=120
+        with pytest.raises(ValueError):
+            cache.get_path(long_key)
+
+    def test_get_path_invalid_key_special_chars(self, cache):
+        with pytest.raises(ValueError):
+            cache.get_path("invalid/key")
+
+    async def test_get_nonexistent_file(self, cache):
+        result = await cache.get("nonexistent")
+        assert result is None
+
+    async def test_get_existing_file(self, cache, temp_dir):
+        key = "test"
+        content = b"Hello, World!"
+        file_path = temp_dir / key
+        file_path.write_bytes(content)
+
+        result = await cache.get(key)
+        assert result == content
+
+    async def test_get_file_read_error(self, cache, temp_dir):
+        key = "test"
+        file_path = temp_dir / key
+        file_path.write_bytes(b"test")
+
+        with patch(
+            "anyio.Path.read_bytes",
+            new_callable=AsyncMock,
+            side_effect=IOError("Read error"),
+        ):
+            result = await cache.get(key)
+            assert result is None
+
+    async def test_set_valid_data(self, cache, temp_dir):
+        key = "test"
+        content = b"Hello, World!"
+
+        result_path = await cache.set(key, content)
+        expected_path = temp_dir / key
+
+        assert result_path == expected_path
+        assert expected_path.exists()
+        assert expected_path.read_bytes() == content
+
+    async def test_set_invalid_key_empty(self, cache):
+        with pytest.raises(ValueError):
+            await cache.set("", b"content")
+
+    async def test_set_invalid_key_too_long(self, cache):
+        long_key = "x" * 121
+        with pytest.raises(ValueError):
+            await cache.set(long_key, b"content")
+
+    async def test_set_invalid_key_special_chars(self, cache):
+        with pytest.raises(ValueError):
+            await cache.set("invalid/key", b"content")
+
+    async def test_set_overwrites_existing(self, cache, temp_dir):
+        key = "test"
+        original_content = b"Original content"
+        new_content = b"New content"
+
+        # Set original content
+        await cache.set(key, original_content)
+        assert (temp_dir / key).read_bytes() == original_content
+
+        # Overwrite with new content
+        await cache.set(key, new_content)
+        assert (temp_dir / key).read_bytes() == new_content
+
+    async def test_set_creates_parent_directory(self, cache, temp_dir):
+        # Test that set method creates parent directories if they don't exist
+        # Note: This test demonstrates a limitation - the current implementation
+        # doesn't support directory paths due to name_validator restrictions
+        key = "subdir_test"  # Using underscore instead of slash
+        content = b"Test content"
+
+        result_path = await cache.set(key, content)
+        expected_path = temp_dir / key
+
+        assert result_path == expected_path
+        assert expected_path.exists()
+        assert expected_path.read_bytes() == content
+
+    async def test_set_with_nested_directories(self, cache, temp_dir):
+        # Test that set method creates deeply nested directories
+        # Note: This test demonstrates a limitation - the current implementation
+        # doesn't support directory paths due to name_validator restrictions
+        key = "level1_level2_level3_test"  # Using underscores instead of slashes
+        content = b"Deeply nested content"
+
+        result_path = await cache.set(key, content)
+        expected_path = temp_dir / key
+
+        assert result_path == expected_path
+        assert expected_path.exists()
+        assert expected_path.read_bytes() == content
+
+    async def test_directory_paths_not_supported(self, cache):
+        # Test that demonstrates the current limitation - directory paths are not supported
+        # due to name_validator forbidding forward slashes
+        with pytest.raises(ValueError, match="Name is invalid"):
+            await cache.set("subdir/test", b"content")
+
+        with pytest.raises(ValueError, match="Name is invalid"):
+            await cache.get("subdir/test")
+
+    async def test_roundtrip_get_set(self, cache):
+        key = "roundtrip"
+        content = b"Roundtrip test content"
+
+        # Set content
+        await cache.set(key, content)
+
+        # Get content back
+        retrieved = await cache.get(key)
+        assert retrieved == content
+
+    async def test_multiple_files(self, cache, temp_dir):
+        files = {
+            "file1": b"Content 1",
+            "file2": b"Content 2",
+            "subdir_file3": b"Content 3",  # Using underscore instead of slash
+        }
+
+        # Set all files
+        for key, content in files.items():
+            await cache.set(key, content)
+
+        # Verify all files exist and have correct content
+        for key, expected_content in files.items():
+            actual_content = await cache.get(key)
+            assert actual_content == expected_content
+
+    async def test_empty_bytes(self, cache):
+        key = "empty"
+        content = b""
+
+        await cache.set(key, content)
+        retrieved = await cache.get(key)
+        assert retrieved == content
+
+    async def test_large_content(self, cache):
+        key = "large"
+        content = b"x" * 10000  # 10KB of data
+
+        await cache.set(key, content)
+        retrieved = await cache.get(key)
+        assert retrieved == content
+
+    @pytest.mark.parametrize(
+        "unicode_text",
+        [
+            "Simple ASCII text",
+            "中文文本测试",
+            "Mixed 中文 and English text",
+            "Emojis: 🎉🚀💻🔥",
+            "Complex: 你好世界! Hello 世界! 🌍 This is 测试 with 中文, emojis 🚀, and English.",
+            "Special chars: ñáéíóú àèìòù çüöä",
+            "Math symbols: ∑∆π∫∞±≤≥≠",
+            "Currency: €£¥$₹₽",
+        ],
+    )
+    async def test_unicode_text_retrieval_integrity(self, cache, unicode_text):
+        # Test that Unicode text is not corrupted during storage and retrieval
+        key = "unicode_integrity"
+        content = unicode_text.encode("utf-8")
+
+        # Store the text
+        await cache.set(key, content)
+
+        # Retrieve and verify
+        retrieved = await cache.get(key)
+        assert retrieved == content
+        assert retrieved.decode("utf-8") == unicode_text
+
+    async def test_key_overwrite_behavior(self, cache):
+        # Test that setting at the same key overwrites whatever was there
+        key = "overwrite_test"
+
+        # Set initial content
+        initial_content = "Initial content".encode("utf-8")
+        await cache.set(key, initial_content)
+
+        # Verify initial content is stored
+        retrieved = await cache.get(key)
+        assert retrieved == initial_content
+        assert retrieved.decode("utf-8") == "Initial content"
+
+        # Overwrite with different content
+        new_content = "New content with 中文 and emojis 🚀".encode("utf-8")
+        await cache.set(key, new_content)
+
+        # Verify the content was overwritten
+        retrieved = await cache.get(key)
+        assert retrieved == new_content
+        assert retrieved.decode("utf-8") == "New content with 中文 and emojis 🚀"
+        assert retrieved != initial_content
+
+        # Overwrite again with empty content
+        empty_content = b""
+        await cache.set(key, empty_content)
+
+        # Verify empty content is stored
+        retrieved = await cache.get(key)
+        assert retrieved == empty_content
+        assert retrieved.decode("utf-8") == ""
+
+    @pytest.mark.parametrize(
+        "invalid_key",
+        [
+            "",
+            "x" * 121,
+            "invalid/key",
+            "invalid\\key",
+            "invalid:key",
+            "invalid*key",
+            "invalid?key",
+            "invalid<key",
+            "invalid>key",
+            "invalid|key",
+        ],
+    )
+    async def test_invalid_keys(self, cache, invalid_key):
+        with pytest.raises(ValueError):
+            cache.get_path(invalid_key)
+
+        with pytest.raises(ValueError):
+            await cache.set(invalid_key, b"content")
+
+
+class TestTemporaryFilesystemCache:
+    def test_temporary_cache_creation(self):
+        """Test that TemporaryFilesystemCache creates a temporary directory."""
+        temp_cache = TemporaryFilesystemCache()
+
+        # Should create a temporary directory
+        assert temp_cache._cache_temp_dir is not None
+        assert Path(temp_cache._cache_temp_dir).exists()
+
+        # Check that the directory name (not full path) starts with the prefix
+        temp_dir_name = Path(temp_cache._cache_temp_dir).name
+        assert temp_dir_name.startswith("kiln_cache_")
+
+        # Should have a FilesystemCache instance
+        assert isinstance(temp_cache.filesystem_cache, FilesystemCache)
+        assert temp_cache.filesystem_cache.cache_dir_path == Path(
+            temp_cache._cache_temp_dir
+        )
+
+    async def test_multiple_instances_share_same_cache(self):
+        """Test that multiple calls to shared() return the same cache instance."""
+        # Get cache instances
+        cache1 = TemporaryFilesystemCache.shared()
+        cache2 = TemporaryFilesystemCache.shared()
+
+        # Should be the same instance
+        assert cache1 is cache2
+
+        # Test that they share the same cache directory
+        assert cache1.cache_dir_path == cache2.cache_dir_path
+
+        # Test that content set in one is available in the other
+        key = "shared_test"
+        content = b"shared content"
+
+        await cache1.set(key, content)
+        retrieved = await cache2.get(key)
+        assert retrieved == content
+
+    def test_cache_directory_naming(self):
+        """Test that the temporary cache directory has the correct naming pattern."""
+        temp_cache = TemporaryFilesystemCache()
+        temp_dir_name = Path(temp_cache._cache_temp_dir).name
+
+        # Should start with the expected prefix
+        assert temp_dir_name.startswith("kiln_cache_")
+
+        # Should be a valid directory name
+        assert len(temp_dir_name) > len("kiln_cache_")
+        assert "/" not in temp_dir_name  # Should not contain path separators

kiln_ai/utils/test_litellm.py

@@ -0,0 +1,206 @@
+import pytest
+
+from kiln_ai.adapters.ml_embedding_model_list import KilnEmbeddingModelProvider
+from kiln_ai.adapters.ml_model_list import KilnModelProvider
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+from kiln_ai.utils.litellm import LitellmProviderInfo, get_litellm_provider_info
+
+
+class TestGetLitellmProviderInfo:
+    """Test cases for get_litellm_provider_info function"""
+
+    @pytest.fixture
+    def sample_model_id(self):
+        """Sample model ID for testing"""
+        return "test-model-id"
+
+    @pytest.fixture
+    def embedding_provider(self, sample_model_id):
+        """Sample KilnEmbeddingModelProvider for testing"""
+        return KilnEmbeddingModelProvider(
+            name=ModelProviderName.openai,
+            model_id=sample_model_id,
+            n_dimensions=1536,
+        )
+
+    @pytest.fixture
+    def model_provider(self, sample_model_id):
+        """Sample KilnModelProvider for testing"""
+        return KilnModelProvider(
+            name=ModelProviderName.openai,
+            model_id=sample_model_id,
+        )
+
+    @pytest.mark.parametrize(
+        "model_id",
+        [
+            None,
+            "",
+        ],
+    )
+    def test_missing_model_id_raises_error(self, model_id):
+        """Test that missing model_id raises ValueError"""
+        provider = KilnModelProvider(
+            name=ModelProviderName.openai,
+            model_id=model_id,
+        )
+
+        with pytest.raises(
+            ValueError, match="Model ID is required for OpenAI compatible models"
+        ):
+            get_litellm_provider_info(provider)
+
+    @pytest.mark.parametrize(
+        "provider_name,expected_litellm_name,expected_is_custom",
+        [
+            (ModelProviderName.openrouter, "openrouter", False),
+            (ModelProviderName.openai, "openai", False),
+            (ModelProviderName.groq, "groq", False),
+            (ModelProviderName.anthropic, "anthropic", False),
+            (ModelProviderName.gemini_api, "gemini", False),
+            (ModelProviderName.fireworks_ai, "fireworks_ai", False),
+            (ModelProviderName.amazon_bedrock, "bedrock", False),
+            (ModelProviderName.azure_openai, "azure", False),
+            (ModelProviderName.huggingface, "huggingface", False),
+            (ModelProviderName.vertex, "vertex_ai", False),
+            (ModelProviderName.together_ai, "together_ai", False),
+            (ModelProviderName.ollama, "openai", True),
+            (ModelProviderName.openai_compatible, "openai", True),
+            (ModelProviderName.kiln_custom_registry, "openai", True),
+            (ModelProviderName.kiln_fine_tune, "openai", True),
+        ],
+    )
+    def test_provider_mappings_with_model_provider(
+        self, provider_name, expected_litellm_name, expected_is_custom, sample_model_id
+    ):
+        """Test provider name mappings for KilnModelProvider"""
+        provider = KilnModelProvider(
+            name=provider_name,
+            model_id=sample_model_id,
+        )
+
+        result = get_litellm_provider_info(provider)
+
+        assert isinstance(result, LitellmProviderInfo)
+        assert result.provider_name == expected_litellm_name
+        assert result.is_custom == expected_is_custom
+        assert result.litellm_model_id == f"{expected_litellm_name}/{sample_model_id}"
+
+    @pytest.mark.parametrize(
+        "provider_name,expected_litellm_name,expected_is_custom",
+        [
+            (ModelProviderName.openrouter, "openrouter", False),
+            (ModelProviderName.openai, "openai", False),
+            (ModelProviderName.groq, "groq", False),
+            (ModelProviderName.anthropic, "anthropic", False),
+            (ModelProviderName.gemini_api, "gemini", False),
+            (ModelProviderName.fireworks_ai, "fireworks_ai", False),
+            (ModelProviderName.amazon_bedrock, "bedrock", False),
+            (ModelProviderName.azure_openai, "azure", False),
+            (ModelProviderName.huggingface, "huggingface", False),
+            (ModelProviderName.vertex, "vertex_ai", False),
+            (ModelProviderName.together_ai, "together_ai", False),
+            (ModelProviderName.ollama, "openai", True),
+            (ModelProviderName.openai_compatible, "openai", True),
+            (ModelProviderName.kiln_custom_registry, "openai", True),
+            (ModelProviderName.kiln_fine_tune, "openai", True),
+        ],
+    )
+    def test_provider_mappings_with_embedding_provider(
+        self, provider_name, expected_litellm_name, expected_is_custom, sample_model_id
+    ):
+        """Test provider name mappings for KilnEmbeddingModelProvider"""
+        provider = KilnEmbeddingModelProvider(
+            name=provider_name,
+            model_id=sample_model_id,
+            n_dimensions=1536,
+        )
+
+        result = get_litellm_provider_info(provider)
+
+        assert isinstance(result, LitellmProviderInfo)
+        assert result.provider_name == expected_litellm_name
+        assert result.is_custom == expected_is_custom
+        assert result.litellm_model_id == f"{expected_litellm_name}/{sample_model_id}"
+
+    def test_custom_providers_use_openai_format(self, sample_model_id):
+        """Test that custom providers use 'openai' as the litellm provider name"""
+        custom_providers = [
+            ModelProviderName.ollama,
+            ModelProviderName.openai_compatible,
+            ModelProviderName.kiln_custom_registry,
+            ModelProviderName.kiln_fine_tune,
+        ]
+
+        for provider_name in custom_providers:
+            provider = KilnModelProvider(
+                name=provider_name,
+                model_id=sample_model_id,
+            )
+
+            result = get_litellm_provider_info(provider)
+
+            assert result.provider_name == "openai"
+            assert result.is_custom is True
+            assert result.litellm_model_id == f"openai/{sample_model_id}"
+
+    def test_non_custom_providers_use_correct_format(self, sample_model_id):
+        """Test that non-custom providers use their actual provider names"""
+        non_custom_providers = [
+            (ModelProviderName.openai, "openai"),
+            (ModelProviderName.anthropic, "anthropic"),
+            (ModelProviderName.groq, "groq"),
+            (ModelProviderName.gemini_api, "gemini"),
+        ]
+
+        for provider_name, expected_name in non_custom_providers:
+            provider = KilnModelProvider(
+                name=provider_name,
+                model_id=sample_model_id,
+            )
+
+            result = get_litellm_provider_info(provider)
+
+            assert result.provider_name == expected_name
+            assert result.is_custom is False
+            assert result.litellm_model_id == f"{expected_name}/{sample_model_id}"
+
+    def test_litellm_model_id_format(self, embedding_provider):
+        """Test that litellm_model_id follows the correct format"""
+        result = get_litellm_provider_info(embedding_provider)
+
+        expected_format = f"{result.provider_name}/{embedding_provider.model_id}"
+        assert result.litellm_model_id == expected_format
+
+    def test_return_type_structure(self, model_provider):
+        """Test that the return type has all expected fields"""
+        result = get_litellm_provider_info(model_provider)
+
+        assert hasattr(result, "provider_name")
+        assert hasattr(result, "is_custom")
+        assert hasattr(result, "litellm_model_id")
+
+        assert isinstance(result.provider_name, str)
+        assert isinstance(result.is_custom, bool)
+        assert isinstance(result.litellm_model_id, str)
+
+    def test_works_with_both_provider_types(self, sample_model_id):
+        """Test that function works with both KilnModelProvider and KilnEmbeddingModelProvider"""
+        model_provider = KilnModelProvider(
+            name=ModelProviderName.openai,
+            model_id=sample_model_id,
+        )
+
+        embedding_provider = KilnEmbeddingModelProvider(
+            name=ModelProviderName.openai,
+            model_id=sample_model_id,
+            n_dimensions=1536,
+        )
+
+        model_result = get_litellm_provider_info(model_provider)
+        embedding_result = get_litellm_provider_info(embedding_provider)
+
+        # Results should be identical for same provider name and model ID
+        assert model_result.provider_name == embedding_result.provider_name
+        assert model_result.is_custom == embedding_result.is_custom
+        assert model_result.litellm_model_id == embedding_result.litellm_model_id