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\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 MultiShotPromptBuilder(BasePromptBuilder):
109    """A prompt builder that includes multiple examples in the prompt."""
110
111    @classmethod
112    def example_count(cls) -> int:
113        """Get the maximum number of examples to include in the prompt.
114
115        Returns:
116            int: The maximum number of examples (default 25).
117        """
118        return 25
119
120    def build_instruction_and_requirements(self) -> str:
121        """Build the instruction and requirements section of the prompt.
122
123        Returns:
124            str: The instruction and requirements sections.
125        """
126        base_prompt = f"# Instruction\n\n{self.task.instruction}\n\n"
127
128        if len(self.task.requirements) > 0:
129            base_prompt += "# Requirements\n\nYour response should respect the following requirements:\n"
130            for i, requirement in enumerate(self.task.requirements):
131                base_prompt += f"{i + 1}) {requirement.instruction}\n"
132            base_prompt += "\n"
133
134        return base_prompt
135
136    def build_base_prompt(self) -> str:
137        """Build a prompt with instruction, requirements, and multiple examples.
138
139        Returns:
140            str: The constructed prompt string with examples.
141        """
142        base_prompt = self.build_instruction_and_requirements()
143
144        valid_examples = self.collect_examples()
145
146        if len(valid_examples) == 0:
147            return base_prompt
148
149        base_prompt += "# Example Outputs\n\n"
150        for i, example in enumerate(valid_examples):
151            base_prompt += self.prompt_section_for_example(i, example)
152
153        return base_prompt
154
155    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
156        # Prefer repaired output if it exists, otherwise use the regular output
157        output = example.repaired_output or example.output
158        return f"## Example {index + 1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
159
160    def collect_examples(self) -> list[TaskRun]:
161        valid_examples: list[TaskRun] = []
162        runs = self.task.runs(readonly=True)
163
164        # first pass, we look for repaired outputs. These are the best examples.
165        for run in runs:
166            if len(valid_examples) >= self.__class__.example_count():
167                break
168            if run.repaired_output is not None:
169                valid_examples.append(run)
170
171        # second pass, we look for high quality outputs (rating based)
172        # Minimum is "high_quality" (4 star in star rating scale), then sort by rating
173        # exclude repaired outputs as they were used above
174        runs_with_rating = [
175            run
176            for run in runs
177            if run.output.rating is not None
178            and run.output.rating.value is not None
179            and run.output.rating.is_high_quality()
180            and run.repaired_output is None
181        ]
182        runs_with_rating.sort(
183            key=lambda x: (x.output.rating and x.output.rating.value) or 0, reverse=True
184        )
185        for run in runs_with_rating:
186            if len(valid_examples) >= self.__class__.example_count():
187                break
188            valid_examples.append(run)
189        return valid_examples
190
191
192class FewShotPromptBuilder(MultiShotPromptBuilder):
193    """A prompt builder that includes a small number of examples in the prompt."""
194
195    @classmethod
196    def example_count(cls) -> int:
197        """Get the maximum number of examples to include in the prompt.
198
199        Returns:
200            int: The maximum number of examples (4).
201        """
202        return 4
203
204
205class CustomExamplePromptBuilder(FewShotPromptBuilder):
206    """A prompt builder that uses custom examples instead of collecting from the dataset."""
207
208    def __init__(self, task: Task, examples: list[PromptExample] | None = None):
209        super().__init__(task)
210        self._custom_examples = examples or []
211
212    def collect_examples(self) -> list[TaskRun]:
213        """Override to return an empty list - we handle examples separately."""
214        return []
215
216    def build_base_prompt(self) -> str:
217        """Build a prompt with instruction, requirements, and custom examples."""
218        base_prompt = self.build_instruction_and_requirements()
219
220        if self._custom_examples:
221            base_prompt += "# Example Outputs\n\n"
222            for i, example in enumerate(self._custom_examples):
223                base_prompt += f"## Example {i + 1}\n\nInput: {example.input}\nOutput: {example.output}\n\n"
224
225        return base_prompt
226
227
228class RepairsPromptBuilder(MultiShotPromptBuilder):
229    """A prompt builder that includes multiple examples in the prompt, including repaired instructions describing what was wrong, and how it was fixed."""
230
231    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
232        if (
233            not example.repaired_output
234            or not example.repair_instructions
235            or not example.repaired_output.output
236        ):
237            return super().prompt_section_for_example(index, example)
238
239        prompt_section = f"## Example {index + 1}\n\nInput: {example.input}\n\n"
240        prompt_section += (
241            f"Initial Output Which Was Insufficient: {example.output.output}\n\n"
242        )
243        prompt_section += f"Instructions On How to Improve the Initial Output: {example.repair_instructions}\n\n"
244        prompt_section += (
245            f"Repaired Output Which is Sufficient: {example.repaired_output.output}\n\n"
246        )
247        return prompt_section
248
249
250def chain_of_thought_prompt(task: Task) -> str:
251    """Standard implementation to build and return the chain of thought prompt string.
252
253    Returns:
254        str: The constructed chain of thought prompt.
255    """
256
257    cot_instruction = task.thinking_instruction
258    if not cot_instruction:
259        cot_instruction = "Think step by step, explaining your reasoning."
260
261    return cot_instruction
262
263
264class SimpleChainOfThoughtPromptBuilder(SimplePromptBuilder):
265    """A prompt builder that includes a chain of thought prompt on top of the simple prompt."""
266
267    def chain_of_thought_prompt(self) -> str | None:
268        return chain_of_thought_prompt(self.task)
269
270
271class FewShotChainOfThoughtPromptBuilder(FewShotPromptBuilder):
272    """A prompt builder that includes a chain of thought prompt on top of the few shot prompt."""
273
274    def chain_of_thought_prompt(self) -> str | None:
275        return chain_of_thought_prompt(self.task)
276
277
278class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
279    """A prompt builder that includes a chain of thought prompt on top of the multi shot prompt."""
280
281    def chain_of_thought_prompt(self) -> str | None:
282        return chain_of_thought_prompt(self.task)
283
284
285class SavedPromptBuilder(BasePromptBuilder):
286    """A prompt builder that looks up a static prompt."""
287
288    def __init__(self, task: Task, prompt_id: str):
289        super().__init__(task)
290        prompt_model = next(
291            (
292                prompt
293                for prompt in task.prompts(readonly=True)
294                if prompt.id == prompt_id
295            ),
296            None,
297        )
298        if not prompt_model:
299            raise ValueError(f"Prompt ID not found: {prompt_id}")
300        self.prompt_model = prompt_model
301
302    def prompt_id(self) -> str | None:
303        return self.prompt_model.id
304
305    def build_base_prompt(self) -> str:
306        """Returns a saved prompt.
307
308        Returns:
309            str: The prompt string.
310        """
311        return self.prompt_model.prompt
312
313    def chain_of_thought_prompt(self) -> str | None:
314        return self.prompt_model.chain_of_thought_instructions
315
316
317class TaskRunConfigPromptBuilder(BasePromptBuilder):
318    """A prompt builder that looks up a static prompt in a task run config."""
319
320    def __init__(self, task: Task, run_config_prompt_id: str):
321        parts = run_config_prompt_id.split("::")
322        if len(parts) != 4:
323            raise ValueError(
324                f"Invalid task run config prompt ID: {run_config_prompt_id}. Expected format: 'task_run_config::[project_id]::[task_id]::[run_config_id]'."
325            )
326
327        task_id = parts[2]
328        if task_id != task.id:
329            raise ValueError(
330                f"Task run config prompt ID: {run_config_prompt_id}. Task ID mismatch. Expected: {task.id}, got: {task_id}."
331            )
332
333        run_config_id = parts[3]
334        run_config = next(
335            (
336                run_config
337                for run_config in task.run_configs(readonly=True)
338                if run_config.id == run_config_id
339            ),
340            None,
341        )
342        if not run_config:
343            raise ValueError(
344                f"Task run config ID not found: {run_config_id} for prompt id {run_config_prompt_id}"
345            )
346        if run_config.prompt is None:
347            raise ValueError(
348                f"Task run config ID {run_config_id} does not have a stored prompt. Used as prompt id {run_config_prompt_id}"
349            )
350
351        # Load the prompt from the model
352        self.prompt = run_config.prompt.prompt
353        self.cot_prompt = run_config.prompt.chain_of_thought_instructions
354        self.id = run_config_prompt_id
355
356        super().__init__(task)
357
358    def prompt_id(self) -> str | None:
359        return self.id
360
361    def build_base_prompt(self) -> str:
362        return self.prompt
363
364    def chain_of_thought_prompt(self) -> str | None:
365        return self.cot_prompt
366
367
368class FineTunePromptBuilder(BasePromptBuilder):
369    """A prompt builder that looks up a fine-tune prompt."""
370
371    def __init__(self, task: Task, nested_fine_tune_id: str):
372        super().__init__(task)
373
374        # IDs are in project_id::task_id::fine_tune_id format
375        self.full_fine_tune_id = nested_fine_tune_id
376        parts = nested_fine_tune_id.split("::")
377        if len(parts) != 3:
378            raise ValueError(
379                f"Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id', got: {nested_fine_tune_id}"
380            )
381        fine_tune_id = parts[2]
382
383        fine_tune_model = next(
384            (
385                fine_tune
386                for fine_tune in task.finetunes(readonly=True)
387                if fine_tune.id == fine_tune_id
388            ),
389            None,
390        )
391        if not fine_tune_model:
392            raise ValueError(f"Fine-tune ID not found: {fine_tune_id}")
393        self.fine_tune_model = fine_tune_model
394
395    def prompt_id(self) -> str | None:
396        return self.full_fine_tune_id
397
398    def build_base_prompt(self) -> str:
399        return self.fine_tune_model.system_message
400
401    def chain_of_thought_prompt(self) -> str | None:
402        return self.fine_tune_model.thinking_instructions
403
404
405# Our UI has some names that are not the same as the class names, which also hint parameters.
406def prompt_builder_from_id(prompt_id: PromptId, task: Task) -> BasePromptBuilder:
407    """Convert a name used in the UI to the corresponding prompt builder class.
408
409    Args:
410        prompt_id (PromptId): The prompt ID.
411
412    Returns:
413        type[BasePromptBuilder]: The corresponding prompt builder class.
414
415    Raises:
416        ValueError: If the UI name is not recognized.
417    """
418
419    # Saved prompts are prefixed with "id::"
420    if prompt_id.startswith("id::"):
421        prompt_id = prompt_id[4:]
422        return SavedPromptBuilder(task, prompt_id)
423
424    # Task run config prompts are prefixed with "task_run_config::"
425    # task_run_config::[project_id]::[task_id]::[run_config_id]
426    if prompt_id.startswith("task_run_config::"):
427        return TaskRunConfigPromptBuilder(task, prompt_id)
428
429    # Fine-tune prompts are prefixed with "fine_tune_prompt::"
430    if prompt_id.startswith("fine_tune_prompt::"):
431        prompt_id = prompt_id[18:]
432        return FineTunePromptBuilder(task, prompt_id)
433
434    # Check if the prompt_id matches any enum value
435    if prompt_id not in [member.value for member in PromptGenerators]:
436        raise ValueError(f"Unknown prompt generator: {prompt_id}")
437    typed_prompt_generator = PromptGenerators(prompt_id)
438
439    match typed_prompt_generator:
440        case PromptGenerators.SIMPLE:
441            return SimplePromptBuilder(task)
442        case PromptGenerators.FEW_SHOT:
443            return FewShotPromptBuilder(task)
444        case PromptGenerators.MULTI_SHOT:
445            return MultiShotPromptBuilder(task)
446        case PromptGenerators.REPAIRS:
447            return RepairsPromptBuilder(task)
448        case PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT:
449            return SimpleChainOfThoughtPromptBuilder(task)
450        case PromptGenerators.FEW_SHOT_CHAIN_OF_THOUGHT:
451            return FewShotChainOfThoughtPromptBuilder(task)
452        case PromptGenerators.MULTI_SHOT_CHAIN_OF_THOUGHT:
453            return MultiShotChainOfThoughtPromptBuilder(task)
454        case _:
455            # Type checking will find missing cases
456            raise_exhaustive_enum_error(typed_prompt_generator)