Merge branch 'main' into mike/update-claude-maintain

Author: tawnymanticore

app/desktop/studio_server/copilot_api.py

@@ -370,14 +370,15 @@ async def create_spec_with_copilot(
         )
 
         # 4. Create TaskRuns for eval, train, and golden datasets
-        task_runs = create_dataset_task_runs(
+        dataset_runs = create_dataset_task_runs(
             all_examples=all_examples,
             reviewed_examples=request.reviewed_examples,
             eval_tag=eval_tag,
             train_tag=train_tag,
             golden_tag=golden_tag,
             spec_name=request.name,
         )
+        task_runs = dataset_runs.task_runs
         for run in task_runs:
             run.parent = task
         models_to_save.extend(task_runs)
@@ -430,6 +431,7 @@ async def create_spec_with_copilot(
             for run in task_runs:
                 run.save_to_file()
                 saved_models.append(run)
+                dataset_runs.save_pending_feedback(run)
 
             spec.save_to_file()
             saved_models.append(spec)

app/desktop/studio_server/test_copilot_api.py

@@ -11,6 +11,7 @@
     RefineSpecApiOutput,
 )
 from app.desktop.studio_server.copilot_api import connect_copilot_api
+from app.desktop.studio_server.utils.copilot_utils import DatasetTaskRuns
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from kiln_ai.datamodel import Project, Task
@@ -417,7 +418,7 @@ def test_create_spec_with_copilot_success(
             ),
             patch(
                 "app.desktop.studio_server.copilot_api.create_dataset_task_runs",
-                return_value=[],
+                return_value=DatasetTaskRuns(),
             ),
             patch(
                 "app.desktop.studio_server.copilot_api.generate_memorable_name",

app/desktop/studio_server/utils/copilot_utils.py

@@ -25,7 +25,7 @@
 )
 from app.desktop.studio_server.utils.response_utils import unwrap_response
 from fastapi import HTTPException
-from kiln_ai.datamodel import TaskRun
+from kiln_ai.datamodel import Feedback, FeedbackSource, TaskRun
 from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
 from kiln_ai.datamodel.task_output import (
     DataSource,
@@ -172,8 +172,12 @@ def create_task_run_from_reviewed(
     tag: str,
     spec_name: str,
     extra_tags: list[str] | None = None,
-) -> TaskRun:
-    """Create a TaskRun from a reviewed example with rating (without parent set)."""
+) -> tuple[TaskRun, str | None]:
+    """Create a TaskRun from a reviewed example with rating (without parent set).
+
+    Returns a (TaskRun, feedback_text) tuple. The caller should create a Feedback
+    child on the TaskRun after saving it, if feedback_text is not None.
+    """
     data_source = DataSource(
         type=DataSourceType.synthetic,
         properties={
@@ -190,7 +194,7 @@ def create_task_run_from_reviewed(
     rating_key = f"named::{spec_name}"
     rating_value = 1.0 if example.user_says_meets_spec else 0.0
 
-    return TaskRun(
+    task_run = TaskRun(
         input=example.input,
         input_source=data_source,
         output=TaskOutput(
@@ -207,9 +211,36 @@ def create_task_run_from_reviewed(
                 },
             ),
         ),
-        user_feedback=example.feedback if example.feedback else None,
         tags=tags,
     )
+    feedback_text = example.feedback if example.feedback else None
+    return task_run, feedback_text
+
+
+class DatasetTaskRuns:
+    """Result of creating dataset task runs, with pending feedback to attach after saving."""
+
+    def __init__(self) -> None:
+        self.task_runs: list[TaskRun] = []
+        self._pending_feedback: dict[str, str] = {}
+
+    def add_run(self, task_run: TaskRun, feedback_text: str | None = None) -> None:
+        self.task_runs.append(task_run)
+        if feedback_text and task_run.id:
+            self._pending_feedback[task_run.id] = feedback_text
+
+    def save_pending_feedback(self, task_run: TaskRun) -> None:
+        """Create Feedback children for a saved TaskRun if it has pending feedback."""
+        if not task_run.id:
+            return
+        feedback_text = self._pending_feedback.get(task_run.id)
+        if feedback_text:
+            fb = Feedback(
+                feedback=feedback_text,
+                source=FeedbackSource.spec_feedback,
+                parent=task_run,
+            )
+            fb.save_to_file()
 
 
 def create_dataset_task_runs(
@@ -219,17 +250,18 @@ def create_dataset_task_runs(
     train_tag: str,
     golden_tag: str,
     spec_name: str,
-) -> list[TaskRun]:
+) -> DatasetTaskRuns:
     """Create TaskRuns for eval, train, and golden datasets.
 
     Samples from all_examples (mutating it) and creates TaskRuns for:
     - Eval dataset
     - Train dataset
     - Golden dataset (reviewed examples + unrated examples to reach MIN_GOLDEN_EXAMPLES)
 
-    Returns TaskRuns without parent set - caller must set parent.
+    Returns DatasetTaskRuns without parent set - caller must set parent and call
+    save_pending_feedback after saving each run.
     """
-    task_runs: list[TaskRun] = []
+    result = DatasetTaskRuns()
 
     # Generate a session tag for all task runs in this batch
     session_id = random.randint(0, 999999999999)
@@ -238,18 +270,17 @@ def create_dataset_task_runs(
 
     # Create TaskRuns for reviewed examples with ratings
     for reviewed in reviewed_examples:
-        task_runs.append(
-            create_task_run_from_reviewed(reviewed, golden_tag, spec_name, extra_tags)
+        task_run, feedback_text = create_task_run_from_reviewed(
+            reviewed, golden_tag, spec_name, extra_tags
         )
+        result.add_run(task_run, feedback_text)
 
     # Create more unrated golden examples from remaining pool if needed
     unrated_golden_count = max(0, MIN_GOLDEN_EXAMPLES - len(reviewed_examples))
     if unrated_golden_count > 0:
         unrated_golden_examples = sample_and_remove(all_examples, unrated_golden_count)
         for example in unrated_golden_examples:
-            task_runs.append(
-                create_task_run_from_sample(example, golden_tag, extra_tags)
-            )
+            result.add_run(create_task_run_from_sample(example, golden_tag, extra_tags))
 
     # Sample half the remaining examples for eval vs train datasets
     example_count = len(all_examples)
@@ -260,10 +291,10 @@ def create_dataset_task_runs(
 
     # Create TaskRuns for eval examples
     for example in eval_examples:
-        task_runs.append(create_task_run_from_sample(example, eval_tag, extra_tags))
+        result.add_run(create_task_run_from_sample(example, eval_tag, extra_tags))
 
     # Create TaskRuns for train examples
     for example in train_examples:
-        task_runs.append(create_task_run_from_sample(example, train_tag, extra_tags))
+        result.add_run(create_task_run_from_sample(example, train_tag, extra_tags))
 
-    return task_runs
+    return result

app/desktop/studio_server/utils/test_copilot_utils.py

@@ -145,7 +145,7 @@ def test_creates_task_run_with_correct_input(self):
             user_says_meets_spec=True,
             feedback="Good example",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
+        task_run, _ = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
         assert task_run.input == "test input"
 
     def test_creates_task_run_with_correct_output(self):
@@ -156,7 +156,7 @@ def test_creates_task_run_with_correct_output(self):
             user_says_meets_spec=True,
             feedback="",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
+        task_run, _ = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
         assert task_run.output.output == "test output"
 
     def test_creates_task_run_with_pass_rating_when_meets_spec(self):
@@ -167,7 +167,7 @@ def test_creates_task_run_with_pass_rating_when_meets_spec(self):
             user_says_meets_spec=True,
             feedback="",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
+        task_run, _ = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
         rating_key = "named::My Spec"
         assert rating_key in task_run.output.rating.requirement_ratings
         assert task_run.output.rating.requirement_ratings[rating_key].value == 1.0
@@ -180,7 +180,7 @@ def test_creates_task_run_with_fail_rating_when_not_meets_spec(self):
             user_says_meets_spec=False,
             feedback="Bad example",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
+        task_run, _ = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
         rating_key = "named::My Spec"
         assert rating_key in task_run.output.rating.requirement_ratings
         assert task_run.output.rating.requirement_ratings[rating_key].value == 0.0
@@ -193,7 +193,7 @@ def test_creates_task_run_with_tag(self):
             user_says_meets_spec=True,
             feedback="",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
+        task_run, _ = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
         assert "golden_tag" in task_run.tags
 
     def test_creates_task_run_with_extra_tags(self):
@@ -204,7 +204,7 @@ def test_creates_task_run_with_extra_tags(self):
             user_says_meets_spec=True,
             feedback="",
         )
-        task_run = create_task_run_from_reviewed(
+        task_run, _ = create_task_run_from_reviewed(
             example, "golden_tag", "My Spec", extra_tags=["session_456"]
         )
         assert "golden_tag" in task_run.tags
@@ -218,34 +218,38 @@ def test_creates_task_run_with_pass_fail_rating_type(self):
             user_says_meets_spec=True,
             feedback="",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
+        task_run, _ = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
         rating_key = "named::My Spec"
         assert (
             task_run.output.rating.requirement_ratings[rating_key].type
             == TaskOutputRatingType.pass_fail
         )
 
-    def test_creates_task_run_with_user_feedback(self):
+    def test_returns_feedback_text_when_present(self):
         example = ReviewedExample(
             input="test input",
             output="test output",
             model_says_meets_spec=True,
             user_says_meets_spec=False,
             feedback="This fails because the output is too vague",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
-        assert task_run.user_feedback == "This fails because the output is too vague"
+        _, feedback_text = create_task_run_from_reviewed(
+            example, "golden_tag", "My Spec"
+        )
+        assert feedback_text == "This fails because the output is too vague"
 
-    def test_creates_task_run_with_no_user_feedback_when_empty(self):
+    def test_returns_none_feedback_when_empty(self):
         example = ReviewedExample(
             input="test input",
             output="test output",
             model_says_meets_spec=True,
             user_says_meets_spec=True,
             feedback="",
         )
-        task_run = create_task_run_from_reviewed(example, "golden_tag", "My Spec")
-        assert task_run.user_feedback is None
+        _, feedback_text = create_task_run_from_reviewed(
+            example, "golden_tag", "My Spec"
+        )
+        assert feedback_text is None
 
 
 class TestCreateDatasetTaskRuns:
@@ -263,7 +267,7 @@ def test_creates_correct_number_of_task_runs(self):
             "train_tag",
             "golden_tag",
             "Test Spec",
-        )
+        ).task_runs
 
         # Should have NUM_SAMPLES_PER_TOPIC * NUM_TOPICS
         expected_count = NUM_SAMPLES_PER_TOPIC * NUM_TOPICS
@@ -291,7 +295,7 @@ def test_includes_reviewed_examples_in_golden_set(self):
             "train_tag",
             "golden_tag",
             "Test Spec",
-        )
+        ).task_runs
 
         # Find the reviewed example in task runs
         reviewed_run = next(
@@ -314,7 +318,7 @@ def test_all_task_runs_have_session_tag(self):
             "train_tag",
             "golden_tag",
             "Test Spec",
-        )
+        ).task_runs
 
         # All task runs should have a session tag
         for task_run in task_runs:
@@ -337,7 +341,7 @@ def test_same_session_tag_for_all_runs(self):
             "train_tag",
             "golden_tag",
             "Test Spec",
-        )
+        ).task_runs
 
         # All task runs should have the same session tag
         session_tags = set()
@@ -362,7 +366,7 @@ def test_eval_examples_have_eval_tag(self):
             "train_tag",
             "golden_tag",
             "Test Spec",
-        )
+        ).task_runs
 
         eval_runs = [tr for tr in task_runs if "eval_tag" in tr.tags]
         num_runs = NUM_SAMPLES_PER_TOPIC * NUM_TOPICS
@@ -383,7 +387,7 @@ def test_train_examples_have_train_tag(self):
             "train_tag",
             "golden_tag",
             "Test Spec",
-        )
+        ).task_runs
 
         train_runs = [tr for tr in task_runs if "train_tag" in tr.tags]
         num_runs = NUM_SAMPLES_PER_TOPIC * NUM_TOPICS
@@ -405,7 +409,7 @@ def test_handles_insufficient_examples(self):
             "train_tag",
             "golden_tag",
             "Test Spec",
-        )
+        ).task_runs
 
         # Should use all available examples
         assert len(task_runs) == 5