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

kiln_ai/adapters/test_provider_tools.py

@@ -1,21 +1,34 @@
-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,
+    openai_compatible_provider_model,
     provider_enabled,
     provider_name_from_id,
     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 +37,53 @@ 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
+
+
+@pytest.fixture
+def mock_shared_config():
+    with patch("kiln_ai.adapters.provider_tools.Config.shared") as mock:
+        config = Mock()
+        config.openai_compatible_providers = [
+            {
+                "name": "test_provider",
+                "base_url": "https://api.test.com",
+                "api_key": "test-key",
+            },
+            {
+                "name": "no_key_provider",
+                "base_url": "https://api.nokey.com",
+            },
+        ]
+        mock.return_value = config
+        yield mock
+
+
 def test_check_provider_warnings_no_warning(mock_config):
     mock_config.return_value = "some_value"
 
@@ -103,6 +163,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 +372,255 @@ 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"}}
+    }
+
+
+def test_openai_compatible_provider_model_success(mock_shared_config):
+    """Test successful creation of an OpenAI compatible provider"""
+    model_id = "test_provider::gpt-4"
+
+    provider = openai_compatible_provider_model(model_id)
+
+    assert provider.name == ModelProviderName.openai_compatible
+    assert provider.provider_options == {
+        "model": "gpt-4",
+        "api_key": "test-key",
+        "openai_api_base": "https://api.test.com",
+    }
+    assert provider.supports_structured_output is False
+    assert provider.supports_data_gen is False
+    assert provider.untested_model is True
+
+
+def test_openai_compatible_provider_model_no_api_key(mock_shared_config):
+    """Test provider creation without API key (should work as some providers don't require it)"""
+    model_id = "no_key_provider::gpt-4"
+
+    provider = openai_compatible_provider_model(model_id)
+
+    assert provider.name == ModelProviderName.openai_compatible
+    assert provider.provider_options == {
+        "model": "gpt-4",
+        "api_key": None,
+        "openai_api_base": "https://api.nokey.com",
+    }
+
+
+def test_openai_compatible_provider_model_invalid_id():
+    """Test handling of invalid model ID format"""
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("invalid-id-format")
+    assert (
+        str(exc_info.value) == "Invalid openai compatible model ID: invalid-id-format"
+    )
+
+
+def test_openai_compatible_provider_model_no_providers(mock_shared_config):
+    """Test handling when no providers are configured"""
+    mock_shared_config.return_value.openai_compatible_providers = None
+
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("test_provider::gpt-4")
+    assert str(exc_info.value) == "OpenAI compatible provider test_provider not found"
+
+
+def test_openai_compatible_provider_model_provider_not_found(mock_shared_config):
+    """Test handling of non-existent provider"""
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("unknown_provider::gpt-4")
+    assert (
+        str(exc_info.value) == "OpenAI compatible provider unknown_provider not found"
+    )
+
+
+def test_openai_compatible_provider_model_no_base_url(mock_shared_config):
+    """Test handling of provider without base URL"""
+    mock_shared_config.return_value.openai_compatible_providers = [
+        {
+            "name": "test_provider",
+            "api_key": "test-key",
+        }
+    ]
+
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("test_provider::gpt-4")
+    assert (
+        str(exc_info.value)
+        == "OpenAI compatible provider test_provider has no base URL"
+    )

kiln_ai/datamodel/__init__.py

@@ -1,3 +1,7 @@
+"""
+See our docs for details about our datamodel: https://kiln-ai.github.io/Kiln/kiln_core_docs/kiln_ai.html
+"""
+
 from __future__ import annotations
 
 import json
@@ -8,7 +12,12 @@ from typing import TYPE_CHECKING, Callable, Dict, List, Type, Union
 
 import jsonschema
 import jsonschema.exceptions
-from pydantic import BaseModel, Field, model_validator
+from pydantic import (
+    BaseModel,
+    Field,
+    ValidationInfo,
+    model_validator,
+)
 from typing_extensions import Self
 
 from kiln_ai.datamodel.json_schema import JsonObjectSchema, schema_from_json_str
@@ -43,9 +52,25 @@ __all__ = [
     "TaskOutputRatingType",
     "TaskRequirement",
     "TaskDeterminism",
+    "strict_mode",
+    "set_strict_mode",
 ]
 
 
+# We want to be hard on ourselves for data completeness generated by the Kiln App, but don't want to make it hard for users to use the datamodel/library.
+# Strict mode enables extra validations that we want to enforce in Kiln App (and any other client that wants best practices), but not in the library (unless they opt in)
+_strict_mode: bool = False
+
+
+def strict_mode() -> bool:
+    return _strict_mode
+
+
+def set_strict_mode(value: bool) -> None:
+    global _strict_mode
+    _strict_mode = value
+
+
 class Priority(IntEnum):
     """Defines priority levels for tasks and requirements, where P0 is highest priority."""
 
