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

opik_optimizer/few_shot_bayesian_optimizer/few_shot_bayesian_optimizer.py

@@ -1,54 +1,108 @@
+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 apply any transformations to the variables either, only the variable name should be included in the format `{{<variable_name>}}`
+    - The few shot examples should include the expected response as the goal is to provide examples of the response.
+    - 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 +120,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 +154,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 +240,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 +248,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:
+                    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 +282,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_for_reporting
+            }
             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 +325,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 +347,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 +362,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 +393,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 +428,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 +542,39 @@ 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)
-
+            for key, value in dataset_item.items():
+                prompt_ = [{
+                    "role": item["role"],
+                    "content": item["content"].replace("{" + key + "}", str(value))
+                } 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_,