kiln-ai 0.11.1__py3-none-any.whl → 0.13.0__py3-none-any.whl

kiln_ai/datamodel/test_example_models.py

@@ -11,7 +11,6 @@ from kiln_ai.datamodel import (
     Finetune,
     Project,
     Task,
-    TaskDeterminism,
     TaskOutput,
     TaskOutputRating,
     TaskOutputRatingType,
@@ -125,7 +124,6 @@ def test_structured_output_workflow(tmp_path):
         name="Structured Output Task",
         parent=project,
         instruction="Generate a JSON object with name and age",
-        determinism=TaskDeterminism.semantic_match,
         output_json_schema=json.dumps(
             {
                 "type": "object",
@@ -142,7 +140,7 @@ def test_structured_output_workflow(tmp_path):
 
     # Create runs
     runs = []
-    for source in DataSourceType:
+    for source in [DataSourceType.human, DataSourceType.synthetic]:
         for _ in range(2):
             task_run = TaskRun(
                 input="Generate info for John Doe",
@@ -157,7 +155,7 @@ def test_structured_output_workflow(tmp_path):
                         "adapter_name": "TestAdapter",
                         "model_name": "GPT-4",
                         "model_provider": "OpenAI",
-                        "prompt_builder_name": "TestPromptBuilder",
+                        "prompt_id": "simple_prompt_builder",
                     },
                 ),
                 parent=task,
@@ -216,9 +214,9 @@ def test_structured_output_workflow(tmp_path):
 
     assert loaded_task.name == "Structured Output Task"
     assert len(loaded_task.requirements) == 2
-    assert len(loaded_task.runs()) == 5
-
     loaded_runs = loaded_task.runs()
+    assert len(loaded_runs) == 5
+
     for task_run in loaded_runs:
         output = task_run.output
         assert output.rating is not None
@@ -472,7 +470,7 @@ def test_valid_synthetic_task_output():
                 "adapter_name": "TestAdapter",
                 "model_name": "GPT-4",
                 "model_provider": "OpenAI",
-                "prompt_builder_name": "TestPromptBuilder",
+                "prompt_id": "simple_prompt_builder",
             },
         ),
     )
@@ -480,7 +478,7 @@ def test_valid_synthetic_task_output():
     assert output.source.properties["adapter_name"] == "TestAdapter"
     assert output.source.properties["model_name"] == "GPT-4"
     assert output.source.properties["model_provider"] == "OpenAI"
-    assert output.source.properties["prompt_builder_name"] == "TestPromptBuilder"
+    assert output.source.properties["prompt_id"] == "simple_prompt_builder"
 
 
 def test_invalid_synthetic_task_output_missing_keys():
@@ -509,23 +507,21 @@ def test_invalid_synthetic_task_output_empty_values():
                     "adapter_name": "TestAdapter",
                     "model_name": "",
                     "model_provider": "OpenAI",
-                    "prompt_builder_name": "TestPromptBuilder",
+                    "prompt_id": "simple_prompt_builder",
                 },
             ),
         )
 
 
 def test_invalid_synthetic_task_output_non_string_values():
-    with pytest.raises(
-        ValidationError, match="'prompt_builder_name' must be of type str"
-    ):
+    with pytest.raises(ValidationError, match="'prompt_id' must be of type str"):
         DataSource(
             type=DataSourceType.synthetic,
             properties={
                 "adapter_name": "TestAdapter",
                 "model_name": "GPT-4",
                 "model_provider": "OpenAI",
-                "prompt_builder_name": 123,
+                "prompt_id": 123,
             },
         )
 

kiln_ai/datamodel/test_json_schema.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel
 from kiln_ai.datamodel.json_schema import (
     JsonObjectSchema,
     schema_from_json_str,
+    string_to_json_key,
     validate_schema,
 )
 
@@ -123,3 +124,25 @@ def test_triangle_schema():
     validate_schema({"a": 1, "b": 2, "c": 3}, json_triangle_schema)
     with pytest.raises(Exception):
         validate_schema({"a": 1, "b": 2, "c": "3"}, json_triangle_schema)
+
+
+@pytest.mark.parametrize(
+    "input_str,expected",
+    [
+        ("hello world", "hello_world"),
+        ("Hello World", "hello_world"),
+        ("hello_world", "hello_world"),
+        ("HELLO WORLD", "hello_world"),
+        ("hello123", "hello123"),
+        ("hello-world", "helloworld"),
+        ("hello!@#$%^&*()world", "helloworld"),
+        ("  hello  world  ", "hello__world"),
+        ("hello__world", "hello__world"),
+        ("", ""),
+        ("!@#$%", ""),
+        ("snake_case_string", "snake_case_string"),
+        ("camelCaseString", "camelcasestring"),
+    ],
+)
+def test_string_to_json_key(input_str: str, expected: str):
+    assert string_to_json_key(input_str) == expected

