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