kiln-ai 0.12.0__py3-none-any.whl → 0.13.2__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,189 @@ 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",
+                    additional_body_options={
+                        # LiteLLM errors without an api_key, even though Ollama doesn't support one.
+                        "api_key": "NA",
+                    },
+                ),
+            )
         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 +219,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] = {
@@ -296,9 +297,12 @@ The model produced the following output for the task:
 
         total_score = 0.0
         total_probability = 0.0
+        top_logprobs_contains_primary_token = False
 
-        # Process all valid scoring tokens
+        # Process all valid scoring tokens from alternatives
         for top_logprob in token_logprob.top_logprobs:
+            if top_logprob.token == token_logprob.token:
+                top_logprobs_contains_primary_token = True
             token_score = self.score_from_token_string(top_logprob.token)
             if token_score is not None:
                 # Convert logprob to probability
@@ -306,9 +310,21 @@ The model produced the following output for the task:
                 total_score += token_score * probability
                 total_probability += probability
 
+        # Weird OpenAI 4o bug - sometimes the primary token is included in the top logprobs, sometimes not.
+        # Add the primary token back in if excluded
+        if not top_logprobs_contains_primary_token:
+            if token_logprob.logprob == -9999.0:
+                # Another "bug" - sometimes the logprob is -9999.0. This seems to happen when the rest of the logprobs are tiny probability.
+                total_score += primary_token_score * 1.0
+                total_probability += 1.0
+            else:
+                probability = math.exp(token_logprob.logprob)
+                total_score += primary_token_score * probability
+                total_probability += probability
+
         if total_probability <= 0.0:
             raise RuntimeError(
-                f"No valid scoring tokens found for {token_logprob.token}. This should never happen. Please file a bug if you see this."
+                f"No valid scoring tokens found for {token_logprob.token}. This should never happen as the token has a valid score (so it must be excluded from top logprobs). Please file a bug if you see this."
             )
 
         # Normalize by total probability of valid tokens (LLM may have wanted to generate other non-rating tokens, these shouldn't lower score of rating tokens)

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
@@ -392,12 +393,13 @@ def test_rating_token_to_score(test_eval_config, test_run_config):
             self.logprob = logprob
 
     class MockTokenLogprob:
-        def __init__(self, token, top_logprobs):
+        def __init__(self, token, top_logprobs, logprob):
             self.token = token
             self.top_logprobs = [MockTopLogprob(t, lp) for t, lp in top_logprobs]
+            self.logprob = logprob
 
     # Test single token case
-    token_logprob = MockTokenLogprob("5", [("5", 0.0)])  # log(1) = 0
+    token_logprob = MockTokenLogprob("5", [("5", 0.0)], logprob=1e-8)  # log(1) = 0
     score = g_eval.rating_token_to_score(token_logprob)
     assert score == 5.0
 
@@ -408,18 +410,22 @@ def test_rating_token_to_score(test_eval_config, test_run_config):
             ("4", math.log(0.6)),  # 60% probability
             ("5", math.log(0.4)),  # 40% probability
         ],
+        logprob=math.log(0.6),
     )
     score = g_eval.rating_token_to_score(token_logprob)
     assert pytest.approx(score) == 4.4  # (4 * 0.6 + 5 * 0.4)
 
     # Test invalid token
-    token_logprob = MockTokenLogprob(":", [(":", 0.0)])
+    token_logprob = MockTokenLogprob(":", [(":", 0.0)], logprob=1e-8)
     assert g_eval.rating_token_to_score(token_logprob) is None
 
-    # Test no valid scoring tokens
-    token_logprob = MockTokenLogprob("5", [])
-    with pytest.raises(RuntimeError, match="No valid scoring tokens found"):
-        g_eval.rating_token_to_score(token_logprob)
+    # Test missing from top logprobs
+    token_logprob = MockTokenLogprob("5", [], logprob=1e-8)
+    assert pytest.approx(g_eval.rating_token_to_score(token_logprob)) == 5.0
+
+    # Test missing from top logprobs, with special case logprob
+    token_logprob = MockTokenLogprob("5", [], logprob=-9999)
+    assert pytest.approx(g_eval.rating_token_to_score(token_logprob)) == 5.0
 
 
 def test_g_eval_system_instruction():

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,
 }

kiln_ai/adapters/fine_tune/fireworks_finetune.py

@@ -132,11 +132,18 @@ class FireworksFinetune(BaseFinetuneAdapter):
                 :60
             ]
         )
-        payload = {
+        payload: dict[str, str | dict[str, str | bool]] = {
             "dataset": f"accounts/{account_id}/datasets/{train_file_id}",
             "displayName": display_name,
             "baseModel": self.datamodel.base_model_id,
         }
+        # Add W&B config if API key is set
+        if Config.shared().wandb_api_key:
+            payload["wandbConfig"] = {
+                "enabled": True,
+                "project": "Kiln_AI",
+                "apiKey": Config.shared().wandb_api_key,
+            }
         hyperparameters = self.create_payload_parameters(self.datamodel.parameters)
         payload.update(hyperparameters)
         headers = {

kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py

@@ -340,6 +340,7 @@ async def test_start_success(
     expected_mode,
     expected_format,
 ):
+    Config.shared().wandb_api_key = "test-api-key"
     mock_task.output_json_schema = output_schema
 
     fireworks_finetune.datamodel.parent = mock_task
@@ -378,6 +379,24 @@ async def test_start_success(
         assert fireworks_finetune.datamodel.structured_output_mode == expected_mode
         assert fireworks_finetune.datamodel.properties["endpoint_version"] == "v2"
 
+        # check mockclent.post call values
+        assert mock_client.post.call_count == 1
+        submit_call_values = mock_client.post.call_args[1]
+        assert submit_call_values["json"]["wandbConfig"] == {
+            "enabled": True,
+            "project": "Kiln_AI",
+            "apiKey": "test-api-key",
+        }
+        assert submit_call_values["json"]["baseModel"] == "llama-v2-7b"
+        assert (
+            submit_call_values["json"]["dataset"]
+            == f"accounts/{Config.shared().fireworks_account_id}/datasets/{mock_dataset_id}"
+        )
+        assert (
+            submit_call_values["json"]["displayName"]
+            == f"Kiln AI fine-tuning [ID:{fireworks_finetune.datamodel.id}][name:{fireworks_finetune.datamodel.name}]"
+        )
+
 
 async def test_start_api_error(
     fireworks_finetune, mock_dataset, mock_task, mock_api_key