opik-optimizer 0.8.1__py3-none-any.whl → 0.9.0rc0__py3-none-any.whl

opik_optimizer/few_shot_bayesian_optimizer/few_shot_bayesian_optimizer.py

@@ -1,54 +1,107 @@
+import json
+import logging
 import random
-from typing import Any, Dict, List, Tuple, Union, Optional, Literal
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import litellm
 import opik
 import optuna
 import optuna.samplers
-import logging
-import json
-from datetime import datetime
-
 from opik import Dataset
-from opik_optimizer.optimization_config import mappers
+from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
+from pydantic import BaseModel
 
-from opik_optimizer.optimization_config.configs import TaskConfig, MetricConfig
 from opik_optimizer import base_optimizer
+from opik_optimizer.optimization_config import mappers
 
-from . import prompt_parameter
-from . import prompt_templates
-from .. import _throttle
-from .. import optimization_result, task_evaluator
-
-import litellm
-
-from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
+from .. import _throttle, optimization_result, task_evaluator, utils
+from ..optimization_config import chat_prompt
+from . import reporting
 
 _limiter = _throttle.get_rate_limiter_for_current_opik_installation()
 
 logger = logging.getLogger(__name__)
 
+FEW_SHOT_EXAMPLE_PLACEHOLDER = "FEW_SHOT_EXAMPLE_PLACEHOLDER"
+SYSTEM_PROMPT_TEMPLATE = f"""
+You are a prompt editor that modifies a message list to support few-shot learning. Your job is to insert a placeholder where few-shot examples can be inserted and generate a reusable string template for formatting those examples.
+
+You will receive a JSON object with the following fields:
+
+- "message_list": a list of messages, each with a role (system, user, or assistant) and a content field.
+- "examples": a list of example pairs, each with input and output fields.
+
+Your task:
+
+- Insert the string "{FEW_SHOT_EXAMPLE_PLACEHOLDER}" into one of the messages in the list. Make sure to:
+    - Insert it at the most logical point for including few-shot examples — typically as part of the system message
+    - Add a section title in XML or markdown format. The examples will be provided as `example_1\nexample_2\n...` with each example following the example template.
+- Analyze the examples to infer a consistent structure, and create a single string few_shot_example_template using the Python .format() style. Make sure to follow the following instructions:
+    - Unless absolutely relevant, do not return an object but instead a string that can be inserted as part of {FEW_SHOT_EXAMPLE_PLACEHOLDER}
+    - Make sure to include the variables as part of this string so we can before string formatting with actual examples. Only variables available in the examples can be used. Do not use anything else, do not apply any transformations to the variables either.
+    - The few shot examples should include the expected response as the goal is to provide examples of the expected output format.
+    - Ensure the format of the few shot examples are consistent with how the model will be called
+
+Return your output as a JSON object with:
+
+- message_list_with_placeholder: the updated list with "FEW_SHOT_EXAMPLE_PLACEHOLDER" inserted.
+- example_template: a string template using the fields provided in the examples (you don't need to use all of them)
+
+Respond only with the JSON object. Do not include any explanation or extra text.
+"""
+
+class FewShotPromptTemplate(BaseModel):
+    message_list_with_placeholder: List[Dict[str, str]]
+    example_template: str
 
 class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
+    """
+    The Few-Shot Bayesian Optimizer can be used to add few-shot examples to prompts. This algorithm
+    employes a two stage pipeline:
+
+    1. We generate a few-shot prompt template that is inserted can be inserted into the prompt 
+       provided
+    2. We use Bayesian Optimization to determine the best examples to include in the prompt.
+
+    This algorithm is best used when you have a well defined task and would like to guide the LLM
+    by providing some examples.
+    """
     def __init__(
         self,
         model: str,
-        project_name: Optional[str] = None,
+        project_name: Optional[str] = "Optimization",
         min_examples: int = 2,
         max_examples: int = 8,
         seed: int = 42,
         n_threads: int = 8,
-        n_initial_prompts: int = 5,
-        n_iterations: int = 10,
         verbose: int = 1,
         **model_kwargs,
     ) -> None:
+        """
+        Args:
+            model: The model to used to evaluate the prompt
+            project_name: Optional project name for tracking
+            min_examples: Minimum number of examples to include
+            max_examples: Maximum number of examples to include
+            seed: Random seed for reproducibility
+            n_threads: Number of threads for parallel evaluation
+            verbose: Controls internal logging/progress bars (0=off, 1=on).
+            **model_kwargs: Additional model parameters
+        """
         super().__init__(model, project_name, **model_kwargs)
         self.min_examples = min_examples
         self.max_examples = max_examples
         self.seed = seed
         self.n_threads = n_threads
-        self.n_initial_prompts = n_initial_prompts
-        self.n_iterations = n_iterations
         self.verbose = verbose
+        if verbose == 0:
+            logger.setLevel(logging.WARNING)
+        elif verbose == 1:
+            logger.setLevel(logging.INFO)
+        elif verbose == 2:
+            logger.setLevel(logging.DEBUG)
+        
         self._opik_client = opik.Opik()
         self.llm_call_counter = 0
         logger.debug(f"Initialized FewShotBayesianOptimizer with model: {model}")
@@ -66,7 +119,6 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         filtered_call_kwargs.pop('n_iterations', None)
         filtered_call_kwargs.pop('min_examples', None)
         filtered_call_kwargs.pop('max_examples', None)
-        filtered_call_kwargs.pop('n_initial_prompts', None)
 
         final_params_for_litellm = opik_litellm_monitor.try_add_opik_monitoring_to_params(filtered_call_kwargs)
 
@@ -101,90 +153,80 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         split_idx = int(len(dataset) * train_ratio)
         return dataset[:split_idx], dataset[split_idx:]
 
