kiln_ai.adapters.repair.repair_task

 1import json
 2
 3from pydantic import BaseModel, Field
 4
 5from kiln_ai.adapters.prompt_builders import (
 6    BasePromptBuilder,
 7    SimplePromptBuilder,
 8    prompt_builder_from_id,
 9)
10from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
11
12
13# We should add evaluator rating
14class RepairTaskInput(BaseModel):
15    original_prompt: str
16    original_input: str
17    original_output: str
18    evaluator_feedback: str = Field(
19        min_length=1,
20        description="Feedback from an evaluator on how to repair the task run.",
21    )
22
23
24class RepairTaskRun(Task, parent_of={}):
25    def __init__(self, original_task: Task):
26        # Keep the typechecker happy
27        tmp_project = Project(name="Repair")
28        super().__init__(
29            name="Repair",
30            parent=tmp_project,
31            description="Repair a task run, given feedback from an evaluator about how the response can be improved.",
32            instruction="You are an assistant which helps improve output from another assistant (original assistant). You'll be provided a task that the original assistant executed (prompt), \
33the input it was given, and the output it generated. An evaluator has determined that the output it generated did not satisfy the task and should be improved. The evaluator will provide \
34feedback describing what should be improved. Your job is to understand the evaluator's feedback and improve the response.",
35            requirements=[
36                TaskRequirement(
37                    name="Follow Eval Feedback",
38                    instruction="The evaluator's feedback is the most important thing to consider. If it conflicts with the original task instruction or prompt, prioritize the evaluator's feedback.",
39                    priority=Priority.p0,
40                )
41            ],
42            input_json_schema=json.dumps(RepairTaskInput.model_json_schema()),
43            output_json_schema=original_task.output_json_schema,
44        )
45
46    @classmethod
47    def _original_prompt(cls, run: TaskRun, task: Task) -> str:
48        # Get the prompt builder id. Need the second check because we used to store this in a prompt_builder_name field, so loading legacy runs will need this.
49        source_properties = (
50            run.output.source.properties
51            if run.output.source and run.output.source.properties
52            else {}
53        )
54        prompt_id = source_properties.get("prompt_id") or source_properties.get(
55            "prompt_builder_name", None
56        )
57        if prompt_id is not None and isinstance(prompt_id, str):
58            try:
59                prompt_builder = prompt_builder_from_id(prompt_id, task)
60            except ValueError:
61                # Unknown/legacy prompt_id — fall through to fallback below
62                prompt_builder = None
63            if isinstance(prompt_builder, BasePromptBuilder):
64                return prompt_builder.build_prompt(include_json_instructions=False)
65
66        # Fallback to simple prompt builder if prompt_id is missing, unknown, or source/properties are absent (e.g. legacy runs)
67        fallback_builder = SimplePromptBuilder(task)
68        return fallback_builder.build_prompt(include_json_instructions=False)
69
70    @classmethod
71    def build_repair_task_input(
72        cls, original_task: Task, task_run: TaskRun, evaluator_feedback: str
73    ) -> RepairTaskInput:
74        original_prompt = cls._original_prompt(task_run, original_task)
75        return RepairTaskInput(
76            original_prompt=original_prompt,
77            original_input=task_run.input,
78            original_output=task_run.output.output,
79            evaluator_feedback=evaluator_feedback,
80        )