kiln-ai 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

kiln_ai/adapters/__init__.py

@@ -12,10 +12,13 @@ The prompt_builders submodule contains classes that build prompts for use with t
 The repair submodule contains an adapter for the repair task.
 
 The parser submodule contains parsers for the output of the AI models.
+
+The eval submodule contains the code for evaluating the performance of a model.
 """
 
 from . import (
     data_gen,
+    eval,
     fine_tune,
     ml_model_list,
     model_adapters,
@@ -30,4 +33,5 @@ __all__ = [
     "ml_model_list",
     "prompt_builders",
     "repair",
+    "eval",
 ]

kiln_ai/adapters/adapter_registry.py

@@ -3,12 +3,11 @@ 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
-from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
-from kiln_ai.adapters.model_adapters.openai_model_adapter import (
-    OpenAICompatibleAdapter,
-    OpenAICompatibleConfig,
+from kiln_ai.adapters.model_adapters.litellm_adapter import (
+    LiteLlmAdapter,
+    LiteLlmConfig,
 )
-from kiln_ai.adapters.provider_tools import core_provider, openai_compatible_config
+from kiln_ai.adapters.provider_tools import core_provider, lite_llm_config
 from kiln_ai.datamodel import PromptId
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
@@ -26,50 +25,185 @@ def adapter_for_task(
 
     match core_provider_name:
         case ModelProviderName.openrouter:
-            return OpenAICompatibleAdapter(
+            return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                config=OpenAICompatibleConfig(
+                config=LiteLlmConfig(
+                    model_name=model_name,
                     base_url=getenv("OPENROUTER_BASE_URL")
                     or "https://openrouter.ai/api/v1",
-                    api_key=Config.shared().open_router_api_key,
-                    model_name=model_name,
                     provider_name=provider,
                     default_headers={
                         "HTTP-Referer": "https://getkiln.ai/openrouter",
                         "X-Title": "KilnAI",
                     },
+                    additional_body_options={
+                        "api_key": Config.shared().open_router_api_key,
+                    },
                 ),
                 prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
             )
         case ModelProviderName.openai:
-            return OpenAICompatibleAdapter(
+            return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                config=OpenAICompatibleConfig(
-                    api_key=Config.shared().open_ai_api_key,
+                config=LiteLlmConfig(
                     model_name=model_name,
                     provider_name=provider,
+                    additional_body_options={
+                        "api_key": Config.shared().open_ai_api_key,
+                    },
                 ),
                 prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
             )
         case ModelProviderName.openai_compatible:
-            config = openai_compatible_config(model_name)
-            return OpenAICompatibleAdapter(
+            config = lite_llm_config(model_name)
+            return LiteLlmAdapter(
                 kiln_task=kiln_task,
                 config=config,
                 prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
             )
-        # Use LangchainAdapter for the rest
         case ModelProviderName.groq:
-            pass
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    additional_body_options={
+                        "api_key": Config.shared().groq_api_key,
+                    },
+                ),
+            )
         case ModelProviderName.amazon_bedrock:
-            pass
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    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:
-            pass
+            ollama_base_url = (
+                Config.shared().ollama_base_url or "http://localhost:11434"
+            )
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    # 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",
+                ),
+            )
         case ModelProviderName.fireworks_ai:
-            pass
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    additional_body_options={
+                        "api_key": Config.shared().fireworks_api_key,
+                    },
+                ),
+            )
+        case ModelProviderName.anthropic:
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    additional_body_options={
+                        "api_key": Config.shared().anthropic_api_key,
+                    },
+                ),
+            )
+        case ModelProviderName.gemini_api:
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    additional_body_options={
+                        "api_key": Config.shared().gemini_api_key,
+                    },
+                ),
+            )
+        case ModelProviderName.vertex:
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    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,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    additional_body_options={
+                        "api_key": Config.shared().together_api_key,
+                    },
+                ),
+            )
+        case ModelProviderName.azure_openai:
+            return LiteLlmAdapter(
+                kiln_task=kiln_task,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    base_url=Config.shared().azure_openai_endpoint,
+                    model_name=model_name,
+                    provider_name=provider,
+                    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,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+                config=LiteLlmConfig(
+                    model_name=model_name,
+                    provider_name=provider,
+                    additional_body_options={
+                        "api_key": Config.shared().huggingface_api_key,
+                    },
+                ),
+            )
         # These are virtual providers that should have mapped to an actual provider in core_provider
         case ModelProviderName.kiln_fine_tune:
             raise ValueError(
@@ -81,12 +215,3 @@ def adapter_for_task(
             )
         case _:
             raise_exhaustive_enum_error(core_provider_name)
-
-    # We use langchain for all others right now, but moving off it as we touch anything.
-    return LangchainAdapter(
-        kiln_task,
-        model_name=model_name,
-        provider=provider,
-        prompt_id=prompt_id,
-        base_adapter_config=base_adapter_config,
-    )

kiln_ai/adapters/eval/__init__.py

@@ -0,0 +1,28 @@
+"""
+# Evals
+
+This module contains the code for evaluating the performance of a model.
+
+The submodules contain:
+
+- BaseEval: each eval technique implements this interface.
+- G-Eval: an eval implementation, that implements G-Eval and LLM as Judge.
+- EvalRunner: a class that runs an full evaluation (many smaller evals jobs). Includes async parallel processing, and the ability to restart where it left off.
+- EvalRegistry: a registry for all eval implementations.
+
+The datamodel for Evals is in the `kiln_ai.datamodel.eval` module.
+"""
+
+from . import (
+    base_eval,
+    eval_runner,
+    g_eval,
+    registry,
+)
+
+__all__ = [
+    "base_eval",
+    "eval_runner",
+    "g_eval",
+    "registry",
+]

kiln_ai/adapters/eval/eval_runner.py

@@ -139,7 +139,10 @@ class EvalRunner:
             for run_config in self.run_configs or []:
                 already_run[eval_config.id][run_config.id] = set()
             for run in eval_config.runs(readonly=True):
-                if run.task_run_config_id is not None:
+                if (
+                    run.task_run_config_id is not None
+                    and run.task_run_config_id in already_run[eval_config.id]
+                ):
                     already_run[eval_config.id][run.task_run_config_id].add(
                         run.dataset_id
                     )

kiln_ai/adapters/eval/g_eval.py

@@ -1,6 +1,8 @@
 import math
 from typing import Dict, List, Tuple
 
+from litellm.types.utils import ChatCompletionTokenLogprob
+
 from kiln_ai.adapters.adapter_registry import adapter_for_task
 from kiln_ai.adapters.eval.base_eval import BaseEval
 from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig, RunOutput
@@ -8,7 +10,6 @@ from kiln_ai.adapters.prompt_builders import PromptGenerators
 from kiln_ai.datamodel import Project, Task, TaskRun
 from kiln_ai.datamodel.eval import EvalConfig, EvalConfigType, EvalScores
 from kiln_ai.datamodel.task import RunConfig
-from openai.types.chat import ChatCompletionTokenLogprob
 
 # all the tokens we score for, and their float scores.
 TOKEN_TO_SCORE_MAP: Dict[str, float] = {

kiln_ai/adapters/eval/test_base_eval.py

@@ -1,6 +1,7 @@
 import json
 
 import pytest
+
 from kiln_ai.adapters.eval.base_eval import BaseEval
 from kiln_ai.datamodel import BasePrompt, DataSource, DataSourceType
 from kiln_ai.datamodel.eval import Eval, EvalConfig, EvalOutputScore

kiln_ai/adapters/eval/test_eval_runner.py

@@ -2,6 +2,7 @@ from typing import Dict
 from unittest.mock import AsyncMock, patch
 
 import pytest
+
 from kiln_ai.adapters.eval.base_eval import BaseEval
 from kiln_ai.adapters.eval.eval_runner import EvalJob, EvalRunner
 from kiln_ai.datamodel import (

kiln_ai/adapters/eval/test_g_eval.py

@@ -2,6 +2,7 @@ import math
 import pickle
 
 import pytest
+
 from kiln_ai.adapters.eval.g_eval import TOKEN_TO_SCORE_MAP, GEval, GEvalTask
 from kiln_ai.adapters.eval.test_g_eval_data import serialized_run_output
 from kiln_ai.adapters.ml_model_list import built_in_models

kiln_ai/adapters/fine_tune/base_finetune.py

@@ -4,7 +4,12 @@ from typing import Literal
 from pydantic import BaseModel
 
 from kiln_ai.adapters.ml_model_list import built_in_models
-from kiln_ai.datamodel import DatasetSplit, FinetuneDataStrategy, FineTuneStatusType
+from kiln_ai.datamodel import (
+    DatasetSplit,
+    FinetuneDataStrategy,
+    FineTuneStatusType,
+    Task,
+)
 from kiln_ai.datamodel import Finetune as FinetuneModel
 from kiln_ai.utils.name_generator import generate_memorable_name
 
@@ -101,7 +106,7 @@ class BaseFinetuneAdapter(ABC):
             train_split_name=train_split_name,
             validation_split_name=validation_split_name,
             parameters=parameters,
-            system_message=system_message,
+            system_message=cls.augment_system_message(system_message, parent_task),
             thinking_instructions=thinking_instructions,
             parent=parent_task,
             data_strategy=data_strategy,
@@ -114,6 +119,15 @@ class BaseFinetuneAdapter(ABC):
 
         return adapter, datamodel
 
+    @classmethod
+    def augment_system_message(cls, system_message: str, task: Task) -> str:
+        """
+        Augment the system message with additional instructions, such as JSON instructions.
+        """
+
+        # Base implementation does nothing, can be overridden by subclasses
+        return system_message
+
     @abstractmethod
     async def _start(self, dataset: DatasetSplit) -> None:
         """

kiln_ai/adapters/fine_tune/finetune_registry.py

@@ -3,9 +3,11 @@ from typing import Type
 from kiln_ai.adapters.fine_tune.base_finetune import BaseFinetuneAdapter
 from kiln_ai.adapters.fine_tune.fireworks_finetune import FireworksFinetune
 from kiln_ai.adapters.fine_tune.openai_finetune import OpenAIFinetune
+from kiln_ai.adapters.fine_tune.together_finetune import TogetherFinetune
 from kiln_ai.adapters.ml_model_list import ModelProviderName
 
 finetune_registry: dict[ModelProviderName, Type[BaseFinetuneAdapter]] = {
     ModelProviderName.openai: OpenAIFinetune,
     ModelProviderName.fireworks_ai: FireworksFinetune,
+    ModelProviderName.together_ai: TogetherFinetune,
 }