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

kiln_ai/datamodel/test_output_rating.py

@@ -1,14 +1,34 @@
+import json
+
 import pytest
 from pydantic import ValidationError
 
-from kiln_ai.datamodel import TaskOutputRating, TaskOutputRatingType
+from kiln_ai.datamodel import RequirementRating, TaskOutputRating, TaskOutputRatingType
 
 
 def test_valid_task_output_rating():
     rating = TaskOutputRating(value=4.0, requirement_ratings={"req1": 5.0, "req2": 3.0})
     assert rating.type == TaskOutputRatingType.five_star
     assert rating.value == 4.0
-    assert rating.requirement_ratings == {"req1": 5.0, "req2": 3.0}
+    dumped = json.loads(rating.model_dump_json())
+    assert dumped["requirement_ratings"] == {
+        "req1": {"type": TaskOutputRatingType.five_star, "value": 5.0},
+        "req2": {"type": TaskOutputRatingType.five_star, "value": 3.0},
+    }
+
+    # new format
+    rating = TaskOutputRating(
+        value=4.0,
+        requirement_ratings={
+            "req1": {"type": TaskOutputRatingType.five_star, "value": 5.0},
+            "req2": {"type": TaskOutputRatingType.five_star, "value": 3.0},
+        },
+    )
+    dumped = json.loads(rating.model_dump_json())
+    assert dumped["requirement_ratings"] == {
+        "req1": {"type": TaskOutputRatingType.five_star, "value": 5.0},
+        "req2": {"type": TaskOutputRatingType.five_star, "value": 3.0},
+    }
 
 
 def test_invalid_rating_type():
@@ -40,34 +60,92 @@ def test_rating_below_range():
         TaskOutputRating(value=0.0)
 
 
-def test_valid_requirement_ratings():
-    rating = TaskOutputRating(
-        value=4.0, requirement_ratings={"req1": 5.0, "req2": 3.0, "req3": 1.0}
+def test_valid_requirement_ratings_old_format():
+    rating = TaskOutputRating.model_validate(
+        {"value": 4.0, "requirement_ratings": {"req1": 5.0, "req2": 3.0, "req3": 1.0}}
+    )
+    dumped = json.loads(rating.model_dump_json())
+    assert dumped["requirement_ratings"] == {
+        "req1": {"type": TaskOutputRatingType.five_star, "value": 5.0},
+        "req2": {"type": TaskOutputRatingType.five_star, "value": 3.0},
+        "req3": {"type": TaskOutputRatingType.five_star, "value": 1.0},
+    }
+
+
+def test_valid_requirement_ratings_new_format():
+    rating = TaskOutputRating.model_validate(
+        {
+            "value": 4.0,
+            "requirement_ratings": {
+                "req1": {"type": TaskOutputRatingType.five_star, "value": 5.0},
+                "req2": {"type": TaskOutputRatingType.five_star, "value": 3.0},
+                "req3": {"type": TaskOutputRatingType.five_star, "value": 1.0},
+            },
+        }
     )
-    assert rating.requirement_ratings == {"req1": 5.0, "req2": 3.0, "req3": 1.0}
+    dumped = json.loads(rating.model_dump_json())
+    assert dumped["requirement_ratings"] == {
+        "req1": {"type": TaskOutputRatingType.five_star, "value": 5.0},
+        "req2": {"type": TaskOutputRatingType.five_star, "value": 3.0},
+        "req3": {"type": TaskOutputRatingType.five_star, "value": 1.0},
+    }
 
 
 def test_invalid_requirement_rating_value():
     with pytest.raises(
         ValidationError,
-        match="Requirement rating for req1 of type five_star must be an integer value",
+        match="Requirement rating for req id: req1 of type five_star must be an integer value",
     ):
         TaskOutputRating(value=4.0, requirement_ratings={"req1": 3.5})
 
+    # new format
+    with pytest.raises(
+        ValidationError,
+        match="Requirement rating for req id: req1 of type five_star must be an integer value",
+    ):
+        TaskOutputRating(
+            value=4.0,
+            requirement_ratings={
+                "req1": {"type": TaskOutputRatingType.five_star, "value": 3.5}
+            },
+        )
+
 
 def test_requirement_rating_out_of_range():
     with pytest.raises(
         ValidationError,
-        match="Requirement rating for req1 of type five_star must be between 1 and 5 stars",
+        match="Requirement rating for req id: req1 of type five_star must be between 1 and 5 stars",
     ):
         TaskOutputRating(value=4.0, requirement_ratings={"req1": 6.0})
 
+    # new format
+    with pytest.raises(
+        ValidationError,
+        match="Requirement rating for req id: req1 of type five_star must be between 1 and 5 stars",
+    ):
+        TaskOutputRating(
+            value=4.0,
+            requirement_ratings={
+                "req1": {"type": TaskOutputRatingType.five_star, "value": 6.0}
+            },
+        )
+
 
 def test_empty_requirement_ratings():
     rating = TaskOutputRating(value=4.0)
     assert rating.requirement_ratings == {}
 
 
+def test_empty_requirement_ratings_integer():
+    rating = TaskOutputRating(
+        value=4,
+        requirement_ratings={
+            "req1": RequirementRating(type=TaskOutputRatingType.five_star, value=5),
+        },
+    )
+    assert rating.requirement_ratings["req1"].value == 5.0
+
+
 def test_invalid_id_type():
     with pytest.raises(ValidationError):
         TaskOutputRating(
@@ -77,13 +155,302 @@ def test_invalid_id_type():
             },
         )
 
