kiln-ai 0.0.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

kiln_ai/datamodel/test_json_schema.py

@@ -0,0 +1,124 @@
+import pytest
+from kiln_ai.datamodel.json_schema import (
+    JsonObjectSchema,
+    schema_from_json_str,
+    validate_schema,
+)
+from pydantic import BaseModel
+
+
+class ExampleModel(BaseModel):
+    x_schema: JsonObjectSchema | None = None
+
+
+json_joke_schema = """{
+  "type": "object",
+  "properties": {
+    "setup": {
+      "description": "The setup of the joke",
+      "title": "Setup",
+      "type": "string"
+    },
+    "punchline": {
+      "description": "The punchline to the joke",
+      "title": "Punchline",
+      "type": "string"
+    },
+    "rating": {
+      "anyOf": [
+        {
+          "type": "integer"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
+      "description": "How funny the joke is, from 1 to 10",
+      "title": "Rating"
+    }
+  },
+  "required": [
+    "setup",
+    "punchline"
+  ]
+}
+"""
+
+
+def test_json_schema():
+    o = ExampleModel(x_schema=json_joke_schema)
+    parsed_schema = schema_from_json_str(o.x_schema)
+    assert parsed_schema is not None
+    assert parsed_schema["type"] == "object"
+    assert parsed_schema["required"] == ["setup", "punchline"]
+    assert parsed_schema["properties"]["setup"]["type"] == "string"
+    assert parsed_schema["properties"]["punchline"]["type"] == "string"
+    assert parsed_schema["properties"]["rating"] is not None
+
+    # Not json schema
+    with pytest.raises(ValueError):
+        o = ExampleModel(x_schema="hello")
+    with pytest.raises(ValueError):
+        o = ExampleModel(x_schema="{'asdf':{}}")
+    with pytest.raises(ValueError):
+        o = ExampleModel(x_schema="{asdf")
+
+
+def test_validate_schema_content():
+    o = {"setup": "asdf", "punchline": "asdf", "rating": 1}
+    validate_schema(o, json_joke_schema)
+    o = {"setup": "asdf"}
+    with pytest.raises(Exception):
+        validate_schema(0, json_joke_schema)
+    o = {"setup": "asdf", "punchline": "asdf"}
+    validate_schema(o, json_joke_schema)
+    o = {"setup": "asdf", "punchline": "asdf", "rating": "1"}
+    with pytest.raises(Exception):
+        validate_schema(o, json_joke_schema)
+
+
+json_triangle_schema = """{
+  "type": "object",
+  "properties": {
+    "a": {
+      "description": "length of side a",
+      "title": "A",
+      "type": "integer"
+    },
+    "b": {
+      "description": "length of side b",
+      "title": "B",
+      "type": "integer"
+    },
+    "c": {
+      "description": "length of side c",
+      "title": "C",
+      "type": "integer"
+    }
+  },
+  "required": [
+    "a",
+    "b",
+    "c"
+  ]
+}
+"""
+
+
+def test_triangle_schema():
+    o = ExampleModel(x_schema=json_joke_schema)
+    parsed_schema = schema_from_json_str(o.x_schema)
+    assert parsed_schema is not None
+
+    o = ExampleModel(x_schema=json_triangle_schema)
+    schema = schema_from_json_str(o.x_schema)
+
+    assert schema is not None
+    assert schema["properties"]["a"]["type"] == "integer"
+    assert schema["properties"]["b"]["type"] == "integer"
+    assert schema["properties"]["c"]["type"] == "integer"
+    assert schema["required"] == ["a", "b", "c"]
+    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)

kiln_ai/datamodel/test_models.py

