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 numbereed 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    def __init__(self, task: Task, examples: list[PromptExample] | None = None):
253        super().__init__(task)
254        self._custom_examples = examples or []
255
256    def collect_examples(self) -> list[TaskRun]:
257        """Override to return an empty list - we handle examples separately."""
258        return []
259
260    def build_base_prompt(self) -> str:
261        """Build a prompt with instruction, requirements, and custom examples."""
262        base_prompt = self.build_instruction_and_requirements()
263
264        if self._custom_examples:
265            base_prompt += "# Example Outputs\n\n"
266            for i, example in enumerate(self._custom_examples):
267                base_prompt += f"## Example {i + 1}\n\nInput: {example.input}\nOutput: {example.output}\n\n"
268
269        return base_prompt
270
271
272class RepairsPromptBuilder(MultiShotPromptBuilder):
273    """A prompt builder that includes multiple examples in the prompt, including repaired instructions describing what was wrong, and how it was fixed."""
274
275    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
276        if (
277            not example.repaired_output
278            or not example.repair_instructions
279            or not example.repaired_output.output
280        ):
281            return super().prompt_section_for_example(index, example)
282
283        prompt_section = f"## Example {index + 1}\n\nInput: {example.input}\n\n"
284        prompt_section += (
285            f"Initial Output Which Was Insufficient: {example.output.output}\n\n"
286        )
287        prompt_section += f"Instructions On How to Improve the Initial Output: {example.repair_instructions}\n\n"
288        prompt_section += (
289            f"Repaired Output Which is Sufficient: {example.repaired_output.output}\n\n"
290        )
291        return prompt_section
292
293
294def chain_of_thought_prompt(task: Task) -> str:
295    """Standard implementation to build and return the chain of thought prompt string.
296
297    Returns:
298        str: The constructed chain of thought prompt.
299    """
300
301    cot_instruction = task.thinking_instruction
302    if not cot_instruction:
303        cot_instruction = "Think step by step, explaining your reasoning."
304
305    return cot_instruction
306
307
308class SimpleChainOfThoughtPromptBuilder(SimplePromptBuilder):
309    """A prompt builder that includes a chain of thought prompt on top of the simple prompt."""
310
311    def chain_of_thought_prompt(self) -> str | None:
312        return chain_of_thought_prompt(self.task)
313
314
315class FewShotChainOfThoughtPromptBuilder(FewShotPromptBuilder):
316    """A prompt builder that includes a chain of thought prompt on top of the few shot prompt."""
317
318    def chain_of_thought_prompt(self) -> str | None:
319        return chain_of_thought_prompt(self.task)
320
321
322class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
323    """A prompt builder that includes a chain of thought prompt on top of the multi shot prompt."""
324
325    def chain_of_thought_prompt(self) -> str | None:
326        return chain_of_thought_prompt(self.task)
327
328
329class SavedPromptBuilder(BasePromptBuilder):
330    """A prompt builder that looks up a static prompt."""
331
332    def __init__(self, task: Task, prompt_id: str):
333        super().__init__(task)
334        prompt_model = next(
335            (
336                prompt
337                for prompt in task.prompts(readonly=True)
338                if prompt.id == prompt_id
339            ),
340            None,
341        )
342        if not prompt_model:
343            raise ValueError(f"Prompt ID not found: {prompt_id}")
344        self.prompt_model = prompt_model
345
346    def prompt_id(self) -> str | None:
347        return self.prompt_model.id
348
349    def build_base_prompt(self) -> str:
350        """Returns a saved prompt.
351
352        Returns:
353            str: The prompt string.
354        """
355        return self.prompt_model.prompt
356
357    def chain_of_thought_prompt(self) -> str | None:
358        return self.prompt_model.chain_of_thought_instructions
359
360
361class TaskRunConfigPromptBuilder(BasePromptBuilder):
362    """A prompt builder that looks up a static prompt in a task run config."""
363
364    def __init__(self, task: Task, run_config_prompt_id: str):
365        parts = run_config_prompt_id.split("::")
366        if len(parts) != 4:
367            raise ValueError(
368                f"Invalid task run config prompt ID: {run_config_prompt_id}. Expected format: 'task_run_config::[project_id]::[task_id]::[run_config_id]'."
369            )
370
371        task_id = parts[2]
372        if task_id != task.id:
373            raise ValueError(
374                f"Task run config prompt ID: {run_config_prompt_id}. Task ID mismatch. Expected: {task.id}, got: {task_id}."
375            )
376
377        run_config_id = parts[3]
378        run_config = next(
379            (
380                run_config
381                for run_config in task.run_configs(readonly=True)
382                if run_config.id == run_config_id
383            ),
384            None,
385        )
386        if not run_config:
387            raise ValueError(
388                f"Task run config ID not found: {run_config_id} for prompt id {run_config_prompt_id}"
389            )
390        if run_config.prompt is None:
391            raise ValueError(
392                f"Task run config ID {run_config_id} does not have a stored prompt. Used as prompt id {run_config_prompt_id}"
393            )
394
395        # Load the prompt from the model
396        self.prompt = run_config.prompt.prompt
397        self.cot_prompt = run_config.prompt.chain_of_thought_instructions
398        self.id = run_config_prompt_id
399
400        super().__init__(task)
401
402    def prompt_id(self) -> str | None:
403        return self.id
404
405    def build_base_prompt(self) -> str:
406        return self.prompt
407
408    def chain_of_thought_prompt(self) -> str | None:
409        return self.cot_prompt
410
411
412class FineTunePromptBuilder(BasePromptBuilder):
413    """A prompt builder that looks up a fine-tune prompt."""
414
415    def __init__(self, task: Task, nested_fine_tune_id: str):
416        super().__init__(task)
417
418        # IDs are in project_id::task_id::fine_tune_id format
419        self.full_fine_tune_id = nested_fine_tune_id
420        parts = nested_fine_tune_id.split("::")
421        if len(parts) != 3:
422            raise ValueError(
423                f"Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id', got: {nested_fine_tune_id}"
424            )
425        fine_tune_id = parts[2]
426
427        fine_tune_model = next(
428            (
429                fine_tune
430                for fine_tune in task.finetunes(readonly=True)
431                if fine_tune.id == fine_tune_id
432            ),
433            None,
434        )
435        if not fine_tune_model:
436            raise ValueError(f"Fine-tune ID not found: {fine_tune_id}")
437        self.fine_tune_model = fine_tune_model
438
439    def prompt_id(self) -> str | None:
440        return self.full_fine_tune_id
441
442    def build_base_prompt(self) -> str:
443        return self.fine_tune_model.system_message
444
445    def chain_of_thought_prompt(self) -> str | None:
446        return self.fine_tune_model.thinking_instructions
447
448
449# Our UI has some names that are not the same as the class names, which also hint parameters.
450def prompt_builder_from_id(prompt_id: PromptId, task: Task) -> BasePromptBuilder:
451    """Convert a name used in the UI to the corresponding prompt builder class.
452
453    Args:
454        prompt_id (PromptId): The prompt ID.
455
456    Returns:
457        type[BasePromptBuilder]: The corresponding prompt builder class.
458
459    Raises:
460        ValueError: If the UI name is not recognized.
461    """
462
463    # Saved prompts are prefixed with "id::"
464    if prompt_id.startswith("id::"):
465        prompt_id = prompt_id[4:]
466        return SavedPromptBuilder(task, prompt_id)
467
468    # Task run config prompts are prefixed with "task_run_config::"
469    # task_run_config::[project_id]::[task_id]::[run_config_id]
470    if prompt_id.startswith("task_run_config::"):
471        return TaskRunConfigPromptBuilder(task, prompt_id)
472
473    # Fine-tune prompts are prefixed with "fine_tune_prompt::"
474    if prompt_id.startswith("fine_tune_prompt::"):
475        prompt_id = prompt_id[18:]
476        return FineTunePromptBuilder(task, prompt_id)
477
478    # Check if the prompt_id matches any enum value
479    if prompt_id not in [member.value for member in PromptGenerators]:
480        raise ValueError(f"Unknown prompt generator: {prompt_id}")
481    typed_prompt_generator = PromptGenerators(prompt_id)
482
483    match typed_prompt_generator:
484        case PromptGenerators.SIMPLE:
485            return SimplePromptBuilder(task)
486        case PromptGenerators.FEW_SHOT:
487            return FewShotPromptBuilder(task)
488        case PromptGenerators.MULTI_SHOT:
489            return MultiShotPromptBuilder(task)
490        case PromptGenerators.REPAIRS:
491            return RepairsPromptBuilder(task)
492        case PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT:
493            return SimpleChainOfThoughtPromptBuilder(task)
494        case PromptGenerators.FEW_SHOT_CHAIN_OF_THOUGHT:
495            return FewShotChainOfThoughtPromptBuilder(task)
496        case PromptGenerators.MULTI_SHOT_CHAIN_OF_THOUGHT:
497            return MultiShotChainOfThoughtPromptBuilder(task)
498        case _:
499            # Type checking will find missing cases
500            raise_exhaustive_enum_error(typed_prompt_generator)