kiln-ai 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

kiln_ai/__init__.py

@@ -0,0 +1,3 @@
+"""
+.. include:: ../README.md
+"""

kiln_ai/adapters/__init__.py

@@ -0,0 +1,23 @@
+"""
+# Adapters
+
+Adapters are used to connect Kiln to external systems, or to add new functionality to Kiln.
+
+BaseAdapter is extensible, and used for adding adapters that provide AI functionality. There's currently a LangChain adapter which provides a bridge to LangChain.
+
+The ml_model_list submodule contains a list of models that can be used for machine learning tasks. More can easily be added, but we keep a list here of models that are known to work well with Kiln's structured data and tool calling systems.
+
+The prompt_builders submodule contains classes that build prompts for use with the AI agents.
+
+The repair submodule contains an adapter for the repair task.
+"""
+
+from . import base_adapter, langchain_adapters, ml_model_list, prompt_builders, repair
+
+__all__ = [
+    "base_adapter",
+    "langchain_adapters",
+    "ml_model_list",
+    "prompt_builders",
+    "repair",
+]

kiln_ai/adapters/base_adapter.py

@@ -25,6 +25,35 @@ class AdapterInfo:
 
 
 class BaseAdapter(metaclass=ABCMeta):
+    """Base class for AI model adapters that handle task execution.
+
+    This abstract class provides the foundation for implementing model-specific adapters
+    that can process tasks with structured or unstructured inputs/outputs. It handles
+    input/output validation, prompt building, and run tracking.
+
+    Attributes:
+        prompt_builder (BasePromptBuilder): Builder for constructing prompts for the model
+        kiln_task (Task): The task configuration and metadata
+        output_schema (dict | None): JSON schema for validating structured outputs
+        input_schema (dict | None): JSON schema for validating structured inputs
+
+    Example:
+        ```python
+        class CustomAdapter(BaseAdapter):
+            async def _run(self, input: Dict | str) -> Dict | str:
+                # Implementation for specific model
+                pass
+
+            def adapter_info(self) -> AdapterInfo:
+                return AdapterInfo(
+                    adapter_name="custom",
+                    model_name="model-1",
+                    model_provider="provider",
+                    prompt_builder_name="simple"
+                )
+        ```
+    """
+
     def __init__(
         self, kiln_task: Task, prompt_builder: BasePromptBuilder | None = None
     ):

kiln_ai/adapters/langchain_adapters.py

@@ -1,11 +1,11 @@
 from typing import Dict
 
-import kiln_ai.datamodel as datamodel
-from kiln_ai.adapters.prompt_builders import SimplePromptBuilder
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.messages.base import BaseMessage
 
+import kiln_ai.datamodel as datamodel
+
 from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder
 from .ml_model_list import langchain_model_from
 
@@ -77,6 +77,7 @@ class LangChainPromptAdapter(BaseAdapter):
             HumanMessage(content=user_msg),
         ]
         response = self.model.invoke(messages)