@@ -0,0 +1,190 @@
+import json
+
+import pytest
+from kiln_ai.datamodel import Priority, Project, Task, TaskDeterminism
+from kiln_ai.datamodel.test_json_schema import json_joke_schema
+from pydantic import ValidationError
+
+
+@pytest.fixture
+def test_project_file(tmp_path):
+    test_file_path = tmp_path / "project.kiln"
+    data = {"v": 1, "name": "Test Project", "model_type": "project"}
+
+    with open(test_file_path, "w") as file:
+        json.dump(data, file, indent=4)
+
+    return test_file_path
+
+
+@pytest.fixture
+def test_task_file(tmp_path):
+    test_file_path = tmp_path / "task.json"
+    data = {
+        "v": 1,
+        "name": "Test Task",
+        "instruction": "Test Instruction",
+        "model_type": "task",
+    }
+
+    with open(test_file_path, "w") as file:
+        json.dump(data, file, indent=4)
+
+    return test_file_path
+
+
+def test_load_from_file(test_project_file):
+    project = Project.load_from_file(test_project_file)
+    assert project.v == 1
+    assert project.name == "Test Project"
+    assert project.path == test_project_file
+
+
+def test_project_init():
+    project = Project(name="test")
+    assert project.name == "test"
+
+
+def test_save_to_file(test_project_file):
+    project = Project(
+        name="Test Project", description="Test Description", path=test_project_file
+    )
+    project.save_to_file()
+
+    with open(test_project_file, "r") as file:
+        data = json.load(file)
+
+    assert data["v"] == 1
+    assert data["name"] == "Test Project"
+    assert data["description"] == "Test Description"
+
+
+def test_task_defaults():
+    task = Task(name="Test Task", instruction="Test Instruction")
+    assert task.description == ""
+    assert task.priority == Priority.p2
+    assert task.determinism == TaskDeterminism.flexible
+
+
+def test_task_serialization(test_project_file):
+    project = Project.load_from_file(test_project_file)
+    task = Task(
+        parent=project,
+        name="Test Task",
+        description="Test Description",
+        determinism=TaskDeterminism.semantic_match,
+        priority=Priority.p0,
+        instruction="Test Base Task Instruction",
+    )
+
+    task.save_to_file()
+
+    parsed_task = Task.all_children_of_parent_path(test_project_file)[0]
+    assert parsed_task.name == "Test Task"
+    assert parsed_task.description == "Test Description"
+    assert parsed_task.instruction == "Test Base Task Instruction"
+    assert parsed_task.determinism == TaskDeterminism.semantic_match
+    assert parsed_task.priority == Priority.p0
+
+
+def test_save_to_file_without_path():
+    project = Project(name="Test Project")
+    with pytest.raises(ValueError):
+        project.save_to_file()
+
+
+def test_name_validation():
+    Project(name="Test Project")
+    Project(name="Te st_Proj- 1234567890")
+    Project(name=("a" * 120))  # longest
+
+    # a string with 120 characters
+
+    with pytest.raises(ValueError):
+        Project(name="Test Project!")
+        Project(name="Test.Project")
+        Project(name=("a" * 121))  # too long
+        Project(name=("a"))  # too short
+
+
+def test_auto_type_name():
+    model = Project(name="Test Project")
+    assert model.model_type == "project"
+
+
+def test_load_tasks(test_project_file):
+    # Set up a project model
+    project = Project.load_from_file(test_project_file)
+
+    # Set up multiple task models under the project
+    task1 = Task(parent=project, name="Task1", instruction="Task 1 instruction")
+    task2 = Task(parent=project, name="Task2", instruction="Task 2 instruction")
+    task3 = Task(parent=project, name="Task3", instruction="Task 3 instruction")
+
+    # Ensure the tasks are saved correctly
+    task1.save_to_file()
+    task2.save_to_file()
+    task3.save_to_file()
+
+    # Load tasks from the project
+    tasks = project.tasks()
+
+    # Verify that all tasks are loaded correctly
+    assert len(tasks) == 3
+    names = [task.name for task in tasks]
+    assert "Task1" in names
+    assert "Task2" in names
+    assert "Task3" in names
+    assert all(task.model_type == "task" for task in tasks)
+    assert all(task.instruction != "" for task in tasks)
+
+
+# verify no error on non-saved model
+def test_load_children_no_path():
+    project = Project(name="Test Project")
+    assert len(project.tasks()) == 0
+
+
+def test_check_model_type(test_project_file, test_task_file):
+    project = Project.load_from_file(test_project_file)
+    task = Task.load_from_file(test_task_file)
+    assert project.model_type == "project"
+    assert task.model_type == "task"
+    assert task.instruction == "Test Instruction"
+
+    with pytest.raises(ValueError):
+        project = Project.load_from_file(test_task_file)
+
+    with pytest.raises(ValueError):
+        task = Task.load_from_file(test_project_file)
+
+
+def test_task_output_schema(tmp_path):
+    path = tmp_path / "task.kiln"
+    task = Task(name="Test Task", path=path, instruction="Test Instruction")
+    task.save_to_file()
+    assert task.output_schema() is None
+    task = Task(
+        name="Test Task",
+        instruction="Test Instruction",
+        output_json_schema=json_joke_schema,
+        input_json_schema=json_joke_schema,
+        path=path,
+    )
+    task.save_to_file()
+    schemas = [task.output_schema(), task.input_schema()]
+    for schema in schemas:
+        assert schema is not None
+        assert schema["properties"]["setup"]["type"] == "string"
+        assert schema["properties"]["punchline"]["type"] == "string"
+        assert schema["properties"]["rating"] is not None
+
+    # Not json schema
+    with pytest.raises(ValidationError):
+        task = Task(name="Test Task", output_json_schema="hello", path=path)
+    with pytest.raises(ValidationError):
+        task = Task(name="Test Task", output_json_schema='{"asdf":{}}', path=path)
+    with pytest.raises(ValidationError):
+        task = Task(name="Test Task", output_json_schema="{'asdf':{}}", path=path)
+    with pytest.raises(ValidationError):
+        task = Task(name="Test Task", input_json_schema="{asdf", path=path)