kiln_ai/datamodel/test_models.py

@@ -385,7 +385,7 @@ def test_task_run_input_source_validation(tmp_path):
     assert task_run.input_source is not None
 
     # Test 3: Creating without input_source should fail when strict mode is on
-    with patch("kiln_ai.datamodel.strict_mode", return_value=True):
+    with patch("kiln_ai.datamodel.task_run.strict_mode", return_value=True):
         with pytest.raises(ValueError) as exc_info:
             task_run = TaskRun(
                 input="test input 3",
@@ -442,7 +442,7 @@ def test_task_output_source_validation(tmp_path):
     assert task_output.source is not None
 
     # Test 3: Creating without source should fail when strict mode is on
-    with patch("kiln_ai.datamodel.strict_mode", return_value=True):
+    with patch("kiln_ai.datamodel.task_output.strict_mode", return_value=True):
         with pytest.raises(ValueError) as exc_info:
             task_output = TaskOutput(
                 output="test output 3",

kiln_ai/datamodel/test_prompt_id.py

@@ -0,0 +1,129 @@
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from kiln_ai.datamodel import (
+    PromptGenerators,
+    PromptId,
+)
+from kiln_ai.datamodel.prompt_id import is_frozen_prompt
+
+
+# Test model to validate the PromptId type
+class ModelTester(BaseModel):
+    prompt_id: PromptId
+
+
+def test_valid_prompt_generator_names():
+    """Test that valid prompt generator names are accepted"""
+    for generator in PromptGenerators:
+        model = ModelTester(prompt_id=generator.value)
+        assert model.prompt_id == generator.value
+
+
+def test_valid_saved_prompt_id():
+    """Test that valid saved prompt IDs are accepted"""
+    valid_id = "id::prompt_789"
+    model = ModelTester(prompt_id=valid_id)
+    assert model.prompt_id == valid_id
+
+
+def test_valid_fine_tune_prompt_id():
+    """Test that valid fine-tune prompt IDs are accepted"""
+    valid_id = "fine_tune_prompt::ft_123456"
+    model = ModelTester(prompt_id=valid_id)
+    assert model.prompt_id == valid_id
+
+
+@pytest.mark.parametrize(
+    "invalid_id",
+    [
+        pytest.param("id::project_123::task_456", id="missing_prompt_id"),
+        pytest.param("id::task_456::prompt_789", id="too_many_parts"),
+        pytest.param("id::", id="empty_parts"),
+    ],
+)
+def test_invalid_saved_prompt_id_format(invalid_id):
+    """Test that invalid saved prompt ID formats are rejected"""
+    with pytest.raises(ValidationError, match="Invalid saved prompt ID"):
+        ModelTester(prompt_id=invalid_id)
+
+
+@pytest.mark.parametrize(
+    "invalid_id,expected_error",
+    [
+        ("fine_tune_prompt::", "Invalid fine-tune prompt ID: fine_tune_prompt::"),
+        ("fine_tune_prompt", "Invalid prompt ID: fine_tune_prompt"),
+    ],
+)
+def test_invalid_fine_tune_prompt_id_format(invalid_id, expected_error):
+    """Test that invalid fine-tune prompt ID formats are rejected"""
+    with pytest.raises(ValidationError, match=expected_error):
+        ModelTester(prompt_id=invalid_id)
+
+
+def test_completely_invalid_formats():
+    """Test that completely invalid formats are rejected"""
+    invalid_ids = [
+        "",  # Empty string
+        "invalid_format",  # Random string
+        "id:wrong_format",  # Almost correct but wrong separator
+        "fine_tune:wrong_format",  # Almost correct but wrong prefix
+        ":::",  # Just separators
+    ]
+
+    for invalid_id in invalid_ids:
+        with pytest.raises(ValidationError, match="Invalid prompt ID"):
+            ModelTester(prompt_id=invalid_id)
+
+
+def test_prompt_generator_case_sensitivity():
+    """Test that prompt generator names are case sensitive"""
+    # Take first generator and modify its case
+    first_generator = next(iter(PromptGenerators)).value
+    wrong_case = first_generator.upper()
+    if wrong_case == first_generator:
+        wrong_case = first_generator.lower()
+
+    with pytest.raises(ValidationError):
+        ModelTester(prompt_id=wrong_case)
+
+
+@pytest.mark.parametrize(
+    "valid_id",
+    [
+        "task_run_config::project_123::task_456::config_123",  # Valid task run config prompt ID
+    ],
+)
+def test_valid_task_run_config_prompt_id(valid_id):
+    """Test that valid eval prompt IDs are accepted"""
+    model = ModelTester(prompt_id=valid_id)
+    assert model.prompt_id == valid_id
+
+
+@pytest.mark.parametrize(
+    "invalid_id,expected_error",
+    [
+        ("task_run_config::", "Invalid task run config prompt ID"),
+        ("task_run_config::p1", "Invalid task run config prompt ID"),
+        ("task_run_config::p1::t1", "Invalid task run config prompt ID"),
+        ("task_run_config::p1::t1::c1::extra", "Invalid task run config prompt ID"),
+    ],
+)
+def test_invalid_eval_prompt_id_format(invalid_id, expected_error):
+    """Test that invalid eval prompt ID formats are rejected"""
+    with pytest.raises(ValidationError, match=expected_error):
+        ModelTester(prompt_id=invalid_id)
+
+
+@pytest.mark.parametrize(
+    "id,should_be_frozen",
+    [
+        ("simple_prompt_builder", False),
+        ("id::prompt_123", True),
+        ("task_run_config::p1::t1", True),
+        ("fine_tune_prompt::ft_123", True),
+    ],
+)
+def test_is_frozen_prompt(id, should_be_frozen):
+    """Test that the is_frozen_prompt function works"""
+    assert is_frozen_prompt(id) == should_be_frozen

kiln_ai/datamodel/test_task.py

@@ -0,0 +1,159 @@
+import pytest
+from pydantic import ValidationError
+
+from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
+from kiln_ai.datamodel.prompt_id import PromptGenerators
+from kiln_ai.datamodel.task import RunConfig, RunConfigProperties, Task, TaskRunConfig
+from kiln_ai.datamodel.task_output import normalize_rating
+
+
+def test_runconfig_valid_creation():
+    task = Task(id="task1", name="Test Task", instruction="Do something")
+
+    config = RunConfig(
+        task=task,
+        model_name="gpt-4",
+        model_provider_name="openai",
+        prompt_id=PromptGenerators.SIMPLE,
+    )
+
+    assert config.task == task
+    assert config.model_name == "gpt-4"
+    assert config.model_provider_name == "openai"
+    assert config.prompt_id == PromptGenerators.SIMPLE  # Check default value
+
+
+def test_runconfig_missing_required_fields():
+    with pytest.raises(ValidationError) as exc_info:
+        RunConfig()
+
+    errors = exc_info.value.errors()
+    assert (
+        len(errors) == 4
+    )  # task, model_name, model_provider_name, and prompt_id are required
+    assert any(error["loc"][0] == "task" for error in errors)
+    assert any(error["loc"][0] == "model_name" for error in errors)
+    assert any(error["loc"][0] == "model_provider_name" for error in errors)
+    assert any(error["loc"][0] == "prompt_id" for error in errors)
+
+
+def test_runconfig_custom_prompt_id():
+    task = Task(id="task1", name="Test Task", instruction="Do something")
+
+    config = RunConfig(
+        task=task,
+        model_name="gpt-4",
+        model_provider_name="openai",
+        prompt_id=PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT,
+    )
+
+    assert config.prompt_id == PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT
+
+
+@pytest.fixture
+def sample_task():
+    return Task(name="Test Task", instruction="Test instruction")
+
+
+@pytest.fixture
+def sample_run_config_props(sample_task):
+    return RunConfigProperties(
+        model_name="gpt-4",
+        model_provider_name="openai",
+        prompt_id=PromptGenerators.SIMPLE,
+    )
+
+
+def test_task_run_config_valid_creation(sample_task, sample_run_config_props):
+    config = TaskRunConfig(
+        name="Test Config",
+        description="Test description",
+        run_config_properties=sample_run_config_props,
+        parent=sample_task,
+    )
+
+    assert config.name == "Test Config"
+    assert config.description == "Test description"
+    assert config.run_config_properties == sample_run_config_props
+    assert config.parent_task() == sample_task
+
+
+def test_task_run_config_minimal_creation(sample_task, sample_run_config_props):
+    # Test creation with only required fields
+    config = TaskRunConfig(
+        name="Test Config",
+        run_config_properties=sample_run_config_props,
+        parent=sample_task,
+    )
+
+    assert config.name == "Test Config"
+    assert config.description is None
+    assert config.run_config_properties == sample_run_config_props
+
+
+def test_task_run_config_missing_required_fields(sample_task):
+    # Test missing name
+    with pytest.raises(ValidationError) as exc_info:
+        TaskRunConfig(
+            run_config_properties=RunConfigProperties(
+                task=sample_task, model_name="gpt-4", model_provider_name="openai"
+            ),
+            parent=sample_task,
+        )
+    assert "Field required" in str(exc_info.value)
+
+    # Test missing run_config
+    with pytest.raises(ValidationError) as exc_info:
+        TaskRunConfig(name="Test Config", parent=sample_task)
+    assert "Field required" in str(exc_info.value)
+
+
+def test_task_run_config_missing_task_in_run_config(sample_task):
+    with pytest.raises(
+        ValidationError, match="Input should be a valid dictionary or instance of Task"
+    ):
+        # Create a run config without a task
+        RunConfig(
+            model_name="gpt-4",
+            model_provider_name="openai",
+            task=None,  # type: ignore
+        )
+
+
+@pytest.mark.parametrize(
+    "rating_type,rating,expected",
+    [
+        (TaskOutputRatingType.five_star, 1, 0),
+        (TaskOutputRatingType.five_star, 2, 0.25),
+        (TaskOutputRatingType.five_star, 3, 0.5),
+        (TaskOutputRatingType.five_star, 4, 0.75),
+        (TaskOutputRatingType.five_star, 5, 1),
+        (TaskOutputRatingType.pass_fail, 0, 0),
+        (TaskOutputRatingType.pass_fail, 1, 1),
+        (TaskOutputRatingType.pass_fail, 0.5, 0.5),
+        (TaskOutputRatingType.pass_fail_critical, -1, 0),
+        (TaskOutputRatingType.pass_fail_critical, 0, 0.5),
+        (TaskOutputRatingType.pass_fail_critical, 1, 1),
+        (TaskOutputRatingType.pass_fail_critical, 0.5, 0.75),
+    ],
+)
+def test_normalize_rating(rating_type, rating, expected):
+    assert normalize_rating(rating, rating_type) == expected
+
+
+@pytest.mark.parametrize(
+    "rating_type,rating",
+    [
+        (TaskOutputRatingType.five_star, 0),
+        (TaskOutputRatingType.five_star, 6),
+        (TaskOutputRatingType.pass_fail, -0.5),
+        (TaskOutputRatingType.pass_fail, 1.5),
+        (TaskOutputRatingType.pass_fail_critical, -1.5),
+        (TaskOutputRatingType.pass_fail_critical, 1.5),
+        (TaskOutputRatingType.custom, 0),
+        (TaskOutputRatingType.custom, 99),
+    ],
+)
+def test_normalize_rating_errors(rating_type, rating):
+    with pytest.raises(ValueError):
+        normalize_rating(rating, rating_type)

kiln_ai/utils/config.py

@@ -78,10 +78,47 @@ class Config:
                 str,
                 env_var="FIREWORKS_ACCOUNT_ID",
             ),
+            "anthropic_api_key": ConfigProperty(
+                str,
+                env_var="ANTHROPIC_API_KEY",
+                sensitive=True,
+            ),
+            "gemini_api_key": ConfigProperty(
+                str,
+                env_var="GEMINI_API_KEY",
+                sensitive=True,
+            ),
             "projects": ConfigProperty(
                 list,
                 default_lambda=lambda: [],
             ),
+            "azure_openai_api_key": ConfigProperty(
+                str,
+                env_var="AZURE_OPENAI_API_KEY",
+                sensitive=True,
+            ),
+            "azure_openai_endpoint": ConfigProperty(
+                str,
+                env_var="AZURE_OPENAI_ENDPOINT",
+            ),
+            "huggingface_api_key": ConfigProperty(
+                str,
+                env_var="HUGGINGFACE_API_KEY",
+                sensitive=True,
+            ),
+            "vertex_project_id": ConfigProperty(
+                str,
+                env_var="VERTEX_PROJECT_ID",
+            ),
+            "vertex_location": ConfigProperty(
+                str,
+                env_var="VERTEX_LOCATION",
+            ),
+            "together_api_key": ConfigProperty(
+                str,
+                env_var="TOGETHERAI_API_KEY",
+                sensitive=True,
+            ),
             "custom_models": ConfigProperty(
                 list,
                 default_lambda=lambda: [],
@@ -142,10 +179,15 @@ class Config:
             raise AttributeError(f"Config has no attribute '{name}'")
 
     @classmethod
-    def settings_path(cls, create=True):
+    def settings_dir(cls, create=True):
         settings_dir = os.path.join(Path.home(), ".kiln_ai")
         if create and not os.path.exists(settings_dir):
             os.makedirs(settings_dir)
+        return settings_dir
+
+    @classmethod
+    def settings_path(cls, create=True):
+        settings_dir = cls.settings_dir(create)
         return os.path.join(settings_dir, "settings.yaml")
 
     @classmethod