kiln-ai 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

kiln_ai/adapters/model_adapters/base_adapter.py

@@ -3,10 +3,7 @@ from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
 from typing import Dict, Tuple
 
-from kiln_ai.adapters.chat.chat_formatter import (
-    ChatFormatter,
-    get_chat_formatter,
-)
+from kiln_ai.adapters.chat.chat_formatter import ChatFormatter, get_chat_formatter
 from kiln_ai.adapters.ml_model_list import (
     KilnModelProvider,
     StructuredOutputMode,
@@ -156,9 +153,16 @@ class BaseAdapter(metaclass=ABCMeta):
                 )
 
         # Validate reasoning content is present (if reasoning)
-        if provider.reasoning_capable and (
-            not parsed_output.intermediate_outputs
-            or "reasoning" not in parsed_output.intermediate_outputs
+        if (
+            provider.reasoning_capable
+            and (
+                not parsed_output.intermediate_outputs
+                or "reasoning" not in parsed_output.intermediate_outputs
+            )
+            and not (
+                provider.reasoning_optional_for_structured_output
+                and self.has_structured_output()
+            )
         ):
             raise RuntimeError(
                 "Reasoning is required for this model, but no reasoning was returned."

kiln_ai/adapters/model_adapters/litellm_adapter.py

@@ -235,7 +235,7 @@ class LiteLlmAdapter(BaseAdapter):
         }
 
     def build_extra_body(self, provider: KilnModelProvider) -> dict[str, Any]:
-        # TODO P1: Don't love having this logic here. But it's a usability improvement
+        # Don't love having this logic here. But it's worth the usability improvement
         # so better to keep it than exclude it. Should figure out how I want to isolate
         # this sort of logic so it's config driven and can be overridden
 
@@ -251,6 +251,15 @@ class LiteLlmAdapter(BaseAdapter):
                 "exclude": False,
             }
 
+        if provider.gemini_reasoning_enabled:
+            extra_body["reasoning"] = {
+                "enabled": True,
+            }
+
+        if provider.name == ModelProviderName.openrouter:
+            # Ask OpenRouter to include usage in the response (cost)
+            extra_body["usage"] = {"include": True}
+
         if provider.anthropic_extended_thinking:
             extra_body["thinking"] = {"type": "enabled", "budget_tokens": 4000}
 
@@ -276,6 +285,10 @@ class LiteLlmAdapter(BaseAdapter):
             # Oddball case, R1 14/8/1.5B fail with this param, even though they support thinking params.
             provider_options["require_parameters"] = False
 
+        # Siliconflow uses a bool flag for thinking, for some models
+        if provider.siliconflow_enable_thinking is not None:
+            extra_body["enable_thinking"] = provider.siliconflow_enable_thinking
+
         if len(provider_options) > 0:
             extra_body["provider"] = provider_options
 
@@ -321,6 +334,10 @@ class LiteLlmAdapter(BaseAdapter):
                 litellm_provider_name = "vertex_ai"
             case ModelProviderName.together_ai:
                 litellm_provider_name = "together_ai"
+            case ModelProviderName.cerebras:
+                litellm_provider_name = "cerebras"
+            case ModelProviderName.siliconflow_cn:
+                is_custom = True
             case ModelProviderName.openai_compatible:
                 is_custom = True
             case ModelProviderName.kiln_custom_registry:
@@ -386,7 +403,12 @@ class LiteLlmAdapter(BaseAdapter):
 
     def usage_from_response(self, response: ModelResponse) -> Usage | None:
         litellm_usage = response.get("usage", None)
+
+        # LiteLLM isn't consistent in how it returns the cost.
         cost = response._hidden_params.get("response_cost", None)
+        if cost is None and litellm_usage:
+            cost = litellm_usage.get("cost", None)
+
         if not litellm_usage and not cost:
             return None
 

kiln_ai/adapters/model_adapters/test_base_adapter.py

@@ -4,7 +4,6 @@ import pytest
 
 from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMode
 from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput
-from kiln_ai.adapters.parsers.request_formatters import request_formatter_from_id
 from kiln_ai.datamodel import Task
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.task import RunConfig, RunConfigProperties
@@ -103,7 +102,7 @@ async def test_model_provider_invalid_provider_model_name(base_task):
     """Test error when model or provider name is missing"""
     # Test with missing model name
     with pytest.raises(ValueError, match="Input should be"):
-        adapter = MockAdapter(
+        MockAdapter(
             run_config=RunConfig(
                 task=base_task,
                 model_name="test_model",

kiln_ai/adapters/model_adapters/test_litellm_adapter.py

@@ -7,9 +7,7 @@ import pytest
 from kiln_ai.adapters.ml_model_list import ModelProviderName, StructuredOutputMode
 from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
 from kiln_ai.adapters.model_adapters.litellm_adapter import LiteLlmAdapter
-from kiln_ai.adapters.model_adapters.litellm_config import (
-    LiteLlmConfig,
-)
+from kiln_ai.adapters.model_adapters.litellm_config import LiteLlmConfig
 from kiln_ai.datamodel import Project, Task, Usage
 from kiln_ai.datamodel.task import RunConfigProperties
 
@@ -242,6 +240,8 @@ def test_tool_call_params_strict(config, mock_task):
         (ModelProviderName.huggingface, "huggingface"),
         (ModelProviderName.vertex, "vertex_ai"),
         (ModelProviderName.together_ai, "together_ai"),
+        # for openai-compatible providers, we expect openai as the provider name
+        (ModelProviderName.siliconflow_cn, "openai"),
     ],
 )
 def test_litellm_model_id_standard_providers(
@@ -352,6 +352,43 @@ def test_litellm_model_id_unknown_provider(config, mock_task):
                 adapter.litellm_model_id()
 
 
+@pytest.mark.parametrize(
+    "provider_name,expected_usage_param",
+    [
+        (ModelProviderName.openrouter, {"usage": {"include": True}}),
+        (ModelProviderName.openai, {}),
+        (ModelProviderName.anthropic, {}),
+        (ModelProviderName.groq, {}),
+    ],
+)
+def test_build_extra_body_openrouter_usage(
+    config, mock_task, provider_name, expected_usage_param
+):
+    """Test build_extra_body includes usage parameter for OpenRouter providers"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Create a mock provider with the specified name and minimal required attributes
+    mock_provider = Mock()
+    mock_provider.name = provider_name
+    mock_provider.thinking_level = None
+    mock_provider.require_openrouter_reasoning = False
+    mock_provider.anthropic_extended_thinking = False
+    mock_provider.r1_openrouter_options = False
+    mock_provider.logprobs_openrouter_options = False
+    mock_provider.openrouter_skip_required_parameters = False
+
+    # Call build_extra_body
+    extra_body = adapter.build_extra_body(mock_provider)
+
+    # Verify the usage parameter is included only for OpenRouter
+    for key, value in expected_usage_param.items():
+        assert extra_body.get(key) == value
+
+    # Verify non-OpenRouter providers don't have the usage parameter
+    if provider_name != ModelProviderName.openrouter:
+        assert "usage" not in extra_body
+
+
 @pytest.mark.asyncio
 async def test_build_completion_kwargs_custom_temperature_top_p(config, mock_task):
     """Test build_completion_kwargs with custom temperature and top_p values"""
@@ -474,6 +511,17 @@ async def test_build_completion_kwargs(
         ({"prompt_tokens": 10}, None, None),
         # Invalid cost type (should be ignored)
         (None, "0.5", None),
+        # Cost in OpenRouter format
+        (
+            litellm.types.utils.Usage(
+                prompt_tokens=10,
+                completion_tokens=20,
+                total_tokens=30,
+                cost=0.5,
+            ),
+            None,
+            Usage(input_tokens=10, output_tokens=20, total_tokens=30, cost=0.5),
+        ),
     ],
 )
 def test_usage_from_response(config, mock_task, litellm_usage, cost, expected_usage):
@@ -504,3 +552,22 @@ def test_usage_from_response(config, mock_task, litellm_usage, cost, expected_us
 
     # Verify the response was queried correctly
     response.get.assert_called_once_with("usage", None)
+
+
+@pytest.mark.parametrize(
+    "enable_thinking",
+    [
+        True,
+        False,
+    ],
+)
+def test_build_extra_body_enable_thinking(config, mock_task, enable_thinking):
+    provider = Mock()
+    provider.name = ModelProviderName.siliconflow_cn
+    provider.siliconflow_enable_thinking = enable_thinking
+
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    extra_body = adapter.build_extra_body(provider)
+
+    assert extra_body["enable_thinking"] == enable_thinking

kiln_ai/adapters/model_adapters/test_structured_output.py

@@ -6,14 +6,8 @@ import pytest
 
 import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.adapter_registry import adapter_for_task
-from kiln_ai.adapters.ml_model_list import (
-    built_in_models,
-)
-from kiln_ai.adapters.model_adapters.base_adapter import (
-    BaseAdapter,
-    RunOutput,
-    Usage,
-)
+from kiln_ai.adapters.ml_model_list import built_in_models
+from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput, Usage
 from kiln_ai.adapters.ollama_tools import ollama_online
 from kiln_ai.adapters.test_prompt_adaptors import get_all_models_and_providers
 from kiln_ai.datamodel import PromptId
@@ -180,8 +174,14 @@ async def run_structured_output_test(tmp_path: Path, model_name: str, provider:
     # Check reasoning models
     assert a._model_provider is not None
     if a._model_provider.reasoning_capable:
-        assert "reasoning" in run.intermediate_outputs
-        assert isinstance(run.intermediate_outputs["reasoning"], str)
+        # some providers have reasoning_capable models that do not return the reasoning
+        # for structured output responses (they provide it only for non-structured output)
+        if a._model_provider.reasoning_optional_for_structured_output:
+            # models may be updated to include the reasoning in the future
+            assert "reasoning" not in run.intermediate_outputs
+        else:
+            assert "reasoning" in run.intermediate_outputs
+            assert isinstance(run.intermediate_outputs["reasoning"], str)
 
 
 def build_structured_input_test_task(tmp_path: Path):
@@ -245,7 +245,7 @@ async def run_structured_input_task_no_validation(
     try:
         run = await a.invoke({"a": 2, "b": 2, "c": 2})
         response = run.output.output
-        return response, a
+        return response, a, run
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
@@ -260,7 +260,7 @@ async def run_structured_input_task(
     provider: str,
     prompt_id: PromptId,
 ):
-    response, a = await run_structured_input_task_no_validation(
+    response, a, run = await run_structured_input_task_no_validation(
         task, model_name, provider, prompt_id
     )
     assert response is not None
@@ -350,7 +350,7 @@ When asked for a final result, this is the format (for an equilateral example):
 """
     task.output_json_schema = json.dumps(triangle_schema)
     task.save_to_file()
-    response, adapter = await run_structured_input_task_no_validation(
+    response, adapter, _ = await run_structured_input_task_no_validation(
         task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
     )
 

kiln_ai/adapters/parsers/parser_registry.py

@@ -1,5 +1,3 @@
-from typing import Type
-
 from kiln_ai.adapters.ml_model_list import ModelParserID
 from kiln_ai.adapters.parsers.base_parser import BaseParser
 from kiln_ai.adapters.parsers.r1_parser import R1ThinkingParser

kiln_ai/adapters/parsers/r1_parser.py

@@ -1,5 +1,4 @@
 from kiln_ai.adapters.parsers.base_parser import BaseParser
-from kiln_ai.adapters.parsers.json_parser import parse_json_string
 from kiln_ai.adapters.run_output import RunOutput
 
 

kiln_ai/adapters/parsers/test_r1_parser.py

@@ -46,7 +46,7 @@ def test_response_with_whitespace(parser):
     assert parsed.output.strip() == "This is the result"
 
 
-def test_empty_thinking_content(parser):
+def test_empty_thinking_content_multiline(parser):
     response = RunOutput(
         output="""
         <think>

kiln_ai/adapters/provider_tools.py

@@ -5,18 +5,13 @@ from typing import Dict, List
 from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     KilnModelProvider,
-    ModelName,
     ModelParserID,
     ModelProviderName,
     StructuredOutputMode,
     built_in_models,
 )
-from kiln_ai.adapters.model_adapters.litellm_config import (
-    LiteLlmConfig,
-)
-from kiln_ai.adapters.ollama_tools import (
-    get_ollama_connection,
-)
+from kiln_ai.adapters.model_adapters.litellm_config import LiteLlmConfig
+from kiln_ai.adapters.ollama_tools import get_ollama_connection
 from kiln_ai.datamodel import Finetune, Task
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.registry import project_from_id
@@ -75,30 +70,24 @@ def builtin_model_from(
     name: str, provider_name: str | None = None
 ) -> KilnModelProvider | None:
     """
-    Gets a model and provider from the built-in list of models.
+    Gets a model provider from the built-in list of models.
 
     Args:
         name: The name of the model to get
         provider_name: Optional specific provider to use (defaults to first available)
 
     Returns:
-        A tuple of (provider, model)
-
-    Raises:
-        ValueError: If the model or provider is not found, or if the provider is misconfigured
+        A KilnModelProvider, or None if not found
     """
-    if name not in ModelName.__members__:
-        return None
-
     # Select the model from built_in_models using the name
-    model = next(filter(lambda m: m.name == name, built_in_models))
+    model = next(filter(lambda m: m.name == name, built_in_models), None)
     if model is None:
-        raise ValueError(f"Model {name} not found")
+        return None
 
-    # If a provider is provided, select the provider from the model's provider_config
+    # If a provider is provided, select the appropriate provider. Otherwise, use the first available.
     provider: KilnModelProvider | None = None
     if model.providers is None or len(model.providers) == 0:
-        raise ValueError(f"Model {name} has no providers")
+        return None
     elif provider_name is None:
         provider = model.providers[0]
     else:
@@ -384,6 +373,10 @@ def provider_name_from_id(id: str) -> str:
                 return "Google Vertex AI"
             case ModelProviderName.together_ai:
                 return "Together AI"
+            case ModelProviderName.siliconflow_cn:
+                return "SiliconFlow"
+            case ModelProviderName.cerebras:
+                return "Cerebras"
             case _:
                 # triggers pyright warning if I miss a case
                 raise_exhaustive_enum_error(enum_id)
@@ -442,4 +435,12 @@ provider_warnings: Dict[ModelProviderName, ModelProviderWarning] = {
         required_config_keys=["together_api_key"],
         message="Attempted to use Together without an API key set. \nGet your API key from https://together.ai/settings/keys",
     ),
+    ModelProviderName.siliconflow_cn: ModelProviderWarning(
+        required_config_keys=["siliconflow_cn_api_key"],
+        message="Attempted to use SiliconFlow without an API key set. \nGet your API key from https://cloud.siliconflow.cn/account/ak",
+    ),
+    ModelProviderName.cerebras: ModelProviderWarning(
+        required_config_keys=["cerebras_api_key"],
+        message="Attempted to use Cerebras without an API key set. \nGet your API key from https://cloud.cerebras.ai/platform",
+    ),
 }

kiln_ai/adapters/remote_config.py

@@ -0,0 +1,113 @@
+import argparse
+import json
+import logging
+import os
+import threading
+from pathlib import Path
+from typing import Any, List
+
+import requests
+from pydantic import ValidationError
+
+from .ml_model_list import KilnModel, KilnModelProvider, built_in_models
+
+logger = logging.getLogger(__name__)
+
+
+def serialize_config(models: List[KilnModel], path: str | Path) -> None:
+    data = {"model_list": [m.model_dump(mode="json") for m in models]}
+    Path(path).write_text(json.dumps(data, indent=2, sort_keys=True))
+
+
+def deserialize_config_at_path(path: str | Path) -> List[KilnModel]:
+    raw = json.loads(Path(path).read_text())
+    return deserialize_config_data(raw)
+
+
+def deserialize_config_data(config_data: Any) -> List[KilnModel]:
+    if not isinstance(config_data, dict):
+        raise ValueError(f"Remote config expected dict, got {type(config_data)}")
+
+    model_list = config_data.get("model_list", None)
+    if not isinstance(model_list, list):
+        raise ValueError(
+            f"Remote config expected list of models, got {type(model_list)}"
+        )
+
+    # We must be careful here, because some of the JSON data may be generated from a forward
+    # version of the code that has newer fields / versions of the fields, that may cause
+    # the current client this code is running on to fail to validate the item into a KilnModel.
+    models = []
+    for model_data in model_list:
+        # We skip any model that fails validation - the models that the client can support
+        # will be pulled from the remote config, but the user will need to update their
+        # client to the latest version to see the newer models that break backwards compatibility.
+        try:
+            providers_list = model_data.get("providers", [])
+
+            providers = []
+            for provider_data in providers_list:
+                try:
+                    provider = KilnModelProvider.model_validate(provider_data)
+                    providers.append(provider)
+                except ValidationError as e:
+                    logger.warning(
+                        "Failed to validate a model provider from remote config. Upgrade Kiln to use this model. Details %s: %s",
+                        provider_data,
+                        e,
+                    )
+
+            # this ensures the model deserialization won't fail because of a bad provider
+            model_data["providers"] = []
+
+            # now we validate the model without its providers
+            model = KilnModel.model_validate(model_data)
+
+            # and we attach back the providers that passed our validation
+            model.providers = providers
+            models.append(model)
+        except ValidationError as e:
+            logger.warning(
+                "Failed to validate a model from remote config. Upgrade Kiln to use this model. Details %s: %s",
+                model_data,
+                e,
+            )
+    return models
+
+
+def load_from_url(url: str) -> List[KilnModel]:
+    response = requests.get(url, timeout=10)
+    response.raise_for_status()
+    data = response.json()
+    return deserialize_config_data(data)
+
+
+def dump_builtin_config(path: str | Path) -> None:
+    serialize_config(built_in_models, path)
+
+
+def load_remote_models(url: str) -> None:
+    if os.environ.get("KILN_SKIP_REMOTE_MODEL_LIST") == "true":
+        return
+
+    def fetch_and_replace() -> None:
+        try:
+            models = load_from_url(url)
+            built_in_models[:] = models
+        except Exception as exc:
+            # Do not crash startup, but surface the issue
+            logger.warning("Failed to fetch remote model list from %s: %s", url, exc)
+
+    thread = threading.Thread(target=fetch_and_replace, daemon=True)
+    thread.start()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("path", help="output path")
+    args = parser.parse_args()
+    dump_builtin_config(args.path)
+
+
+if __name__ == "__main__":
+    main()

kiln_ai/adapters/repair/repair_task.py

@@ -1,17 +1,12 @@
 import json
-from typing import Type
 
 from pydantic import BaseModel, Field
 
-from kiln_ai.adapters.prompt_builders import (
-    BasePromptBuilder,
-    SavedPromptBuilder,
-    prompt_builder_from_id,
-)
+from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_from_id
 from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
 
 
-# TODO add evaluator rating
+# We should add evaluator rating
 class RepairTaskInput(BaseModel):
     original_prompt: str
     original_input: str

kiln_ai/adapters/test_adapter_registry.py

@@ -16,6 +16,7 @@ def mock_config():
     with patch("kiln_ai.adapters.adapter_registry.Config") as mock:
         mock.shared.return_value.open_ai_api_key = "test-openai-key"
         mock.shared.return_value.open_router_api_key = "test-openrouter-key"
+        mock.shared.return_value.siliconflow_cn_api_key = "test-siliconflow-key"
         yield mock
 
 
@@ -85,6 +86,33 @@ def test_openrouter_adapter_creation(mock_config, basic_task):
     }
 
 
+def test_siliconflow_adapter_creation(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        run_config_properties=RunConfigProperties(
+            model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+            model_provider_name=ModelProviderName.siliconflow_cn,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        ),
+    )
+
+    assert isinstance(adapter, LiteLlmAdapter)
+    assert (
+        adapter.config.run_config_properties.model_name
+        == "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"
+    )
+    assert adapter.config.additional_body_options == {"api_key": "test-siliconflow-key"}
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.siliconflow_cn
+    )
+    assert adapter.config.default_headers == {
+        "HTTP-Referer": "https://getkiln.ai/siliconflow",
+        "X-Title": "KilnAI",
+    }
+
+
 @pytest.mark.parametrize(
     "provider",
     [
@@ -109,7 +137,7 @@ def test_openai_compatible_adapter_creation(mock_config, basic_task, provider):
     assert adapter.run_config.model_name == "test-model"
 
 
-# TODO should run for all cases
+# We should run for all cases
 def test_custom_prompt_builder(mock_config, basic_task):
     adapter = adapter_for_task(
         kiln_task=basic_task,
@@ -124,7 +152,7 @@ def test_custom_prompt_builder(mock_config, basic_task):
     assert adapter.run_config.prompt_id == "simple_chain_of_thought_prompt_builder"
 
 
-# TODO should run for all cases
+# We should run for all cases
 def test_tags_passed_through(mock_config, basic_task):
     tags = ["test-tag-1", "test-tag-2"]
     adapter = adapter_for_task(

kiln_ai/adapters/test_ml_model_list.py

@@ -2,6 +2,7 @@ import pytest
 
 from kiln_ai.adapters.ml_model_list import (
     ModelName,
+    built_in_models,
     default_structured_output_mode_for_model_provider,
     get_model_by_name,
 )
@@ -156,3 +157,32 @@ class TestDefaultStructuredOutputModeForModelProvider:
             provider=first_provider.name,
         )
         assert result == first_provider.structured_output_mode
+
+
+def test_uncensored():
+    """Test that uncensored is set correctly"""
+    model = get_model_by_name(ModelName.grok_3_mini)
+    for provider in model.providers:
+        assert provider.uncensored
+        assert not provider.suggested_for_uncensored_data_gen
+
+    model = get_model_by_name(ModelName.gpt_4_1_nano)
+    for provider in model.providers:
+        assert not provider.uncensored
+        assert not provider.suggested_for_uncensored_data_gen
+
+    model = get_model_by_name(ModelName.grok_4)
+    for provider in model.providers:
+        assert provider.uncensored
+        assert provider.suggested_for_uncensored_data_gen
+
+
+def test_no_reasoning_for_structured_output():
+    """Test that no reasoning is returned for structured output"""
+    # get all models
+    for model in built_in_models:
+        for provider in model.providers:
+            if provider.reasoning_optional_for_structured_output is not None:
+                assert provider.reasoning_capable, (
+                    f"{model.name} {provider.name} has reasoning_optional_for_structured_output but is not reasoning capable. This field should only be defined for models that are reasoning capable."
+                )

kiln_ai/adapters/test_prompt_adaptors.py

@@ -13,10 +13,6 @@ from kiln_ai.adapters.model_adapters.litellm_adapter import (
     LiteLlmConfig,
 )
 from kiln_ai.adapters.ollama_tools import ollama_online
-from kiln_ai.adapters.prompt_builders import (
-    BasePromptBuilder,
-    SimpleChainOfThoughtPromptBuilder,
-)
 from kiln_ai.datamodel import PromptId
 from kiln_ai.datamodel.task import RunConfigProperties