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

kiln_ai/datamodel/basemodel.py

@@ -0,0 +1,372 @@
+import json
+import uuid
+from abc import ABCMeta
+from builtins import classmethod
+from datetime import datetime
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Optional,
+    Self,
+    Type,
+    TypeVar,
+)
+
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.formatting import snake_case
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationError,
+    computed_field,
+    model_validator,
+)
+from pydantic_core import ErrorDetails
+
+# ID is a 12 digit random integer string.
+# Should be unique per item, at least inside the context of a parent/child relationship.
+# Use integers to make it easier to type for a search function.
+# Allow none, even though we generate it, because we clear it in the REST API if the object is ephemeral (not persisted to disk)
+ID_FIELD = Field(default_factory=lambda: str(uuid.uuid4().int)[:12])
+ID_TYPE = Optional[str]
+T = TypeVar("T", bound="KilnBaseModel")
+PT = TypeVar("PT", bound="KilnParentedModel")
+
+
+class KilnBaseModel(BaseModel):
+    model_config = ConfigDict(validate_assignment=True)
+
+    v: int = Field(default=1)  # schema_version
+    id: ID_TYPE = ID_FIELD
+    path: Optional[Path] = Field(default=None)
+    created_at: datetime = Field(default_factory=datetime.now)
+    created_by: str = Field(default_factory=lambda: Config.shared().user_id)
+
+    @computed_field()
+    def model_type(self) -> str:
+        return self.type_name()
+
+    # if changing the model name, should keep the original name here for parsing old files
+    @classmethod
+    def type_name(cls) -> str:
+        return snake_case(cls.__name__)
+
+    # used as /obj_folder/base_filename.kiln
+    @classmethod
+    def base_filename(cls) -> str:
+        return cls.type_name() + ".kiln"
+
+    @classmethod
+    def load_from_folder(cls: Type[T], folderPath: Path) -> T:
+        path = folderPath / cls.base_filename()
+        return cls.load_from_file(path)
+
+    @classmethod
+    def load_from_file(cls: Type[T], path: Path) -> T:
+        with open(path, "r") as file:
+            file_data = file.read()
+            # TODO P2 perf: parsing the JSON twice here.
+            # Once for model_type, once for model. Can't call model_validate with parsed json because enum types break; they get strings instead of enums.
+            parsed_json = json.loads(file_data)
+            m = cls.model_validate_json(file_data, strict=True)
+            if not isinstance(m, cls):
+                raise ValueError(f"Loaded model is not of type {cls.__name__}")
+            file_data = None
+        m.path = path
+        if m.v > m.max_schema_version():
+            raise ValueError(
+                f"Cannot load from file because the schema version is higher than the current version. Upgrade kiln to the latest version. "
+                f"Class: {m.__class__.__name__}, id: {getattr(m, 'id', None)}, path: {path}, "
+                f"version: {m.v}, max version: {m.max_schema_version()}"
+            )
+        if parsed_json["model_type"] != cls.type_name():
+            raise ValueError(
+                f"Cannot load from file because the model type is incorrect. Expected {cls.type_name()}, got {parsed_json['model_type']}. "
+                f"Class: {m.__class__.__name__}, id: {getattr(m, 'id', None)}, path: {path}, "
+                f"version: {m.v}, max version: {m.max_schema_version()}"
+            )
+        return m
+
+    def save_to_file(self) -> None:
+        path = self.build_path()
+        if path is None:
+            raise ValueError(
+                f"Cannot save to file because 'path' is not set. Class: {self.__class__.__name__}, "
+                f"id: {getattr(self, 'id', None)}, path: {path}"
+            )
+        path.parent.mkdir(parents=True, exist_ok=True)
+        json_data = self.model_dump_json(indent=2, exclude={"path"})
+        with open(path, "w") as file:
+            file.write(json_data)
+        # save the path so even if something like name changes, the file doesn't move
+        self.path = path
+
+    def build_path(self) -> Path | None:
+        if self.path is not None:
+            return self.path
+        return None
+
+    # increment for breaking changes
+    def max_schema_version(self) -> int:
+        return 1
+
+
+class KilnParentedModel(KilnBaseModel, metaclass=ABCMeta):
+    _parent: KilnBaseModel | None = None
+
+    # workaround to tell typechecker that we support the parent property, even though it's not a stock property
+    if TYPE_CHECKING:
+        parent: KilnBaseModel  # type: ignore
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        if "parent" in data:
+            self.parent = data["parent"]
+
+    @property
+    def parent(self) -> Optional[KilnBaseModel]:
+        if self._parent is not None:
+            return self._parent
+        # lazy load parent from path
+        if self.path is None:
+            return None
+        # TODO: this only works with base_filename. If we every support custom names, we need to change this.
+        parent_path = (
+            self.path.parent.parent.parent
+            / self.__class__.parent_type().base_filename()
+        )
+        if parent_path is None:
+            return None
+        self._parent = self.__class__.parent_type().load_from_file(parent_path)
+        return self._parent
+
+    @parent.setter
+    def parent(self, value: Optional[KilnBaseModel]):
+        if value is not None:
+            expected_parent_type = self.__class__.parent_type()
+            if not isinstance(value, expected_parent_type):
+                raise ValueError(
+                    f"Parent must be of type {expected_parent_type}, but was {type(value)}"
+                )
+        self._parent = value
+
+    # Dynamically implemented by KilnParentModel method injection
+    @classmethod
+    def relationship_name(cls) -> str:
+        raise NotImplementedError("Relationship name must be implemented")
+
+    # Dynamically implemented by KilnParentModel method injection
+    @classmethod
+    def parent_type(cls) -> Type[KilnBaseModel]:
+        raise NotImplementedError("Parent type must be implemented")
+
+    @model_validator(mode="after")
+    def check_parent_type(self) -> Self:
+        if self._parent is not None:
+            expected_parent_type = self.__class__.parent_type()
+            if not isinstance(self._parent, expected_parent_type):
+                raise ValueError(
+                    f"Parent must be of type {expected_parent_type}, but was {type(self._parent)}"
+                )
+        return self
+
+    def build_child_dirname(self) -> Path:
+        # Default implementation for readable folder names.
+        # {id} - {name}/{type}.kiln
+        if self.id is None:
+            # consider generating an ID here. But if it's been cleared, we've already used this without one so raise for now.
+            raise ValueError("ID is not set - can not save or build path")
+        path = self.id
+        name = getattr(self, "name", None)
+        if name is not None:
+            path = f"{path} - {name[:32]}"
+        return Path(path)
+
+    def build_path(self) -> Path | None:
+        # if specifically loaded from an existing path, keep that no matter what
+        # this ensures the file structure is easy to use with git/version control
+        # and that changes to things like name (which impacts default path) don't leave dangling files
+        if self.path is not None:
+            return self.path
+        # Build a path under parent_folder/relationship/file.kiln
+        if self.parent is None:
+            return None
+        parent_path = self.parent.build_path()
+        if parent_path is None:
+            return None
+        parent_folder = parent_path.parent
+        if parent_folder is None:
+            return None
+        return (
+            parent_folder
+            / self.__class__.relationship_name()
+            / self.build_child_dirname()
+            / self.__class__.base_filename()
+        )
+
+    @classmethod
+    def all_children_of_parent_path(
+        cls: Type[PT], parent_path: Path | None
+    ) -> list[PT]:
+        if parent_path is None:
+            # children are disk based. If not saved, they don't exist
+            return []
+
+        # Determine the parent folder
+        if parent_path.is_file():
+            parent_folder = parent_path.parent
+        else:
+            parent_folder = parent_path
+
+        parent = cls.parent_type().load_from_file(parent_path)
+        if parent is None:
+            raise ValueError("Parent must be set to load children")
+
+        # Ignore type error: this is abstract base class, but children must implement relationship_name
+        relationship_folder = parent_folder / Path(cls.relationship_name())  # type: ignore
+
+        if not relationship_folder.exists() or not relationship_folder.is_dir():
+            return []
+
+        # Collect all /relationship/{id}/{base_filename.kiln} files in the relationship folder
+        children = []
+        for child_file in relationship_folder.glob(f"**/{cls.base_filename()}"):
+            child = cls.load_from_file(child_file)
+            children.append(child)
+
+        return children
+
+
+# Parent create methods for all child relationships
+# You must pass in parent_of in the subclass definition, defining the child relationships
+class KilnParentModel(KilnBaseModel, metaclass=ABCMeta):
+    @classmethod
+    def _create_child_method(
+        cls, relationship_name: str, child_class: Type[KilnParentedModel]
+    ):
+        def child_method(self) -> list[child_class]:
+            return child_class.all_children_of_parent_path(self.path)
+
+        child_method.__name__ = relationship_name
+        child_method.__annotations__ = {"return": List[child_class]}
+        setattr(cls, relationship_name, child_method)
+
+    @classmethod
+    def _create_parent_methods(
+        cls, targetCls: Type[KilnParentedModel], relationship_name: str
+    ):
+        def parent_class_method() -> Type[KilnParentModel]:
+            return cls
+
+        parent_class_method.__name__ = "parent_type"
+        parent_class_method.__annotations__ = {"return": Type[KilnParentModel]}
+        setattr(targetCls, "parent_type", parent_class_method)
+
+        def relationship_name_method() -> str:
+            return relationship_name
+
+        relationship_name_method.__name__ = "relationship_name"
+        relationship_name_method.__annotations__ = {"return": str}
+        setattr(targetCls, "relationship_name", relationship_name_method)
+
+    @classmethod
+    def __init_subclass__(cls, parent_of: Dict[str, Type[KilnParentedModel]], **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls._parent_of = parent_of
+        for relationship_name, child_class in parent_of.items():
+            cls._create_child_method(relationship_name, child_class)
+            cls._create_parent_methods(child_class, relationship_name)
+
+    @classmethod
+    def validate_and_save_with_subrelations(
+        cls,
+        data: Dict[str, Any],
+        path: Path | None = None,
+        parent: KilnBaseModel | None = None,
+    ):
+        # Validate first, then save. Don't want error half way through, and partly persisted
+        # TODO P2: save to tmp dir, then move atomically. But need to merge directories so later.
+        cls._validate_nested(data, save=False, path=path, parent=parent)
+        instance = cls._validate_nested(data, save=True, path=path, parent=parent)
+        return instance
+
+    @classmethod
+    def _validate_nested(
+        cls,
+        data: Dict[str, Any],
+        save: bool = False,
+        parent: KilnBaseModel | None = None,
+        path: Path | None = None,
+    ):
+        # Collect all validation errors so we can report them all at once
+        validation_errors = []
+
+        try:
+            instance = cls.model_validate(data, strict=True)
+            if path is not None:
+                instance.path = path
+            if parent is not None and isinstance(instance, KilnParentedModel):
+                instance.parent = parent
+            if save:
+                instance.save_to_file()
+        except ValidationError as e:
+            instance = None
+            for suberror in e.errors():
+                validation_errors.append(suberror)
+
+        for key, value_list in data.items():
+            if key in cls._parent_of:
+                parent_type = cls._parent_of[key]
+                if not isinstance(value_list, list):
+                    raise ValueError(
+                        f"Expected a list for {key}, but got {type(value_list)}"
+                    )
+                for value_index, value in enumerate(value_list):
+                    try:
+                        if issubclass(parent_type, KilnParentModel):
+                            kwargs = {"data": value, "save": save}
+                            if instance is not None:
+                                kwargs["parent"] = instance
+                            parent_type._validate_nested(**kwargs)
+                        elif issubclass(parent_type, KilnParentedModel):
+                            # Root node
+                            subinstance = parent_type.model_validate(value, strict=True)
+                            if instance is not None:
+                                subinstance.parent = instance
+                            if save:
+                                subinstance.save_to_file()
+                        else:
+                            raise ValueError(
+                                f"Invalid type {parent_type}. Should be KilnBaseModel based."
+                            )
+                    except ValidationError as e:
+                        for suberror in e.errors():
+                            cls._append_loc(suberror, key, value_index)
+                            validation_errors.append(suberror)
+
+        if len(validation_errors) > 0:
+            raise ValidationError.from_exception_data(
+                title=f"Validation failed for {cls.__name__}",
+                line_errors=validation_errors,
+                input_type="json",
+            )
+
+        return instance
+
+    @classmethod
+    def _append_loc(
+        cls, error: ErrorDetails, current_loc: str, value_index: int | None = None
+    ):
+        orig_loc = error["loc"] if "loc" in error else None
+        new_loc: list[str | int] = [current_loc]
+        if value_index is not None:
+            new_loc.append(value_index)
+        if isinstance(orig_loc, tuple):
+            new_loc.extend(list(orig_loc))
+        elif isinstance(orig_loc, list):
+            new_loc.extend(orig_loc)
+        error["loc"] = tuple(new_loc)

kiln_ai/datamodel/json_schema.py

@@ -0,0 +1,45 @@
+import json
+from typing import Annotated, Dict
+
+import jsonschema
+import jsonschema.exceptions
+import jsonschema.validators
+from pydantic import AfterValidator
+
+JsonObjectSchema = Annotated[
+    str,
+    AfterValidator(lambda v: _check_json_schema(v)),
+]
+
+
+def _check_json_schema(v: str) -> str:
+    # parsing returns needed errors
+    schema_from_json_str(v)
+    return v
+
+
+def validate_schema(instance: Dict, schema_str: str) -> None:
+    schema = schema_from_json_str(schema_str)
+    v = jsonschema.Draft202012Validator(schema)
+    return v.validate(instance)
+
+
+def schema_from_json_str(v: str) -> Dict:
+    try:
+        parsed = json.loads(v)
+        jsonschema.Draft202012Validator.check_schema(parsed)
+        if not isinstance(parsed, dict):
+            raise ValueError(f"JSON schema must be a dict, not {type(parsed)}")
+        if (
+            "type" not in parsed
+            or parsed["type"] != "object"
+            or "properties" not in parsed
+        ):
+            raise ValueError(f"JSON schema must be an object with properties: {v}")
+        return parsed
+    except jsonschema.exceptions.SchemaError as e:
+        raise ValueError(f"Invalid JSON schema: {v} \n{e}")
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Invalid JSON: {v}\n {e}")
+    except Exception as e:
+        raise ValueError(f"Unexpected error parsing JSON schema: {v}\n {e}")

kiln_ai/datamodel/test_basemodel.py

@@ -0,0 +1,277 @@
+import datetime
+import json
+from pathlib import Path
+from typing import Optional
+
+import pytest
+from kiln_ai.datamodel.basemodel import KilnBaseModel, KilnParentedModel
+
+
+@pytest.fixture
+def test_base_file(tmp_path) -> Path:
+    test_file_path = tmp_path / "test_model.json"
+    data = {"v": 1, "model_type": "kiln_base_model"}
+
+    with open(test_file_path, "w") as file:
+        json.dump(data, file, indent=4)
+
+    return test_file_path
+
+
+@pytest.fixture
+def test_base_parented_file(tmp_path) -> Path:
+    test_file_path = tmp_path / "test_model.json"
+    data = {"v": 1, "model_type": "base_parent_example"}
+
+    with open(test_file_path, "w") as file:
+        json.dump(data, file, indent=4)
+
+    return test_file_path
+
+
+@pytest.fixture
+def test_newer_file(tmp_path) -> Path:
+    test_file_path = tmp_path / "test_model_sv.json"
+    data = {"v": 99}
+
+    with open(test_file_path, "w") as file:
+        json.dump(data, file, indent=4)
+
+    return test_file_path
+
+
+def test_load_from_file(test_base_file):
+    model = KilnBaseModel.load_from_file(test_base_file)
+    assert model.v == 1
+    assert model.path == test_base_file
+
+
+def test_save_to_file(test_base_file):
+    model = KilnBaseModel(path=test_base_file)
+    model.save_to_file()
+
+    with open(test_base_file, "r") as file:
+        data = json.load(file)
+
+    assert data["v"] == 1
+    assert data["model_type"] == "kiln_base_model"
+
+
+def test_save_to_file_without_path():
+    model = KilnBaseModel()
+    with pytest.raises(ValueError):
+        model.save_to_file()
+
+
+def test_max_schema_version(test_newer_file):
+    with pytest.raises(ValueError):
+        KilnBaseModel.load_from_file(test_newer_file)
+
+
+def test_type_name():
+    model = KilnBaseModel()
+    assert model.model_type == "kiln_base_model"
+
+
+def test_created_atby():
+    model = KilnBaseModel()
+    assert model.created_at is not None
+    # Check it's within 2 seconds of now
+    now = datetime.datetime.now()
+    assert abs((model.created_at - now).total_seconds()) < 2
+
+    # Created by
+    assert len(model.created_by) > 0
+    # assert model.created_by == "scosman"
+
+
+# Instance of the parented model for abstract methods
+class NamedParentedModel(KilnParentedModel):
+    @classmethod
+    def relationship_name(cls) -> str:
+        return "tests"
+
+    @classmethod
+    def parent_type(cls):
+        return KilnBaseModel
+
+
+def test_parented_model_path_gen(tmp_path):
+    parent = KilnBaseModel(path=tmp_path)
+    assert parent.id is not None
+    child = NamedParentedModel(parent=parent)
+    child_path = child.build_path()
+    assert child_path.name == "named_parented_model.kiln"
+    assert child_path.parent.name == child.id
+    assert child_path.parent.parent.name == "tests"
+    assert child_path.parent.parent.parent == tmp_path.parent
+
+
+class BaseParentExample(KilnBaseModel):
+    name: Optional[str] = None
+
+
+# Instance of the parented model for abstract methods, with default name builder
+class DefaultParentedModel(KilnParentedModel):
+    name: Optional[str] = None
+
+    @classmethod
+    def relationship_name(self):
+        return "children"
+
+    @classmethod
+    def parent_type(cls):
+        return BaseParentExample
+
+
+def test_build_default_child_filename(tmp_path):
+    parent = BaseParentExample(path=tmp_path)
+    child = DefaultParentedModel(parent=parent)
+    child_path = child.build_path()
+    assert child_path.name == "default_parented_model.kiln"
+    assert child_path.parent.name == child.id
+    assert child_path.parent.parent.name == "children"
+    assert child_path.parent.parent.parent == tmp_path.parent
+    # now with name
+    child = DefaultParentedModel(parent=parent, name="Name")
+    child_path = child.build_path()
+    assert child_path.name == "default_parented_model.kiln"
+    assert child_path.parent.name == child.id + " - Name"
+    assert child_path.parent.parent.name == "children"
+    assert child_path.parent.parent.parent == tmp_path.parent
+
+
+def test_serialize_child(tmp_path):
+    parent = BaseParentExample(path=tmp_path)
+    child = DefaultParentedModel(parent=parent, name="Name")
+
+    expected_path = child.build_path()
+    assert child.path is None
+    child.save_to_file()
+
+    # ensure we save exact path
+    assert child.path is not None
+    assert child.path == expected_path
+
+    # should have made the directory, and saved the file
+    with open(child.path, "r") as file:
+        data = json.load(file)
+
+    assert data["v"] == 1
+    assert data["name"] == "Name"
+    assert data["model_type"] == "default_parented_model"
+    assert len(data["id"]) == 12
+    assert child.path.parent.name == child.id + " - Name"
+    assert child.path.parent.parent.name == "children"
+    assert child.path.parent.parent.parent == tmp_path.parent
+
+    # change name, see it serializes, but path stays the same
+    child.name = "Name2"
+    child.save_to_file()
+    assert child.path == expected_path
+    with open(child.path, "r") as file:
+        data = json.load(file)
+    assert data["name"] == "Name2"
+
+
+def test_save_to_set_location(tmp_path):
+    # Keeps the OG path if parent and path are both set
+    parent = BaseParentExample(path=tmp_path)
+    child_path = tmp_path.parent / "child.kiln"
+    child = DefaultParentedModel(path=child_path, parent=parent, name="Name")
+    assert child.build_path() == child_path
+
+    # check file created at child_path, not the default smart path
+    assert not child_path.exists()
+    child.save_to_file()
+    assert child_path.exists()
+
+    # if we don't set the path, use the parent + smartpath
+    child2 = DefaultParentedModel(parent=parent, name="Name2")
+    assert child2.build_path().parent.name == child2.id + " - Name2"
+    assert child2.build_path().parent.parent.name == "children"
+    assert child2.build_path().parent.parent.parent == tmp_path.parent
+
+
+def test_parent_without_path():
+    # no path from parent or direct path
+    parent = BaseParentExample()
+    child = DefaultParentedModel(parent=parent, name="Name")
+    with pytest.raises(ValueError):
+        child.save_to_file()
+
+
+def test_parent_wrong_type():
+    # DefaultParentedModel is parented to BaseParentExample, not KilnBaseModel
+    parent = KilnBaseModel()
+    with pytest.raises(ValueError):
+        DefaultParentedModel(parent=parent, name="Name")
+
+
+def test_load_children(test_base_parented_file):
+    # Set up parent and children models
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+
+    child1 = DefaultParentedModel(parent=parent, name="Child1")
+    child2 = DefaultParentedModel(parent=parent, name="Child2")
+    child3 = DefaultParentedModel(parent=parent, name="Child3")
+
+    # Ensure the children are saved correctly
+    child1.save_to_file()
+    child2.save_to_file()
+    child3.save_to_file()
+
+    # Load children from parent path
+    children = DefaultParentedModel.all_children_of_parent_path(test_base_parented_file)
+
+    # Verify that all children are loaded correctly
+    assert len(children) == 3
+    names = [child.name for child in children]
+    assert "Child1" in names
+    assert "Child2" in names
+    assert "Child3" in names
+    assert all(child.model_type == "default_parented_model" for child in children)
+
+
+def test_base_filename():
+    model = DefaultParentedModel(name="Test")
+    assert model.base_filename() == "default_parented_model.kiln"
+    model = NamedParentedModel(name="Test")
+    assert model.base_filename() == "named_parented_model.kiln"
+    assert NamedParentedModel.base_filename() == "named_parented_model.kiln"
+
+
+def test_load_from_folder(test_base_parented_file):
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+    child1 = DefaultParentedModel(parent=parent, name="Child1")
+    child1.save_to_file()
+
+    loaded_child1 = DefaultParentedModel.load_from_folder(child1.path.parent)
+    assert loaded_child1.name == "Child1"
+
+
+def test_lazy_load_parent(tmp_path):
+    # Create a parent
+    parent = BaseParentExample(
+        name="Parent", path=(tmp_path / BaseParentExample.base_filename())
+    )
+    parent.save_to_file()
+
+    # Create a child
+    child = DefaultParentedModel(parent=parent, name="Child")
+    child.save_to_file()
+
+    # Load the child by path
+    loaded_child = DefaultParentedModel.load_from_file(child.path)
+
+    # Access the parent to trigger lazy loading
+    loaded_parent = loaded_child.parent
+
+    # Verify that the parent is now loaded and correct
+    assert loaded_parent is not None
+    assert loaded_parent.name == "Parent"
+    assert loaded_parent.path == parent.path
+
+    # Verify that the _parent attribute is now set
+    assert hasattr(loaded_child, "_parent")
+    assert loaded_child._parent is loaded_parent