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

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/adapters/adapter_registry.py

@@ -1,5 +1,3 @@
-from os import getenv
-
 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
@@ -9,11 +7,47 @@ from kiln_ai.adapters.model_adapters.litellm_adapter import (
 )
 from kiln_ai.adapters.provider_tools import (
     core_provider,
-    lite_llm_config_for_openai_compatible,
+    lite_llm_core_config_for_provider,
 )
 from kiln_ai.datamodel.task import RunConfigProperties
-from kiln_ai.utils.config import Config
-from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+
+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(
@@ -21,225 +55,8 @@ def adapter_for_task(
     run_config_properties: RunConfigProperties,
     base_adapter_config: AdapterConfig | None = None,
 ) -> BaseAdapter:
-    # Get the provider to run. For things like the fine-tune provider, we want to run the underlying provider
-    core_provider_name = core_provider(
-        run_config_properties.model_name, run_config_properties.model_provider_name
+    return LiteLlmAdapter(
+        kiln_task=kiln_task,
+        config=litellm_core_provider_config(run_config_properties),
+        base_adapter_config=base_adapter_config,
     )
-
-    match core_provider_name:
-        case ModelProviderName.openrouter:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    base_url=getenv("OPENROUTER_BASE_URL")
-                    or "https://openrouter.ai/api/v1",
-                    default_headers={
-                        "HTTP-Referer": "https://getkiln.ai/openrouter",
-                        "X-Title": "KilnAI",
-                    },
-                    additional_body_options={
-                        "api_key": Config.shared().open_router_api_key,
-                    },
-                ),
-                base_adapter_config=base_adapter_config,
-            )
-        case ModelProviderName.siliconflow_cn:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    base_url=getenv("SILICONFLOW_BASE_URL")
-                    or "https://api.siliconflow.cn/v1",
-                    default_headers={
-                        "HTTP-Referer": "https://kiln.tech/siliconflow",
-                        "X-Title": "KilnAI",
-                    },
-                    additional_body_options={
-                        "api_key": Config.shared().siliconflow_cn_api_key,
-                    },
-                ),
-                base_adapter_config=base_adapter_config,
-            )
-        case ModelProviderName.openai:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().open_ai_api_key,
-                    },
-                ),
-                base_adapter_config=base_adapter_config,
-            )
-        case ModelProviderName.openai_compatible:
-            config = lite_llm_config_for_openai_compatible(run_config_properties)
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                config=config,
-                base_adapter_config=base_adapter_config,
-            )
-        case ModelProviderName.groq:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().groq_api_key,
-                    },
-                ),
-            )
-        case ModelProviderName.amazon_bedrock:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "aws_access_key_id": Config.shared().bedrock_access_key,
-                        "aws_secret_access_key": Config.shared().bedrock_secret_key,
-                        # The only region that's widely supported for bedrock
-                        "aws_region_name": "us-west-2",
-                    },
-                ),
-            )
-        case ModelProviderName.ollama:
-            ollama_base_url = (
-                Config.shared().ollama_base_url or "http://localhost:11434"
-            )
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    # Set the Ollama base URL for 2 reasons:
-                    # 1. To use the correct base URL
-                    # 2. We use Ollama's OpenAI compatible API (/v1), and don't just let litellm use the Ollama API. We use more advanced features like json_schema.
-                    base_url=ollama_base_url + "/v1",
-                    additional_body_options={
-                        # LiteLLM errors without an api_key, even though Ollama doesn't support one.
-                        "api_key": "NA",
-                    },
-                ),
-            )
-        case ModelProviderName.docker_model_runner:
-            docker_base_url = (
-                Config.shared().docker_model_runner_base_url
-                or "http://localhost:12434/engines/llama.cpp"
-            )
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    # Docker Model Runner uses OpenAI-compatible API at /v1 endpoint
-                    base_url=docker_base_url + "/v1",
-                    additional_body_options={
-                        # LiteLLM errors without an api_key, even though Docker Model Runner doesn't require one.
-                        "api_key": "DMR",
-                    },
-                ),
-            )
-        case ModelProviderName.fireworks_ai:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().fireworks_api_key,
-                    },
-                ),
-            )
-        case ModelProviderName.anthropic:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().anthropic_api_key,
-                    },
-                ),
-            )
-        case ModelProviderName.gemini_api:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().gemini_api_key,
-                    },
-                ),
-            )
-        case ModelProviderName.vertex:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "vertex_project": Config.shared().vertex_project_id,
-                        "vertex_location": Config.shared().vertex_location,
-                    },
-                ),
-            )
-        case ModelProviderName.together_ai:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().together_api_key,
-                    },
-                ),
-            )
-        case ModelProviderName.azure_openai:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    base_url=Config.shared().azure_openai_endpoint,
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().azure_openai_api_key,
-                        "api_version": "2025-02-01-preview",
-                    },
-                ),
-            )
-        case ModelProviderName.huggingface:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().huggingface_api_key,
-                    },
-                ),
-            )
-        case ModelProviderName.cerebras:
-            return LiteLlmAdapter(
-                kiln_task=kiln_task,
-                base_adapter_config=base_adapter_config,
-                config=LiteLlmConfig(
-                    run_config_properties=run_config_properties,
-                    additional_body_options={
-                        "api_key": Config.shared().cerebras_api_key,
-                    },
-                ),
-            )
-        # These are virtual providers that should have mapped to an actual provider in core_provider
-        case ModelProviderName.kiln_fine_tune:
-            raise ValueError(
-                "Fine tune is not a supported core provider. It should map to an actual provider."
-            )
-        case ModelProviderName.kiln_custom_registry:
-            raise ValueError(
-                "Custom openai compatible provider is not a supported core provider. It should map to an actual provider."
-            )
-        case _:
-            raise_exhaustive_enum_error(core_provider_name)

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/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/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/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/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/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/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", {})