kiln_ai.adapters.prompt_builders

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