kiln_ai/datamodel/test_nested_save.py

@@ -0,0 +1,205 @@
+import pytest
+from kiln_ai.datamodel.basemodel import KilnParentedModel, KilnParentModel
+from pydantic import Field, ValidationError
+
+
+class ModelC(KilnParentedModel):
+    code: str = Field(..., pattern=r"^[A-Z]{3}$")
+
+    @classmethod
+    def relationship_name(cls) -> str:
+        return "cs"
+
+    @classmethod
+    def parent_type(cls):
+        return ModelB
+
+
+class ModelB(KilnParentedModel, KilnParentModel, parent_of={"cs": ModelC}):
+    value: int = Field(..., ge=0)
+
+    @classmethod
+    def relationship_name(cls) -> str:
+        return "bs"
+
+    @classmethod
+    def parent_type(cls):
+        return ModelA
+
+
+# Define the hierarchy
+class ModelA(KilnParentModel, parent_of={"bs": ModelB}):
+    name: str = Field(..., min_length=3)
+
+
+def test_validation_error_in_c_level():
+    data = {
+        "name": "Root",
+        "bs": [
+            {
+                "value": 10,
+                "cs": [
+                    {"code": "ABC"},
+                    {"code": "DEF"},
+                    {"code": "invalid"},  # This should cause a validation error
+                ],
+            }
+        ],
+    }
+
+    with pytest.raises(ValidationError) as exc_info:
+        ModelA.validate_and_save_with_subrelations(data)
+
+    assert "String should match pattern" in str(exc_info.value)
+
+
+def test_persist_three_level_hierarchy(tmp_path):
+    # Set up temporary paths
+    root_path = tmp_path / "model_a.kiln"
+
+    data = {
+        "name": "Root",
+        "bs": [
+            {"value": 10, "cs": [{"code": "ABC"}, {"code": "DEF"}]},
+            {"value": 20, "cs": [{"code": "XYZ"}]},
+        ],
+    }
+
+    instance = ModelA.validate_and_save_with_subrelations(data, path=root_path)
+
+    assert isinstance(instance, ModelA)
+    assert instance.name == "Root"
+    assert instance.path == root_path
+    assert len(instance.bs()) == 2
+
+    # Load the instance back from the file to double-check
+    instance = ModelA.load_from_file(root_path)
+
+    bs = instance.bs()
+    assert len(bs) == 2
+
+    # Check for the existence of both expected B models
+    b_values = [b.value for b in bs]
+    assert 10 in b_values
+    assert 20 in b_values
+
+    # Find the B models by their values
+    b10 = next(b for b in bs if b.value == 10)
+    b20 = next(b for b in bs if b.value == 20)
+
+    assert len(b10.cs()) == 2
+    assert len(b20.cs()) == 1
+
+    # Check C models for b10
+    c_codes_b10 = [c.code for c in b10.cs()]
+    assert "ABC" in c_codes_b10
+    assert "DEF" in c_codes_b10
+
+    # Check C model for b20
+    c_codes_b20 = [c.code for c in b20.cs()]
+    assert "XYZ" in c_codes_b20
+
+    # Check that all objects have their parent set correctly
+    assert all(b.parent == instance for b in bs)
+    assert all(c.parent.id == b10.id for c in b10.cs())
+    assert all(c.parent.id == b20.id for c in b20.cs())
+
+
+def test_persist_model_a_without_children(tmp_path):
+    # Set up temporary path
+    root_path = tmp_path / "model_a_no_children.kiln"
+
+    data = {"name": "RootNoChildren"}
+
+    instance = ModelA.validate_and_save_with_subrelations(data, path=root_path)
+
+    assert isinstance(instance, ModelA)
+    assert instance.name == "RootNoChildren"
+    assert instance.path == root_path
+    assert len(instance.bs()) == 0
+
+    # Verify that the file was created
+    assert root_path.exists()
+
+    # Load the instance back from the file to double-check
+    loaded_instance = ModelA.load_from_file(root_path)
+    assert loaded_instance.name == "RootNoChildren"
+    assert len(loaded_instance.bs()) == 0
+
+
+def test_validate_without_saving(tmp_path):
+    data = {
+        "name": "ValidateOnly",
+        "bs": [
+            {"value": 30, "cs": [{"code": "GHI"}, {"code": "JKL"}]},
+            {"value": 40, "cs": [{"code": "MNO"}]},
+        ],
+    }
+
+    # Validate the data without saving
+    ModelA._validate_nested(data, save=False)
+
+    data = {
+        "name": "ValidateOnly",
+        "bs": [
+            {"value": 30, "cs": [{"code": "GHI"}, {"code": "JKL"}]},
+            {"value": 40, "cs": [{"code": 123}]},
+        ],
+    }
+
+    with pytest.raises(ValidationError):
+        ModelA._validate_nested(data, save=False)
+
+
+def test_validation_error_in_multiple_levels():
+    data = {
+        "missing_name": "Root",
+        "bs": [
+            {
+                "value": -1,
+                "cs": [
+                    {"code": "ABC"},
+                    {"code": "DEF"},
+                    {"code": "invalid"},
+                ],
+            }
+        ],
+    }
+
+    with pytest.raises(ValidationError) as exc_info:
+        ModelA.validate_and_save_with_subrelations(data)
+
+    assert len(exc_info.value.errors()) == 3
+
+    first = exc_info.value.errors()[0]
+    assert "Field required" in first["msg"]
+    assert first["loc"] == ("name",)
+
+    second = exc_info.value.errors()[1]
+    assert "Input should be greater than or equal to 0" in second["msg"]
+    assert second["loc"] == ("bs", 0, "value")
+
+    third = exc_info.value.errors()[2]
+    assert "String should match pattern" in third["msg"]
+    assert third["loc"] == ("bs", 0, "cs", 2, "code")
+
+
+def test_validation_error_in_c_level_length():
+    data = {
+        "name": "Root",
+        "bs": [
+            {
+                "value": 10,
+                "cs": [
+                    {"code": "ABC"},
+                    {"code": "DEF"},
+                    {"code": "GE"},  # This should cause a validation error
+                ],
+            }
+        ],
+    }
+
+    with pytest.raises(ValidationError) as exc_info:
+        ModelA.validate_and_save_with_subrelations(data)
+
+    assert "String should match pattern" in str(exc_info.value)

