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

kiln_ai/tools/test_tool_registry.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from unittest.mock import Mock
 
 import pytest
@@ -8,6 +9,7 @@ from kiln_ai.datamodel.task import Task
 from kiln_ai.datamodel.tool_id import (
     MCP_LOCAL_TOOL_ID_PREFIX,
     MCP_REMOTE_TOOL_ID_PREFIX,
+    RAG_TOOL_ID_PREFIX,
     KilnBuiltInToolId,
     _check_tool_id,
     mcp_server_and_tool_name_from_id,
@@ -143,7 +145,7 @@ class TestToolRegistry:
         tool_id = f"{MCP_LOCAL_TOOL_ID_PREFIX}test_server::test_tool"
         with pytest.raises(
             ValueError,
-            match="Unable to resolve tool from id.*Requires a parent project/task",
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
         ):
             tool_from_id(tool_id, task=None)
 
@@ -181,6 +183,93 @@ class TestToolRegistry:
             ):
                 tool_from_id(tool_id, task=mock_task)
 
+    def test_tool_from_id_rag_tool_success(self):
+        """Test that tool_from_id works with RAG tool IDs."""
+        # Create mock RAG config
+        from unittest.mock import patch
+
+        with (
+            patch("kiln_ai.tools.tool_registry.RagConfig") as mock_rag_config_class,
+            patch("kiln_ai.tools.rag_tools.RagTool") as mock_rag_tool_class,
+        ):
+            # Setup mock RAG config
+            mock_rag_config = Mock()
+            mock_rag_config.id = "test_rag_config"
+            mock_rag_config_class.from_id_and_parent_path.return_value = mock_rag_config
+
+            # Setup mock RAG tool
+            mock_rag_tool = Mock()
+            mock_rag_tool_class.return_value = mock_rag_tool
+
+            # Create mock project
+            mock_project = Mock(spec=Project)
+            mock_project.id = "test_project_id"
+            mock_project.path = Path("/test/path")
+
+            # Create mock task with parent project
+            mock_task = Mock(spec=Task)
+            mock_task.parent_project.return_value = mock_project
+
+            # Test with RAG tool ID
+            tool_id = f"{RAG_TOOL_ID_PREFIX}test_rag_config"
+            tool = tool_from_id(tool_id, task=mock_task)
+
+            # Verify the tool is RagTool
+            assert tool == mock_rag_tool
+            mock_rag_config_class.from_id_and_parent_path.assert_called_once_with(
+                "test_rag_config", Path("/test/path")
+            )
+            mock_rag_tool_class.assert_called_once_with(tool_id, mock_rag_config)
+
+    def test_tool_from_id_rag_tool_no_task(self):
+        """Test that RAG tool ID without task raises ValueError."""
+        tool_id = f"{RAG_TOOL_ID_PREFIX}test_rag_config"
+
+        with pytest.raises(
+            ValueError,
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
+        ):
+            tool_from_id(tool_id, task=None)
+
+    def test_tool_from_id_rag_tool_no_project(self):
+        """Test that RAG tool ID with task but no project raises ValueError."""
+        # Create mock task without parent project
+        mock_task = Mock(spec=Task)
+        mock_task.parent_project.return_value = None
+
+        tool_id = f"{RAG_TOOL_ID_PREFIX}test_rag_config"
+
+        with pytest.raises(
+            ValueError,
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
+        ):
+            tool_from_id(tool_id, task=mock_task)
+
+    def test_tool_from_id_rag_config_not_found(self):
+        """Test that RAG tool ID with missing RAG config raises ValueError."""
+        from unittest.mock import patch
+
+        with patch("kiln_ai.tools.tool_registry.RagConfig") as mock_rag_config_class:
+            # Setup mock to return None (config not found)
+            mock_rag_config_class.from_id_and_parent_path.return_value = None
+
+            # Create mock project
+            mock_project = Mock(spec=Project)
+            mock_project.id = "test_project_id"
+            mock_project.path = Path("/test/path")
+
+            # Create mock task with parent project
+            mock_task = Mock(spec=Task)
+            mock_task.parent_project.return_value = mock_project
+
+            tool_id = f"{RAG_TOOL_ID_PREFIX}missing_rag_config"
+
+            with pytest.raises(
+                ValueError,
+                match="RAG config not found: missing_rag_config in project test_project_id for tool",
+            ):
+                tool_from_id(tool_id, task=mock_task)
+
     def test_all_built_in_tools_are_registered(self):
         """Test that all KilnBuiltInToolId enum members are handled by the registry."""
         for tool_id in KilnBuiltInToolId:
