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

kiln_ai/adapters/provider_tools.py

@@ -1,20 +1,24 @@
 from dataclasses import dataclass
-from typing import Dict, List, NoReturn
+from typing import Dict, List
 
 from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     KilnModelProvider,
     ModelName,
     ModelProviderName,
+    StructuredOutputMode,
     built_in_models,
 )
+from kiln_ai.adapters.model_adapters.openai_compatible_config import (
+    OpenAICompatibleConfig,
+)
 from kiln_ai.adapters.ollama_tools import (
     get_ollama_connection,
 )
 from kiln_ai.datamodel import Finetune, Task
 from kiln_ai.datamodel.registry import project_from_id
-
-from ..utils.config import Config
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
 async def provider_enabled(provider_name: ModelProviderName) -> bool:
@@ -61,7 +65,7 @@ def check_provider_warnings(provider_name: ModelProviderName):
             raise ValueError(warning_check.message)
 
 
-async def builtin_model_from(
+def builtin_model_from(
     name: str, provider_name: str | None = None
 ) -> KilnModelProvider | None:
     """
@@ -102,7 +106,47 @@ async def builtin_model_from(
     return provider
 
 
-async def kiln_model_provider_from(
+def core_provider(model_id: str, provider_name: ModelProviderName) -> ModelProviderName:
+    """
+    Get the provider that should be run.
+
+    Some provider IDs are wrappers (fine-tunes, custom models). This maps these to runnable providers (openai, ollama, etc)
+    """
+
+    # Custom models map to the underlying provider
+    if provider_name is ModelProviderName.kiln_custom_registry:
+        provider_name, _ = parse_custom_model_id(model_id)
+        return provider_name
+
+    # Fine-tune provider maps to an underlying provider
+    if provider_name is ModelProviderName.kiln_fine_tune:
+        finetune = finetune_from_id(model_id)
+        if finetune.provider not in ModelProviderName.__members__:
+            raise ValueError(
+                f"Finetune {model_id} has no underlying provider {finetune.provider}"
+            )
+        return ModelProviderName(finetune.provider)
+
+    return provider_name
+
+
+def parse_custom_model_id(
+    model_id: str,
+) -> tuple[ModelProviderName, str]:
+    if "::" not in model_id:
+        raise ValueError(f"Invalid custom model ID: {model_id}")
+
+    # For custom registry, get the provider name and model name from the model id
+    provider_name = model_id.split("::", 1)[0]
+    model_name = model_id.split("::", 1)[1]
+
+    if provider_name not in ModelProviderName.__members__:
+        raise ValueError(f"Invalid provider name: {provider_name}")
+
+    return ModelProviderName(provider_name), model_name
+
+
+def kiln_model_provider_from(
     name: str, provider_name: str | None = None
 ) -> KilnModelProvider:
     if provider_name == ModelProviderName.kiln_fine_tune:
@@ -111,14 +155,13 @@ async def kiln_model_provider_from(
     if provider_name == ModelProviderName.openai_compatible:
         return openai_compatible_provider_model(name)
 
-    built_in_model = await builtin_model_from(name, provider_name)
+    built_in_model = builtin_model_from(name, provider_name)
     if built_in_model:
         return built_in_model
 
     # For custom registry, get the provider name and model name from the model id
     if provider_name == ModelProviderName.kiln_custom_registry:
-        provider_name = name.split("::", 1)[0]
-        name = name.split("::", 1)[1]
+        provider_name, name = parse_custom_model_id(name)
 
     # Custom/untested model. Set untested, and build a ModelProvider at runtime
     if provider_name is None:
@@ -136,12 +179,9 @@ async def kiln_model_provider_from(
     )
 
 
-finetune_cache: dict[str, KilnModelProvider] = {}
-
-
-def openai_compatible_provider_model(
+def openai_compatible_config(
     model_id: str,
-) -> KilnModelProvider:
+) -> OpenAICompatibleConfig:
     try:
         openai_provider_name, model_id = model_id.split("::")
     except Exception:
@@ -165,12 +205,21 @@ def openai_compatible_provider_model(
             f"OpenAI compatible provider {openai_provider_name} has no base URL"
         )
 
+    return OpenAICompatibleConfig(
+        api_key=api_key,
+        model_name=model_id,
+        provider_name=ModelProviderName.openai_compatible,
+        base_url=base_url,
+    )
+
+
+def openai_compatible_provider_model(
+    model_id: str,
+) -> KilnModelProvider:
     return KilnModelProvider(
         name=ModelProviderName.openai_compatible,
         provider_options={
             "model": model_id,
-            "api_key": api_key,
-            "openai_api_base": base_url,
         },
         supports_structured_output=False,
         supports_data_gen=False,
@@ -178,9 +227,10 @@ def openai_compatible_provider_model(
     )
 
 
-def finetune_provider_model(
-    model_id: str,
-) -> KilnModelProvider:
+finetune_cache: dict[str, Finetune] = {}
+
+
+def finetune_from_id(model_id: str) -> Finetune:
     if model_id in finetune_cache:
         return finetune_cache[model_id]
 
@@ -202,6 +252,15 @@ def finetune_provider_model(
             f"Fine tune {fine_tune_id} not completed. Refresh it's status in the fine-tune tab."
         )
 
+    finetune_cache[model_id] = fine_tune
+    return fine_tune
+
+
+def finetune_provider_model(
+    model_id: str,
+) -> KilnModelProvider:
+    fine_tune = finetune_from_id(model_id)
+
     provider = ModelProviderName[fine_tune.provider]
     model_provider = KilnModelProvider(
         name=provider,
@@ -210,18 +269,18 @@ def finetune_provider_model(
         },
     )
 
-    # TODO: Don't love this abstraction/logic.
-    if fine_tune.provider == ModelProviderName.fireworks_ai:
-        # Fireworks finetunes are trained with json, not tool calling (which is LC default format)
-        model_provider.adapter_options = {
-            "langchain": {
-                "with_structured_output_options": {
-                    "method": "json_mode",
-                }
-            }
-        }
-
-    finetune_cache[model_id] = model_provider
+    if fine_tune.structured_output_mode is not None:
+        # If we know the model was trained with specific output mode, set it
+        model_provider.structured_output_mode = fine_tune.structured_output_mode
+    else:
+        # Some early adopters won't have structured_output_mode set on their fine-tunes.
+        # We know that OpenAI uses json_schema, and Fireworks (only other provider) use json_mode.
+        # This can be removed in the future
+        if provider == ModelProviderName.openai:
+            model_provider.structured_output_mode = StructuredOutputMode.json_schema
+        else:
+            model_provider.structured_output_mode = StructuredOutputMode.json_mode
+
     return model_provider
 
 
@@ -274,7 +333,7 @@ def provider_name_from_id(id: str) -> str:
                 return "OpenAI Compatible"
             case _:
                 # triggers pyright warning if I miss a case
-                raise_exhaustive_error(enum_id)
+                raise_exhaustive_enum_error(enum_id)
 
     return "Unknown provider: " + id
 
@@ -316,16 +375,12 @@ def provider_options_for_custom_model(
             )
         case _:
             # triggers pyright warning if I miss a case
-            raise_exhaustive_error(enum_id)
+            raise_exhaustive_enum_error(enum_id)
 
     # Won't reach this, type checking will catch missed values
     return {"model": model_name}
 
 
-def raise_exhaustive_error(value: NoReturn) -> NoReturn:
-    raise ValueError(f"Unhandled enum value: {value}")
-
-
 @dataclass
 class ModelProviderWarning:
     required_config_keys: List[str]

kiln_ai/adapters/repair/repair_task.py

@@ -3,7 +3,11 @@ from typing import Type
 
 from pydantic import BaseModel, Field
 
-from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_registry
+from kiln_ai.adapters.prompt_builders import (
+    BasePromptBuilder,
+    SavedPromptBuilder,
+    prompt_builder_from_id,
+)
 from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
 
 
@@ -42,24 +46,19 @@ feedback describing what should be improved. Your job is to understand the evalu
 
     @classmethod
     def _original_prompt(cls, run: TaskRun, task: Task) -> str:
-        prompt_builder_class: Type[BasePromptBuilder] | None = None
-        prompt_builder_name = (
-            run.output.source.properties.get("prompt_builder_name", None)
-            if run.output.source
-            else None
-        )
-        if prompt_builder_name is not None and isinstance(prompt_builder_name, str):
-            prompt_builder_class = prompt_builder_registry.get(
-                prompt_builder_name, None
-            )
-        if prompt_builder_class is None:
-            raise ValueError(f"No prompt builder found for name: {prompt_builder_name}")
-        prompt_builder = prompt_builder_class(task=task)
-        if not isinstance(prompt_builder, BasePromptBuilder):
-            raise ValueError(
-                f"Prompt builder {prompt_builder_name} is not a valid prompt builder"
-            )
-        return prompt_builder.build_prompt()
+        if run.output.source is None or run.output.source.properties is None:
+            raise ValueError("No source properties found")
+
+        # Get the prompt builder id. Need the second check because we used to store this in a prompt_builder_name field, so loading legacy runs will need this.
+        prompt_id = run.output.source.properties.get(
+            "prompt_id"
+        ) or run.output.source.properties.get("prompt_builder_name", None)
+        if prompt_id is not None and isinstance(prompt_id, str):
+            prompt_builder = prompt_builder_from_id(prompt_id, task)
+            if isinstance(prompt_builder, BasePromptBuilder):
+                return prompt_builder.build_prompt(include_json_instructions=False)
+
+        raise ValueError(f"Prompt builder '{prompt_id}' is not a valid prompt builder")
 
     @classmethod
     def build_repair_task_input(

kiln_ai/adapters/repair/test_repair_task.py

@@ -6,8 +6,8 @@ import pytest
 from pydantic import ValidationError
 
 from kiln_ai.adapters.adapter_registry import adapter_for_task
-from kiln_ai.adapters.base_adapter import RunOutput
-from kiln_ai.adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.model_adapters.base_adapter import RunOutput
+from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
 from kiln_ai.adapters.repair.repair_task import (
     RepairTaskInput,
     RepairTaskRun,
@@ -95,7 +95,7 @@ def sample_task_run(sample_task):
                     "model_name": "gpt_4o",
                     "model_provider": "openai",
                     "adapter_name": "langchain_adapter",
-                    "prompt_builder_name": "simple_prompt_builder",
+                    "prompt_id": "simple_prompt_builder",
                 },
             ),
         ),
@@ -201,7 +201,7 @@ async def test_live_run(sample_task, sample_task_run, sample_repair_data):
         "adapter_name": "kiln_langchain_adapter",
         "model_name": "llama_3_1_8b",
         "model_provider": "groq",
-        "prompt_builder_name": "simple_prompt_builder",
+        "prompt_id": "simple_prompt_builder",
     }
 
 
@@ -223,7 +223,7 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         )
 
         adapter = adapter_for_task(
-            repair_task, model_name="llama_3_1_8b", provider="groq"
+            repair_task, model_name="llama_3_1_8b", provider="ollama"
         )
 
         run = await adapter.invoke(repair_task_input.model_dump())
@@ -237,8 +237,8 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     assert run.output.source.properties == {
         "adapter_name": "kiln_langchain_adapter",
         "model_name": "llama_3_1_8b",
-        "model_provider": "groq",
-        "prompt_builder_name": "simple_prompt_builder",
+        "model_provider": "ollama",
+        "prompt_id": "simple_prompt_builder",
     }
     assert run.input_source.type == DataSourceType.human
     assert "created_by" in run.input_source.properties

kiln_ai/adapters/run_output.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+from typing import Dict
+
+from openai.types.chat.chat_completion import ChoiceLogprobs
+
+
+@dataclass
+class RunOutput:
+    output: Dict | str
+    intermediate_outputs: Dict[str, str] | None
+    output_logprobs: ChoiceLogprobs | None = None

kiln_ai/adapters/test_adapter_registry.py

@@ -0,0 +1,177 @@
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai import datamodel
+from kiln_ai.adapters.adapter_registry import adapter_for_task
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
+from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.model_adapters.openai_model_adapter import OpenAICompatibleAdapter
+from kiln_ai.adapters.prompt_builders import BasePromptBuilder
+from kiln_ai.adapters.provider_tools import kiln_model_provider_from
+
+
+@pytest.fixture
+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"
+        yield mock
+
+
+@pytest.fixture
+def basic_task():
+    return datamodel.Task(
+        task_id="test-task",
+        task_type="test",
+        input_text="test input",
+        name="test-task",
+        instruction="test-task",
+    )
+
+
+@pytest.fixture
+def mock_finetune_from_id():
+    with patch("kiln_ai.adapters.provider_tools.finetune_from_id") as mock:
+        mock.return_value.provider = ModelProviderName.openai
+        mock.return_value.fine_tune_model_id = "test-model"
+        yield mock
+
+
+def test_openai_adapter_creation(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task, model_name="gpt-4", provider=ModelProviderName.openai
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.model_name == "gpt-4"
+    assert adapter.config.api_key == "test-openai-key"
+    assert adapter.config.provider_name == ModelProviderName.openai
+    assert adapter.config.base_url is None  # OpenAI url is default
+    assert adapter.config.default_headers is None
+
+
+def test_openrouter_adapter_creation(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="anthropic/claude-3-opus",
+        provider=ModelProviderName.openrouter,
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.model_name == "anthropic/claude-3-opus"
+    assert adapter.config.api_key == "test-openrouter-key"
+    assert adapter.config.provider_name == ModelProviderName.openrouter
+    assert adapter.config.base_url == "https://openrouter.ai/api/v1"
+    assert adapter.config.default_headers == {
+        "HTTP-Referer": "https://getkiln.ai/openrouter",
+        "X-Title": "KilnAI",
+    }
+
+
+@pytest.mark.parametrize(
+    "provider",
+    [
+        ModelProviderName.groq,
+        ModelProviderName.amazon_bedrock,
+        ModelProviderName.ollama,
+        ModelProviderName.fireworks_ai,
+    ],
+)
+def test_langchain_adapter_creation(mock_config, basic_task, provider):
+    adapter = adapter_for_task(
+        kiln_task=basic_task, model_name="test-model", provider=provider
+    )
+
+    assert isinstance(adapter, LangchainAdapter)
+    assert adapter.run_config.model_name == "test-model"
+
+
+# TODO should run for all cases
+def test_custom_prompt_builder(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="gpt-4",
+        provider=ModelProviderName.openai,
+        prompt_id="simple_chain_of_thought_prompt_builder",
+    )
+
+    assert adapter.run_config.prompt_id == "simple_chain_of_thought_prompt_builder"
+
+
+# TODO 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_task=basic_task,
+        model_name="gpt-4",
+        provider=ModelProviderName.openai,
+        base_adapter_config=AdapterConfig(
+            default_tags=tags,
+        ),
+    )
+
+    assert adapter.base_adapter_config.default_tags == tags
+
+
+def test_invalid_provider(mock_config, basic_task):
+    with pytest.raises(ValueError, match="Unhandled enum value"):
+        adapter_for_task(
+            kiln_task=basic_task, model_name="test-model", provider="invalid"
+        )
+
+
+@patch("kiln_ai.adapters.adapter_registry.openai_compatible_config")
+def test_openai_compatible_adapter(mock_compatible_config, mock_config, basic_task):
+    mock_compatible_config.return_value.model_name = "test-model"
+    mock_compatible_config.return_value.api_key = "test-key"
+    mock_compatible_config.return_value.base_url = "https://test.com/v1"
+    mock_compatible_config.return_value.provider_name = "CustomProvider99"
+
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="provider::test-model",
+        provider=ModelProviderName.openai_compatible,
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    mock_compatible_config.assert_called_once_with("provider::test-model")
+    assert adapter.config.model_name == "test-model"
+    assert adapter.config.api_key == "test-key"
+    assert adapter.config.base_url == "https://test.com/v1"
+    assert adapter.config.provider_name == "CustomProvider99"
+
+
+def test_custom_openai_compatible_provider(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="openai::test-model",
+        provider=ModelProviderName.kiln_custom_registry,
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.model_name == "openai::test-model"
+    assert adapter.config.api_key == "test-openai-key"
+    assert adapter.config.base_url is None  # openai is none
+    assert adapter.config.provider_name == ModelProviderName.kiln_custom_registry
+
+
+async def test_fine_tune_provider(mock_config, basic_task, mock_finetune_from_id):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="proj::task::tune",
+        provider=ModelProviderName.kiln_fine_tune,
+    )
+
+    mock_finetune_from_id.assert_called_once_with("proj::task::tune")
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.provider_name == ModelProviderName.kiln_fine_tune
+    # Kiln model name here, but the underlying openai model id below
+    assert adapter.config.model_name == "proj::task::tune"
+
+    provider = kiln_model_provider_from(
+        "proj::task::tune", provider_name=ModelProviderName.kiln_fine_tune
+    )
+    # The actual model name from the fine tune object
+    assert provider.provider_options["model"] == "test-model"

kiln_ai/adapters/test_generate_docs.py

@@ -0,0 +1,69 @@
+import logging
+from typing import List
+
+import pytest
+
+from libs.core.kiln_ai.adapters.ml_model_list import KilnModelProvider, built_in_models
+from libs.core.kiln_ai.adapters.provider_tools import provider_name_from_id
+
+logger = logging.getLogger(__name__)
+
+
+def _all_providers_support(providers: List[KilnModelProvider], attribute: str) -> bool:
+    """Check if all providers support a given feature"""
+    return all(getattr(provider, attribute) for provider in providers)
+
+
+def _any_providers_support(providers: List[KilnModelProvider], attribute: str) -> bool:
+    """Check if any providers support a given feature"""
+    return any(getattr(provider, attribute) for provider in providers)
+
+
+def _get_support_status(providers: List[KilnModelProvider], attribute: str) -> str:
+    """Get the support status for a feature"""
+    if _all_providers_support(providers, attribute):
+        return "✅︎"
+    elif _any_providers_support(providers, attribute):
+        return "✅︎ (some providers)"
+    return ""
+
+
+def _has_finetune_support(providers: List[KilnModelProvider]) -> str:
+    """Check if any provider supports fine-tuning"""
+    return "✅︎" if any(p.provider_finetune_id for p in providers) else ""
+
+
+@pytest.mark.paid(reason="Marking as paid so it isn't run by default")
+def test_generate_model_table():
+    """Generate a markdown table of all models and their capabilities"""
+
+    # Table header
+    table = [
+        "| Model Name | Providers | Structured Output | Reasoning | Synthetic Data | API Fine-Tuneable |",
+        "|------------|-----------|-------------------|-----------|----------------|-------------------|",
+    ]
+
+    for model in built_in_models:
+        provider_names = ", ".join(
+            sorted(provider_name_from_id(p.name.value) for p in model.providers)
+        )
+        structured_output = _get_support_status(
+            model.providers, "supports_structured_output"
+        )
+        reasoning = _get_support_status(model.providers, "reasoning_capable")
+        data_gen = _get_support_status(model.providers, "supports_data_gen")
+        finetune = _has_finetune_support(model.providers)
+
+        row = f"| {model.friendly_name} | {provider_names} | {structured_output} | {reasoning} | {data_gen} | {finetune} |"
+        table.append(row)
+
+    # Print the table (useful for documentation)
+    logger.info("\nModel Capability Matrix:\n")
+    logger.info("\n".join(table))
+
+    # Basic assertions to ensure the table is well-formed
+    assert len(table) > 2, "Table should have header and at least one row"
+    assert all("|" in row for row in table), "All rows should be properly formatted"
+    assert len(table[0].split("|")) == len(table[1].split("|")), (
+        "Header and separator should have same number of columns"
+    )

kiln_ai/adapters/test_ollama_tools.py

@@ -10,7 +10,6 @@ from kiln_ai.adapters.ollama_tools import (
 def test_parse_ollama_tags_no_models():
     json_response = '{"models":[{"name":"scosman_net","model":"scosman_net:latest"},{"name":"phi3.5:latest","model":"phi3.5:latest","modified_at":"2024-10-02T12:04:35.191519822-04:00","size":2176178843,"digest":"61819fb370a3c1a9be6694869331e5f85f867a079e9271d66cb223acb81d04ba","details":{"parent_model":"","format":"gguf","family":"phi3","families":["phi3"],"parameter_size":"3.8B","quantization_level":"Q4_0"}},{"name":"gemma2:2b","model":"gemma2:2b","modified_at":"2024-09-09T16:46:38.64348929-04:00","size":1629518495,"digest":"8ccf136fdd5298f3ffe2d69862750ea7fb56555fa4d5b18c04e3fa4d82ee09d7","details":{"parent_model":"","format":"gguf","family":"gemma2","families":["gemma2"],"parameter_size":"2.6B","quantization_level":"Q4_0"}},{"name":"llama3.1:latest","model":"llama3.1:latest","modified_at":"2024-09-01T17:19:43.481523695-04:00","size":4661230720,"digest":"f66fc8dc39ea206e03ff6764fcc696b1b4dfb693f0b6ef751731dd4e6269046e","details":{"parent_model":"","format":"gguf","family":"llama","families":["llama"],"parameter_size":"8.0B","quantization_level":"Q4_0"}}]}'
     tags = json.loads(json_response)
-    print(json.dumps(tags, indent=2))
     conn = parse_ollama_tags(tags)
     assert "phi3.5:latest" in conn.supported_models
     assert "gemma2:2b" in conn.supported_models