@@ -60,30 +85,71 @@ class TaskOutputRatingType(str, Enum):
     """Defines the types of rating systems available for task outputs."""
 
     five_star = "five_star"
+    pass_fail = "pass_fail"
+    pass_fail_critical = "pass_fail_critical"
     custom = "custom"
 
 
+class RequirementRating(BaseModel):
+    """Rating for a specific requirement within a task output."""
+
+    value: float = Field(
+        description="The rating value. Interpretation depends on rating type"
+    )
+    type: TaskOutputRatingType = Field(description="The type of rating")
+
+
 class TaskOutputRating(KilnBaseModel):
     """
     A rating for a task output, including an overall rating and ratings for each requirement.
 
-    Only supports five star ratings for now, but extensible for custom values.
+    Supports:
+    - five_star: 1-5 star ratings
+    - pass_fail: boolean pass/fail (1.0 = pass, 0.0 = fail)
+    - pass_fail_critical: tri-state (1.0 = pass, 0.0 = fail, -1.0 = critical fail)
     """
 
     type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
     value: float | None = Field(
-        description="The overall rating value (typically 1-5 stars).",
+        description="The rating value. Interpretation depends on rating type:\n- five_star: 1-5 stars\n- pass_fail: 1.0 (pass) or 0.0 (fail)\n- pass_fail_critical: 1.0 (pass), 0.0 (fail), or -1.0 (critical fail)",
         default=None,
     )