kiln_ai/datamodel/test_output_rating.py

@@ -0,0 +1,88 @@
+import pytest
+from kiln_ai.datamodel import TaskOutputRating, TaskOutputRatingType
+from pydantic import ValidationError
+
+
+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}
+
+
+def test_invalid_rating_type():
+    with pytest.raises(ValidationError, match="Input should be"):
+        TaskOutputRating(type="invalid_type", value=4.0)
+
+
+def test_invalid_rating_value():
+    with pytest.raises(
+        ValidationError,
+        match="Overall rating of type five_star must be an integer value",
+    ):
+        TaskOutputRating(value=3.5)
+
+
+def test_rating_out_of_range():
+    with pytest.raises(
+        ValidationError,
+        match="Overall rating of type five_star must be between 1 and 5 stars",
+    ):
+        TaskOutputRating(value=6.0)
+
+
+def test_rating_below_range():
+    with pytest.raises(
+        ValidationError,
+        match="Overall rating of type five_star must be between 1 and 5 stars",
+    ):
+        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}
+    )
+    assert rating.requirement_ratings == {"req1": 5.0, "req2": 3.0, "req3": 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",
+    ):
+        TaskOutputRating(value=4.0, requirement_ratings={"req1": 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",
+    ):
+        TaskOutputRating(value=4.0, requirement_ratings={"req1": 6.0})
+
+
+def test_empty_requirement_ratings():
+    rating = TaskOutputRating(value=4.0)
+    assert rating.requirement_ratings == {}
+
+
+def test_invalid_id_type():
+    with pytest.raises(ValidationError):
+        TaskOutputRating(
+            value=4.0,
+            requirement_ratings={
+                123: 4.0  # Assuming ID_TYPE is str
+            },
+        )
+
+
+def test_valid_custom_rating():
+    rating = TaskOutputRating(
+        type=TaskOutputRatingType.custom,
+        value=31.459,
+        requirement_ratings={"req1": 42.0, "req2": 3.14},
+    )
+    assert rating.type == TaskOutputRatingType.custom
+    assert rating.value == 31.459
+    assert rating.requirement_ratings == {"req1": 42.0, "req2": 3.14}