kiln_ai.adapters.prompt_builders

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