kiln-ai 0.20.1__tar.gz → 0.21.0__tar.gz

{kiln_ai-0.20.1 → kiln_ai-0.21.0}/.gitignore

@@ -8,6 +8,7 @@ __pycache__/
 .coverage
 **/.venv
 **/*.egg-info
+node_modules/
 
 libs/core/docs
 libs/core/build

{kiln_ai-0.20.1 → kiln_ai-0.21.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiln-ai
-Version: 0.20.1
+Version: 0.21.0
 Summary: Kiln AI
 Project-URL: Homepage, https://kiln.tech
 Project-URL: Repository, https://github.com/Kiln-AI/kiln
@@ -15,15 +15,21 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
+Requires-Dist: anyio>=4.10.0
 Requires-Dist: boto3>=1.37.10
 Requires-Dist: coverage>=7.6.4
 Requires-Dist: exceptiongroup>=1.0.0; python_version < '3.11'
 Requires-Dist: google-cloud-aiplatform>=1.84.0
+Requires-Dist: google-genai>=1.21.1
 Requires-Dist: jsonschema>=4.23.0
+Requires-Dist: lancedb>=0.24.2
 Requires-Dist: litellm>=1.72.6
+Requires-Dist: llama-index-vector-stores-lancedb>=0.3.3
+Requires-Dist: llama-index>=0.13.3
 Requires-Dist: openai>=1.53.0
 Requires-Dist: pdoc>=15.0.0
 Requires-Dist: pydantic>=2.9.2
+Requires-Dist: pypdf>=6.0.0
 Requires-Dist: pytest-benchmark>=5.1.0
 Requires-Dist: pytest-cov>=6.0.0
 Requires-Dist: pyyaml>=6.0.2

{kiln_ai-0.20.1 → kiln_ai-0.21.0}/kiln_ai/adapters/__init__.py

@@ -18,9 +18,12 @@ The eval submodule contains the code for evaluating the performance of a model.
 
 from . import (
     chat,
+    chunkers,
     data_gen,
     eval,
+    extractors,
     fine_tune,
+    ml_embedding_model_list,
     ml_model_list,
     model_adapters,
     prompt_builders,
@@ -29,9 +32,12 @@ from . import (
 
 __all__ = [
     "chat",
+    "chunkers",
     "data_gen",
     "eval",
+    "extractors",
     "fine_tune",
+    "ml_embedding_model_list",
     "ml_model_list",
     "model_adapters",
     "prompt_builders",

kiln_ai-0.21.0/kiln_ai/adapters/adapter_registry.py

@@ -0,0 +1,62 @@
+from kiln_ai import datamodel
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig, BaseAdapter
+from kiln_ai.adapters.model_adapters.litellm_adapter import (
+    LiteLlmAdapter,
+    LiteLlmConfig,
+)
+from kiln_ai.adapters.provider_tools import (
+    core_provider,
+    lite_llm_core_config_for_provider,
+)
+from kiln_ai.datamodel.task import RunConfigProperties
+
+
+def litellm_core_provider_config(
+    run_config_properties: RunConfigProperties,
+) -> LiteLlmConfig:
+    # For things like the fine-tune provider, we want to run the underlying provider (e.g. openai)
+    core_provider_name = core_provider(
+        run_config_properties.model_name, run_config_properties.model_provider_name
+    )
+
+    # For OpenAI compatible providers, we want to retrieve the underlying provider and update the run config properties to match
+    openai_compatible_provider_name = None
+    if run_config_properties.model_provider_name == ModelProviderName.openai_compatible:
+        model_id = run_config_properties.model_name
+        try:
+            openai_compatible_provider_name, model_id = model_id.split("::")
+        except Exception:
+            raise ValueError(f"Invalid openai compatible model ID: {model_id}")
+
+        # Update a copy of the run config properties to use the openai compatible provider
+        updated_run_config_properties = run_config_properties.model_copy(deep=True)
+        updated_run_config_properties.model_name = model_id
+        run_config_properties = updated_run_config_properties
+
+    config = lite_llm_core_config_for_provider(
+        core_provider_name, openai_compatible_provider_name
+    )
+    if config is None:
+        raise ValueError(
+            "Fine tune or custom openai compatible provider is not a core provider. The underlying provider should be used when requesting the adapter litellm config instead."
+        )
+
+    return LiteLlmConfig(
+        run_config_properties=run_config_properties,
+        base_url=config.base_url,
+        default_headers=config.default_headers,
+        additional_body_options=config.additional_body_options or {},
+    )
+
+
+def adapter_for_task(
+    kiln_task: datamodel.Task,
+    run_config_properties: RunConfigProperties,
+    base_adapter_config: AdapterConfig | None = None,
+) -> BaseAdapter:
+    return LiteLlmAdapter(
+        kiln_task=kiln_task,
+        config=litellm_core_provider_config(run_config_properties),
+        base_adapter_config=base_adapter_config,
+    )

kiln_ai-0.21.0/kiln_ai/adapters/chunkers/__init__.py

@@ -0,0 +1,13 @@
+"""
+Chunkers for processing different document types.
+
+This package provides a framework for chunking text into smaller chunks.
+"""
+
+from . import base_chunker, chunker_registry, fixed_window_chunker
+
+__all__ = [
+    "base_chunker",
+    "chunker_registry",
+    "fixed_window_chunker",
+]

kiln_ai-0.21.0/kiln_ai/adapters/chunkers/base_chunker.py

@@ -0,0 +1,42 @@
+import logging
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel, Field
+
+from kiln_ai.adapters.chunkers.helpers import clean_up_text
+from kiln_ai.datamodel.chunk import ChunkerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TextChunk(BaseModel):
+    text: str = Field(description="The text of the chunk.")
+
+
+class ChunkingResult(BaseModel):
+    chunks: list[TextChunk] = Field(description="The chunks of the text.")
+
+
+class BaseChunker(ABC):
+    """
+    Base class for all chunkers.
+
+    Should be subclassed by each chunker.
+    """
+
+    def __init__(self, chunker_config: ChunkerConfig):
+        self.chunker_config = chunker_config
+
+    async def chunk(self, text: str) -> ChunkingResult:
+        if not text:
+            return ChunkingResult(chunks=[])
+
+        sanitized_text = clean_up_text(text)
+        if not sanitized_text:
+            return ChunkingResult(chunks=[])
+
+        return await self._chunk(sanitized_text)
+
+    @abstractmethod
+    async def _chunk(self, text: str) -> ChunkingResult:
+        pass

kiln_ai-0.21.0/kiln_ai/adapters/chunkers/chunker_registry.py

@@ -0,0 +1,16 @@
+from kiln_ai.adapters.chunkers.base_chunker import BaseChunker
+from kiln_ai.adapters.chunkers.fixed_window_chunker import FixedWindowChunker
+from kiln_ai.datamodel.chunk import ChunkerConfig, ChunkerType
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+
+def chunker_adapter_from_type(
+    chunker_type: ChunkerType,
+    chunker_config: ChunkerConfig,
+) -> BaseChunker:
+    match chunker_type:
+        case ChunkerType.FIXED_WINDOW:
+            return FixedWindowChunker(chunker_config)
+        case _:
+            # type checking will catch missing cases
+            raise_exhaustive_enum_error(chunker_type)

kiln_ai-0.21.0/kiln_ai/adapters/chunkers/fixed_window_chunker.py

@@ -0,0 +1,39 @@
+from typing import List
+
+from llama_index.core.text_splitter import SentenceSplitter
+
+from kiln_ai.adapters.chunkers.base_chunker import (
+    BaseChunker,
+    ChunkingResult,
+    TextChunk,
+)
+from kiln_ai.datamodel.chunk import ChunkerConfig, ChunkerType
+
+
+class FixedWindowChunker(BaseChunker):
+    def __init__(self, chunker_config: ChunkerConfig):
+        if chunker_config.chunker_type != ChunkerType.FIXED_WINDOW:
+            raise ValueError("Chunker type must be FIXED_WINDOW")
+
+        chunk_size = chunker_config.chunk_size()
+        if chunk_size is None:
+            raise ValueError("Chunk size must be set")
+
+        chunk_overlap = chunker_config.chunk_overlap()
+        if chunk_overlap is None:
+            raise ValueError("Chunk overlap must be set")
+
+        super().__init__(chunker_config)
+        self.splitter = SentenceSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+        )
+
+    async def _chunk(self, text: str) -> ChunkingResult:
+        sentences = self.splitter.split_text(text)
+
+        chunks: List[TextChunk] = []
+        for sentence in sentences:
+            chunks.append(TextChunk(text=sentence))
+
+        return ChunkingResult(chunks=chunks)

kiln_ai-0.21.0/kiln_ai/adapters/chunkers/helpers.py

@@ -0,0 +1,23 @@
+import re
+
+
+def clean_up_text(text: str) -> str:
+    """
+    Clean up text by limiting consecutive newlines and consecutive whitespace. Models sometimes send a lot of those.
+    It seems to happen more when the transcription is done at low temperature.
+
+    - Replaces 6+ consecutive newlines with exactly 6 newlines
+    - Replaces 50+ consecutive spaces with exactly 50 spaces
+    - Leaves 1-5 consecutive newlines unchanged
+    - Leaves 1-49 consecutive spaces unchanged
+    """
+    max_consecutive_newlines = 6
+    max_consecutive_whitespace = 50
+
+    # Replace 6+ consecutive newlines with exactly 6 newlines
+    text = re.sub(r"\n{6,}", "\n" * max_consecutive_newlines, text)
+
+    # Replace 50+ consecutive spaces with exactly 50 spaces
+    text = re.sub(r" {50,}", " " * max_consecutive_whitespace, text)
+
+    return text.strip()

kiln_ai-0.21.0/kiln_ai/adapters/chunkers/test_base_chunker.py

@@ -0,0 +1,63 @@
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai.adapters.chunkers.base_chunker import (
+    BaseChunker,
+    ChunkingResult,
+    TextChunk,
+)
+from kiln_ai.adapters.chunkers.helpers import clean_up_text
+from kiln_ai.datamodel.chunk import ChunkerConfig, ChunkerType
+
+
+@pytest.fixture
+def config() -> ChunkerConfig:
+    return ChunkerConfig(
+        name="test-chunker",
+        chunker_type=ChunkerType.FIXED_WINDOW,
+        properties={"chunk_size": 100, "chunk_overlap": 10},
+    )
+
+
+class WhitespaceChunker(BaseChunker):
+    async def _chunk(self, text: str) -> ChunkingResult:
+        return ChunkingResult(chunks=[TextChunk(text=chunk) for chunk in text.split()])
+
+
+@pytest.fixture
+def chunker(config: ChunkerConfig) -> WhitespaceChunker:
+    return WhitespaceChunker(config)
+
+
+async def test_base_chunker_chunk_empty_text(chunker: WhitespaceChunker):
+    assert await chunker.chunk("") == ChunkingResult(chunks=[])
+
+
+async def test_base_chunker_concrete_chunker(chunker: WhitespaceChunker):
+    output = await chunker.chunk("Hello, world!")
+    assert len(output.chunks) == 2
+
+
+async def test_base_chunker_calls_clean_up_text(chunker: WhitespaceChunker):
+    with patch(
+        "kiln_ai.adapters.chunkers.base_chunker.clean_up_text"
+    ) as mock_clean_up_text:
+        mock_clean_up_text.side_effect = clean_up_text
+        await chunker.chunk("Hello, world!")
+        mock_clean_up_text.assert_called_once_with("Hello, world!")
+
+
+async def test_base_chunker_empty_text(chunker: WhitespaceChunker):
+    chunks = await chunker.chunk("")
+    assert chunks == ChunkingResult(chunks=[])
+
+
+async def test_base_chunker_empty_text_after_clean_up(chunker: WhitespaceChunker):
+    with patch(
+        "kiln_ai.adapters.chunkers.base_chunker.clean_up_text"
+    ) as mock_clean_up_text:
+        mock_clean_up_text.side_effect = clean_up_text
+        chunks = await chunker.chunk("\n\n   ")
+        mock_clean_up_text.assert_called_once_with("\n\n   ")
+        assert chunks == ChunkingResult(chunks=[])

kiln_ai-0.21.0/kiln_ai/adapters/chunkers/test_chunker_registry.py

@@ -0,0 +1,28 @@
+import pytest
+
+from kiln_ai.adapters.chunkers.chunker_registry import chunker_adapter_from_type
+from kiln_ai.adapters.chunkers.fixed_window_chunker import FixedWindowChunker
+from kiln_ai.datamodel.chunk import ChunkerConfig, ChunkerType
+
+
+def test_chunker_adapter_from_type():
+    chunker = chunker_adapter_from_type(
+        ChunkerType.FIXED_WINDOW,
+        ChunkerConfig(
+            name="test-chunker",
+            chunker_type=ChunkerType.FIXED_WINDOW,
+            properties={
+                # do not use these values in production!
+                "chunk_size": 5555,
+                "chunk_overlap": 1111,
+            },
+        ),
+    )
+    assert isinstance(chunker, FixedWindowChunker)
+    assert chunker.chunker_config.chunk_size() == 5555
+    assert chunker.chunker_config.chunk_overlap() == 1111
+
+
+def test_chunker_adapter_from_type_invalid():
+    with pytest.raises(ValueError):
+        chunker_adapter_from_type("invalid-type", {})