kiln_ai.adapters.prompt_builders

  1import json
  2from abc import ABCMeta, abstractmethod
  3from typing import Dict
  4
  5from kiln_ai.datamodel import Task, TaskRun
  6from kiln_ai.utils.formatting import snake_case
  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    @abstractmethod
 24    def build_prompt(self) -> str:
 25        """Build and return the complete prompt string.
 26
 27        Returns:
 28            str: The constructed prompt.
 29        """
 30        pass
 31
 32    @classmethod
 33    def prompt_builder_name(cls) -> str:
 34        """Returns the name of the prompt builder, to be used for persisting into the datastore.
 35
 36        Default implementation gets the name of the prompt builder in snake case. If you change the class name, you should override this so prior saved data is compatible.
 37
 38        Returns:
 39            str: The prompt builder name in snake_case format.
 40        """
 41        return snake_case(cls.__name__)
 42
 43    def build_user_message(self, input: Dict | str) -> str:
 44        """Build a user message from the input.
 45
 46        Args:
 47            input (Union[Dict, str]): The input to format into a message.
 48
 49        Returns:
 50            str: The formatted user message.
 51        """
 52        if isinstance(input, Dict):
 53            return f"The input is:\n{json.dumps(input, indent=2)}"
 54
 55        return f"The input is:\n{input}"
 56
 57    def chain_of_thought_prompt(self) -> str | None:
 58        """Build and return the chain of thought prompt string.
 59
 60        Returns:
 61            str: The constructed chain of thought prompt.
 62        """
 63        return None
 64
 65    def build_prompt_for_ui(self) -> str:
 66        """Build a prompt for the UI. It includes additional instructions (like chain of thought), even if they are passed to the model in stages.
 67
 68        Designed for end-user consumption, not for model consumption.
 69
 70        Returns:
 71            str: The constructed prompt string.
 72        """
 73        base_prompt = self.build_prompt()
 74        cot_prompt = self.chain_of_thought_prompt()
 75        if cot_prompt:
 76            base_prompt += "\n# Thinking Instructions\n\n" + cot_prompt
 77        return base_prompt
 78
 79
 80class SimplePromptBuilder(BasePromptBuilder):
 81    """A basic prompt builder that combines task instruction with requirements."""
 82
 83    def build_prompt(self) -> str:
 84        """Build a simple prompt with instruction and requirements.
 85
 86        Returns:
 87            str: The constructed prompt string.
 88        """
 89        base_prompt = self.task.instruction
 90
 91        # TODO: this is just a quick version. Formatting and best practices TBD
 92        if len(self.task.requirements) > 0:
 93            base_prompt += (
 94                "\n\nYour response should respect the following requirements:\n"
 95            )
 96            # iterate requirements, formatting them in numbereed list like 1) task.instruction\n2)...
 97            for i, requirement in enumerate(self.task.requirements):
 98                base_prompt += f"{i+1}) {requirement.instruction}\n"
 99