+    # new format
+    with pytest.raises(ValidationError):
+        TaskOutputRating(
+            value=4.0,
+            requirement_ratings={
+                123: {"type": TaskOutputRatingType.five_star, "value": 4.0}
+            },
+        )
+
 
 def test_valid_custom_rating():
     rating = TaskOutputRating(
         type=TaskOutputRatingType.custom,
         value=31.459,
-        requirement_ratings={"req1": 42.0, "req2": 3.14},
+        requirement_ratings={
+            "req1": {"type": TaskOutputRatingType.custom, "value": 42.0},
+            "req2": {"type": TaskOutputRatingType.custom, "value": 3.14},
+        },
     )
     assert rating.type == TaskOutputRatingType.custom
     assert rating.value == 31.459
-    assert rating.requirement_ratings == {"req1": 42.0, "req2": 3.14}
+    dumped = json.loads(rating.model_dump_json())
+    assert dumped["requirement_ratings"] == {
+        "req1": {"type": TaskOutputRatingType.custom, "value": 42.0},
+        "req2": {"type": TaskOutputRatingType.custom, "value": 3.14},
+    }
+
+
+# We upgraded the format of requirement_ratings to be a dict of RequirementRating objects from a dict of floats
+def test_task_output_rating_format_upgrade():
+    # Test old format (dict of floats)
+    old_format = {
+        "type": "five_star",
+        "value": 4.0,
+        "requirement_ratings": {"req1": 5.0, "req2": 3.0},
+    }
+
+    rating = TaskOutputRating.model_validate(old_format)
+
+    # Verify the upgrade worked
+    assert isinstance(rating.requirement_ratings["req1"], RequirementRating)
+    assert rating.requirement_ratings["req1"].value == 5.0
+    assert rating.requirement_ratings["req1"].type == TaskOutputRatingType.five_star
+    assert rating.requirement_ratings["req2"].value == 3.0
+    assert rating.requirement_ratings["req2"].type == TaskOutputRatingType.five_star
+
+    # Verify the json dump is new format
+    json_dump = json.loads(rating.model_dump_json())
+    assert json_dump["requirement_ratings"]["req1"]["type"] == "five_star"
+    assert json_dump["requirement_ratings"]["req1"]["value"] == 5.0
+    assert json_dump["requirement_ratings"]["req2"]["type"] == "five_star"
+    assert json_dump["requirement_ratings"]["req2"]["value"] == 3.0
+
+    # Test new format (dict of RequirementRating)
+    new_format = {
+        "type": "five_star",
+        "value": 4.0,
+        "requirement_ratings": {
+            "req1": {"value": 5.0, "type": "five_star"},
+            "req2": {"value": 3.0, "type": "five_star"},
+        },
+    }
+
+    rating = TaskOutputRating.model_validate(new_format)
+
+    # Verify new format works as expected
+    assert isinstance(rating.requirement_ratings["req1"], RequirementRating)
+    assert rating.requirement_ratings["req1"].value == 5.0
+    assert rating.requirement_ratings["req1"].type == TaskOutputRatingType.five_star
+
+    # Verify the json dump is new format
+    json_dump = json.loads(rating.model_dump_json())
+    assert json_dump["requirement_ratings"]["req1"]["type"] == "five_star"
+    assert json_dump["requirement_ratings"]["req1"]["value"] == 5.0
+    assert json_dump["requirement_ratings"]["req2"]["type"] == "five_star"
+    assert json_dump["requirement_ratings"]["req2"]["value"] == 3.0
+
+    # Test mixed format (should fail)
+    mixed_format = {
+        "type": "five_star",
+        "value": 4.0,
+        "requirement_ratings": {
+            "req1": 5.0,
+            "req2": {"value": 3.0, "type": "five_star"},
+        },
+    }
+
+    with pytest.raises(ValidationError):
+        TaskOutputRating.model_validate(mixed_format)
+
+    # Test empty requirement_ratings
+    empty_format = {"type": "five_star", "value": 4.0, "requirement_ratings": {}}
+
+    rating = TaskOutputRating.model_validate(empty_format)
+    assert rating.requirement_ratings == {}
+
+
+def test_valid_pass_fail_rating():
+    rating = TaskOutputRating(
+        type=TaskOutputRatingType.pass_fail,
+        value=1.0,
+        requirement_ratings={
+            "req1": {"type": TaskOutputRatingType.pass_fail, "value": 1.0},
+            "req2": {"type": TaskOutputRatingType.pass_fail, "value": 0.0},
+        },
+    )
+    assert rating.type == TaskOutputRatingType.pass_fail
+    assert rating.value == 1.0
+    dumped = json.loads(rating.model_dump_json())
+    assert dumped["requirement_ratings"] == {
+        "req1": {"type": TaskOutputRatingType.pass_fail, "value": 1.0},
+        "req2": {"type": TaskOutputRatingType.pass_fail, "value": 0.0},
+    }
+
+
+def test_invalid_pass_fail_rating_value():
+    with pytest.raises(
+        ValidationError,
+        match="Overall rating of type pass_fail must be an integer value",
+    ):
+        TaskOutputRating(type=TaskOutputRatingType.pass_fail, value=0.5)
+
+    with pytest.raises(
+        ValidationError,
+        match="Requirement rating for req id: req1 of type pass_fail must be an integer value",
+    ):
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail,
+            value=1.0,
+            requirement_ratings={
+                "req1": {"type": TaskOutputRatingType.pass_fail, "value": 0.5}
+            },
+        )
+
+
+def test_pass_fail_rating_out_of_range():
+    with pytest.raises(
+        ValidationError,
+        match="Overall rating of type pass_fail must be 0 \\(fail\\) or 1 \\(pass\\)",
+    ):
+        TaskOutputRating(type=TaskOutputRatingType.pass_fail, value=2.0)
+
+    with pytest.raises(
+        ValidationError,
+        match="Requirement rating for req id: req1 of type pass_fail must be 0 \\(fail\\) or 1 \\(pass\\)",
+    ):
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail,
+            value=1.0,
+            requirement_ratings={
+                "req1": {"type": TaskOutputRatingType.pass_fail, "value": 2.0}
+            },
+        )
+
+
+def test_valid_pass_fail_critical_rating():
+    rating = TaskOutputRating(
+        type=TaskOutputRatingType.pass_fail_critical,
+        value=1.0,
+        requirement_ratings={
+            "req1": {"type": TaskOutputRatingType.pass_fail_critical, "value": 1.0},
+            "req2": {"type": TaskOutputRatingType.pass_fail_critical, "value": 0.0},
+            "req3": {"type": TaskOutputRatingType.pass_fail_critical, "value": -1.0},
+        },
+    )
+    assert rating.type == TaskOutputRatingType.pass_fail_critical
+    assert rating.value == 1.0
+    dumped = json.loads(rating.model_dump_json())
+    assert dumped["requirement_ratings"] == {
+        "req1": {"type": TaskOutputRatingType.pass_fail_critical, "value": 1.0},
+        "req2": {"type": TaskOutputRatingType.pass_fail_critical, "value": 0.0},
+        "req3": {"type": TaskOutputRatingType.pass_fail_critical, "value": -1.0},
+    }
+
+
+def test_invalid_pass_fail_critical_rating_value():
+    with pytest.raises(
+        ValidationError,
+        match="Overall rating of type pass_fail_critical must be an integer value",
+    ):
+        TaskOutputRating(type=TaskOutputRatingType.pass_fail_critical, value=0.5)
+
+    with pytest.raises(
+        ValidationError,
+        match="Requirement rating for req id: req1 of type pass_fail_critical must be an integer value",
+    ):
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail_critical,
+            value=1.0,
+            requirement_ratings={
+                "req1": {"type": TaskOutputRatingType.pass_fail_critical, "value": 0.5}
+            },
+        )
+
+
+def test_pass_fail_critical_rating_out_of_range():
+    with pytest.raises(
+        ValidationError,
+        match="Overall rating of type pass_fail_critical must be -1 \\(critical fail\\), 0 \\(fail\\), or 1 \\(pass\\)",
+    ):
+        TaskOutputRating(type=TaskOutputRatingType.pass_fail_critical, value=2.0)
+
+    with pytest.raises(
+        ValidationError,
+        match="Requirement rating for req id: req1 of type pass_fail_critical must be -1 \\(critical fail\\), 0 \\(fail\\), or 1 \\(pass\\)",
+    ):
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail_critical,
+            value=1.0,
+            requirement_ratings={
+                "req1": {"type": TaskOutputRatingType.pass_fail_critical, "value": 2.0}
+            },
+        )
+
+
+def test_is_high_quality():
+    # Test five_star ratings
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.five_star, value=5.0
+        ).is_high_quality()
+        is True
+    )
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.five_star, value=4.0
+        ).is_high_quality()
+        is True
+    )
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.five_star, value=3.0
+        ).is_high_quality()
+        is False
+    )
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.five_star, value=2.0
+        ).is_high_quality()
+        is False
+    )
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.five_star, value=1.0
+        ).is_high_quality()
+        is False
+    )
+
+    # Test pass_fail ratings
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail, value=1.0
+        ).is_high_quality()
+        is True
+    )
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail, value=0.0
+        ).is_high_quality()
+        is False
+    )
+
+    # Test pass_fail_critical ratings
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail_critical, value=1.0
+        ).is_high_quality()
+        is True
+    )
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail_critical, value=0.0
+        ).is_high_quality()
+        is False
+    )
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.pass_fail_critical, value=-1.0
+        ).is_high_quality()
+        is False
+    )
+
+    # Test custom ratings (should always return False)
+    assert (
+        TaskOutputRating(
+            type=TaskOutputRatingType.custom, value=100.0
+        ).is_high_quality()
+        is False
+    )
+    assert (
+        TaskOutputRating(type=TaskOutputRatingType.custom, value=0.0).is_high_quality()
+        is False
+    )
+
+    # Test None value
+    assert (
+        TaskOutputRating(type=TaskOutputRatingType.custom, value=None).is_high_quality()
+        is False
+    )

