Skip to content

Prompt Optimization with Textual Gradients (ProTeGi)

About

Prompt Optimization with Textual Gradients (ProTeGi) starts from an initial prompt and iteratively updates using feedback from the language model, called "gradients". The initial prompt is scored using the validation set and the failures are collected. The failures are described to the language model, and the language model is asked to provide multiple feedbacks (gradients) to describe how the prompt should be improved. The language model then uses each of these gradients to generate a number of new prompts, and these new prompts are scored. Then the language model generates variations The top scoring prompt is retained and becomes the initial prompt for the next iteration. The iterations continue until the maximum iteration depth is reached or the score threshold is exceeded.

A variation of this using "greedy" search instead keeps all prompts at each iteration, leading to larger trees of prompts.

Usage

The ProtegiOptimizer requires a description of the failures after each step. You must provide this feedback in your evaluator by capturing errors and saving them in the prompt object's errors attribute.

[!IMPORTANT] Important Note Your evaluator function **MUST** save any errors to the prompt object's errors attribute. Otherwise the optimization will fail.

from lagnchain_openai import ChatOpenAI
from prompt_optimizer import PredictionError, Prompt
from prompt_optimizer.optimizers import ProtegiOptimizer

# Simple QA validation set
validation_set = [
    {"question": "What is the capital of France?", "answer": "Paris"},
    {"question": "What is the largest planet in our solar system?", "answer": "Jupiter"},
    {"question": "What is the smallest planet in our solar system?", "answer": "Mercury"},
    {"question": "What is the longest river in the world?", "answer": "Nile"},
    {"question": "What is the smallest river in the world?", "answer": "Reprua River"},
]

# A langchain ChatModel for generating new prompts
client = ChatOpenAI(model="gpt-5", temperature=0.7)

# Evaluator function
def evaluator(prompt: Prompt, validation_set: list[dict]) -> list[str]:
    """Prompt evaluator function."""
    # Run the prompt through the AI system
    predictions = []
    num_correct = 0
    agent = get_agent()
    for row in validation_set:
        question = row["input"]
        messages = [{"role": "system", "content": prompt.content}, {"role": "user", "content": question}]
        response = agent.invoke(messages)
        prediction = response.content.strip()
        predictions.append(prediction)

        # Reward exact matches and collect errors
        actual = row["target"]
        if actual == prediction:
            num_correct += 1
        else:
            num_correct += 0
            # Save prediction error - Required for ProtegiOptimizer
            error = PredictionError(input=question, prediction=prediction, actual=actual, feedback=None)
            prompt.errors.append(error)

    # Compute the score
    score = num_correct / len(validation_set)

    return score

# Initialize the optimizer
baseline_prompt = "Answer the user's questions to the best of your ability."
optimizer = ProtegiOptimizer(
    client=client,
    seed_prompts=[baseline_prompt],
    validation_set=validation_set,
    max_depth=3,
    evaluator=evaluator,
)

# Run the optimization
optimized_prompt = optimizer.run()

Citation