100        return base_prompt
101
102
103class MultiShotPromptBuilder(BasePromptBuilder):
104    """A prompt builder that includes multiple examples in the prompt."""
105
106    @classmethod
107    def example_count(cls) -> int:
108        """Get the maximum number of examples to include in the prompt.
109
110        Returns:
111            int: The maximum number of examples (default 25).
112        """
113        return 25
114
115    def build_prompt(self) -> str:
116        """Build a prompt with instruction, requirements, and multiple examples.
117
118        Returns:
119            str: The constructed prompt string with examples.
120        """
121        base_prompt = f"# Instruction\n\n{ self.task.instruction }\n\n"
122
123        if len(self.task.requirements) > 0:
124            base_prompt += "# Requirements\n\nYour response should respect the following requirements:\n"
125            for i, requirement in enumerate(self.task.requirements):
126                base_prompt += f"{i+1}) {requirement.instruction}\n"
127            base_prompt += "\n"
128
129        valid_examples = self.collect_examples()
130
131        if len(valid_examples) == 0:
132            return base_prompt
133
134        base_prompt += "# Example Outputs\n\n"
135        for i, example in enumerate(valid_examples):
136            base_prompt += self.prompt_section_for_example(i, example)
137
138        return base_prompt
139
140    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
141        # Prefer repaired output if it exists, otherwise use the regular output
142        output = example.repaired_output or example.output
143        return f"## Example {index+1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
144
145    def collect_examples(self) -> list[TaskRun]:
146        valid_examples: list[TaskRun] = []
147        runs = self.task.runs()
148
149        # first pass, we look for repaired outputs. These are the best examples.
150        for run in runs:
151            if len(valid_examples) >= self.__class__.example_count():
152                break
153            if run.repaired_output is not None:
154                valid_examples.append(run)
155
156        # second pass, we look for high quality outputs (rating based)
157        # Minimum is "high_quality" (4 star in star rating scale), then sort by rating
158        # exclude repaired outputs as they were used above
159        runs_with_rating = [
160            run
161            for run in runs
162            if run.output.rating is not None
163            and run.output.rating.value is not None
164            and run.output.rating.is_high_quality()
165            and run.repaired_output is None
166        ]
167        runs_with_rating.sort(
168            key=lambda x: (x.output.rating and x.output.rating.value) or 0, reverse=True
169        )
170        for run in runs_with_rating:
171            if len(valid_examples) >= self.__class__.example_count():
172                break
173            valid_examples.append(run)
174        return valid_examples
175
176
177class FewShotPromptBuilder(MultiShotPromptBuilder):
178    """A prompt builder that includes a small number of examples in the prompt."""
179
180    @classmethod
181    def example_count(cls) -> int:
182        """Get the maximum number of examples to include in the prompt.
183
184        Returns:
185            int: The maximum number of examples (4).
186        """
187        return 4
188
189
190class RepairsPromptBuilder(MultiShotPromptBuilder):
191    """A prompt builder that includes multiple examples in the prompt, including repaired instructions describing what was wrong, and how it was fixed."""
192
193    def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
194        if (
195            not example.repaired_output
196            or not example.repair_instructions
197            or not example.repaired_output.output
198        ):
199            return super().prompt_section_for_example(index, example)
200
201        prompt_section = f"## Example {index+1}\n\nInput: {example.input}\n\n"
202        prompt_section += (
203            f"Initial Output Which Was Insufficient: {example.output.output}\n\n"
204        )
205        prompt_section += f"Instructions On How to Improve the Initial Output: {example.repair_instructions}\n\n"
206        prompt_section += (
207            f"Repaired Output Which is Sufficient: {example.repaired_output.output}\n\n"
208        )
209        return prompt_section
210
211
212def chain_of_thought_prompt(task: Task) -> str | None:
213    """Standard implementation to build and return the chain of thought prompt string.
214
215    Returns:
216        str: The constructed chain of thought prompt.
217    """
218
219    cot_instruction = task.thinking_instruction
220    if not cot_instruction:
221        cot_instruction = "Think step by step, explaining your reasoning."
222
223    return cot_instruction
224
225
226class SimpleChainOfThoughtPromptBuilder(SimplePromptBuilder):
227    """A prompt builder that includes a chain of thought prompt on top of the simple prompt."""
228
229    def chain_of_thought_prompt(self) -> str | None:
230        return chain_of_thought_prompt(self.task)
231
232
233class FewShotChainOfThoughtPromptBuilder(FewShotPromptBuilder):
234    """A prompt builder that includes a chain of thought prompt on top of the few shot prompt."""
235
236    def chain_of_thought_prompt(self) -> str | None:
237        return chain_of_thought_prompt(self.task)
238
239
240class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
241    """A prompt builder that includes a chain of thought prompt on top of the multi shot prompt."""
242
243    def chain_of_thought_prompt(self) -> str | None:
244        return chain_of_thought_prompt(self.task)
245
246
247prompt_builder_registry = {
248    "simple_prompt_builder": SimplePromptBuilder,
249    "multi_shot_prompt_builder": MultiShotPromptBuilder,
250    "few_shot_prompt_builder": FewShotPromptBuilder,
251    "repairs_prompt_builder": RepairsPromptBuilder,
252    "simple_chain_of_thought_prompt_builder": SimpleChainOfThoughtPromptBuilder,
253    "few_shot_chain_of_thought_prompt_builder": FewShotChainOfThoughtPromptBuilder,
254    "multi_shot_chain_of_thought_prompt_builder": MultiShotChainOfThoughtPromptBuilder,
255}
256
257
258# Our UI has some names that are not the same as the class names, which also hint parameters.
259def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
260    """Convert a name used in the UI to the corresponding prompt builder class.
261
262    Args:
263        ui_name (str): The UI name for the prompt builder type.
264
265    Returns:
266        type[BasePromptBuilder]: The corresponding prompt builder class.
267
268    Raises:
269        ValueError: If the UI name is not recognized.
270    """
271    match ui_name:
272        case "basic":
273            return SimplePromptBuilder
274        case "few_shot":
275            return FewShotPromptBuilder
276        case "many_shot":
277            return MultiShotPromptBuilder
278        case "repairs":
279            return RepairsPromptBuilder
280        case "simple_chain_of_thought":
281            return SimpleChainOfThoughtPromptBuilder
282        case "few_shot_chain_of_thought":
283            return FewShotChainOfThoughtPromptBuilder
284        case "multi_shot_chain_of_thought":
285            return MultiShotChainOfThoughtPromptBuilder
286        case _:
287            raise ValueError(f"Unknown prompt builder: {ui_name}")