-    def _optimize_prompt(
+    def _create_fewshot_prompt_template(
         self,
-        dataset: Union[str, Dataset],
-        metric_config: MetricConfig,
-        task_config: TaskConfig,
+        model: str,
+        prompt: chat_prompt.ChatPrompt,
+        few_shot_examples: List[Dict[str, Any]]
+    ) -> FewShotPromptTemplate:
+        """
+        During this step we update the system prompt to include few-shot examples.
+        """
+        user_message = {
+            "message_list": prompt.formatted_messages,
+            "examples": few_shot_examples
+        }
+        
+        messages: List[Dict[str, str]] = [
+            {"role": "system", "content": SYSTEM_PROMPT_TEMPLATE},
+            {"role": "user", "content": json.dumps(user_message)},
+        ]
+        
+        logger.debug(f"fewshot_prompt_template - Calling LLM with: {messages}")
+        response = self._call_model(
+            model,
+            messages,
+            self.seed,
+            self.model_kwargs
+        )
+        logger.debug(f"fewshot_prompt_template - LLM response: {response}")
+
+        try:
+            res = utils.json_to_dict(response["choices"][0]["message"]["content"])
+            return FewShotPromptTemplate(
+                message_list_with_placeholder=res["message_list_with_placeholder"],
+                example_template=res["example_template"]
+            )
+        except Exception as e:
+            logger.error(f"Failed to compute few-shot prompt template: {e} - response: {response}")
+            raise
+
+    def _run_optimization(
+        self,
+        fewshot_prompt_template: FewShotPromptTemplate,
+        dataset: Dataset,
+        metric: Callable,
         n_trials: int = 10,
+        baseline_score: Optional[float] = None,
         optimization_id: Optional[str] = None,
         experiment_config: Optional[Dict] = None,
-        n_samples: int = None,
+        n_samples: Optional[int] = None,
     ) -> optimization_result.OptimizationResult:
+        reporting.start_optimization_run(verbose=self.verbose)
+
         random.seed(self.seed)
         self.llm_call_counter = 0
-
-        if not task_config.use_chat_prompt:
-            raise ValueError(
-                "Few-shot Bayesian optimization is only supported for chat prompts."
-            )
-
-        opik_dataset: opik.Dataset = dataset
-
+        
         # Load the dataset
-        if isinstance(dataset, str):
-            opik_dataset = self._opik_client.get_dataset(dataset)
-            dataset_items = opik_dataset.get_items()
-        else:
-            opik_dataset = dataset
-            dataset_items = opik_dataset.get_items()
-
+        dataset_items = dataset.get_items()
+        all_dataset_item_ids = [item["id"] for item in dataset_items]
+        eval_dataset_item_ids = all_dataset_item_ids
+        if n_samples is not None and n_samples < len(dataset_items):
+            eval_dataset_item_ids = random.sample(all_dataset_item_ids, n_samples)
+        
+        # Define the experiment configuration
         experiment_config = experiment_config or {}
         base_experiment_config = {  # Base config for reuse
             **experiment_config,
             **{
                 "optimizer": self.__class__.__name__,
-                "metric": metric_config.metric.name,
-                "dataset": opik_dataset.name,
+                "metric": metric.__name__,
+                "dataset": dataset.name,
                 "configuration": {},
             },
         }
 
-        # Evaluate Initial (Zero-Shot) Prompt
-        logger.info("Evaluating initial (zero-shot) prompt...")
-        initial_instruction = task_config.instruction_prompt
-        zero_shot_param = prompt_parameter.ChatPromptParameter(
-            name="zero_shot_prompt",
-            instruction=initial_instruction,
-            task_input_parameters=task_config.input_dataset_fields,
-            task_output_parameter=task_config.output_dataset_field,
-            demo_examples=[],  # No examples
-        )
-        zero_shot_llm_task = self._build_task_from_prompt_template(
-            zero_shot_param.as_template()
-        )
-
-        initial_eval_config = base_experiment_config.copy()
-        initial_eval_config["configuration"]["prompt"] = initial_instruction
-        initial_eval_config["configuration"]["n_examples"] = 0
-
-        # Determine dataset item IDs for evaluation (initial and trials)
-        all_dataset_item_ids = [item["id"] for item in dataset_items]
-        eval_dataset_item_ids = all_dataset_item_ids
-        if n_samples is not None and n_samples < len(all_dataset_item_ids):
-            eval_dataset_item_ids = random.sample(all_dataset_item_ids, n_samples)
-            logger.info(f"Using {n_samples} samples for evaluations.")
-        else:
-            logger.info(
-                f"Using all {len(all_dataset_item_ids)} samples for evaluations."
-            )
-
-        initial_score = task_evaluator.evaluate(
-            dataset=opik_dataset,
-            dataset_item_ids=eval_dataset_item_ids,
-            metric_config=metric_config,
-            evaluated_task=zero_shot_llm_task,
-            num_threads=self.n_threads,
-            project_name=self.project_name,
-            experiment_config=initial_eval_config,
-            optimization_id=optimization_id,
-            verbose=self.verbose,
-        )
-        logger.info(f"Initial (zero-shot) score: {initial_score:.4f}")
-
         # Start Optuna Study
-        logger.info("Starting Optuna study for Few-Shot Bayesian Optimization...")
-
         def optimization_objective(trial: optuna.Trial) -> float:
             n_examples = trial.suggest_int(
                 "n_examples", self.min_examples, self.max_examples
@@ -197,7 +239,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             ]
             trial.set_user_attr("example_indices", example_indices)
 
-            instruction = task_config.instruction_prompt
+            # Process few shot examples
             demo_examples = [dataset_items[idx] for idx in example_indices]
 
             processed_demo_examples = []
@@ -205,21 +247,29 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
                 processed_example = {}
                 for key, value in example.items():
                     processed_example[key] = str(value)
-                processed_demo_examples.append(processed_example)
-
-            param = prompt_parameter.ChatPromptParameter(
-                name=f"trial_{trial.number}_prompt",
-                instruction=instruction,
-                task_input_parameters=task_config.input_dataset_fields,
-                task_output_parameter=task_config.output_dataset_field,
-                demo_examples=processed_demo_examples,
+
+                try:
+                    processed_demo_examples.append(
+                        fewshot_prompt_template.example_template.format(**processed_example)
+                    )
+                except Exception as e:
+                    logger.error(f"Failed to format fewshot prompt template {fewshot_prompt_template} with example: {processed_example} ")
+                    raise
+            few_shot_examples = "\n\n".join(processed_demo_examples)
+            
+            llm_task = self._build_task_from_messages(
+                messages=fewshot_prompt_template.message_list_with_placeholder,
+                few_shot_examples=few_shot_examples
             )
 
-            llm_task = self._build_task_from_prompt_template(param.as_template())
+            messages_for_reporting = [{
+                "role": item["role"],
+                "content": item["content"].replace(FEW_SHOT_EXAMPLE_PLACEHOLDER, few_shot_examples)
+            } for item in fewshot_prompt_template.message_list_with_placeholder]
 
             # Log trial config
             trial_config = base_experiment_config.copy()
-            trial_config["configuration"]["prompt"] = instruction  # Base instruction
+            trial_config["configuration"]["prompt"] = messages_for_reporting  # Base instruction
             trial_config["configuration"][
                 "examples"
             ] = processed_demo_examples  # Log stringified examples
@@ -231,21 +281,30 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             )
             logger.debug(f"Evaluating trial {trial.number}...")
 
-            score = task_evaluator.evaluate(
-                dataset=opik_dataset,
-                dataset_item_ids=eval_dataset_item_ids,
-                metric_config=metric_config,
-                evaluated_task=llm_task,
-                num_threads=self.n_threads,
-                project_name=self.project_name,
-                experiment_config=trial_config,
-                optimization_id=optimization_id,
-                verbose=self.verbose,
-            )
+            with reporting.start_optimization_trial(trial.number, n_trials, verbose=self.verbose) as trial_reporter:
+                trial_reporter.start_trial(messages_for_reporting)
+                score = task_evaluator.evaluate(
+                    dataset=dataset,
+                    dataset_item_ids=eval_dataset_item_ids,
+                    metric=metric,
+                    evaluated_task=llm_task,
+                    num_threads=self.n_threads,
+                    project_name=self.project_name,
+                    experiment_config=trial_config,
+                    optimization_id=optimization_id,
+                    verbose=self.verbose,
+                )
+                trial_reporter.set_score(baseline_score, score)
             logger.debug(f"Trial {trial.number} score: {score:.4f}")
 
+            # Trial results
+            trial_config = {
+                "demo_examples": demo_examples,
+                "message_list_with_placeholder": fewshot_prompt_template.message_list_with_placeholder,
+                "message_list": messages
+            }
             trial.set_user_attr("score", score)
-            trial.set_user_attr("param", param)
+            trial.set_user_attr("config", trial_config)
             return score
 
         # Configure Optuna Logging
@@ -265,29 +324,18 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         sampler = optuna.samplers.TPESampler(seed=self.seed)
         study = optuna.create_study(direction="maximize", sampler=sampler)
         
-        study.optimize(optimization_objective, n_trials=n_trials, show_progress_bar=(self.verbose >= 1))
-        logger.info("Optuna study finished.")
-
+        study.optimize(
+            optimization_objective,
+            n_trials=n_trials,
+            show_progress_bar=False
+        )
+        
         optuna_history_processed = []
         for trial_idx, trial in enumerate(study.trials):
             if trial.state == optuna.trial.TrialState.COMPLETE:
-                param_obj: Optional[prompt_parameter.ChatPromptParameter] = trial.user_attrs.get("param")
-                prompt_cand_display = None # Default to None
-                if param_obj and hasattr(param_obj, 'as_template') and callable(param_obj.as_template):
-                    try:
-                        # .format() on ChatPromptTemplate returns the list of messages
-                        chat_messages_for_history = param_obj.as_template().format()
-                        prompt_cand_display = json.dumps(chat_messages_for_history) 
-                    except Exception as e_param_format:
-                        logger.warning(f"Trial {trial.number}: Error formatting prompt from param_obj: {e_param_format}")
-                        prompt_cand_display = "Error: Could not format prompt content."
-                elif not param_obj:
-                    logger.warning(f"Trial {trial.number}: 'param' object not found in user_attrs.")
-                    prompt_cand_display = "Error: Prompt data missing in trial."
-                else:
-                    logger.warning(f"Trial {trial.number}: 'param' object is not of expected type or lacks methods.")
-                    prompt_cand_display = "Error: Invalid prompt data structure in trial."
-
+                trial_config = trial.user_attrs.get("config", {})
+                prompt_cand_display = trial_config.get('message_list') # Default to None
+                
                 score_val = trial.value # This can be None if trial failed to produce a score
                 duration_val = None
                 if trial.datetime_complete and trial.datetime_start:
@@ -298,16 +346,13 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
                     "timestamp": trial.datetime_start.isoformat() if trial.datetime_start else datetime.now().isoformat(),
                     "prompt_candidate": prompt_cand_display,
                     "parameters_used": { 
-                        "optuna_params": trial.params, 
+                        "optuna_params": trial.user_attrs.get("config", {}), 
                         "example_indices": trial.user_attrs.get("example_indices", []) # Default to empty list
                     },
                     "scores": [{
-                        "metric_name": metric_config.metric.name, 
+                        "metric_name": metric.__name__, 
                         "score": score_val, # Can be None
-                        "opik_evaluation_id": None # TODO
                     }],
-                    "tokens_used": None, # TODO
-                    "cost": None, # TODO
                     "duration_seconds": duration_val,
                 }
                 optuna_history_processed.append(iter_detail)