kiln_ai/utils/config.py

@@ -2,7 +2,7 @@ import getpass
 import os
 import threading
 from pathlib import Path
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import yaml
 
@@ -15,12 +15,14 @@ class ConfigProperty:
         env_var: Optional[str] = None,
         default_lambda: Optional[Callable[[], Any]] = None,
         sensitive: bool = False,
+        sensitive_keys: Optional[List[str]] = None,
     ):
         self.type = type_
         self.default = default
         self.env_var = env_var
         self.default_lambda = default_lambda
         self.sensitive = sensitive
+        self.sensitive_keys = sensitive_keys
 
 
 class Config:
@@ -80,6 +82,15 @@ class Config:
                 list,
                 default_lambda=lambda: [],
             ),
+            "custom_models": ConfigProperty(
+                list,
+                default_lambda=lambda: [],
+            ),
+            "openai_compatible_providers": ConfigProperty(
+                list,
+                default_lambda=lambda: [],
+                sensitive_keys=["api_key"],
+            ),
         }
         self._settings = self.load_settings()
 
@@ -145,15 +156,27 @@ class Config:
             settings = yaml.safe_load(f.read()) or {}
         return settings
 
-    def settings(self, hide_sensitive=False):
-        if hide_sensitive:
-            return {
-                k: "[hidden]"
-                if k in self._properties and self._properties[k].sensitive
-                else v
-                for k, v in self._settings.items()
-            }
-        return self._settings
+    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
+        if not hide_sensitive:
+            return self._settings
+
+        settings = {
+            k: "[hidden]"
+            if k in self._properties and self._properties[k].sensitive
+            else v
+            for k, v in self._settings.items()
+        }
+        # Hide sensitive keys in lists. Could generalize this if we every have more types, but right not it's only needed for root elements of lists
+        for key, value in settings.items():
+            if key in self._properties and self._properties[key].sensitive_keys:
+                sensitive_keys = self._properties[key].sensitive_keys or []
+                for sensitive_key in sensitive_keys:
+                    if isinstance(value, list):
+                        for item in value:
+                            if sensitive_key in item:
+                                item[sensitive_key] = "[hidden]"
+
+        return settings
 
     def save_setting(self, name: str, value: Any):
         self.update_settings({name: value})