+
         if self.has_structured_output():
             if (
                 not isinstance(response, dict)

kiln_ai/adapters/ml_model_list.py

@@ -14,8 +14,18 @@ from pydantic import BaseModel
 
 from ..utils.config import Config
 
+"""
+Provides model configuration and management for various LLM providers and models.
+This module handles the integration with different AI model providers and their respective models,
+including configuration, validation, and instantiation of language models.
+"""
+
 
 class ModelProviderName(str, Enum):
+    """
+    Enumeration of supported AI model providers.
+    """
+
     openai = "openai"
     groq = "groq"
     amazon_bedrock = "amazon_bedrock"
@@ -24,6 +34,10 @@ class ModelProviderName(str, Enum):
 
 
 class ModelFamily(str, Enum):
+    """
+    Enumeration of supported model families/architectures.
+    """
+
     gpt = "gpt"
     llama = "llama"
     phi = "phi"
@@ -33,6 +47,11 @@ class ModelFamily(str, Enum):
 
 # Where models have instruct and raw versions, instruct is default and raw is specified
 class ModelName(str, Enum):
+    """
+    Enumeration of specific model versions supported by the system.
+    Where models have instruct and raw versions, instruct is default and raw is specified.
+    """
+
     llama_3_1_8b = "llama_3_1_8b"
     llama_3_1_70b = "llama_3_1_70b"
     llama_3_1_405b = "llama_3_1_405b"
@@ -47,13 +66,32 @@ class ModelName(str, Enum):
 
 
 class KilnModelProvider(BaseModel):
+    """
+    Configuration for a specific model provider.
+
+    Attributes:
+        name: The provider's identifier
+        supports_structured_output: Whether the provider supports structured output formats
+        provider_options: Additional provider-specific configuration options
+    """
+
     name: ModelProviderName
-    # Allow overriding the model level setting
     supports_structured_output: bool = True
     provider_options: Dict = {}
 
 
 class KilnModel(BaseModel):
+    """
+    Configuration for a specific AI model.
+
+    Attributes:
+        family: The model's architecture family
+        name: The model's identifier
+        friendly_name: Human-readable name for the model
+        providers: List of providers that offer this model
+        supports_structured_output: Whether the model supports structured output formats
+    """
+
     family: str
     name: str
     friendly_name: str
@@ -292,6 +330,18 @@ built_in_models: List[KilnModel] = [
 
 
 def provider_name_from_id(id: str) -> str:
+    """
+    Converts a provider ID to its human-readable name.
+
+    Args:
+        id: The provider identifier string
+
+    Returns:
+        The human-readable name of the provider
+
+    Raises:
+        ValueError: If the provider ID is invalid or unhandled
+    """
     if id in ModelProviderName.__members__:
         enum_id = ModelProviderName(id)
         match enum_id:
@@ -350,6 +400,15 @@ def get_config_value(key: str):
 
 
 def check_provider_warnings(provider_name: ModelProviderName):
+    """
+    Validates that required configuration is present for a given provider.
+
+    Args:
+        provider_name: The provider to check
+
+    Raises:
+        ValueError: If required configuration keys are missing
+    """
     warning_check = provider_warnings.get(provider_name)
     if warning_check is None:
         return
@@ -359,6 +418,19 @@ def check_provider_warnings(provider_name: ModelProviderName):
 
 
 def langchain_model_from(name: str, provider_name: str | None = None) -> BaseChatModel:
+    """
+    Creates a LangChain chat model instance for the specified model and provider.
+
+    Args:
+        name: The name of the model to instantiate
+        provider_name: Optional specific provider to use (defaults to first available)
+
+    Returns:
+        A configured LangChain chat model instance
+
+    Raises:
+        ValueError: If the model/provider combination is invalid or misconfigured
+    """
     if name not in ModelName.__members__:
         raise ValueError(f"Invalid name: {name}")
 
@@ -413,7 +485,7 @@ def langchain_model_from(name: str, provider_name: str | None = None) -> BaseCha
             openai_api_key=api_key,  # type: ignore[arg-type]
             openai_api_base=base_url,  # type: ignore[arg-type]
             default_headers={
-                "HTTP-Referer": "https://kiln-ai.com/openrouter",
+                "HTTP-Referer": "https://getkiln.ai/openrouter",
                 "X-Title": "KilnAI",
             },
         )
@@ -421,14 +493,27 @@ def langchain_model_from(name: str, provider_name: str | None = None) -> BaseCha
         raise ValueError(f"Invalid model or provider: {name} - {provider_name}")
 
 
-def ollama_base_url():
+def ollama_base_url() -> str:
+    """
+    Gets the base URL for Ollama API connections.
+
+    Returns:
+        The base URL to use for Ollama API calls, using environment variable if set
+        or falling back to localhost default
+    """
     env_base_url = os.getenv("OLLAMA_BASE_URL")
     if env_base_url is not None:
         return env_base_url
     return "http://localhost:11434"
 
 
-async def ollama_online():
+async def ollama_online() -> bool:
+    """
+    Checks if the Ollama service is available and responding.
+
+    Returns:
+        True if Ollama is available and responding, False otherwise
+    """
     try:
         httpx.get(ollama_base_url() + "/api/tags")
     except httpx.RequestError:

kiln_ai/adapters/prompt_builders.py

@@ -2,25 +2,53 @@ import json
 from abc import ABCMeta, abstractmethod
 from typing import Dict
 
-from kiln_ai.datamodel import Task
+from kiln_ai.datamodel import Task, TaskRun
 from kiln_ai.utils.formatting import snake_case
 
 
 class BasePromptBuilder(metaclass=ABCMeta):
+    """Base class for building prompts from tasks.
+
+    Provides the core interface and basic functionality for prompt builders.
+    """
+
     def __init__(self, task: Task):
+        """Initialize the prompt builder with a task.
+
+        Args:
+            task (Task): The task containing instructions and requirements.
+        """
         self.task = task
 
     @abstractmethod
     def build_prompt(self) -> str:
+        """Build and return the complete prompt string.
+
+        Returns:
+            str: The constructed prompt.
+        """
         pass
 
-    # override to change the name of the prompt builder (if changing class names)
     @classmethod
     def prompt_builder_name(cls) -> str:
+        """Returns the name of the prompt builder, to be used for persisting into the datastore.
+
+        Default implementation gets the name of the prompt builder in snake case. If you change the class name, you should override this so prior saved data is compatible.
+
+        Returns:
+            str: The prompt builder name in snake_case format.
+        """
         return snake_case(cls.__name__)
 
-    # Can be overridden to add more information to the user message
     def build_user_message(self, input: Dict | str) -> str:
+        """Build a user message from the input.
+
+        Args:
+            input (Union[Dict, str]): The input to format into a message.
+
+        Returns:
+            str: The formatted user message.
+        """
         if isinstance(input, Dict):
             return f"The input is:\n{json.dumps(input, indent=2)}"
 
@@ -28,7 +56,14 @@ class BasePromptBuilder(metaclass=ABCMeta):
 
 
 class SimplePromptBuilder(BasePromptBuilder):
+    """A basic prompt builder that combines task instruction with requirements."""
+
     def build_prompt(self) -> str:
+        """Build a simple prompt with instruction and requirements.
+
+        Returns:
+            str: The constructed prompt string.
+        """
         base_prompt = self.task.instruction
 
         # TODO: this is just a quick version. Formatting and best practices TBD
@@ -44,11 +79,23 @@ class SimplePromptBuilder(BasePromptBuilder):
 
 
 class MultiShotPromptBuilder(BasePromptBuilder):
+    """A prompt builder that includes multiple examples in the prompt."""
+
     @classmethod
     def example_count(cls) -> int:
+        """Get the maximum number of examples to include in the prompt.
+
+        Returns:
+            int: The maximum number of examples (default 25).
+        """
         return 25
 
     def build_prompt(self) -> str:
+        """Build a prompt with instruction, requirements, and multiple examples.
+
+        Returns:
+            str: The constructed prompt string with examples.
+        """
         base_prompt = f"# Instruction\n\n{ self.task.instruction }\n\n"
 
         if len(self.task.requirements) > 0:
@@ -57,7 +104,24 @@ class MultiShotPromptBuilder(BasePromptBuilder):
                 base_prompt += f"{i+1}) {requirement.instruction}\n"
             base_prompt += "\n"
 
-        valid_examples: list[tuple[str, str]] = []
+        valid_examples = self.collect_examples()
+
+        if len(valid_examples) == 0:
+            return base_prompt
+
+        base_prompt += "# Example Outputs\n\n"
+        for i, example in enumerate(valid_examples):
+            base_prompt += self.prompt_section_for_example(i, example)
+
+        return base_prompt
+
+    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
+        # Prefer repaired output if it exists, otherwise use the regular output
+        output = example.repaired_output or example.output
+        return f"## Example {index+1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
+
+    def collect_examples(self) -> list[TaskRun]:
+        valid_examples: list[TaskRun] = []
         runs = self.task.runs()
 
         # first pass, we look for repaired outputs. These are the best examples.
@@ -65,7 +129,7 @@ class MultiShotPromptBuilder(BasePromptBuilder):
             if len(valid_examples) >= self.__class__.example_count():
                 break
             if run.repaired_output is not None:
-                valid_examples.append((run.input, run.repaired_output.output))
+                valid_examples.append(run)
 
         # second pass, we look for high quality outputs (rating based)
         # Minimum is "high_quality" (4 star in star rating scale), then sort by rating
@@ -84,33 +148,66 @@ class MultiShotPromptBuilder(BasePromptBuilder):
         for run in runs_with_rating:
             if len(valid_examples) >= self.__class__.example_count():
                 break
-            valid_examples.append((run.input, run.output.output))
-
-        if len(valid_examples) > 0:
-            base_prompt += "# Example Outputs\n\n"
-            for i, example in enumerate(valid_examples):
-                base_prompt += (
-                    f"## Example {i+1}\n\nInput: {example[0]}\nOutput: {example[1]}\n\n"
-                )
-
-        return base_prompt
+            valid_examples.append(run)
+        return valid_examples
 
 
 class FewShotPromptBuilder(MultiShotPromptBuilder):
+    """A prompt builder that includes a small number of examples in the prompt."""
+
     @classmethod
     def example_count(cls) -> int:
+        """Get the maximum number of examples to include in the prompt.
+
+        Returns:
+            int: The maximum number of examples (4).
+        """
         return 4
 
 
+class RepairsPromptBuilder(MultiShotPromptBuilder):
+    """A prompt builder that includes multiple examples in the prompt, including repaired instructions describing what was wrong, and how it was fixed."""
+
+    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
+        if (
+            not example.repaired_output
+            or not example.repair_instructions
+            or not example.repaired_output.output
+        ):
+            return super().prompt_section_for_example(index, example)
+
+        prompt_section = f"## Example {index+1}\n\nInput: {example.input}\n\n"
+        prompt_section += (
+            f"Initial Output Which Was Insufficient: {example.output.output}\n\n"
+        )
+        prompt_section += f"Instructions On How to Improve the Initial Output: {example.repair_instructions}\n\n"
+        prompt_section += (
+            f"Repaired Output Which is Sufficient: {example.repaired_output.output}\n\n"
+        )
+        return prompt_section
+
+
 prompt_builder_registry = {
     "simple_prompt_builder": SimplePromptBuilder,
     "multi_shot_prompt_builder": MultiShotPromptBuilder,
     "few_shot_prompt_builder": FewShotPromptBuilder,
+    "repairs_prompt_builder": RepairsPromptBuilder,
 }
 
 
 # Our UI has some names that are not the same as the class names, which also hint parameters.
 def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
+    """Convert a name used in the UI to the corresponding prompt builder class.
+
+    Args:
+        ui_name (str): The UI name for the prompt builder type.
+
+    Returns:
+        type[BasePromptBuilder]: The corresponding prompt builder class.
+
+    Raises:
+        ValueError: If the UI name is not recognized.
+    """
     match ui_name:
         case "basic":
             return SimplePromptBuilder
@@ -118,5 +215,7 @@ def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
             return FewShotPromptBuilder
         case "many_shot":
             return MultiShotPromptBuilder
+        case "repairs":
+            return RepairsPromptBuilder
         case _:
             raise ValueError(f"Unknown prompt builder: {ui_name}")

kiln_ai/adapters/repair/__init__.py

@@ -0,0 +1,11 @@
+"""
+# Repair
+
+A task to repair a task output. Used to fix up task outputs that don't meet the requirements.
+
+Not intended to be persisted into a project file.
+"""
+
+from . import repair_task
+
+__all__ = ["repair_task"]

kiln_ai/adapters/repair/repair_task.py

@@ -1,9 +1,10 @@
 import json
 from typing import Type
 
+from pydantic import BaseModel, Field
+
 from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_registry
 from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
-from pydantic import BaseModel, Field
 
 
 # TODO add evaluator rating

kiln_ai/adapters/repair/test_repair_task.py

@@ -3,6 +3,8 @@ import os
 from unittest.mock import AsyncMock, patch
 
 import pytest
+from pydantic import ValidationError
+
 from kiln_ai.adapters.langchain_adapters import (
     LangChainPromptAdapter,
 )
@@ -19,7 +21,6 @@ from kiln_ai.datamodel import (
     TaskRequirement,
     TaskRun,
 )
-from pydantic import ValidationError
 
 json_joke_schema = """{
   "type": "object",

kiln_ai/adapters/test_langchain_adapter.py

@@ -1,6 +1,7 @@
+from langchain_groq import ChatGroq
+
 from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
 from kiln_ai.adapters.test_prompt_adaptors import build_test_task
-from langchain_groq import ChatGroq
 
 
 def test_langchain_adapter_munge_response(tmp_path):

kiln_ai/adapters/test_ml_model_list.py

@@ -2,7 +2,7 @@ from unittest.mock import patch
 
 import pytest
 
-from libs.core.kiln_ai.adapters.ml_model_list import (
+from kiln_ai.adapters.ml_model_list import (
     ModelProviderName,
     check_provider_warnings,
     provider_name_from_id,
@@ -12,7 +12,7 @@ from libs.core.kiln_ai.adapters.ml_model_list import (
 
 @pytest.fixture
 def mock_config():
-    with patch("libs.core.kiln_ai.adapters.ml_model_list.get_config_value") as mock:
+    with patch("kiln_ai.adapters.ml_model_list.get_config_value") as mock:
         yield mock
 
 

kiln_ai/adapters/test_prompt_adaptors.py

@@ -1,11 +1,12 @@
 import os
 from pathlib import Path
 
-import kiln_ai.datamodel as datamodel
 import pytest
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+
+import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
 from kiln_ai.adapters.ml_model_list import built_in_models, ollama_online
-from langchain_core.language_models.fake_chat_models import FakeListChatModel
 
 
 @pytest.mark.paid

kiln_ai/adapters/test_prompt_builders.py

@@ -1,10 +1,12 @@
 import json
 
 import pytest
+
 from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter
 from kiln_ai.adapters.prompt_builders import (
     FewShotPromptBuilder,
     MultiShotPromptBuilder,
+    RepairsPromptBuilder,
     SimplePromptBuilder,
     prompt_builder_from_ui_name,
 )
@@ -71,7 +73,8 @@ def test_simple_prompt_builder_structured_output(tmp_path):
     assert input not in prompt
 
 
-def test_multi_shot_prompt_builder(tmp_path):
+@pytest.fixture
+def task_with_examples(tmp_path):
     # Create a project and task hierarchy
     project = Project(name="Test Project", path=(tmp_path / "test_project.kiln"))
     project.save_to_file()
@@ -192,9 +195,12 @@ def test_multi_shot_prompt_builder(tmp_path):
     )
     e3.save_to_file()
     check_example_outputs(task, 3)
+    return task
+
 
+def test_multi_shot_prompt_builder(task_with_examples):
     # Verify the order of examples
-    prompt_builder = MultiShotPromptBuilder(task=task)
+    prompt_builder = MultiShotPromptBuilder(task=task_with_examples)
     prompt = prompt_builder.build_prompt()
     assert "Why did the cow cross the road?" in prompt
     assert prompt.index("Why did the cow cross the road?") < prompt.index(
@@ -299,12 +305,14 @@ def check_example_outputs(task: Task, count: int):
 def test_prompt_builder_name():
     assert SimplePromptBuilder.prompt_builder_name() == "simple_prompt_builder"
     assert MultiShotPromptBuilder.prompt_builder_name() == "multi_shot_prompt_builder"
+    assert RepairsPromptBuilder.prompt_builder_name() == "repairs_prompt_builder"
 
 
 def test_prompt_builder_from_ui_name():
     assert prompt_builder_from_ui_name("basic") == SimplePromptBuilder
     assert prompt_builder_from_ui_name("few_shot") == FewShotPromptBuilder
     assert prompt_builder_from_ui_name("many_shot") == MultiShotPromptBuilder
+    assert prompt_builder_from_ui_name("repairs") == RepairsPromptBuilder
 
     with pytest.raises(ValueError, match="Unknown prompt builder: invalid_name"):
         prompt_builder_from_ui_name("invalid_name")
@@ -313,3 +321,18 @@ def test_prompt_builder_from_ui_name():
 def test_example_count():
     assert FewShotPromptBuilder.example_count() == 4
     assert MultiShotPromptBuilder.example_count() == 25
+
+
+def test_repair_multi_shot_prompt_builder(task_with_examples):
+    # Verify the order of examples
+    prompt_builder = RepairsPromptBuilder(task=task_with_examples)
+    prompt = prompt_builder.build_prompt()
+    assert (
+        'Repaired Output Which is Sufficient: {"joke": "Why did the cow cross the road? To get to the udder side!"}'
+        in prompt
+    )
+    assert "Instructions On How to Improve the Initial Output: Fix the joke" in prompt
+    assert (
+        'Initial Output Which Was Insufficient: {"joke": "Moo I am a cow joke."}'
+        in prompt
+    )

kiln_ai/adapters/test_saving_adapter_results.py

@@ -1,6 +1,7 @@
 from unittest.mock import patch
 
 import pytest
+
 from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter
 from kiln_ai.datamodel import (
     DataSource,

kiln_ai/adapters/test_structured_output.py

@@ -3,8 +3,9 @@ from typing import Dict
 
 import jsonschema
 import jsonschema.exceptions
-import kiln_ai.datamodel as datamodel
 import pytest
+
+import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter
 from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
 from kiln_ai.adapters.ml_model_list import (