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

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:
+    dir = os.path.dirname(path)
+    if sys.platform.startswith("darwin"):
+        subprocess.run(["open", dir], check=True)
+    elif sys.platform.startswith("win"):
+        os.startfile(dir)  # type: ignore[attr-defined]
+    else:
+        subprocess.run(["xdg-open", 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

@@ -17,11 +17,11 @@ from typing import (
 )
 
 from openai.types.chat import (
+    ChatCompletionContentPartTextParam,
     ChatCompletionDeveloperMessageParam,
     ChatCompletionFunctionMessageParam,
     ChatCompletionMessageToolCallParam,
     ChatCompletionSystemMessageParam,
-    ChatCompletionToolMessageParam,
     ChatCompletionUserMessageParam,
 )
 from openai.types.chat.chat_completion_assistant_message_param import (
@@ -84,11 +84,28 @@ class ChatCompletionAssistantMessageParamWrapper(TypedDict, total=False):
     """The tool calls generated by the model, such as function calls."""
 
 
+class ChatCompletionToolMessageParamWrapper(TypedDict, total=False):
+    content: Required[Union[str, Iterable[ChatCompletionContentPartTextParam]]]
+    """The contents of the tool message."""
+
+    role: Required[Literal["tool"]]
+    """The role of the messages author, in this case `tool`."""
+
+    tool_call_id: Required[str]
+    """Tool call that this message is responding to."""
+
+    kiln_task_tool_data: Optional[str]
+    """The data for the Kiln task tool that this message is responding to.
+    
+    Formatted as `<project_id>:::<tool_id>:::<task_id>:::<run_id>`
+    """
+
+
 ChatCompletionMessageParam: TypeAlias = Union[
     ChatCompletionDeveloperMessageParam,
     ChatCompletionSystemMessageParam,
     ChatCompletionUserMessageParam,
     ChatCompletionAssistantMessageParamWrapper,
-    ChatCompletionToolMessageParam,
+    ChatCompletionToolMessageParamWrapper,
     ChatCompletionFunctionMessageParam,
 ]

kiln_ai/utils/pdf_utils.py

@@ -0,0 +1,59 @@
+"""
+Utilities for working with PDF files.
+"""
+
+import asyncio
+import tempfile
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import AsyncGenerator
+
+import pypdfium2
+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
+
+
+async def convert_pdf_to_images(pdf_path: Path, output_dir: Path) -> list[Path]:
+    image_paths = []
+
+    # note: doing this in a thread causes a segfault - but this is slow and blocking
+    # so we should try to find a better way
+    pdf = pypdfium2.PdfDocument(pdf_path)
+    try:
+        for idx, page in enumerate(pdf):
+            await asyncio.sleep(0)
+            # scale=2 is legible for ~A4 pages (research papers, etc.) - lower than this is blurry
+            bitmap = page.render(scale=2).to_pil()
+            target_path = output_dir / f"img-{pdf_path.name}-{idx}.png"
+            bitmap.save(target_path)
+            image_paths.append(target_path)
+
+        return image_paths
+    finally:
+        pdf.close()