kiln-ai 0.7.0__py3-none-any.whl → 0.7.1__py3-none-any.whl

kiln_ai/adapters/ml_model_list.py

@@ -22,6 +22,7 @@ class ModelProviderName(str, Enum):
     openrouter = "openrouter"
     fireworks_ai = "fireworks_ai"
     kiln_fine_tune = "kiln_fine_tune"
+    kiln_custom_registry = "kiln_custom_registry"
 
 
 class ModelFamily(str, Enum):
@@ -54,6 +55,7 @@ class ModelName(str, Enum):
     llama_3_2_3b = "llama_3_2_3b"
     llama_3_2_11b = "llama_3_2_11b"
     llama_3_2_90b = "llama_3_2_90b"
+    llama_3_3_70b = "llama_3_3_70b"
     gpt_4o_mini = "gpt_4o_mini"
     gpt_4o = "gpt_4o"
     phi_3_5 = "phi_3_5"
@@ -502,6 +504,38 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # Llama 3.3 70B
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_3_70b,
+        friendly_name="Llama 3.3 70B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "meta-llama/llama-3.3-70b-instruct"},
+                # Openrouter not supporing tools yet. Once they do probably can remove. JSON mode sometimes works, but not consistently.
+                supports_structured_output=False,
+                supports_data_gen=False,
+                adapter_options={
+                    "langchain": {
+                        "with_structured_output_options": {"method": "json_mode"}
+                    }
+                },
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={"model": "llama3.3"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.fireworks_ai,
+                # Finetuning not live yet
+                # provider_finetune_id="accounts/fireworks/models/llama-v3p3-70b-instruct",
+                provider_options={
+                    "model": "accounts/fireworks/models/llama-v3p3-70b-instruct"
+                },
+            ),
+        ],
+    ),
     # Phi 3.5
     KilnModel(
         family=ModelFamily.phi,
@@ -598,18 +632,6 @@ built_in_models: List[KilnModel] = [
         name=ModelName.mixtral_8x7b,
         friendly_name="Mixtral 8x7B",
         providers=[
-            KilnModelProvider(
-                name=ModelProviderName.fireworks_ai,
-                provider_options={
-                    "model": "accounts/fireworks/models/mixtral-8x7b-instruct-hf",
-                },
-                provider_finetune_id="accounts/fireworks/models/mixtral-8x7b-instruct-hf",
-                adapter_options={
-                    "langchain": {
-                        "with_structured_output_options": {"method": "json_mode"}
-                    }
-                },
-            ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 provider_options={"model": "mistralai/mixtral-8x7b-instruct"},

kiln_ai/adapters/ollama_tools.py

@@ -6,6 +6,7 @@ import requests
 from pydantic import BaseModel, Field
 
 from kiln_ai.adapters.ml_model_list import ModelProviderName, built_in_models
+from kiln_ai.utils.config import Config
 
 
 def ollama_base_url() -> str:
@@ -16,9 +17,9 @@ def ollama_base_url() -> str:
         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
+    config_base_url = Config.shared().ollama_base_url
+    if config_base_url:
+        return config_base_url
     return "http://localhost:11434"
 
 

kiln_ai/adapters/provider_tools.py

@@ -11,6 +11,7 @@ from kiln_ai.adapters.ml_model_list import (
 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
@@ -111,6 +112,11 @@ async def kiln_model_provider_from(
     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]
+
     # Custom/untested model. Set untested, and build a ModelProvider at runtime
     if provider_name is None:
         raise ValueError("Provider name is required for custom models")
@@ -143,10 +149,10 @@ def finetune_provider_model(
     project = project_from_id(project_id)
     if project is None:
         raise ValueError(f"Project {project_id} not found")
-    task = next((t for t in project.tasks() if t.id == task_id), None)
+    task = Task.from_id_and_parent_path(task_id, project.path)
     if task is None:
         raise ValueError(f"Task {task_id} not found")
-    fine_tune = next((f for f in task.finetunes() if f.id == fine_tune_id), None)
+    fine_tune = Finetune.from_id_and_parent_path(fine_tune_id, task.path)
     if fine_tune is None:
         raise ValueError(f"Fine tune {fine_tune_id} not found")
     if fine_tune.fine_tune_model_id is None:
@@ -220,6 +226,8 @@ def provider_name_from_id(id: str) -> str:
                 return "Fine Tuned Models"
             case ModelProviderName.fireworks_ai:
                 return "Fireworks AI"
+            case ModelProviderName.kiln_custom_registry:
+                return "Custom Models"
             case _:
                 # triggers pyright warning if I miss a case
                 raise_exhaustive_error(enum_id)
@@ -233,6 +241,7 @@ def provider_options_for_custom_model(
     """
     Generated model provider options for a custom model. Each has their own format/options.
     """
+
     if provider_name not in ModelProviderName.__members__:
         raise ValueError(f"Invalid provider name: {provider_name}")
 
@@ -249,6 +258,10 @@ def provider_options_for_custom_model(
             | ModelProviderName.groq
         ):
             return {"model": model_name}
+        case ModelProviderName.kiln_custom_registry:
+            raise ValueError(
+                "Custom models from registry should be parsed into provider/model before calling this."
+            )
         case ModelProviderName.kiln_fine_tune:
             raise ValueError(
                 "Fine tuned models should populate provider options via another path"

kiln_ai/adapters/repair/repair_task.py

@@ -43,8 +43,10 @@ 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
+        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(

kiln_ai/adapters/test_langchain_adapter.py

@@ -1,12 +1,20 @@
+import os
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import pytest
+from langchain_aws import ChatBedrockConverse
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_fireworks import ChatFireworks
 from langchain_groq import ChatGroq
+from langchain_ollama import ChatOllama
+from langchain_openai import ChatOpenAI
 
 from kiln_ai.adapters.langchain_adapters import (
     LangchainAdapter,
     get_structured_output_options,
+    langchain_model_from_provider,
 )
+from kiln_ai.adapters.ml_model_list import KilnModelProvider, ModelProviderName
 from kiln_ai.adapters.prompt_builders import SimpleChainOfThoughtPromptBuilder
 from kiln_ai.adapters.test_prompt_adaptors import build_test_task
 
@@ -150,3 +158,178 @@ async def test_get_structured_output_options():
     ):
         options = await get_structured_output_options("model_name", "provider")
         assert options == {}
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_openai():
+    provider = KilnModelProvider(
+        name=ModelProviderName.openai, provider_options={"model": "gpt-4"}
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.open_ai_api_key = "test_key"
+        model = await langchain_model_from_provider(provider, "gpt-4")
+        assert isinstance(model, ChatOpenAI)
+        assert model.model_name == "gpt-4"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_groq():
+    provider = KilnModelProvider(
+        name=ModelProviderName.groq, provider_options={"model": "mixtral-8x7b"}
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.groq_api_key = "test_key"
+        model = await langchain_model_from_provider(provider, "mixtral-8x7b")
+        assert isinstance(model, ChatGroq)
+        assert model.model_name == "mixtral-8x7b"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_bedrock():
+    provider = KilnModelProvider(
+        name=ModelProviderName.amazon_bedrock,
+        provider_options={"model": "anthropic.claude-v2", "region_name": "us-east-1"},
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.bedrock_access_key = "test_access"
+        mock_config.return_value.bedrock_secret_key = "test_secret"
+        model = await langchain_model_from_provider(provider, "anthropic.claude-v2")
+        assert isinstance(model, ChatBedrockConverse)
+        assert os.environ.get("AWS_ACCESS_KEY_ID") == "test_access"
+        assert os.environ.get("AWS_SECRET_ACCESS_KEY") == "test_secret"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_fireworks():
+    provider = KilnModelProvider(
+        name=ModelProviderName.fireworks_ai, provider_options={"model": "mixtral-8x7b"}
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.fireworks_api_key = "test_key"
+        model = await langchain_model_from_provider(provider, "mixtral-8x7b")
+        assert isinstance(model, ChatFireworks)
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_ollama():
+    provider = KilnModelProvider(
+        name=ModelProviderName.ollama,
+        provider_options={"model": "llama2", "model_aliases": ["llama2-uncensored"]},
+    )
+
+    mock_connection = MagicMock()
+    with (
+        patch(
+            "kiln_ai.adapters.langchain_adapters.get_ollama_connection",
+            return_value=AsyncMock(return_value=mock_connection),
+        ),
+        patch(
+            "kiln_ai.adapters.langchain_adapters.ollama_model_installed",
+            return_value=True,
+        ),
+        patch(
+            "kiln_ai.adapters.langchain_adapters.ollama_base_url",
+            return_value="http://localhost:11434",
+        ),
+    ):
+        model = await langchain_model_from_provider(provider, "llama2")
+        assert isinstance(model, ChatOllama)
+        assert model.model == "llama2"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_invalid():
+    provider = KilnModelProvider.model_construct(
+        name="invalid_provider", provider_options={}
+    )
+
+    with pytest.raises(ValueError, match="Invalid model or provider"):
+        await langchain_model_from_provider(provider, "test_model")
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter_model_caching(tmp_path):
+    task = build_test_task(tmp_path)
+    custom_model = ChatGroq(model="mixtral-8x7b", groq_api_key="test")
+
+    adapter = LangchainAdapter(kiln_task=task, custom_model=custom_model)
+
+    # First call should return the cached model
+    model1 = await adapter.model()
+    assert model1 is custom_model
+
+    # Second call should return the same cached instance
+    model2 = await adapter.model()
+    assert model2 is model1
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter_model_structured_output(tmp_path):
+    task = build_test_task(tmp_path)
+    task.output_json_schema = """
+    {
+        "type": "object",
+        "properties": {
+            "count": {"type": "integer"}
+        }
+    }
+    """
+
+    mock_model = MagicMock()
+    mock_model.with_structured_output = MagicMock(return_value="structured_model")
+
+    adapter = LangchainAdapter(
+        kiln_task=task, model_name="test_model", provider="test_provider"
+    )
+
+    with (
+        patch(
+            "kiln_ai.adapters.langchain_adapters.langchain_model_from",
+            AsyncMock(return_value=mock_model),
+        ),
+        patch(
+            "kiln_ai.adapters.langchain_adapters.get_structured_output_options",
+            AsyncMock(return_value={"option1": "value1"}),
+        ),
+    ):
+        model = await adapter.model()
+
+        # Verify the model was configured with structured output
+        mock_model.with_structured_output.assert_called_once_with(
+            {
+                "type": "object",
+                "properties": {"count": {"type": "integer"}},
+                "title": "task_response",
+                "description": "A response from the task",
+            },
+            include_raw=True,
+            option1="value1",
+        )
+        assert model == "structured_model"
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter_model_no_structured_output_support(tmp_path):
+    task = build_test_task(tmp_path)
+    task.output_json_schema = (
+        '{"type": "object", "properties": {"count": {"type": "integer"}}}'
+    )
+
+    mock_model = MagicMock()
+    # Remove with_structured_output method
+    del mock_model.with_structured_output
+
+    adapter = LangchainAdapter(
+        kiln_task=task, model_name="test_model", provider="test_provider"
+    )
+
+    with patch(
+        "kiln_ai.adapters.langchain_adapters.langchain_model_from",
+        AsyncMock(return_value=mock_model),
+    ):
+        with pytest.raises(ValueError, match="does not support structured output"):
+            await adapter.model()

kiln_ai/adapters/test_provider_tools.py

@@ -1,14 +1,18 @@
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
 from kiln_ai.adapters.ml_model_list import (
+    KilnModel,
     ModelName,
     ModelProviderName,
 )
 from kiln_ai.adapters.ollama_tools import OllamaConnection
 from kiln_ai.adapters.provider_tools import (
+    builtin_model_from,
     check_provider_warnings,
+    finetune_cache,
+    finetune_provider_model,
     get_model_and_provider,
     kiln_model_provider_from,
     provider_enabled,
@@ -16,6 +20,14 @@ from kiln_ai.adapters.provider_tools import (
     provider_options_for_custom_model,
     provider_warnings,
 )
+from kiln_ai.datamodel import Finetune, Task
+
+
+@pytest.fixture(autouse=True)
+def clear_finetune_cache():
+    """Clear the finetune provider model cache before each test"""
+    finetune_cache.clear()
+    yield
 
 
 @pytest.fixture
@@ -24,6 +36,34 @@ def mock_config():
         yield mock
 
 
+@pytest.fixture
+def mock_project():
+    with patch("kiln_ai.adapters.provider_tools.project_from_id") as mock:
+        project = Mock()
+        project.path = "/fake/path"
+        mock.return_value = project
+        yield mock
+
+
+@pytest.fixture
+def mock_task():
+    with patch("kiln_ai.datamodel.Task.from_id_and_parent_path") as mock:
+        task = Mock(spec=Task)
+        task.path = "/fake/path/task"
+        mock.return_value = task
+        yield mock
+
+
+@pytest.fixture
+def mock_finetune():
+    with patch("kiln_ai.datamodel.Finetune.from_id_and_parent_path") as mock:
+        finetune = Mock(spec=Finetune)
+        finetune.provider = ModelProviderName.openai
+        finetune.fine_tune_model_id = "ft:gpt-3.5-turbo:custom:model-123"
+        mock.return_value = finetune
+        yield mock
+
+
 def test_check_provider_warnings_no_warning(mock_config):
     mock_config.return_value = "some_value"
 
@@ -103,6 +143,8 @@ def test_provider_name_from_id_case_sensitivity():
         (ModelProviderName.ollama, "Ollama"),
         (ModelProviderName.openai, "OpenAI"),
         (ModelProviderName.fireworks_ai, "Fireworks AI"),
+        (ModelProviderName.kiln_fine_tune, "Fine Tuned Models"),
+        (ModelProviderName.kiln_custom_registry, "Custom Models"),
     ],
 )
 def test_provider_name_from_id_parametrized(provider_id, expected_name):
@@ -310,3 +352,180 @@ def test_provider_options_for_custom_model_invalid_enum():
     """Test handling of invalid enum value"""
     with pytest.raises(ValueError):
         provider_options_for_custom_model("model_name", "invalid_enum_value")
+
+
+@pytest.mark.asyncio
+async def test_kiln_model_provider_from_custom_registry(mock_config):
+    # Mock config to pass provider warnings check
+    mock_config.return_value = "fake-api-key"
+
+    # Test with a custom registry model ID in format "provider::model_name"
+    provider = await kiln_model_provider_from(
+        "openai::gpt-4-turbo", ModelProviderName.kiln_custom_registry
+    )
+
+    assert provider.name == ModelProviderName.openai
+    assert provider.supports_structured_output is False
+    assert provider.supports_data_gen is False
+    assert provider.untested_model is True
+    assert provider.provider_options == {"model": "gpt-4-turbo"}
+
+
+@pytest.mark.asyncio
+async def test_builtin_model_from_invalid_model():
+    """Test that an invalid model name returns None"""
+    result = await builtin_model_from("non_existent_model")
+    assert result is None
+
+
+@pytest.mark.asyncio
+async def test_builtin_model_from_valid_model_default_provider(mock_config):
+    """Test getting a valid model with default provider"""
+    mock_config.return_value = "fake-api-key"
+
+    provider = await builtin_model_from(ModelName.phi_3_5)
+
+    assert provider is not None
+    assert provider.name == ModelProviderName.ollama
+    assert provider.provider_options["model"] == "phi3.5"
+
+
+@pytest.mark.asyncio
+async def test_builtin_model_from_valid_model_specific_provider(mock_config):
+    """Test getting a valid model with specific provider"""
+    mock_config.return_value = "fake-api-key"
+
+    provider = await builtin_model_from(
+        ModelName.llama_3_1_70b, provider_name=ModelProviderName.groq
+    )
+
+    assert provider is not None
+    assert provider.name == ModelProviderName.groq
+    assert provider.provider_options["model"] == "llama-3.1-70b-versatile"
+
+
+@pytest.mark.asyncio
+async def test_builtin_model_from_invalid_provider(mock_config):
+    """Test that requesting an invalid provider returns None"""
+    mock_config.return_value = "fake-api-key"
+
+    provider = await builtin_model_from(
+        ModelName.phi_3_5, provider_name="invalid_provider"
+    )
+
+    assert provider is None
+
+
+@pytest.mark.asyncio
+async def test_builtin_model_from_model_no_providers():
+    """Test handling of a model with no providers"""
+    with patch("kiln_ai.adapters.provider_tools.built_in_models") as mock_models:
+        # Create a mock model with no providers
+        mock_model = KilnModel(
+            name=ModelName.phi_3_5,
+            friendly_name="Test Model",
+            providers=[],
+            family="test_family",
+        )
+        mock_models.__iter__.return_value = [mock_model]
+
+        with pytest.raises(ValueError) as exc_info:
+            await builtin_model_from(ModelName.phi_3_5)
+
+        assert str(exc_info.value) == f"Model {ModelName.phi_3_5} has no providers"
+
+
+@pytest.mark.asyncio
+async def test_builtin_model_from_provider_warning_check(mock_config):
+    """Test that provider warnings are checked"""
+    # Make the config check fail
+    mock_config.return_value = None
+
+    with pytest.raises(ValueError) as exc_info:
+        await builtin_model_from(ModelName.llama_3_1_70b, ModelProviderName.groq)
+
+    assert provider_warnings[ModelProviderName.groq].message in str(exc_info.value)
+
+
+def test_finetune_provider_model_success(mock_project, mock_task, mock_finetune):
+    """Test successful creation of a fine-tuned model provider"""
+    model_id = "project-123::task-456::finetune-789"
+
+    provider = finetune_provider_model(model_id)
+
+    assert provider.name == ModelProviderName.openai
+    assert provider.provider_options == {"model": "ft:gpt-3.5-turbo:custom:model-123"}
+
+    # Test cache
+    cached_provider = finetune_provider_model(model_id)
+    assert cached_provider is provider
+
+
+def test_finetune_provider_model_invalid_id():
+    """Test handling of invalid model ID format"""
+    with pytest.raises(ValueError) as exc_info:
+        finetune_provider_model("invalid-id-format")
+    assert str(exc_info.value) == "Invalid fine tune ID: invalid-id-format"
+
+
+def test_finetune_provider_model_project_not_found(mock_project):
+    """Test handling of non-existent project"""
+    mock_project.return_value = None
+
+    with pytest.raises(ValueError) as exc_info:
+        finetune_provider_model("project-123::task-456::finetune-789")
+    assert str(exc_info.value) == "Project project-123 not found"
+
+
+def test_finetune_provider_model_task_not_found(mock_project, mock_task):
+    """Test handling of non-existent task"""
+    mock_task.return_value = None
+
+    with pytest.raises(ValueError) as exc_info:
+        finetune_provider_model("project-123::task-456::finetune-789")
+    assert str(exc_info.value) == "Task task-456 not found"
+
+
+def test_finetune_provider_model_finetune_not_found(
+    mock_project, mock_task, mock_finetune
+):
+    """Test handling of non-existent fine-tune"""
+    mock_finetune.return_value = None
+
+    with pytest.raises(ValueError) as exc_info:
+        finetune_provider_model("project-123::task-456::finetune-789")
+    assert str(exc_info.value) == "Fine tune finetune-789 not found"
+
+
+def test_finetune_provider_model_incomplete_finetune(
+    mock_project, mock_task, mock_finetune
+):
+    """Test handling of incomplete fine-tune"""
+    finetune = Mock(spec=Finetune)
+    finetune.fine_tune_model_id = None
+    mock_finetune.return_value = finetune
+
+    with pytest.raises(ValueError) as exc_info:
+        finetune_provider_model("project-123::task-456::finetune-789")
+    assert (
+        str(exc_info.value)
+        == "Fine tune finetune-789 not completed. Refresh it's status in the fine-tune tab."
+    )
+
+
+def test_finetune_provider_model_fireworks_provider(
+    mock_project, mock_task, mock_finetune
+):
+    """Test creation of Fireworks AI provider with specific adapter options"""
+    finetune = Mock(spec=Finetune)
+    finetune.provider = ModelProviderName.fireworks_ai
+    finetune.fine_tune_model_id = "fireworks-model-123"
+    mock_finetune.return_value = finetune
+
+    provider = finetune_provider_model("project-123::task-456::finetune-789")
+
+    assert provider.name == ModelProviderName.fireworks_ai
+    assert provider.provider_options == {"model": "fireworks-model-123"}
+    assert provider.adapter_options == {
+        "langchain": {"with_structured_output_options": {"method": "json_mode"}}
+    }