@@ -406,7 +495,7 @@ class TestToolRegistry:
 
         with pytest.raises(
             ValueError,
-            match="Unable to resolve tool from id.*Requires a parent project/task",
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
         ):
             tool_from_id(mcp_tool_id, task=None)
 

kiln_ai/tools/tool_registry.py

@@ -1,9 +1,12 @@
+from kiln_ai.datamodel.rag import RagConfig
 from kiln_ai.datamodel.task import Task
 from kiln_ai.datamodel.tool_id import (
     MCP_LOCAL_TOOL_ID_PREFIX,
     MCP_REMOTE_TOOL_ID_PREFIX,
+    RAG_TOOL_ID_PREFIX,
     KilnBuiltInToolId,
     mcp_server_and_tool_name_from_id,
+    rag_config_id_from_id,
 )
 from kiln_ai.tools.base_tool import KilnToolInterface
 from kiln_ai.tools.built_in_tools.math_tools import (
@@ -60,5 +63,23 @@ def tool_from_id(tool_id: str, task: Task | None = None) -> KilnToolInterface:
             )
 
         return MCPServerTool(server, tool_name)
+    elif tool_id.startswith(RAG_TOOL_ID_PREFIX):
+        project = task.parent_project() if task is not None else None
+        if project is None:
+            raise ValueError(
+                f"Unable to resolve tool from id: {tool_id}. Requires a parent project/task."
+            )
+
+        rag_config_id = rag_config_id_from_id(tool_id)
+        rag_config = RagConfig.from_id_and_parent_path(rag_config_id, project.path)
+        if rag_config is None:
+            raise ValueError(
+                f"RAG config not found: {rag_config_id} in project {project.id} for tool {tool_id}"
+            )
+
+        # Lazy import to avoid circular dependency
+        from kiln_ai.tools.rag_tools import RagTool
+
+        return RagTool(tool_id, rag_config)
 
     raise ValueError(f"Tool ID {tool_id} not found in tool registry")

kiln_ai/utils/__init__.py