@@ -316,33 +361,30 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
 
         best_trial = study.best_trial
         best_score = best_trial.value
-        best_n_examples = best_trial.params["n_examples"]
         best_example_indices = best_trial.user_attrs.get("example_indices", [])
-        best_param: prompt_parameter.ChatPromptParameter = best_trial.user_attrs[
-            "param"
-        ]
 
-        chat_messages_list = best_param.as_template().format()
-        main_prompt_string = best_param.instruction
+        reporting.display_result(
+            initial_score=baseline_score,
+            best_score=best_score,
+            best_prompt=best_trial.user_attrs["config"]["message_list"],
+            verbose=self.verbose
+        )
 
         return optimization_result.OptimizationResult(
             optimizer=self.__class__.__name__,
-            prompt=main_prompt_string,
+            prompt=best_trial.user_attrs["config"]["message_list"],
             score=best_score,
-            metric_name=metric_config.metric.name,
+            metric_name=metric.__name__,
             details={
-                "prompt_type": "chat" if task_config.use_chat_prompt else "non-chat",
-                "chat_messages": chat_messages_list,
-                "prompt_parameter": best_param,
-                "n_examples": best_n_examples,
+                "chat_messages": best_trial.user_attrs["config"]["message_list"],
+                "prompt_parameter": best_trial.user_attrs["config"],
+                #"n_examples": best_n_examples,
                 "example_indices": best_example_indices,
                 "trial_number": best_trial.number,
-                "initial_score": initial_score,
                 "total_trials": n_trials,
                 "rounds": [],
                 "stopped_early": False,
-                "metric_config": metric_config.model_dump(),
-                "task_config": task_config.model_dump(),
+                "metric_name": metric.__name__,
                 "model": self.model,
                 "temperature": self.model_kwargs.get("temperature"),
             },
@@ -350,20 +392,32 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             llm_calls=self.llm_call_counter
         )
 
