Merge pull request #1176 from Kiln-AI/leonard/kil-490-refactor-unnest-taskruns

Author: leonardmq

app/web_ui/src/lib/api_schema.d.ts

@@ -7048,11 +7048,6 @@ export interface components {
          *
          *     Contains the input used, its source, the output produced, and optional
          *     repair information if the output needed correction.
-         *
-         *     Can be nested under another TaskRun; nested runs are stored as child runs
-         *     in a "runs" subfolder (same relationship name as Task's runs).
-         *
-         *     Accepts both Task and TaskRun as parents (polymorphic).
          */
         "TaskRun-Input": {
             /**
@@ -7108,18 +7103,18 @@ export interface components {
              * @description The trace of the task run in OpenAI format. This is the list of messages that were sent to/from the model.
              */
             trace?: (components["schemas"]["ChatCompletionDeveloperMessageParam"] | components["schemas"]["ChatCompletionSystemMessageParam"] | components["schemas"]["ChatCompletionUserMessageParam-Input"] | components["schemas"]["ChatCompletionAssistantMessageParamWrapper-Input"] | components["schemas"]["ChatCompletionToolMessageParamWrapper"] | components["schemas"]["ChatCompletionFunctionMessageParam"])[] | null;
+            /**
+             * Parent Task Run Id
+             * @description The ID of the parent task run. This is the ID of the task run that contains this task run.
+             */
+            parent_task_run_id?: string | null;
         };
         /**
          * TaskRun
          * @description Represents a single execution of a Task.
          *
          *     Contains the input used, its source, the output produced, and optional
          *     repair information if the output needed correction.
-         *
-         *     Can be nested under another TaskRun; nested runs are stored as child runs
-         *     in a "runs" subfolder (same relationship name as Task's runs).
-         *
-         *     Accepts both Task and TaskRun as parents (polymorphic).
          */
         "TaskRun-Output": {
             /**
@@ -7174,6 +7169,11 @@ export interface components {
              * @description The trace of the task run in OpenAI format. This is the list of messages that were sent to/from the model.
              */
             trace?: (components["schemas"]["ChatCompletionDeveloperMessageParam"] | components["schemas"]["ChatCompletionSystemMessageParam"] | components["schemas"]["ChatCompletionUserMessageParam-Output"] | components["schemas"]["ChatCompletionAssistantMessageParamWrapper-Output"] | components["schemas"]["ChatCompletionToolMessageParamWrapper"] | components["schemas"]["ChatCompletionFunctionMessageParam"])[] | null;
+            /**
+             * Parent Task Run Id
+             * @description The ID of the parent task run. This is the ID of the task run that contains this task run.
+             */
+            parent_task_run_id?: string | null;
             /** Model Type */
             readonly model_type: string;
         };

libs/core/kiln_ai/adapters/model_adapters/base_adapter.py

@@ -605,8 +605,18 @@ def generate_run(
                 properties={"created_by": Config.shared().user_id},
             )
 
+        parent_task_run_id: str | None = None
+        if parent_task_run is not None:
+            if parent_task_run.id is None:
+                raise ValueError(
+                    "parent_task_run must be persisted before using as parent: save the parent "
+                    "TaskRun (e.g. save_to_file()) so it has a stable id."
+                )
+            parent_task_run_id = parent_task_run.id
+
         return TaskRun(
-            parent=parent_task_run if parent_task_run is not None else self.task,
+            parent=self.task,
+            parent_task_run_id=parent_task_run_id,
             input=input_str,
             input_source=input_source,
             output=new_output,

libs/core/kiln_ai/adapters/model_adapters/test_saving_adapter_results.py

@@ -3,7 +3,7 @@
 import pytest
 
 from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput
-from kiln_ai.datamodel import DataSource, DataSourceType, Project, Task, TaskRun, Usage
+from kiln_ai.datamodel import DataSource, DataSourceType, Project, Task, Usage
 from kiln_ai.datamodel.datamodel_enums import InputType
 from kiln_ai.datamodel.run_config import KilnAgentRunConfigProperties
 from kiln_ai.utils.config import Config
@@ -427,15 +427,15 @@ def test_properties_for_task_output_custom_values(test_task):
     assert output.source.properties["top_p"] == 0.9
 
 
-def test_generate_run_with_parent_task_run_sets_parent(test_task, adapter):
-    """Test that generate_run with parent_task_run uses it as parent instead of the task."""
+def test_generate_run_with_parent_task_run_sets_parent_task_run_id(test_task, adapter):
     prior_run = adapter.generate_run(
         input="prior input",
         input_source=None,
         run_output=RunOutput(output="prior output", intermediate_outputs=None),
     )
     prior_run.save_to_file()
     assert prior_run.id is not None
+    assert prior_run.parent_task_run_id is None
 
     new_run = adapter.generate_run(
         input="new input",
@@ -444,41 +444,55 @@ def test_generate_run_with_parent_task_run_sets_parent(test_task, adapter):
         parent_task_run=prior_run,
     )
 
-    assert new_run.parent == prior_run
-
+    assert new_run.parent_task_run_id == prior_run.id
     new_run.save_to_file()
 
-    reloaded_prior_run = TaskRun.load_from_file(prior_run.path)
-    child_runs = reloaded_prior_run.runs()
-    assert len(child_runs) == 1
-    assert child_runs[0].output.output == "new output"
-
-    # The task should only have the prior run as a direct child
     reloaded_task = Task.load_from_file(test_task.path)
     task_runs = reloaded_task.runs()
-    assert len(task_runs) == 1
-    assert task_runs[0].id == prior_run.id
+    assert len(task_runs) == 2
+    by_id = {r.id: r for r in task_runs}
+    assert by_id[prior_run.id].parent_task_run_id is None
+    assert by_id[new_run.id].parent_task_run_id == prior_run.id
+    assert by_id[new_run.id].output.output == "new output"
 
 
 def test_generate_run_without_parent_task_run_defaults_to_task(test_task, adapter):
-    """Test that generate_run without parent_task_run defaults to using the task as parent."""
+    """Test that generate_run without parent_task_run leaves parent_task_run_id unset."""
     run = adapter.generate_run(
         input="input",
         input_source=None,
         run_output=RunOutput(output="output", intermediate_outputs=None),
     )
+    assert run.parent_task_run_id is None
     assert run.parent == test_task
 
 
+def test_generate_run_with_unsaved_parent_task_run_raises(adapter):
+    prior_run = adapter.generate_run(
+        input="prior input",
+        input_source=None,
+        run_output=RunOutput(output="prior output", intermediate_outputs=None),
+    )
+    prior_run.id = None
+
+    with pytest.raises(ValueError, match="parent_task_run must be persisted"):
+        adapter.generate_run(
+            input="new input",
+            input_source=None,
+            run_output=RunOutput(output="new output", intermediate_outputs=None),
+            parent_task_run=prior_run,
+        )
+
+
 @pytest.mark.asyncio
-async def test_invoke_with_parent_task_run_saves_as_child(test_task, adapter):
-    """Test that invoke with parent_task_run saves the new run as a child of that run."""
+async def test_invoke_with_parent_task_run_saves_under_task_with_link(
+    test_task, adapter
+):
     trace = [
         {"role": "user", "content": "Hello"},
         {"role": "assistant", "content": "Hi there!"},
     ]
 
-    # Create and save a prior run to act as parent
     prior_run = adapter.generate_run(
         input="Hello",
         input_source=None,
@@ -531,16 +545,11 @@ async def test_invoke_with_parent_task_run_saves_as_child(test_task, adapter):
         )
 
     assert new_run.id is not None
-    assert new_run.parent == prior_run
-
-    # The prior run should have the new run as a child
-    reloaded_prior_run = TaskRun.load_from_file(prior_run.path)
-    child_runs = reloaded_prior_run.runs()
-    assert len(child_runs) == 1
-    assert child_runs[0].output.output == "More details!"
+    assert new_run.parent_task_run_id == prior_run.id
 
-    # The task should only have the prior run as a direct child
     reloaded_task = Task.load_from_file(test_task.path)
     task_runs = reloaded_task.runs()
-    assert len(task_runs) == 1
-    assert task_runs[0].id == prior_run.id
+    assert len(task_runs) == 2
+    by_id = {r.id: r for r in task_runs}
+    assert by_id[new_run.id].output.output == "More details!"
+    assert by_id[new_run.id].parent_task_run_id == prior_run.id

libs/core/kiln_ai/datamodel/basemodel.py

@@ -545,36 +545,10 @@ def relationship_name(cls) -> str:
     def parent_type(cls) -> Type[KilnBaseModel]:
         raise NotImplementedError("Parent type must be implemented")
 
-    @classmethod
-    def _parent_types(cls) -> List[Type["KilnBaseModel"]] | None:
-        """Return accepted parent types. This must be implemented by the subclass if
-        the model can have multiple parent types.
-
-        Return None (default) to use the single parent_type() check.
-        Override and return a list of parent types for models that can be nested
-        under more than one parent type (e.g. TaskRun can be nested under Task or TaskRun).
-        """
-        return None
-
-    def _check_parent_type(
-        self,
-        expected_parent_types: List[Type[KilnBaseModel]] | None = None,
-    ) -> Self:
+    @model_validator(mode="after")
+    def check_parent_type(self) -> Self:
         cached_parent = self.cached_parent()
-        if cached_parent is None:
-            return self
-
-        # some models support having multiple parent types, so we allow overriding the expected parent
-        if expected_parent_types is not None:
-            if not any(
-                isinstance(cached_parent, expected_parent_type)
-                for expected_parent_type in expected_parent_types
-            ):
-                raise ValueError(
-                    f"Parent must be one of {expected_parent_types}, but was {type(cached_parent)}"
-                )
-        else:
-            # default case where we expect a single parent type to be valid
+        if cached_parent is not None:
             expected_parent_type = self.__class__.parent_type()
             if not isinstance(cached_parent, expected_parent_type):
                 raise ValueError(
@@ -583,11 +557,6 @@ def _check_parent_type(
 
         return self
 
-    @model_validator(mode="after")
-    def check_parent_type(self) -> Self:
-        """Default validation for parent type. Can be overridden by subclasses - for example if the parent is polymorphic."""
-        return self._check_parent_type()
-
     def build_child_dirname(self) -> Path:
         # Default implementation for readable folder names.
         # {id} - {name}/{type}.kiln
@@ -642,39 +611,10 @@ def iterate_children_paths_of_parent_path(cls: Type[PT], parent_path: Path | Non
         else:
             parent_folder = parent_path
 
-        if not parent_path.exists():
+        parent = cls.parent_type().load_from_file(parent_path)
+        if parent is None:
             raise ValueError("Parent must be set to load children")
 
-        # Validate the parent file's declared type so we fail fast when the caller
-        # passes a wrong path.  For polymorphic children (e.g. TaskRun) the
-        # subclass overrides _accepted_parent_types() to broaden the check to all
-        # accepted parent types
-        parent_types_override = cls._parent_types()
-        if parent_types_override is None:
-            # Default: single expected parent type — original behaviour
-            parent = cls.parent_type().load_from_file(parent_path)
-            if parent is None:
-                raise ValueError("Parent must be set to load children")
-        else:
-            # Polymorphic parent: read only the model_type field to avoid a full load.
-            with open(parent_path, "r", encoding="utf-8") as fh:
-                actual_parent_type_name = json.loads(fh.read()).get("model_type", "")
-            parent_type_names = {t.type_name() for t in parent_types_override}
-            if actual_parent_type_name not in parent_type_names:
-                raise ValueError(
-                    f"Parent model_type '{actual_parent_type_name}' is not one of "
-                    f"{parent_type_names}"
-                )
-
-            parent_type = next(
-                t
-                for t in parent_types_override
-                if t.type_name() == actual_parent_type_name
-            )
-            parent = 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
 

libs/core/kiln_ai/datamodel/task.py

@@ -203,21 +203,3 @@ def parent_project(self) -> Union["Project", None]:
         if self.parent is None or self.parent.__class__.__name__ != "Project":
             return None
         return self.parent  # type: ignore
-
-    def find_task_run_by_id_dfs(
-        self, task_run_id: str, readonly: bool = False
-    ) -> TaskRun | None:
-        """
-        Find a task run by id in the entire task run tree. This is an expensive DFS
-        traversal of the file system so do not use too willy nilly.
-
-        If you already know the root task run, you can use the same method on
-        the root TaskRun instead - that will save a bunch of subtree traversals.
-        """
-        stack: List[TaskRun] = list(self.runs(readonly=readonly))
-        while stack:
-            run = stack.pop()
-            if run.id == task_run_id:
-                return run
-            stack.extend(run.runs(readonly=readonly))
-        return None