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

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

@@ -6,6 +6,9 @@ from typing import Any, Callable, Dict, List, Optional
 
 import yaml
 
+# Configuration keys
+MCP_SECRETS_KEY = "mcp_secrets"
+
 
 class ConfigProperty:
     def __init__(
@@ -54,6 +57,10 @@ class Config:
                 str,
                 env_var="OLLAMA_BASE_URL",
             ),
+            "docker_model_runner_base_url": ConfigProperty(
+                str,
+                env_var="DOCKER_MODEL_RUNNER_BASE_URL",
+            ),
             "bedrock_access_key": ConfigProperty(
                 str,
                 env_var="AWS_ACCESS_KEY_ID",
@@ -147,6 +154,21 @@ class Config:
                 env_var="CEREBRAS_API_KEY",
                 sensitive=True,
             ),
+            "enable_demo_tools": ConfigProperty(
+                bool,
+                env_var="ENABLE_DEMO_TOOLS",
+                default=False,
+            ),
+            # Allow the user to set the path to lookup MCP server commands, like npx.
+            "custom_mcp_path": ConfigProperty(
+                str,
+                env_var="CUSTOM_MCP_PATH",
+            ),
+            # Allow the user to set secrets for MCP servers, the key is mcp_server_id::key_name
+            MCP_SECRETS_KEY: ConfigProperty(
+                dict[str, str],
+                sensitive=True,
+            ),
         }
         self._lock = threading.Lock()
         self._settings = self.load_settings()
@@ -199,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)

kiln_ai/utils/open_ai_types.py

@@ -0,0 +1,94 @@
+"""
+Wrapper for OpenAI types to make them compatible with Pydantic.
+
+Pydantic doesn't support Iterable[T] well, so we use List[T] instead for tool_calls,
+https://github.com/pydantic/pydantic/issues/9541
+
+Otherwise we are using OpenAI SDK types directly.
+"""
+
+from typing import (
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    TypeAlias,
+    Union,
+)
+
+from openai.types.chat import (
+    ChatCompletionDeveloperMessageParam,
+    ChatCompletionFunctionMessageParam,
+    ChatCompletionMessageToolCallParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.chat.chat_completion_assistant_message_param import (
+    Audio,
+    ContentArrayOfContentPart,
+    FunctionCall,
+)
+from typing_extensions import Required, TypedDict
+
+
+class ChatCompletionAssistantMessageParamWrapper(TypedDict, total=False):
+    """
+    Almost exact copy of ChatCompletionAssistantMessageParam, but two changes.
+
+    First change: List[T] instead of Iterable[T] for tool_calls. Addresses pydantic issue.
+    https://github.com/pydantic/pydantic/issues/9541
+
+    Second change: Add reasoning_content to the message. A LiteLLM property for reasoning data.
+    """
+
+    role: Required[Literal["assistant"]]
+    """The role of the messages author, in this case `assistant`."""
+
+    audio: Optional[Audio]
+    """Data about a previous audio response from the model.
+
+    [Learn more](https://platform.openai.com/docs/guides/audio).
+    """
+
+    content: Union[str, Iterable[ContentArrayOfContentPart], None]
+    """The contents of the assistant message.
+
+    Required unless `tool_calls` or `function_call` is specified.
+    """
+
+    reasoning_content: Optional[str]
+    """The reasoning content of the assistant message. 
+    
+    A LiteLLM property for reasoning data: https://docs.litellm.ai/docs/reasoning_content
+    """
+
+    function_call: Optional[FunctionCall]
+    """Deprecated and replaced by `tool_calls`.
+
+    The name and arguments of a function that should be called, as generated by the
+    model.
+    """
+
+    name: str
+    """An optional name for the participant.
+
+    Provides the model information to differentiate between participants of the same
+    role.
+    """
+
+    refusal: Optional[str]
+    """The refusal message by the assistant."""
+
+    tool_calls: List[ChatCompletionMessageToolCallParam]
+    """The tool calls generated by the model, such as function calls."""
+
+
+ChatCompletionMessageParam: TypeAlias = Union[
+    ChatCompletionDeveloperMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionUserMessageParam,
+    ChatCompletionAssistantMessageParamWrapper,
+    ChatCompletionToolMessageParam,
+    ChatCompletionFunctionMessageParam,
+]

kiln_ai/utils/pdf_utils.py

@@ -0,0 +1,38 @@
+"""
+Utilities for working with PDF files.
+"""
+
+import asyncio
+import tempfile
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import AsyncGenerator
+
+from pypdf import PdfReader, PdfWriter
+
+
+@asynccontextmanager
+async def split_pdf_into_pages(pdf_path: Path) -> AsyncGenerator[list[Path], None]:
+    with tempfile.TemporaryDirectory(prefix="kiln_pdf_pages_") as temp_dir:
+        page_paths = []
+
+        with open(pdf_path, "rb") as file:
+            # Reader init can be heavy; offload to thread
+            pdf_reader = await asyncio.to_thread(PdfReader, file)
+
+            for page_num in range(len(pdf_reader.pages)):
+                await asyncio.sleep(0)
+                pdf_writer = PdfWriter()
+                pdf_writer.add_page(pdf_reader.pages[page_num])
+
+                # Create temporary file for this page
+                page_filename = f"page_{page_num + 1}.pdf"
+                page_path = Path(temp_dir) / page_filename
+
+                with open(page_path, "wb") as page_file:
+                    # Writing/compression can be expensive; offload to thread
+                    await asyncio.to_thread(pdf_writer.write, page_file)
+
+                page_paths.append(page_path)
+
+        yield page_paths

kiln_ai/utils/project_utils.py

@@ -0,0 +1,17 @@
+from kiln_ai.datamodel.project import Project
+from kiln_ai.utils.config import Config
+
+
+def project_from_id(project_id: str) -> Project | None:
+    project_paths = Config.shared().projects
+    if project_paths is not None:
+        for project_path in project_paths:
+            try:
+                project = Project.load_from_file(project_path)
+                if project.id == project_id:
+                    return project
+            except Exception:
+                # deleted files are possible continue with the rest
+                continue
+
+    return None