-    requirement_ratings: Dict[ID_TYPE, float] = Field(
+    requirement_ratings: Dict[ID_TYPE, RequirementRating] = Field(
         default={},
-        description="The ratings of the requirements of the task. The keys are the ids of the requirements. The values are the ratings (typically 1-5 stars).",
+        description="The ratings of the requirements of the task.",
     )
 
+    # Previously we stored rating values as a dict of floats, but now we store them as RequirementRating objects.
+    @model_validator(mode="before")
+    def upgrade_old_format(cls, data: dict) -> dict:
+        if not isinstance(data, dict):
+            return data
+
+        # Check if we have the old format (dict of floats)
+        req_ratings = data.get("requirement_ratings", {})
+        if req_ratings and all(
+            isinstance(v, (int, float)) for v in req_ratings.values()
+        ):
+            # Convert each float to a RequirementRating object
+            # all ratings are five star at the point we used this format
+            data["requirement_ratings"] = {
+                k: {"value": v, "type": TaskOutputRatingType.five_star}
+                for k, v in req_ratings.items()
+            }
+
+        return data
+
     # Used to select high quality outputs for example selection (MultiShotPromptBuilder, etc)
     def is_high_quality(self) -> bool:
+        if self.value is None:
+            return False
+
         if self.type == TaskOutputRatingType.five_star:
-            return self.value is not None and self.value >= 4
+            return self.value >= 4
+        elif self.type == TaskOutputRatingType.pass_fail:
+            return self.value == 1.0
+        elif self.type == TaskOutputRatingType.pass_fail_critical:
+            return self.value == 1.0
         return False
 
     @model_validator(mode="after")
@@ -91,24 +157,61 @@ class TaskOutputRating(KilnBaseModel):
         if self.type not in TaskOutputRatingType:
             raise ValueError(f"Invalid rating type: {self.type}")
 
-        if self.type == TaskOutputRatingType.five_star:
-            if self.value is not None:
-                self._validate_five_star(self.value, "overall rating")
-            for req_id, req_rating in self.requirement_ratings.items():
-                self._validate_five_star(req_rating, f"requirement rating for {req_id}")
+        # Overall rating is optional
+        if self.value is not None:
+            self._validate_rating(self.type, self.value, "overall rating")
+
+        for req_id, req_rating in self.requirement_ratings.items():
+            self._validate_rating(
+                req_rating.type,
+                req_rating.value,
+                f"requirement rating for req ID: {req_id}",
+            )
 
         return self
 
-    def _validate_five_star(self, rating: float, rating_name: str) -> None:
-        if not isinstance(rating, float) or not rating.is_integer():
+    def _validate_rating(
+        self, type: TaskOutputRatingType, rating: float | None, rating_name: str
+    ) -> None:
+        if type == TaskOutputRatingType.five_star:
+            self._validate_five_star(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail:
+            self._validate_pass_fail(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail_critical:
+            self._validate_pass_fail_critical(rating, rating_name)
+
+    def _validate_five_star(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
             raise ValueError(
-                f"{rating_name.capitalize()} of type five_star must be an integer value (1.0, 2.0, 3.0, 4.0, or 5.0)"
+                f"{rating_name.capitalize()} of type five_star must be an integer value (1-5)"
             )
         if rating < 1 or rating > 5:
             raise ValueError(
                 f"{rating_name.capitalize()} of type five_star must be between 1 and 5 stars"
             )
 
+    def _validate_pass_fail(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be an integer value (0 or 1)"
+            )
+        if rating not in [0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be 0 (fail) or 1 (pass)"
+            )
+
+    def _validate_pass_fail_critical(
+        self, rating: float | None, rating_name: str
+    ) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be an integer value (-1, 0, or 1)"
+            )
+        if rating not in [-1, 0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be -1 (critical fail), 0 (fail), or 1 (pass)"
+            )
+
 
 class TaskOutput(KilnBaseModel):
     """
@@ -121,8 +224,9 @@ class TaskOutput(KilnBaseModel):
     output: str = Field(
         description="The output of the task. JSON formatted for structured output, plaintext for unstructured output."
     )
-    source: DataSource = Field(
-        description="The source of the output: human or synthetic."
+    source: DataSource | None = Field(
+        description="The source of the output: human or synthetic.",
+        default=None,
     )
     rating: TaskOutputRating | None = Field(
         default=None, description="The rating of the output"
@@ -139,6 +243,18 @@ class TaskOutput(KilnBaseModel):
                 raise ValueError(f"Output does not match task output schema: {e}")
         return self
 
+    @model_validator(mode="after")
+    def validate_output_source(self, info: ValidationInfo) -> Self:
+        # On strict mode and not loaded from file, we validate output_source is not None.
+        # We want to be able to load any data, even if it's not perfect. But we want to create perfect data when adding new data.
+        if not strict_mode():
+            return self
+        if self.loaded_from_file(info):
+            return self
+        if self.source is None:
+            raise ValueError("Output source is required when strict mode is enabled")
+        return self
+
 
 class FineTuneStatusType(str, Enum):
     """
@@ -326,8 +442,8 @@ class TaskRun(KilnParentedModel):
     input: str = Field(
         description="The inputs to the task. JSON formatted for structured input, plaintext for unstructured input."
     )
-    input_source: DataSource = Field(
-        description="The source of the input: human or synthetic."
+    input_source: DataSource | None = Field(
+        default=None, description="The source of the input: human or synthetic."
     )
 
     output: TaskOutput = Field(description="The output of the task run.")
@@ -343,6 +459,10 @@ class TaskRun(KilnParentedModel):
         default=None,
         description="Intermediate outputs from the task run. Keys are the names of the intermediate output steps (cot=chain of thought, etc), values are the output data.",
     )
+    tags: List[str] = Field(
+        default=[],
+        description="Tags for the task run. Tags are used to categorize task runs for filtering and reporting.",
+    )
 
     def parent_task(self) -> Task | None:
         if not isinstance(self.parent, Task):
@@ -392,6 +512,28 @@ class TaskRun(KilnParentedModel):
             )
         return self
 
+    @model_validator(mode="after")
+    def validate_input_source(self, info: ValidationInfo) -> Self:
+        # On strict mode and not loaded from file, we validate input_source is not None.
+        # We want to be able to load any data, even if it's not perfect. But we want to create perfect data when adding new data.
+        if not strict_mode():
+            return self
+        if self.loaded_from_file(info):
+            return self
+        if self.input_source is None:
+            raise ValueError("input_source is required when strict mode is enabled")
+        return self
+
+    @model_validator(mode="after")
+    def validate_tags(self) -> Self:
+        for tag in self.tags:
+            if not tag:
+                raise ValueError("Tags cannot be empty strings")
+            if " " in tag:
+                raise ValueError("Tags cannot contain spaces. Try underscores.")
+
+        return self
+
 
 # Define the type alias for clarity
 DatasetFilter = Callable[[TaskRun], bool]
@@ -552,7 +694,7 @@ class TaskRequirement(BaseModel):
     Defines a specific requirement that should be met by task outputs.
 
     Includes an identifier, name, description, instruction for meeting the requirement,
-    and priority level.
+    priority level, and rating type (five_star, pass_fail, pass_fail_critical, custom).
     """
 
     id: ID_TYPE = ID_FIELD
@@ -560,6 +702,7 @@ class TaskRequirement(BaseModel):
     description: str | None = Field(default=None)
     instruction: str = Field(min_length=1)
     priority: Priority = Field(default=Priority.p2)
+    type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
 
 
 class TaskDeterminism(str, Enum):