-    def optimize_prompt(
+    def optimize_prompt( # type: ignore
         self,
-        dataset: Union[str, Dataset],
-        metric_config: MetricConfig,
-        task_config: TaskConfig,
+        prompt: chat_prompt.ChatPrompt,
+        dataset: Dataset,
+        metric: Callable,
         n_trials: int = 10,
         experiment_config: Optional[Dict] = None,
-        n_samples: int = None,
+        n_samples: Optional[int] = None,
     ) -> optimization_result.OptimizationResult:
+        """
+        Args:
+            prompt: The prompt to optimize
+            dataset: Opik Dataset to optimize on
+            metric: Metric function to evaluate on
+            n_trials: Number of trials for Bayesian Optimization
+            experiment_config: Optional configuration for the experiment, useful to log additional metadata
+            n_samples: Optional number of items to test in the dataset
+        
+        Returns:
+            OptimizationResult: Result of the optimization
+        """
         optimization = None
         try:
             optimization = self._opik_client.create_optimization(
                 dataset_name=dataset.name,
-                objective_name=metric_config.metric.name,
+                objective_name=metric.__name__,
                 metadata={"optimizer": self.__class__.__name__},
             )
         except Exception:
@@ -373,72 +427,109 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             optimization = None
 
         try:
-            result = self._optimize_prompt(
-                optimization_id=optimization.id if optimization is not None else None,
+            # Start experiment reporting
+            reporting.display_header("Few-Shot Bayesian Optimizer", verbose=self.verbose)
+            reporting.display_configuration(
+                prompt.formatted_messages,
+                optimizer_config={
+                    "optimizer": self.__class__.__name__,
+                    "metric": metric.__name__,
+                    "n_trials": n_trials,
+                    "n_samples": n_samples
+                },
+                verbose=self.verbose
+            )
+
+            utils.disable_experiment_reporting()
+
+            # Step 1. Compute the baseline evaluation
+            with reporting.display_evaluation(message="First we will establish the baseline performance:", verbose=self.verbose) as eval_report:
+                baseline_score = self.evaluate_prompt(
+                    prompt=prompt,
+                    dataset=dataset,
+                    metric=metric,
+                    n_samples=n_samples,
+                    optimization_id=optimization.id if optimization is not None else None
+                )
+
+                eval_report.set_score(baseline_score)
+            
+            # Step 2. Create the few-shot prompt template
+            with reporting.creation_few_shot_prompt_template(verbose=self.verbose) as fewshot_template_report:
+                fewshot_template = self._create_fewshot_prompt_template(
+                    model=self.model,
+                    prompt=prompt,
+                    few_shot_examples=[{k: v for k, v in item.items() if k != 'id'} 
+                                        for item in dataset.get_items(nb_samples=10)]
+                )
+
+                fewshot_template_report.set_fewshot_template(fewshot_template)
+
+            # Step 3. Start the optimization process
+            result = self._run_optimization(
+                fewshot_prompt_template=fewshot_template,
                 dataset=dataset,
-                metric_config=metric_config,
-                task_config=task_config,
-                n_trials=n_trials,
+                metric=metric,
+                optimization_id=optimization.id if optimization is not None else None,
                 experiment_config=experiment_config,
+                n_trials=n_trials,
+                baseline_score=baseline_score,
                 n_samples=n_samples,
             )
             if optimization:
                 self.update_optimization(optimization, status="completed")
+
+            utils.enable_experiment_reporting()
             return result
         except Exception as e:
             if optimization:
                 self.update_optimization(optimization, status="cancelled")
             logger.error(f"FewShotBayesian optimization failed: {e}", exc_info=True)
+            utils.enable_experiment_reporting()
             raise e
 
     def evaluate_prompt(
         self,
-        prompt: List[Dict[Literal["role", "content"], str]],
+        prompt: chat_prompt.ChatPrompt,
         dataset: opik.Dataset,
-        metric_config: MetricConfig,
-        task_config: Optional[TaskConfig] = None,
+        metric: Callable,
         dataset_item_ids: Optional[List[str]] = None,
         experiment_config: Optional[Dict] = None,
-        n_samples: int = None,
+        optimization_id: Optional[str] = None,
+        n_samples: Optional[int] = None,
     ) -> float:
-
-        if isinstance(prompt, str):
-            if task_config is None:
-                raise ValueError(
-                    "To use a string prompt, please pass in task_config to evaluate_prompt()"
-                )
-
-            questions = {
-                field: ("{{%s}}" % field) for field in task_config.input_dataset_fields
-            }
-            prompt = [
-                {"role": "system", "content": prompt},
-                {"role": "user", "content": json.dumps(questions)},
-            ]
-
+        """
+        Args:
+            prompt: The prompt to evaluate
+            dataset: Opik Dataset to evaluate the prompt on
+            metric: Metric function to evaluate on, should have the arguments `dataset_item` and `llm_output`
+            dataset_item_ids: Optional list of dataset item IDs to evaluate
+            experiment_config: Optional configuration for the experiment
+            optimization_id: Optional ID of the optimization
+            n_samples: Optional number of items to test in the dataset
+        Returns:
+            float: The evaluation score
+        """
         # Ensure prompt is correctly formatted
         if not all(
             isinstance(item, dict) and "role" in item and "content" in item
-            for item in prompt
+            for item in prompt.formatted_messages
         ):
             raise ValueError(
                 "A ChatPrompt must be a list of dictionaries with 'role' and 'content' keys."
             )
 
-        template = prompt_templates.ChatPromptTemplate(
-            prompt, validate_placeholders=False
-        )
-        llm_task = self._build_task_from_prompt_template(template)
+        llm_task = self._build_task_from_messages(prompt.formatted_messages)
 
         experiment_config = experiment_config or {}
         experiment_config = {
             **experiment_config,
             **{
                 "optimizer": self.__class__.__name__,
-                "metric": metric_config.metric.name,
+                "metric": metric.__name__,
                 "dataset": dataset.name,
                 "configuration": {
-                    "examples": prompt,
+                    "prompt": prompt.formatted_messages,
                 },
             },
         }
@@ -450,27 +541,38 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             all_ids = [dataset_item["id"] for dataset_item in dataset.get_items()]
             dataset_item_ids = random.sample(all_ids, n_samples)
 
-        logger.debug(f"Starting FewShotBayesian evaluation...")
+        logger.debug("Starting FewShotBayesian evaluation...")
         score = task_evaluator.evaluate(
             dataset=dataset,
             dataset_item_ids=dataset_item_ids,
-            metric_config=metric_config,
+            metric=metric,
             evaluated_task=llm_task,
             num_threads=self.n_threads,
             project_name=self.project_name,
             experiment_config=experiment_config,
+            optimization_id=optimization_id,
             verbose=self.verbose,
         )
         logger.debug(f"Evaluation score: {score:.4f}")
 
         return score
 
-    def _build_task_from_prompt_template(
-        self, template: prompt_templates.ChatPromptTemplate
+
+    def _build_task_from_messages(
+        self, messages: List[Dict[str, str]], few_shot_examples: Optional[str] = None
     ):
         def llm_task(dataset_item: Dict[str, Any]) -> Dict[str, Any]:
-            prompt_ = template.format(**dataset_item)
-
+            prompt_ = [{
+                "role": item["role"],
+                "content": item["content"].format(**dataset_item)
+            } for item in messages]
+
+            if few_shot_examples:
+                prompt_ = [{
+                    "role": item["role"],
+                    "content": item["content"].replace(FEW_SHOT_EXAMPLE_PLACEHOLDER, few_shot_examples)
+                } for item in prompt_]
+                
             response = self._call_model(
                 model=self.model,
                 messages=prompt_,
@@ -482,4 +584,4 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
                 mappers.EVALUATED_LLM_TASK_OUTPUT: response.choices[0].message.content
             }
 
-        return llm_task
+        return llm_task, messages