@misc{pryzant2023automaticpromptoptimizationgradient,
    title={Automatic Prompt Optimization with "Gradient Descent" and Beam Search}, 
    author={Reid Pryzant and Dan Iter and Jerry Li and Yin Tat Lee and Chenguang Zhu and Michael Zeng},
    year={2023},
    eprint={2305.03495},
    archivePrefix={arXiv},
    primaryClass={cs.CL},
    url={https://arxiv.org/abs/2305.03495}, 
}

Source

ProtegiOptimizer

Bases: BaseOptimizer

ProTeGi Optimizer.

Based on ProTeGi with Successive Rejects.

@misc{pryzant2023automaticpromptoptimizationgradient,
    title={Automatic Prompt Optimization with "Gradient Descent" and Beam Search},
    author={Reid Pryzant and Dan Iter and Jerry Li and Yin Tat Lee and Chenguang Zhu and Michael Zeng},
    year={2023},
    eprint={2305.03495},
    archivePrefix={arXiv},
    primaryClass={cs.CL},
    url={https://arxiv.org/abs/2305.03495},
}
Source code in src/prompt_optimizer/optimizers/protegi.py 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
class ProtegiOptimizer(BaseOptimizer):
    """
    ProTeGi Optimizer.

    Based on ProTeGi with Successive Rejects.

    ```
    @misc{pryzant2023automaticpromptoptimizationgradient,
        title={Automatic Prompt Optimization with "Gradient Descent" and Beam Search},
        author={Reid Pryzant and Dan Iter and Jerry Li and Yin Tat Lee and Chenguang Zhu and Michael Zeng},
        year={2023},
        eprint={2305.03495},
        archivePrefix={arXiv},
        primaryClass={cs.CL},
        url={https://arxiv.org/abs/2305.03495},
    }
    ```
    """

    def __init__(
        self,
        *,
        client: ClientType,
        seed_prompts: list[Prompt],
        validation_set: ValidationSetType,
        max_depth: int,
        evaluator: Callable[[Prompt, ValidationSetType], ScoreType],
        output_path: Optional[Union[str, Path]] = None,
        num_feedbacks: int = 3,
        steps_per_gradient: int = 3,
        num_resample: int = 3,
        search_mode: Literal["greedy", "beam"] = "beam",
        score_threshold: Optional[Union[float, int]] = None,
        **kwargs,
    ):
        """
        Initialize the ProTeGi optimizer.

        Args:
            client (ClientType):
                Language model client to use for prompt generation and feedback.
            seed_prompts (list[Prompt]):
                List of prompts to seed generation.
            validation_set (ValidationSetType):
                Set of examples to evaluate the prompt on.
            max_depth (int):
                Maximum iteration depth for prompt generation.
            evaluator (Callable[[Prompt, ValidationSetType], ScoreType]):
                Function that takes a prompt and the validation data and returns a score.
            output_path (Union[str, Path], optional):
                Path to store run results. Should be a .jsonl file path.
                If None, no outputs will be written to disk. Defaults to None.
            num_feedbacks (int, optional):
                Number of feedbacks to generate per prompt. Defaults to 3.
            steps_per_gradient (int, optional):
                Number of new prompts to generate per feedback. Defaults to 3.
            num_resample (int, optional):
                Number of Monte Carlo rewrites per new prompt generated from feedback. The paper recommends
                setting this equal to steps_per_gradient. Defaults to 3.
            search_mode (Literal["greedy", "beam"], optional):
                Mode for filtering prompt candidates after each step. "greedy" keeps all prompts from the previous step.
                "beam" keeps only the highest scoring prompt from the previous step. Defaults to "beam".
            score_threshold (float, optional):
                Threshold for early convergence. If a prompt exceeds this score after any iteration, the optimization loop
                immediately ends. If set to None, the optimization loop will not terminate early. Defaults to None.
            kwargs:
                Additional keyword arguments.

        """
        super().__init__(
            client=client,
            seed_prompts=seed_prompts,
            validation_set=validation_set,
            max_depth=max_depth,
            evaluator=evaluator,
            output_path=output_path,
        )
        self.num_feedbacks = num_feedbacks
        self.steps_per_gradient = steps_per_gradient
        self.num_resample = num_resample
        self.search_mode = search_mode
        self.score_threshold = score_threshold

    def _extract_responses(self, content: str) -> list[str]:
        """
        Extract the responses between <START> and <END>.

        Args:
            content (str): Output string from an LLM generation request.

        Returns:
            list[str]: List of all responses within <START> and </?END> or <START> and </?START>.

        """
        pattern = r"<START>(.*?)(?:<\/?END>|<\/?START>)"
        matches = re.findall(pattern, content, flags=re.DOTALL)
        return matches

    def _generate(self, metaprompt_template: str, template_kwargs: dict) -> str:
        """
        Generate a completion for a given template and kwargs and parse the results.

        Args:
            metaprompt_template (str): Template for the metaprompt.
            template_kwargs (dict): Key word arguments to fill the template values.
            kwargs: Additional kwargs to pass to the OpenAI client.completions.create (e.g. temperature)

        Returns:
            list[str]: The parsed generation results.

        """
        metaprompt = metaprompt_template.format(**template_kwargs)
        input = [{"role": "user", "content": metaprompt}]
        raw_response = self.client.invoke(input=input)
        response = raw_response.content.strip()
        return response

    def generate_prompt_candidates(self, *, prompts: list[Prompt], **kwargs) -> list[Prompt]:
        """Generate prompt candidates using gradients."""
        prompt_candidates = []
        for prompt in track(prompts, description="Generating prompt candidates", transient=True):
            # Build error string
            error_string = "\n\n".join(
                [
                    f"Input: {error.input}\nPrediction: {error.prediction}\nFeedback: {error.feedback}" 
                    if error.feedback is not None 
                    else f"Input: {error.input}\nPrediction: {error.prediction}\nActual: {error.actual}" 
                    for error in prompt.errors
                ]
            )

            # Generate gradients
            template_kwargs = {
                "prompt": prompt.content,
                "error_string": error_string,
                "num_feedbacks": self.num_feedbacks,
                "steps_per_gradient": self.steps_per_gradient,
            }
            raw_gradients = self._generate(metaprompt_template=GENERATE_GRADIENT_PROMPT_TEMPLATE, template_kwargs=template_kwargs)
            gradients = self._extract_responses(raw_gradients)
            gradients = gradients[: self.num_feedbacks]
            # Generate prompts for each gradient
            for gradient in gradients:
                template_kwargs.update({"gradient": gradient})
                raw_new_prompts = self._generate(metaprompt_template=INCORPORTATE_GRADIENT_PROMPT_TEMPLATE, template_kwargs=template_kwargs)
                new_prompts = self._extract_responses(raw_new_prompts)
                new_prompts = new_prompts[: self.steps_per_gradient]
                metadata = {"_origin_prompt": prompt.content, "_gradient": gradient, "_resampled": False}
                new_prompt_candidates = [Prompt(content=new_prompt, metadata=metadata) for new_prompt in new_prompts]

                # Resample new prompts
                for new_prompt in new_prompts:
                    varied_prompts = [
                        self._generate(metaprompt_template=RESAMPLING_PROMPT_TEMPLATE, template_kwargs={"prompt": new_prompt})
                        for _ in range(self.num_resample)
                    ]
                    metadata = {"_origin_prompt": new_prompt, "_gradient": None, "_resampled": True}
                    varied_prompts = [Prompt(content=new_prompt) for new_prompt in varied_prompts]
                    new_prompt_candidates.extend(varied_prompts)

                # Save prompts to prompt candidates
                prompt_candidates.extend(new_prompt_candidates)

        # Add back the initial prompts to the pool
        prompt_candidates = prompts + prompt_candidates

        return prompt_candidates

    def _get_best_prompt(self, prompts: list[Prompt]):
        """Get the highest scoring prompt."""
        if any(prompt.score is None for prompt in prompts):
            raise ValueError("All prompts must be scored before calling this function.")
        return max(prompts, key=lambda x: x.score)

    def select_prompt_candidates(self, *, prompts: list[Prompt], validation_set: ValidationSetType) -> list[Prompt]:
        """Select prompt candidates according to the search mode."""
        self._score_prompts(prompts=prompts, validation_set=validation_set)
        if self.search_mode == "greedy":
            return prompts
        elif self.search_mode == "beam":
            return [self._get_best_prompt(prompts=prompts)]

    def check_early_convergence(self, *, all_prompts: list[list[Prompt]]):
        """Check if the early convergence criteria is met."""
        if self.score_threshold is None:
            return False

        # Flatten all iterations
        prompts = sum(all_prompts, start=[])

        # Check if early convergence criteria is met
        highest_score = max(prompts, key=lambda x: x.score).score
        if highest_score >= self.score_threshold:
            return True
        return False

    def select_best_prompt(self, *, all_prompts: list[list[Prompt]]) -> Prompt:
        """Select the top scoring prompt."""
        # Flatten all iterations
        prompts = sum(all_prompts, start=[])

        # Select the single prompt with the highest score
        best_prompt = self._get_best_prompt(prompts=prompts)
        logger.info(f"Best score: {best_prompt.score:.3f}")
        return best_prompt

__init__(*, client, seed_prompts, validation_set, max_depth, evaluator, output_path=None, num_feedbacks=3, steps_per_gradient=3, num_resample=3, search_mode='beam', score_threshold=None, **kwargs)

Initialize the ProTeGi optimizer.

Parameters:

Name Type Description Default
client ClientType

Language model client to use for prompt generation and feedback.

 required
seed_prompts list[Prompt]

List of prompts to seed generation.

 required
validation_set ValidationSetType

Set of examples to evaluate the prompt on.

 required
max_depth int

Maximum iteration depth for prompt generation.

 required
evaluator Callable[[Prompt, ValidationSetType], ScoreType]

Function that takes a prompt and the validation data and returns a score.

 required
output_path Union[str, Path]

Path to store run results. Should be a .jsonl file path. If None, no outputs will be written to disk. Defaults to None.

 None
num_feedbacks int

Number of feedbacks to generate per prompt. Defaults to 3.

 3
steps_per_gradient int

Number of new prompts to generate per feedback. Defaults to 3.

 3
num_resample int

Number of Monte Carlo rewrites per new prompt generated from feedback. The paper recommends setting this equal to steps_per_gradient. Defaults to 3.

 3
search_mode Literal['greedy', 'beam']

Mode for filtering prompt candidates after each step. "greedy" keeps all prompts from the previous step. "beam" keeps only the highest scoring prompt from the previous step. Defaults to "beam".

 'beam'
score_threshold float

Threshold for early convergence. If a prompt exceeds this score after any iteration, the optimization loop immediately ends. If set to None, the optimization loop will not terminate early. Defaults to None.

 None
kwargs

Additional keyword arguments.

 {}
Source code in src/prompt_optimizer/optimizers/protegi.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def __init__(
    self,
    *,
    client: ClientType,
    seed_prompts: list[Prompt],
    validation_set: ValidationSetType,
    max_depth: int,
    evaluator: Callable[[Prompt, ValidationSetType], ScoreType],
    output_path: Optional[Union[str, Path]] = None,
    num_feedbacks: int = 3,
    steps_per_gradient: int = 3,
    num_resample: int = 3,
    search_mode: Literal["greedy", "beam"] = "beam",
    score_threshold: Optional[Union[float, int]] = None,
    **kwargs,
):
    """
    Initialize the ProTeGi optimizer.

    Args:
        client (ClientType):
            Language model client to use for prompt generation and feedback.
        seed_prompts (list[Prompt]):
            List of prompts to seed generation.
        validation_set (ValidationSetType):
            Set of examples to evaluate the prompt on.
        max_depth (int):
            Maximum iteration depth for prompt generation.
        evaluator (Callable[[Prompt, ValidationSetType], ScoreType]):
            Function that takes a prompt and the validation data and returns a score.
        output_path (Union[str, Path], optional):
            Path to store run results. Should be a .jsonl file path.
            If None, no outputs will be written to disk. Defaults to None.
        num_feedbacks (int, optional):
            Number of feedbacks to generate per prompt. Defaults to 3.
        steps_per_gradient (int, optional):
            Number of new prompts to generate per feedback. Defaults to 3.
        num_resample (int, optional):
            Number of Monte Carlo rewrites per new prompt generated from feedback. The paper recommends
            setting this equal to steps_per_gradient. Defaults to 3.
        search_mode (Literal["greedy", "beam"], optional):
            Mode for filtering prompt candidates after each step. "greedy" keeps all prompts from the previous step.
            "beam" keeps only the highest scoring prompt from the previous step. Defaults to "beam".
        score_threshold (float, optional):
            Threshold for early convergence. If a prompt exceeds this score after any iteration, the optimization loop
            immediately ends. If set to None, the optimization loop will not terminate early. Defaults to None.
        kwargs:
            Additional keyword arguments.

    """
    super().__init__(
        client=client,
        seed_prompts=seed_prompts,
        validation_set=validation_set,
        max_depth=max_depth,
        evaluator=evaluator,
        output_path=output_path,
    )
    self.num_feedbacks = num_feedbacks
    self.steps_per_gradient = steps_per_gradient
    self.num_resample = num_resample
    self.search_mode = search_mode
    self.score_threshold = score_threshold

check_early_convergence(*, all_prompts)

Check if the early convergence criteria is met.

Source code in src/prompt_optimizer/optimizers/protegi.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
def check_early_convergence(self, *, all_prompts: list[list[Prompt]]):
    """Check if the early convergence criteria is met."""
    if self.score_threshold is None:
        return False

    # Flatten all iterations
    prompts = sum(all_prompts, start=[])

    # Check if early convergence criteria is met
    highest_score = max(prompts, key=lambda x: x.score).score
    if highest_score >= self.score_threshold:
        return True
    return False

generate_prompt_candidates(*, prompts, **kwargs)

Generate prompt candidates using gradients.

Source code in src/prompt_optimizer/optimizers/protegi.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def generate_prompt_candidates(self, *, prompts: list[Prompt], **kwargs) -> list[Prompt]:
    """Generate prompt candidates using gradients."""
    prompt_candidates = []
    for prompt in track(prompts, description="Generating prompt candidates", transient=True):
        # Build error string
        error_string = "\n\n".join(
            [
                f"Input: {error.input}\nPrediction: {error.prediction}\nFeedback: {error.feedback}" 
                if error.feedback is not None 
                else f"Input: {error.input}\nPrediction: {error.prediction}\nActual: {error.actual}" 
                for error in prompt.errors
            ]
        )

        # Generate gradients
        template_kwargs = {
            "prompt": prompt.content,
            "error_string": error_string,
            "num_feedbacks": self.num_feedbacks,
            "steps_per_gradient": self.steps_per_gradient,
        }
        raw_gradients = self._generate(metaprompt_template=GENERATE_GRADIENT_PROMPT_TEMPLATE, template_kwargs=template_kwargs)
        gradients = self._extract_responses(raw_gradients)
        gradients = gradients[: self.num_feedbacks]
        # Generate prompts for each gradient
        for gradient in gradients:
            template_kwargs.update({"gradient": gradient})
            raw_new_prompts = self._generate(metaprompt_template=INCORPORTATE_GRADIENT_PROMPT_TEMPLATE, template_kwargs=template_kwargs)
            new_prompts = self._extract_responses(raw_new_prompts)
            new_prompts = new_prompts[: self.steps_per_gradient]
            metadata = {"_origin_prompt": prompt.content, "_gradient": gradient, "_resampled": False}
            new_prompt_candidates = [Prompt(content=new_prompt, metadata=metadata) for new_prompt in new_prompts]

            # Resample new prompts
            for new_prompt in new_prompts:
                varied_prompts = [
                    self._generate(metaprompt_template=RESAMPLING_PROMPT_TEMPLATE, template_kwargs={"prompt": new_prompt})
                    for _ in range(self.num_resample)
                ]
                metadata = {"_origin_prompt": new_prompt, "_gradient": None, "_resampled": True}
                varied_prompts = [Prompt(content=new_prompt) for new_prompt in varied_prompts]
                new_prompt_candidates.extend(varied_prompts)

            # Save prompts to prompt candidates
            prompt_candidates.extend(new_prompt_candidates)

    # Add back the initial prompts to the pool
    prompt_candidates = prompts + prompt_candidates

    return prompt_candidates

get_all_prompts(include_candidates=False)

Get all the prompts from the latest training run.

The default behavior returns a list of lists, where each internal list contains the retained candidates after one iteration step. Setting include_candidates to True will also include all generated candidate prompts.

Parameters:

Name Type Description Default
include_candidates bool

Whether to include all the candidate prompts in the output. If True, candidate prompts from each iteration will be included. Defaults to False.

 False

Returns:

Type Description
list[list[Prompt]]

list[list[Prompt]]: List of lists where each list contains the prompts from each iteration. E.g. list[0] contains prompts from the first iteration, list[1] the second, etc. If include_candidates is False, each inner list contains only the retained prompts at each iteration. If include_candidates is True, each inner list contains all candidate prompts at each iteration, including those that were discarded.
Source code in src/prompt_optimizer/optimizers/base.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def get_all_prompts(self, include_candidates: bool = False) -> list[list[Prompt]]:
    """
    Get all the prompts from the latest training run.

    The default behavior returns a list of lists, where each internal list contains the
    retained candidates after one iteration step.
    Setting include_candidates to True will also include all generated candidate prompts.

    Args:
        include_candidates (bool, optional):
            Whether to include all the candidate prompts in the output.
            If True, candidate prompts from each iteration will be included.
            Defaults to False.

    Returns:
        list[list[Prompt]]:
            List of lists where each list contains the prompts from each iteration.
            E.g. list[0] contains prompts from the first iteration, list[1] the second, etc.
            If include_candidates is False, each inner list contains only the retained prompts at each iteration.
            If include_candidates is True, each inner list contains all candidate prompts at each iteration,
            including those that were discarded.

    """
    # Decide whether to include candidates
    if include_candidates:
        all_prompts = self._g
    else:
        all_prompts = self._p

    return all_prompts

run()

Run the optimization pipeline.

Source code in src/prompt_optimizer/optimizers/base.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def run(self) -> Prompt:
    """Run the optimization pipeline."""
    # Score seed_prompts
    self.seed_prompts = self._score_prompts(self.seed_prompts, self.validation_set)

    # Initialize objects
    self._p = [self.seed_prompts]
    self._g = [self.seed_prompts]

    # Iterate until max depth
    for t in track(range(1, self.max_depth + 1), description="Step", total=self.max_depth):
        # Generate prompt candidates
        g_t = self.generate_prompt_candidates(prompts=self._p[t - 1], validation_set=self.validation_set)
        self._g.append(g_t)
        # Select prompt candidates
        p_t = self.select_prompt_candidates(prompts=self._g[t], validation_set=self.validation_set)
        self._p.append(p_t)
        # Check for early convergence
        if self.check_early_convergence(all_prompts=self._p):
            break

    # Save prompts if requested
    self.save_prompts(output_path=self.output_path)

    # Return best prompt
    return self.select_best_prompt(all_prompts=self._p)

save_prompts(output_path)

Save prompts in jsonl format.

Source code in src/prompt_optimizer/optimizers/base.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def save_prompts(self, output_path: Optional[Union[str, Path]]):
    """Save prompts in jsonl format."""
    # Exit if no output path is set
    if self.output_path is None:
        return

    # Get and deduplicate prompts
    prompts = sum(self._p, start=[])
    prompts = list(set(prompts))

    # Save the prompts to the file
    lines = [prompt.model_dump_json() for prompt in prompts]
    with open(output_path, "w") as f:
        for line in lines:
            f.write(line)
            f.write("\n")

select_best_prompt(*, all_prompts)

Select the top scoring prompt.

Source code in src/prompt_optimizer/optimizers/protegi.py 
239
240
241
242
243
244
245
246
247
def select_best_prompt(self, *, all_prompts: list[list[Prompt]]) -> Prompt:
    """Select the top scoring prompt."""
    # Flatten all iterations
    prompts = sum(all_prompts, start=[])

    # Select the single prompt with the highest score
    best_prompt = self._get_best_prompt(prompts=prompts)
    logger.info(f"Best score: {best_prompt.score:.3f}")
    return best_prompt

select_prompt_candidates(*, prompts, validation_set)

Select prompt candidates according to the search mode.

Source code in src/prompt_optimizer/optimizers/protegi.py 
217
218
219
220
221
222
223
def select_prompt_candidates(self, *, prompts: list[Prompt], validation_set: ValidationSetType) -> list[Prompt]:
    """Select prompt candidates according to the search mode."""
    self._score_prompts(prompts=prompts, validation_set=validation_set)
    if self.search_mode == "greedy":
        return prompts
    elif self.search_mode == "beam":
        return [self._get_best_prompt(prompts=prompts)]