kiln_ai/utils/test_config.py

@@ -27,6 +27,7 @@ def config_with_yaml(mock_yaml_file):
                 ),
                 "int_property": ConfigProperty(int, default=0),
                 "empty_property": ConfigProperty(str),
+                "list_of_objects": ConfigProperty(list, default=[]),
             }
         )
 
@@ -251,3 +252,50 @@ def test_stale_values_bug(config_with_yaml):
     # Simulate updating the settings file with set_settings
     config_with_yaml.update_settings({"example_property": "third_value"})
     assert config_with_yaml.example_property == "third_value"
+
+
+async def test_openai_compatible_providers():
+    config = Config.shared()
+    assert config.openai_compatible_providers == []
+
+    new_settings = [
+        {
+            "name": "provider1",
+            "url": "https://provider1.com",
+            "api_key": "password1",
+        },
+        {
+            "name": "provider2",
+            "url": "https://provider2.com",
+        },
+    ]
+    config.save_setting("openai_compatible_providers", new_settings)
+    assert config.openai_compatible_providers == new_settings
+
+    # Test that sensitive keys are hidden
+    settings = config.settings(hide_sensitive=True)
+    assert settings["openai_compatible_providers"] == [
+        {"name": "provider1", "url": "https://provider1.com", "api_key": "[hidden]"},
+        {"name": "provider2", "url": "https://provider2.com"},
+    ]
+
+
+def test_yaml_persistence_structured_data(config_with_yaml, mock_yaml_file):
+    # Set a value
+    new_settings = [
+        {
+            "name": "provider1",
+            "url": "https://provider1.com",
+            "api_key": "password1",
+        },
+        {
+            "name": "provider2",
+            "url": "https://provider2.com",
+        },
+    ]
+    config_with_yaml.list_of_objects = new_settings
+
+    # Check that the value was saved to the YAML file
+    with open(mock_yaml_file, "r") as f:
+        saved_settings = yaml.safe_load(f)
+    assert saved_settings["list_of_objects"] == new_settings