@@ -5,8 +5,11 @@ Misc utilities used in the kiln_ai library.
 """
 
 from . import config, formatting
+from .lock import AsyncLockManager, shared_async_lock_manager
 
 __all__ = [
+    "AsyncLockManager",
     "config",
     "formatting",
+    "shared_async_lock_manager",
 ]

kiln_ai/utils/async_job_runner.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from dataclasses import dataclass
-from typing import AsyncGenerator, Awaitable, Callable, List, TypeVar
+from typing import AsyncGenerator, Awaitable, Callable, Generic, List, TypeVar
 
 logger = logging.getLogger(__name__)
 
@@ -15,29 +15,66 @@ class Progress:
     errors: int
 
 
-class AsyncJobRunner:
-    def __init__(self, concurrency: int = 1):
+class AsyncJobRunnerObserver(Generic[T]):
+    async def on_error(self, job: T, error: Exception):
+        """
+        Called when a job raises an unhandled exception.
+        """
+        pass
+
+    async def on_success(self, job: T):
+        """
+        Called when a job completes successfully.
+        """
+        pass
+
+    async def on_job_start(self, job: T):
+        """
+        Called when a job starts.
+        """
+        pass
+
+
+class AsyncJobRunner(Generic[T]):
+    def __init__(
+        self,
+        jobs: List[T],
+        run_job_fn: Callable[[T], Awaitable[bool]],
+        concurrency: int = 1,
+        observers: List[AsyncJobRunnerObserver[T]] | None = None,
+    ):
         if concurrency < 1:
             raise ValueError("concurrency must be ≥ 1")
         self.concurrency = concurrency
+        self.jobs = jobs
+        self.run_job_fn = run_job_fn
+        self.observers = observers or []
 
-    async def run(
-        self,
-        jobs: List[T],
-        run_job: Callable[[T], Awaitable[bool]],
-    ) -> AsyncGenerator[Progress, None]:
+    async def notify_error(self, job: T, error: Exception):
+        for observer in self.observers:
+            await observer.on_error(job, error)
+
+    async def notify_success(self, job: T):
+        for observer in self.observers:
+            await observer.on_success(job)
+
+    async def notify_job_start(self, job: T):
+        for observer in self.observers:
+            await observer.on_job_start(job)
+
+    async def run(self) -> AsyncGenerator[Progress, None]:
         """
         Runs the jobs with parallel workers and yields progress updates.
         """
         complete = 0
         errors = 0
-        total = len(jobs)
+        total = len(self.jobs)
 
         # Send initial status
         yield Progress(complete=complete, total=total, errors=errors)
 
         worker_queue: asyncio.Queue[T] = asyncio.Queue()
-        for job in jobs:
+        for job in self.jobs:
             worker_queue.put_nowait(job)
 
         # simple status queue to return progress. True=success, False=error
@@ -46,7 +83,7 @@ class AsyncJobRunner:
         workers = []
         for _ in range(self.concurrency):
             task = asyncio.create_task(
-                self._run_worker(worker_queue, status_queue, run_job),
+                self._run_worker(worker_queue, status_queue, self.run_job_fn),
             )
             workers.append(task)
 
@@ -64,7 +101,11 @@ class AsyncJobRunner:
                     else:
                         errors += 1
 
-                    yield Progress(complete=complete, total=total, errors=errors)
+                    yield Progress(
+                        complete=complete,
+                        total=total,
+                        errors=errors,
+                    )
                 except asyncio.TimeoutError:
                     # Timeout is expected, just continue to recheck worker status
                     # Don't love this but beats sentinels for reliability
@@ -82,7 +123,7 @@ class AsyncJobRunner:
         self,
         worker_queue: asyncio.Queue[T],
         status_queue: asyncio.Queue[bool],
-        run_job: Callable[[T], Awaitable[bool]],
+        run_job_fn: Callable[[T], Awaitable[bool]],
     ):
         while True:
             try:
@@ -92,13 +133,17 @@ class AsyncJobRunner:
                 break
 
             try:
-                success = await run_job(job)
-            except Exception:
+                await self.notify_job_start(job)
+                result = await run_job_fn(job)
+                if result:
+                    await self.notify_success(job)
+            except Exception as e:
                 logger.error("Job failed to complete", exc_info=True)
-                success = False
+                await self.notify_error(job, e)
+                result = False
 
             try:
-                await status_queue.put(success)
+                await status_queue.put(result)
             except Exception:
                 logger.error("Failed to enqueue status for job", exc_info=True)
             finally:

kiln_ai/utils/config.py

@@ -221,14 +221,14 @@ class Config:
             raise AttributeError(f"Config has no attribute '{name}'")
 
     @classmethod
-    def settings_dir(cls, create=True):
+    def settings_dir(cls, create=True) -> str:
         settings_dir = os.path.join(Path.home(), ".kiln_ai")
         if create and not os.path.exists(settings_dir):
             os.makedirs(settings_dir)
         return settings_dir
 
     @classmethod
-    def settings_path(cls, create=True):
+    def settings_path(cls, create=True) -> str:
         settings_dir = cls.settings_dir(create)
         return os.path.join(settings_dir, "settings.yaml")
 

kiln_ai/utils/env.py

@@ -0,0 +1,15 @@
+import os
+from contextlib import contextmanager
+
+
+@contextmanager
+def temporary_env(var_name: str, value: str):
+    old_value = os.environ.get(var_name)
+    os.environ[var_name] = value
+    try:
+        yield
+    finally:
+        if old_value is None:
+            os.environ.pop(var_name, None)  # remove if it did not exist before
+        else:
+            os.environ[var_name] = old_value

kiln_ai/utils/filesystem.py

@@ -0,0 +1,14 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def open_folder(path: str | Path) -> None:
+    log_dir = os.path.dirname(path)
+    if sys.platform.startswith("darwin"):
+        subprocess.run(["open", log_dir], check=True)
+    elif sys.platform.startswith("win"):
+        os.startfile(log_dir)  # type: ignore[attr-defined]
+    else:
+        subprocess.run(["xdg-open", log_dir], check=True)

kiln_ai/utils/filesystem_cache.py

@@ -0,0 +1,60 @@
+import logging
+import tempfile
+from pathlib import Path
+
+import anyio
+
+from kiln_ai.datamodel.basemodel import name_validator
+
+logger = logging.getLogger(__name__)
+
+
+class FilesystemCache:
+    def __init__(self, path: Path):
+        self.cache_dir_path = path
+
+    def validate_key(self, key: str) -> None:
+        # throws if invalid
+        name_validator(min_length=1, max_length=120)(key)
+
+    def get_path(self, key: str) -> Path:
+        self.validate_key(key)
+        return self.cache_dir_path / key
+
+    async def get(self, key: str) -> bytes | None:
+        # check if the file exists - don't need to validate the key
+        # worst case we just return None
+        if not self.get_path(key).exists():
+            return None
+
+        # we don't want to raise because of internal cache corruption issues
+        try:
+            return await anyio.Path(self.get_path(key)).read_bytes()
+        except Exception:
+            logger.error(f"Error reading file {self.get_path(key)}", exc_info=True)
+            return None
+
+    async def set(self, key: str, value: bytes) -> Path:
+        logger.debug(f"Caching {key} at {self.get_path(key)}")
+        self.validate_key(key)
+        path = self.get_path(key)
+        await anyio.Path(path).write_bytes(value)
+        return path
+
+
+class TemporaryFilesystemCache:
+    _shared_instance = None
+
+    def __init__(self):
+        self._cache_temp_dir = tempfile.mkdtemp(prefix="kiln_cache_")
+        self.filesystem_cache = FilesystemCache(path=Path(self._cache_temp_dir))
+
+        logger.debug(
+            f"Created temporary filesystem cache directory: {self._cache_temp_dir}"
+        )
+
+    @classmethod
+    def shared(cls) -> FilesystemCache:
+        if cls._shared_instance is None:
+            cls._shared_instance = cls()
+        return cls._shared_instance.filesystem_cache

kiln_ai/utils/litellm.py

@@ -0,0 +1,94 @@
+from dataclasses import dataclass
+
+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.exhaustive_error import raise_exhaustive_enum_error
+
+
+@dataclass
+class LitellmProviderInfo:
+    # The name of the provider, as it appears in litellm
+    provider_name: str
+    # Whether the provider is custom - e.g. custom models, ollama, fine tunes, and custom registry models
+    is_custom: bool
+    # The model ID slug to use in litellm
+    litellm_model_id: str
+
+
+def get_litellm_provider_info(
+    model_provider: KilnEmbeddingModelProvider | KilnModelProvider,
+) -> LitellmProviderInfo:
+    """
+    Maps a Kiln model provider to a litellm provider.
+
+    Args:
+        model_provider: The model provider to get litellm provider info for
+
+    Returns:
+        LitellmProviderInfo containing the provider name and whether it's custom
+    """
+    if not model_provider.model_id:
+        raise ValueError("Model ID is required for OpenAI compatible models")
+
+    litellm_provider_name: str | None = None
+    is_custom = False
+    match model_provider.name:
+        case ModelProviderName.openrouter:
+            litellm_provider_name = "openrouter"
+        case ModelProviderName.openai:
+            litellm_provider_name = "openai"
+        case ModelProviderName.groq:
+            litellm_provider_name = "groq"
+        case ModelProviderName.anthropic:
+            litellm_provider_name = "anthropic"
+        case ModelProviderName.ollama:
+            # We don't let litellm use the Ollama API and muck with our requests. We use Ollama's OpenAI compatible API.
+            # This is because we're setting detailed features like response_format=json_schema and want lower level control.
+            is_custom = True
+        case ModelProviderName.docker_model_runner:
+            # Docker Model Runner uses OpenAI-compatible API, similar to Ollama
+            # We want direct control over the requests for features like response_format=json_schema
+            is_custom = True
+        case ModelProviderName.gemini_api:
+            litellm_provider_name = "gemini"
+        case ModelProviderName.fireworks_ai:
+            litellm_provider_name = "fireworks_ai"
+        case ModelProviderName.amazon_bedrock:
+            litellm_provider_name = "bedrock"
+        case ModelProviderName.azure_openai:
+            litellm_provider_name = "azure"
+        case ModelProviderName.huggingface:
+            litellm_provider_name = "huggingface"
+        case ModelProviderName.vertex:
+            litellm_provider_name = "vertex_ai"
+        case ModelProviderName.together_ai:
+            litellm_provider_name = "together_ai"
+        case ModelProviderName.cerebras:
+            litellm_provider_name = "cerebras"
+        case ModelProviderName.siliconflow_cn:
+            is_custom = True
+        case ModelProviderName.openai_compatible:
+            is_custom = True
+        case ModelProviderName.kiln_custom_registry:
+            is_custom = True
+        case ModelProviderName.kiln_fine_tune:
+            is_custom = True
+        case _:
+            raise_exhaustive_enum_error(model_provider.name)
+
+    if is_custom:
+        # Use openai as it's only used for format, not url
+        litellm_provider_name = "openai"
+
+    # Shouldn't be possible but keep type checker happy
+    if litellm_provider_name is None:
+        raise ValueError(
+            f"Provider name could not lookup valid litellm provider ID {model_provider.model_id}"
+        )
+
+    return LitellmProviderInfo(
+        provider_name=litellm_provider_name,
+        is_custom=is_custom,
+        litellm_model_id=f"{litellm_provider_name}/{model_provider.model_id}",
+    )

kiln_ai/utils/lock.py

@@ -0,0 +1,100 @@
+import asyncio
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Dict, Hashable
+
+
+@dataclass
+class _Entry:
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    waiters: int = 0  # tasks waiting to acquire
+    holders: int = 0  # 0 or 1 for a mutex
+
+
+class AsyncLockManager:
+    """
+    A per-key asyncio lock manager that automatically cleans up locks when they're no longer needed.
+
+    Usage:
+        locks = AsyncLockManager()
+
+        async with locks.acquire("user:123"):
+            # critical section for "user:123"
+            ...
+
+    The manager removes a key when there are no holders and no waiters.
+    """
+
+    def __init__(self) -> None:
+        # Protects the _locks dict and bookkeeping counters.
+        self._mu = asyncio.Lock()
+        self._locks: Dict[Hashable, _Entry] = {}
+
+    @asynccontextmanager
+    async def acquire(self, key: Hashable, *, timeout: float | None = None):
+        """
+        Acquire the lock for `key` as an async context manager.
+
+        - `timeout`: optional seconds to wait; raises TimeoutError on expiry.
+        """
+        # Phase 1: register as a waiter and get/create the entry (under manager mutex).
+        async with self._mu:
+            entry = self._locks.get(key)
+            if entry is None:
+                entry = self._locks[key] = _Entry()
+            entry.waiters += 1
+
+        # Phase 2: wait on the per-key lock (outside manager mutex).
+        try:
+            if timeout is None:
+                await entry.lock.acquire()
+            else:
+                # Manual timeout to keep compatibility across Python versions.
+                await asyncio.wait_for(entry.lock.acquire(), timeout=timeout)
+
+            # Phase 3: update counters: became a holder.
+            async with self._mu:
+                entry.waiters -= 1
+                entry.holders += 1
+
+            try:
+                yield  # critical section
+            finally:
+                # Phase 4: release holder and maybe cleanup.
+                entry.lock.release()
+                async with self._mu:
+                    entry.holders -= 1
+                    # Remove the entry if fully idle.
+                    if entry.waiters == 0 and entry.holders == 0:
+                        # Double-check we still point to same object (paranoia/race safety).
+                        if self._locks.get(key) is entry:
+                            del self._locks[key]
+
+        except asyncio.TimeoutError:
+            # Timed out while waiting; undo waiter count and maybe cleanup.
+            async with self._mu:
+                entry.waiters -= 1
+                if entry.waiters == 0 and entry.holders == 0:
+                    if self._locks.get(key) is entry:
+                        del self._locks[key]
+            raise
+        except asyncio.CancelledError:
+            # Cancelled while waiting; same cleanup as timeout.
+            async with self._mu:
+                entry.waiters -= 1
+                if entry.waiters == 0 and entry.holders == 0:
+                    if self._locks.get(key) is entry:
+                        del self._locks[key]
+            raise
+
+    # Optional: expose a snapshot for metrics/debugging
+    async def snapshot(self) -> Dict[Hashable, dict]:
+        async with self._mu:
+            return {
+                k: {"waiters": e.waiters, "holders": e.holders}
+                for k, e in self._locks.items()
+            }
+
+
+# callers should use this global instance instead of creating their own
+shared_async_lock_manager = AsyncLockManager()

kiln_ai/utils/mime_type.py

@@ -0,0 +1,38 @@
+import mimetypes
+
+
+def guess_mime_type(filename: str) -> str | None:
+    filename_normalized = filename.lower()
+
+    # we override the mimetypes.guess_type for some common cases
+    # because it does not handle them correctly
+    if filename_normalized.endswith(".mov"):
+        return "video/quicktime"
+    elif filename_normalized.endswith(".mp3"):
+        return "audio/mpeg"
+    elif filename_normalized.endswith(".wav"):
+        return "audio/wav"
+    elif filename_normalized.endswith(".mp4"):
+        return "video/mp4"
+
+    mime_type, _ = mimetypes.guess_type(filename_normalized)
+    return mime_type
+
+
+def guess_extension(mime_type: str) -> str | None:
+    mapping = {
+        "application/pdf": ".pdf",
+        "image/png": ".png",
+        "video/mp4": ".mp4",
+        "audio/ogg": ".ogg",
+        "text/markdown": ".md",
+        "text/plain": ".txt",
+        "text/html": ".html",
+        "text/csv": ".csv",
+        "image/jpeg": ".jpeg",
+        "image/jpg": ".jpeg",
+        "audio/mpeg": ".mp3",
+        "audio/wav": ".wav",
+        "video/quicktime": ".mov",
+    }
+    return mapping.get(mime_type)