kiln_ai.adapters.prompt_builders

  1from abc import ABCMeta, abstractmethod
  2from dataclasses import dataclass
  3
  4from kiln_ai.datamodel import PromptGenerators, PromptId, Task, TaskRun
  5from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
  6
  7
  8@dataclass
  9class PromptExample:
 10    """A simple input/output example for use in prompt building."""
 11
 12    input: str
 13    output: str
 14
 15
 16class BasePromptBuilder(metaclass=ABCMeta):
 17    """Base class for building prompts from tasks.
 18
 19    Provides the core interface and basic functionality for prompt builders.
 20    """
 21
 22    def __init__(self, task: Task):
 23        """Initialize the prompt builder with a task.
 24
 25        Args:
 26            task (Task): The task containing instructions and requirements.
 27        """
 28        self.task = task
 29
 30    def prompt_id(self) -> str | None:
 31        """Returns the ID of the prompt, scoped to this builder.
 32
 33        Returns:
 34            str | None: The ID of the prompt, or None if not set.
 35        """
 36        return None
 37
 38    def build_prompt(self, include_json_instructions) -> str:
 39        """Build and return the complete prompt string.
 40
 41        Returns:
 42            str: The constructed prompt.
 43        """
 44        prompt = self.build_base_prompt()
 45
 46        if include_json_instructions and self.task.output_schema():
 47            prompt = (
 48                prompt
 49                + f"\n\n# Format Instructions\n\nReturn a JSON object conforming to the following schema:\n```\n{self.task.output_schema()}\n```"
 50            )
 51
 52        return prompt
 53
 54    @abstractmethod
 55    def build_base_prompt(self) -> str:
 56        """Build and return the complete prompt string.
 57
 58        Returns:
 59            str: The constructed prompt.
 60        """
 61        pass
 62
 63    def chain_of_thought_prompt(self) -> str | None:
 64        """Build and return the chain of thought prompt string.
 65
 66        Returns:
 67            str: The constructed chain of thought prompt.
 68        """
 69        return None
 70
 71    def build_prompt_for_ui(self) -> str:
 72        """Build a prompt for the UI. It includes additional instructions (like chain of thought), even if they are passed to the model in stages.
 73
 74        Designed for end-user consumption, not for model consumption.
 75
 76        Returns:
 77            str: The constructed prompt string.
 78        """
 79        base_prompt = self.build_prompt(include_json_instructions=False)
 80        cot_prompt = self.chain_of_thought_prompt()
 81        if cot_prompt:
 82            base_prompt += "\n# Thinking Instructions\n\n" + cot_prompt
 83        return base_prompt
 84
 85
 86class SimplePromptBuilder(BasePromptBuilder):
 87    """A basic prompt builder that combines task instruction with requirements."""
 88
 89    def build_base_prompt(self) -> str:
 90        """Build a simple prompt with instruction and requirements.
 91
 92        Returns:
 93            str: The constructed prompt string.
 94        """
 95        base_prompt = self.task.instruction
 96
 97        if len(self.task.requirements) > 0:
 98            base_prompt += (
 99                "\n\nYour response should respect the following requirements:\n"
100            )
101            # iterate requirements, formatting them in numbereed list like 1) task.instruction\n2)...
102            for i, requirement in enumerate(self.task.requirements):
103                base_prompt += f"{i + 1}) {requirement.instruction}\n"
104
105        return base_prompt
106
107
108class ShortPromptBuilder(BasePromptBuilder):
109    """A prompt builder that includes a the base prompt but excludes the requirements."""
110
111    def build_base_prompt(self) -> str:
112        """Build a short prompt with just the base prompt, no requirements.
113
114        Returns:
115            str: The constructed prompt string.
116        """
117        return self.task.instruction
118
119
120class MultiShotPromptBuilder(BasePromptBuilder):
121    """A prompt builder that includes multiple examples in the prompt."""
122
123    @classmethod
124    def example_count(cls) -> int:
125        """Get the maximum number of examples to include in the prompt.
126
127        Returns:
128            int: The maximum number of examples (default 25).
129        """
130        return 25
131
132    def build_instruction_and_requirements(self) -> str:
133        """Build the instruction and requirements section of the prompt.
134
135        Returns:
136            str: The instruction and requirements sections.
137        """
138        base_prompt = f"# Instruction\n\n{self.task.instruction}\n\n"
139
140        if len(self.task.requirements) > 0:
141            base_prompt += "# Requirements\n\nYour response should respect the following requirements:\n"
142            for i, requirement in enumerate(self.task.requirements):
143                base_prompt += f"{i + 1}) {requirement.instruction}\n"
144            base_prompt += "\n"
145
146        return base_prompt
147
148    def build_base_prompt(self) -> str:
149        """Build a prompt with instruction, requirements, and multiple examples.
150
151        Returns:
152            str: The constructed prompt string with examples.
153        """
154        base_prompt = self.build_instruction_and_requirements()
155
156        valid_examples = self.collect_examples()
157
158        if len(valid_examples) == 0:
159            return base_prompt
160
161        base_prompt += "# Example Outputs\n\n"
162        for i, example in enumerate(valid_examples):
163            base_prompt += self.prompt_section_for_example(i, example)
164
165        return base_prompt
166
167    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
168        # Prefer repaired output if it exists, otherwise use the regular output
169        output = example.repaired_output or example.output
170        return f"## Example {index + 1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
171
172    def collect_examples(self) -> list[TaskRun]:
173        valid_examples: list[TaskRun] = []
174        runs = self.task.runs(readonly=True)
175
176        # first pass, we look for repaired outputs. These are the best examples.
177        for run in runs:
178            if len(valid_examples) >= self.__class__.example_count():
179                break
180            if run.repaired_output is not None:
181                valid_examples.append(run)
182
183        # second pass, we look for high quality outputs (rating based)
184        # Minimum is "high_quality" (4 star in star rating scale), then sort by rating
185        # exclude repaired outputs as they were used above
186        runs_with_rating = [
187            run
188            for run in runs
189            if run.output.rating is not None
190            and run.output.rating.value is not None
191            and run.output.rating.is_high_quality()
192            and run.repaired_output is None
193        ]
194        runs_with_rating.sort(
195            key=lambda x: (x.output.rating and x.output.rating.value) or 0, reverse=True
196        )
197        for run in runs_with_rating:
198            if len(valid_examples) >= self.__class__.example_count():
199                break
200            valid_examples.append(run)
201        return valid_examples
202
203
204class FewShotPromptBuilder(MultiShotPromptBuilder):
205    """A prompt builder that includes a small number of examples in the prompt."""
206
207    @classmethod
208    def example_count(cls) -> int:
209        """Get the maximum number of examples to include in the prompt.
210
211        Returns:
212            int: The maximum number of examples (4).
213        """
214        return 4
215
216
217class CustomExamplePromptBuilder(FewShotPromptBuilder):
218    """A prompt builder that uses custom examples instead of collecting from the dataset."""
219
220    def __init__(self, task: Task, examples: list[PromptExample] | None = None):
221        super().__init__(task)
222        self._custom_examples = examples or []
223
224    def collect_examples(self) -> list[TaskRun]:
225        """Override to return an empty list - we handle examples separately."""
226        return []
227
228    def build_base_prompt(self) -> str:
229        """Build a prompt with instruction, requirements, and custom examples."""
230        base_prompt = self.build_instruction_and_requirements()
231
232        if self._custom_examples:
233            base_prompt += "# Example Outputs\n\n"
234            for i, example in enumerate(self._custom_examples):
235                base_prompt += f"## Example {i + 1}\n\nInput: {example.input}\nOutput: {example.output}\n\n"
236
237        return base_prompt
238
239
240class RepairsPromptBuilder(MultiShotPromptBuilder):
241    """A prompt builder that includes multiple examples in the prompt, including repaired instructions describing what was wrong, and how it was fixed."""
242
243    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
244        if (
245            not example.repaired_output
246            or not example.repair_instructions
247            or not example.repaired_output.output
248        ):
249            return super().prompt_section_for_example(index, example)
250
251        prompt_section = f"## Example {index + 1}\n\nInput: {example.input}\n\n"
252        prompt_section += (
253            f"Initial Output Which Was Insufficient: {example.output.output}\n\n"
254        )
255        prompt_section += f"Instructions On How to Improve the Initial Output: {example.repair_instructions}\n\n"
256        prompt_section += (
257            f"Repaired Output Which is Sufficient: {example.repaired_output.output}\n\n"
258        )
259        return prompt_section
260
261
262def chain_of_thought_prompt(task: Task) -> str:
263    """Standard implementation to build and return the chain of thought prompt string.
264
265    Returns:
266        str: The constructed chain of thought prompt.
267    """
268
269    cot_instruction = task.thinking_instruction
270    if not cot_instruction:
271        cot_instruction = "Think step by step, explaining your reasoning."
272
273    return cot_instruction
274
275
276class SimpleChainOfThoughtPromptBuilder(SimplePromptBuilder):
277    """A prompt builder that includes a chain of thought prompt on top of the simple prompt."""
278
279    def chain_of_thought_prompt(self) -> str | None:
280        return chain_of_thought_prompt(self.task)
281
282
283class FewShotChainOfThoughtPromptBuilder(FewShotPromptBuilder):
284    """A prompt builder that includes a chain of thought prompt on top of the few shot prompt."""
285
286    def chain_of_thought_prompt(self) -> str | None:
287        return chain_of_thought_prompt(self.task)
288
289
290class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
291    """A prompt builder that includes a chain of thought prompt on top of the multi shot prompt."""
292
293    def chain_of_thought_prompt(self) -> str | None:
294        return chain_of_thought_prompt(self.task)
295
296
297class SavedPromptBuilder(BasePromptBuilder):
298    """A prompt builder that looks up a static prompt."""
299
300    def __init__(self, task: Task, prompt_id: str):
301        super().__init__(task)
302        prompt_model = next(
303            (
304                prompt
305                for prompt in task.prompts(readonly=True)
306                if prompt.id == prompt_id
307            ),
308            None,
309        )
310        if not prompt_model:
311            raise ValueError(f"Prompt ID not found: {prompt_id}")
312        self.prompt_model = prompt_model
313
314    def prompt_id(self) -> str | None:
315        return self.prompt_model.id
316
317    def build_base_prompt(self) -> str:
318        """Returns a saved prompt.
319
320        Returns:
321            str: The prompt string.
322        """
323        return self.prompt_model.prompt
324
325    def chain_of_thought_prompt(self) -> str | None:
326        return self.prompt_model.chain_of_thought_instructions
327
328
329class TaskRunConfigPromptBuilder(BasePromptBuilder):
330    """A prompt builder that looks up a static prompt in a task run config."""
331
332    def __init__(self, task: Task, run_config_prompt_id: str):
333        parts = run_config_prompt_id.split("::")
334        if len(parts) != 4:
335            raise ValueError(
336                f"Invalid task run config prompt ID: {run_config_prompt_id}. Expected format: 'task_run_config::[project_id]::[task_id]::[run_config_id]'."
337            )
338
339        task_id = parts[2]
340        if task_id != task.id:
341            raise ValueError(
342                f"Task run config prompt ID: {run_config_prompt_id}. Task ID mismatch. Expected: {task.id}, got: {task_id}."
343            )
344
345        run_config_id = parts[3]
346        run_config = next(
347            (
348                run_config
349                for run_config in task.run_configs(readonly=True)
350                if run_config.id == run_config_id
351            ),
352            None,
353        )
354        if not run_config:
355            raise ValueError(
356                f"Task run config ID not found: {run_config_id} for prompt id {run_config_prompt_id}"
357            )
358        if run_config.prompt is None:
359            raise ValueError(
360                f"Task run config ID {run_config_id} does not have a stored prompt. Used as prompt id {run_config_prompt_id}"
361            )
362
363        # Load the prompt from the model
364        self.prompt = run_config.prompt.prompt
365        self.cot_prompt = run_config.prompt.chain_of_thought_instructions
366        self.id = run_config_prompt_id
367
368        super().__init__(task)
369
370    def prompt_id(self) -> str | None:
371        return self.id
372
373    def build_base_prompt(self) -> str:
374        return self.prompt
375
376    def chain_of_thought_prompt(self) -> str | None:
377        return self.cot_prompt
378
379
380class FineTunePromptBuilder(BasePromptBuilder):
381    """A prompt builder that looks up a fine-tune prompt."""
382
383    def __init__(self, task: Task, nested_fine_tune_id: str):
384        super().__init__(task)
385
386        # IDs are in project_id::task_id::fine_tune_id format
387        self.full_fine_tune_id = nested_fine_tune_id
388        parts = nested_fine_tune_id.split("::")
389        if len(parts) != 3:
390            raise ValueError(
391                f"Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id', got: {nested_fine_tune_id}"
392            )
393        fine_tune_id = parts[2]
394
395        fine_tune_model = next(
396            (
397                fine_tune
398                for fine_tune in task.finetunes(readonly=True)
399                if fine_tune.id == fine_tune_id
400            ),
401            None,
402        )
403        if not fine_tune_model:
404            raise ValueError(f"Fine-tune ID not found: {fine_tune_id}")
405        self.fine_tune_model = fine_tune_model
406
407    def prompt_id(self) -> str | None:
408        return self.full_fine_tune_id
409
410    def build_base_prompt(self) -> str:
411        return self.fine_tune_model.system_message
412
413    def chain_of_thought_prompt(self) -> str | None:
414        return self.fine_tune_model.thinking_instructions
415
416
417# Our UI has some names that are not the same as the class names, which also hint parameters.
418def prompt_builder_from_id(prompt_id: PromptId, task: Task) -> BasePromptBuilder:
419    """Convert a name used in the UI to the corresponding prompt builder class.
420
421    Args:
422        prompt_id (PromptId): The prompt ID.
423
424    Returns:
425        type[BasePromptBuilder]: The corresponding prompt builder class.
426
427    Raises:
428        ValueError: If the UI name is not recognized.
429    """
430
431    # Saved prompts are prefixed with "id::"
432    if prompt_id.startswith("id::"):
433        prompt_id = prompt_id[4:]
434        return SavedPromptBuilder(task, prompt_id)
435
436    # Task run config prompts are prefixed with "task_run_config::"
437    # task_run_config::[project_id]::[task_id]::[run_config_id]
438    if prompt_id.startswith("task_run_config::"):
439        return TaskRunConfigPromptBuilder(task, prompt_id)
440
441    # Fine-tune prompts are prefixed with "fine_tune_prompt::"
442    if prompt_id.startswith("fine_tune_prompt::"):
443        prompt_id = prompt_id[18:]
444        return FineTunePromptBuilder(task, prompt_id)
445
446    # Check if the prompt_id matches any enum value
447    if prompt_id not in [member.value for member in PromptGenerators]:
448        raise ValueError(f"Unknown prompt generator: {prompt_id}")
449    typed_prompt_generator = PromptGenerators(prompt_id)
450
451    match typed_prompt_generator:
452        case PromptGenerators.SIMPLE:
453            return SimplePromptBuilder(task)
454        case PromptGenerators.SHORT:
455            return ShortPromptBuilder(task)
456        case PromptGenerators.FEW_SHOT:
457            return FewShotPromptBuilder(task)
458        case PromptGenerators.MULTI_SHOT:
459            return MultiShotPromptBuilder(task)
460        case PromptGenerators.REPAIRS:
461            return RepairsPromptBuilder(task)
462        case PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT:
463            return SimpleChainOfThoughtPromptBuilder(task)
464        case PromptGenerators.FEW_SHOT_CHAIN_OF_THOUGHT:
465            return FewShotChainOfThoughtPromptBuilder(task)
466        case PromptGenerators.MULTI_SHOT_CHAIN_OF_THOUGHT:
467            return MultiShotChainOfThoughtPromptBuilder(task)
468        case _:
469            # Type checking will find missing cases
470            raise_exhaustive_enum_error(typed_prompt_generator)