Skip to content

Commit

Permalink
Add input validation
Browse files Browse the repository at this point in the history
  • Loading branch information
scosman committed Sep 11, 2024
1 parent 9a7a02e commit 89bf6b9
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 8 deletions.
1 change: 0 additions & 1 deletion libs/core/kiln_ai/datamodel/basemodel.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
BaseModel,
ConfigDict,
Field,
ValidationError,
computed_field,
model_validator,
)
Expand Down
21 changes: 21 additions & 0 deletions libs/core/kiln_ai/datamodel/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,27 @@ def parent_type(cls):
def outputs(self) -> list[ExampleOutput]:
return ExampleOutput.all_children_of_parent_path(self.path)

@model_validator(mode="after")
def validate_input_format(self) -> Self:
task = self.parent
if task is None:
# don't validate this relationship until we have a path or parent. Give them time to build it (but will catch it before saving)
return self
if not isinstance(task, Task):
raise ValueError(
"ExampleOutput's parent Example must have a valid parent Task"
)

# validate output
if task.input_json_schema is not None:
try:
validate_schema(json.loads(self.input), task.input_json_schema)
except json.JSONDecodeError:
raise ValueError("Input is not a valid JSON object")
except jsonschema.exceptions.ValidationError as e:
raise ValueError(f"Input does not match task input schema: {e}")
return self


class TaskRequirement(KilnParentedModel):
name: str = NAME_FIELD
Expand Down
58 changes: 51 additions & 7 deletions libs/core/kiln_ai/datamodel/test_example_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,12 @@
from pydantic import ValidationError


def test_example_model_validation():
def test_example_model_validation(tmp_path):
# Valid example
task = Task(name="Test Task", path=tmp_path / Task.base_filename())
task.save_to_file()
valid_example = Example(
path="/test/path",
parent=task,
input="Test input",
source=ExampleSource.human,
source_properties={"creator": "John Doe"},
Expand All @@ -30,29 +32,31 @@ def test_example_model_validation():
# Invalid source
with pytest.raises(ValidationError):
Example(
path="/test/path",
parent=task,
input="Test input",
source="invalid_source",
source_properties={},
)

# Missing required field
with pytest.raises(ValidationError):
Example(path="/test/path", source=ExampleSource.human, source_properties={})
Example(parent=task, source=ExampleSource.human, source_properties={})

# Invalid source_properties type
with pytest.raises(ValidationError):
Example(
path="/test/path",
parent=task,
input="Test input",
source=ExampleSource.human,
source_properties="invalid",
)


def test_example_relationship():
def test_example_relationship(tmp_path):
task = Task(name="Test Task", path=tmp_path / Task.base_filename())
task.save_to_file()
example = Example(
path="/test/path",
parent=task,
input="Test input",
source=ExampleSource.human,
source_properties={},
Expand Down Expand Up @@ -303,3 +307,43 @@ def test_example_output_schema_validation(tmp_path):
parent=example,
)
output.save_to_file()


def test_example_input_schema_validation(tmp_path):
# Create a project and task hierarchy
project = Project(name="Test Project", path=(tmp_path / "test_project"))
project.save_to_file()
task = Task(
name="Test Task",
parent=project,
input_json_schema=json.dumps(
{
"type": "object",
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
"required": ["name", "age"],
}
),
)
task.save_to_file()

# Create an example with a valid input schema
valid_example = Example(
input='{"name": "John Doe", "age": 30}',
source=ExampleSource.human,
parent=task,
)
valid_example.save_to_file()

# Changing to invalid input
with pytest.raises(ValueError):
valid_example.input = '{"name": "John Doe", "age": "thirty"}'
valid_example.save_to_file()

# Invalid case: input does not match task input schema
with pytest.raises(ValueError):
example = Example(
input='{"name": "John Doe", "age": "thirty"}',
source=ExampleSource.human,
parent=task,
)
example.save_to_file()

0 comments on commit 89bf6b9

Please sign in to comment.