kiln_ai.adapters.repair.repair_task

 1import json
 2from typing import Type
 3
 4from pydantic import BaseModel, Field
 5
 6from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_registry
 7from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
 8
 9
10# TODO add evaluator rating
11class RepairTaskInput(BaseModel):
12    original_prompt: str
13    original_input: str
14    original_output: str
15    evaluator_feedback: str = Field(
16        min_length=1,
17        description="Feedback from an evaluator on how to repair the task run.",
18    )
19
20
21class RepairTaskRun(Task, parent_of={}):
22    def __init__(self, original_task: Task):
23        # Keep the typechecker happy
24        tmp_project = Project(name="Repair")
25        super().__init__(
26            name="Repair",
27            parent=tmp_project,
28            description="Repair a task run, given feedback from an evaluator about how the response can be improved.",
29            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), \
30the 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 \
31feedback describing what should be improved. Your job is to understand the evaluator's feedback and improve the response.",
32            requirements=[
33                TaskRequirement(
34                    name="Follow Eval Feedback",
35                    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.",
36                    priority=Priority.p0,
37                )
38            ],
39            input_json_schema=json.dumps(RepairTaskInput.model_json_schema()),
40            output_json_schema=original_task.output_json_schema,
41        )
42
43    @classmethod
44    def _original_prompt(cls, run: TaskRun, task: Task) -> str:
45        prompt_builder_class: Type[BasePromptBuilder] | None = None
46        prompt_builder_name = (
47            run.output.source.properties.get("prompt_builder_name", None)
48            if run.output.source
49            else None
50        )
51        if prompt_builder_name is not None and isinstance(prompt_builder_name, str):
52            prompt_builder_class = prompt_builder_registry.get(
53                prompt_builder_name, None
54            )
55        if prompt_builder_class is None:
56            raise ValueError(f"No prompt builder found for name: {prompt_builder_name}")
57        prompt_builder = prompt_builder_class(task=task)
58        if not isinstance(prompt_builder, BasePromptBuilder):
59            raise ValueError(
60                f"Prompt builder {prompt_builder_name} is not a valid prompt builder"
61            )
62        return prompt_builder.build_prompt()
63
64    @classmethod
65    def build_repair_task_input(
66        cls, original_task: Task, task_run: TaskRun, evaluator_feedback: str
67    ) -> RepairTaskInput:
68        original_prompt = cls._original_prompt(task_run, original_task)
69        return RepairTaskInput(
70            original_prompt=original_prompt,
71            original_input=task_run.input,
72            original_output=task_run.output.output,
73            evaluator_feedback=evaluator_feedback,
74        )