opik-optimizer 0.9.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

opik_optimizer/meta_prompt_optimizer/meta_prompt_optimizer.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import os
-from typing import Any, Callable, Dict, List, Optional, overload
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type
 
 import litellm
 import opik
@@ -13,11 +13,13 @@ from opik.environment import get_tqdm_for_current_environment
 from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
 
 from opik_optimizer import task_evaluator
+from opik_optimizer import utils
 
 from .. import _throttle
 from ..base_optimizer import BaseOptimizer, OptimizationRound
 from ..optimization_config import chat_prompt, mappers
 from ..optimization_result import OptimizationResult
+from ..optimizable_agent import OptimizableAgent
 from . import reporting
 
 tqdm = get_tqdm_for_current_environment()
@@ -35,17 +37,18 @@ _rate_limiter = _throttle.get_rate_limiter_for_current_opik_installation()
 class MetaPromptOptimizer(BaseOptimizer):
     """
     The Meta-Prompt Optimizer uses meta-prompting to improve prompts based on examples and performance.
-    
+
     This algorithm is best used when you have a prompt and would like to make sure it follows best
     practices.
     """
+
     # --- Constants for Default Configuration ---
     DEFAULT_ROUNDS = 3
     DEFAULT_PROMPTS_PER_ROUND = 4
 
     # --- Reasoning System Prompt ---
     _REASONING_SYSTEM_PROMPT = """You are an expert prompt engineer. Your task is to improve prompts for any type of task.
-        
+
         Focus on making the prompt more effective by:
         1. Being clear and specific about what is expected
         2. Providing necessary context and constraints
@@ -83,11 +86,10 @@ class MetaPromptOptimizer(BaseOptimizer):
         rounds: int = DEFAULT_ROUNDS,
         num_prompts_per_round: int = DEFAULT_PROMPTS_PER_ROUND,
         num_threads: int = 12,
-        project_name: str = "Optimization",
         verbose: int = 1,
         enable_context: bool = True,
-        **model_kwargs,
-    ):
+        **model_kwargs: Any,
+    ) -> None:
         """
         Args:
             model: The model to use for evaluation
@@ -95,18 +97,16 @@ class MetaPromptOptimizer(BaseOptimizer):
             rounds: Number of optimization rounds
             num_prompts_per_round: Number of prompts to generate per round
             num_threads: Number of threads for parallel evaluation
-            project_name: Optional project name for tracking
             verbose: Controls internal logging/progress bars (0=off, 1=on).
             enable_context: Whether to include task-specific context (metrics, examples) in the reasoning prompt.
             **model_kwargs: Additional model parameters
         """
-        super().__init__(model=model, project_name=project_name, **model_kwargs)
+        super().__init__(model=model, verbose=verbose, **model_kwargs)
         self.reasoning_model = reasoning_model if reasoning_model is not None else model
         self.rounds = rounds
         self.num_prompts_per_round = num_prompts_per_round
         self.num_threads = num_threads
-        self.verbose = verbose
-        self.dataset = None
+        self.dataset: Optional[Dataset] = None
         self._opik_client = opik_client.get_client_cached()
         self.llm_call_counter = 0
         self.enable_context = enable_context
@@ -120,6 +120,7 @@ class MetaPromptOptimizer(BaseOptimizer):
     @_throttle.rate_limited(_rate_limiter)
     def _call_model(
         self,
+        project_name: str,
         messages: List[Dict[str, str]],
         is_reasoning: bool = False,
         optimization_id: Optional[str] = None,
@@ -133,12 +134,18 @@ class MetaPromptOptimizer(BaseOptimizer):
             base_max_tokens = getattr(self, "max_tokens", 1000)
 
             # Use potentially different settings for reasoning calls
-            reasoning_temperature = base_temperature # Keep same temp unless specified otherwise
+            reasoning_temperature = (
+                base_temperature  # Keep same temp unless specified otherwise
+            )
             # Increase max_tokens for reasoning to ensure JSON fits, unless already high
-            reasoning_max_tokens = max(base_max_tokens, 3000) if is_reasoning else base_max_tokens 
+            reasoning_max_tokens = (
+                max(base_max_tokens, 3000) if is_reasoning else base_max_tokens
+            )
 
             llm_config_params = {
-                "temperature": reasoning_temperature if is_reasoning else base_temperature,
+                "temperature": (
+                    reasoning_temperature if is_reasoning else base_temperature
+                ),
                 "max_tokens": reasoning_max_tokens,
                 "top_p": getattr(self, "top_p", 1.0),
                 "frequency_penalty": getattr(self, "frequency_penalty", 0.0),
@@ -146,12 +153,12 @@ class MetaPromptOptimizer(BaseOptimizer):
             }
 
             # Prepare metadata that we want to be part of the LLM call context.
-            metadata_for_opik = {}
-            if self.project_name:
+            metadata_for_opik: Dict[str, Any] = {}
+            if project_name:
                 metadata_for_opik["project_name"] = (
-                    self.project_name
-                )  # Top-level for general use
-                metadata_for_opik["opik"] = {"project_name": self.project_name}
+                    project_name  # Top-level for general use
+                )
+                metadata_for_opik["opik"] = {"project_name": project_name}
 
             if optimization_id:
                 # Also add to opik-specific structure if project_name was added
@@ -181,10 +188,10 @@ class MetaPromptOptimizer(BaseOptimizer):
             )
 
             response = litellm.completion(
-                model=model_to_use, 
-                messages=messages, 
+                model=model_to_use,
+                messages=messages,
                 num_retries=6,
-                **final_call_params
+                **final_call_params,
             )
             return response.choices[0].message.content
         except litellm.exceptions.RateLimitError as e:
@@ -197,27 +204,26 @@ class MetaPromptOptimizer(BaseOptimizer):
             logger.error(f"LiteLLM Context Window Exceeded Error: {e}")
             # Log prompt length if possible? Needs access to prompt_for_llm here.
             raise
-        except Exception as e:
-            logger.error(
-                f"Error calling model '{model_to_use}': {type(e).__name__} - {e}"
-            )
+        except Exception:
+            # logger.error(
+            #    f"Error calling model '{model_to_use}': {type(e).__name__} - {e}"
+            # )
             raise
 
-     # type: ignore
-    def evaluate_prompt(
+    def _evaluate_prompt(
         self,
         prompt: chat_prompt.ChatPrompt,
         dataset: opik.Dataset,
         metric: Callable,
-        use_full_dataset: bool = True,
-        experiment_config: Optional[Dict] = None,
         n_samples: Optional[int] = None,
+        dataset_item_ids: Optional[List[str]] = None,
+        experiment_config: Optional[Dict] = None,
+        use_full_dataset: bool = True,
         optimization_id: Optional[str] = None,
-        verbose: int = 1,
+        **kwargs: Any,
     ) -> float:
         """
         Args:
-            prompt: The prompt to evaluate
             dataset: Opik Dataset to evaluate the prompt on
             metric: Metric functions
             use_full_dataset: Whether to use the full dataset or a subset
@@ -249,16 +255,18 @@ class MetaPromptOptimizer(BaseOptimizer):
         else:
             subset_size = None  # Use all items for final checks
             logger.debug("Using full dataset for evaluation")
-        
+
         experiment_config = experiment_config or {}
         experiment_config = {
             **experiment_config,
             **{
                 "optimizer": self.__class__.__name__,
-                "metric": metric.__name__,
+                "agent_class": self.agent_class.__name__,
+                "agent_config": prompt.to_dict(),
+                "metric": getattr(metric, "__name__", str(metric)),
                 "dataset": dataset.name,
                 "configuration": {
-                    "prompt": prompt.formatted_messages,
+                    "prompt": prompt.get_messages(),
                     "n_samples": subset_size,
                     "use_full_dataset": use_full_dataset,
                 },
@@ -269,25 +277,33 @@ class MetaPromptOptimizer(BaseOptimizer):
 
         def llm_task(dataset_item: Dict[str, Any]) -> Dict[str, str]:
             # --- Step 1: Prepare the prompt for the LLM ---
-            messages = [{
-                "role": item["role"],
-                "content": item["content"].format(**dataset_item)
-            } for item in prompt.formatted_messages]
+            # messages = [
+            #    {
+            #        "role": item["role"],
+            #        "content": item["content"].format(**dataset_item),
+            #    }
+            #    for item in prompt.get_messages()
+            # ]
+            # Step 1: create the agent
+            new_prompt = prompt.copy()
+            messages = new_prompt.get_messages(dataset_item)
+            new_prompt.set_messages(messages)
+            agent = self.agent_class(new_prompt)
 
             # --- Step 2: Call the model ---
             try:
-                logger.debug(f"Calling LLM with prompt length: {sum(len(msg['content']) for msg in messages)}")
-                raw_model_output = self._call_model(
-                    messages=messages,
-                    is_reasoning=False,
-                    optimization_id=optimization_id,
+                logger.debug(
+                    f"Calling LLM with prompt length: {sum(len(msg['content']) for msg in messages)}"
                 )
+                raw_model_output = agent.invoke(messages)
                 logger.debug(f"LLM raw response length: {len(raw_model_output)}")
                 logger.debug(f"LLM raw output: {raw_model_output}")
             except Exception as e:
                 logger.error(f"Error calling model with prompt: {e}")
                 logger.error(f"Failed prompt: {messages}")
-                logger.error(f"Prompt length: {sum(len(msg['content']) for msg in messages)}")
+                logger.error(
+                    f"Prompt length: {sum(len(msg['content']) for msg in messages)}"
+                )
                 raise
 
             # --- Step 3: Clean the model's output before metric evaluation ---
@@ -300,14 +316,15 @@ class MetaPromptOptimizer(BaseOptimizer):
 
         # Use dataset's get_items with limit for sampling
         logger.debug(
-            f"Starting evaluation with {subset_size if subset_size else 'all'} samples for metric: {metric.__name__}"
+            f"Starting evaluation with {subset_size if subset_size else 'all'} samples for metric: {getattr(metric, '__name__', str(metric))}"
         )
         score = task_evaluator.evaluate(
             dataset=dataset,
             metric=metric,
             evaluated_task=llm_task,
+            dataset_item_ids=dataset_item_ids,
             num_threads=self.num_threads,
-            project_name=self.project_name,
+            project_name=self.agent_class.project_name,
             n_samples=subset_size,  # Use subset_size for trials, None for full dataset
             experiment_config=experiment_config,
             optimization_id=optimization_id,
@@ -316,7 +333,7 @@ class MetaPromptOptimizer(BaseOptimizer):
         logger.debug(f"Evaluation score: {score:.4f}")
         return score
 
-    def optimize_prompt( # type: ignore[override]
+    def optimize_prompt(
         self,
         prompt: chat_prompt.ChatPrompt,
         dataset: Dataset,
@@ -324,13 +341,13 @@ class MetaPromptOptimizer(BaseOptimizer):
         experiment_config: Optional[Dict] = None,
         n_samples: Optional[int] = None,
         auto_continue: bool = False,
-        **kwargs,
+        agent_class: Optional[Type[OptimizableAgent]] = None,
+        **kwargs: Any,
     ) -> OptimizationResult:
         """
         Optimize a prompt using meta-reasoning.
 
         Args:
-            prompt: The prompt to optimize
             dataset: The dataset to evaluate against
             metric: The metric to use for evaluation
             experiment_config: A dictionary to log with the experiments
@@ -341,8 +358,27 @@ class MetaPromptOptimizer(BaseOptimizer):
         Returns:
             OptimizationResult: Structured result containing optimization details
         """
-        reporting.display_header(self.__class__.__name__, verbose=self.verbose)
-        
+        if not isinstance(prompt, chat_prompt.ChatPrompt):
+            raise ValueError("Prompt must be a ChatPrompt object")
+
+        if not isinstance(dataset, Dataset):
+            raise ValueError("Dataset must be a Dataset object")
+
+        if not callable(metric):
+            raise ValueError(
+                "Metric must be a function that takes `dataset_item` and `llm_output` as arguments."
+            )
+
+        if prompt.model is None:
+            prompt.model = self.model
+        if prompt.model_kwargs is None:
+            prompt.model_kwargs = self.model_kwargs
+
+        if agent_class is None:
+            self.agent_class = utils.create_litellm_agent_class(prompt)
+        else:
+            self.agent_class = agent_class
+
         total_items = len(dataset.get_items())
         if n_samples is not None and n_samples > total_items:
             logger.warning(
@@ -350,21 +386,11 @@ class MetaPromptOptimizer(BaseOptimizer):
             )
             n_samples = None
 
-        reporting.display_configuration(
-            messages=prompt.formatted_messages,
-            optimizer_config={
-                "optimizer": self.__class__.__name__,
-                "n_samples": n_samples,
-                "auto_continue": auto_continue
-            },
-            verbose=self.verbose
-        )
-
         optimization = None
         try:
             optimization = self._opik_client.create_optimization(
                 dataset_name=dataset.name,
-                objective_name=metric.__name__,
+                objective_name=getattr(metric, "__name__", str(metric)),
                 metadata={"optimizer": self.__class__.__name__},
             )
             logger.debug(f"Created optimization with ID: {optimization.id}")
@@ -374,9 +400,26 @@ class MetaPromptOptimizer(BaseOptimizer):
             )
             optimization = None
 
+        reporting.display_header(
+            algorithm=self.__class__.__name__,
+            optimization_id=optimization.id if optimization is not None else None,
+            dataset_id=dataset.id,
+            verbose=self.verbose,
+        )
+        reporting.display_configuration(
+            messages=prompt.get_messages(),
+            optimizer_config={
+                "optimizer": self.__class__.__name__,
+                "n_samples": n_samples,
+                "auto_continue": auto_continue,
+            },
+            verbose=self.verbose,
+        )
+
         try:
+            optimization_id = optimization.id if optimization is not None else None
             result = self._optimize_prompt(
-                optimization_id=optimization.id if optimization is not None else None,
+                optimization_id=optimization_id,
                 prompt=prompt,
                 dataset=dataset,
                 metric=metric,
@@ -398,30 +441,33 @@ class MetaPromptOptimizer(BaseOptimizer):
 
     def _optimize_prompt(
         self,
-        optimization_id: str,
+        optimization_id: Optional[str],
         prompt: chat_prompt.ChatPrompt,
         dataset: Dataset,
         metric: Callable,
         experiment_config: Optional[Dict],
-        n_samples: int,
+        n_samples: Optional[int],
         auto_continue: bool,
-        **kwargs,
+        **kwargs: Any,
     ) -> OptimizationResult:
         self.auto_continue = auto_continue
         self.dataset = dataset
         self.prompt = prompt
-        self.llm_call_counter = 0 # Reset counter for run
+        self.llm_call_counter = 0  # Reset counter for run
+        initial_prompt = prompt
 
-        current_prompt = prompt.formatted_messages
+        current_prompt = prompt
         experiment_config = experiment_config or {}
         experiment_config = {
             **experiment_config,
             **{
                 "optimizer": self.__class__.__name__,
-                "metric": metric.__name__,
-                "dataset": self.dataset.name,
+                "agent_class": self.agent_class.__name__,
+                "agent_config": prompt.to_dict(),
+                "metric": getattr(metric, "__name__", str(metric)),
+                "dataset": dataset.name,
                 "configuration": {
-                    "prompt": current_prompt,
+                    "prompt": prompt.get_messages(),
                     "rounds": self.rounds,
                     "num_prompts_per_round": self.num_prompts_per_round,
                 },
@@ -429,8 +475,8 @@ class MetaPromptOptimizer(BaseOptimizer):
         }
 
         with reporting.display_evaluation(verbose=self.verbose) as baseline_reporter:
-            initial_score = self.evaluate_prompt(
-                prompt=prompt,
+            initial_score = self._evaluate_prompt(
+                prompt,
                 optimization_id=optimization_id,
                 dataset=dataset,
                 metric=metric,
@@ -441,20 +487,22 @@ class MetaPromptOptimizer(BaseOptimizer):
             )
             best_score = initial_score
             best_prompt = current_prompt
-            rounds = []
+            rounds: List[OptimizationRound] = []
 
             baseline_reporter.set_score(initial_score)
 
         reporting.display_optimization_start_message(verbose=self.verbose)
-        with reporting.display_round_progress(self.rounds, verbose=self.verbose) as round_reporter:
+        with reporting.display_round_progress(
+            self.rounds, verbose=self.verbose
+        ) as round_reporter:
             for round_num in range(self.rounds):
-                
                 round_reporter.round_start(round_num)
                 previous_best_score = best_score
-                
+
                 # Step 1. Create a set of candidate prompts
                 try:
                     candidate_prompts = self._generate_candidate_prompts(
+                        project_name=self.agent_class.project_name,
                         current_prompt=best_prompt,
                         best_score=best_score,
                         round_num=round_num,
@@ -467,14 +515,19 @@ class MetaPromptOptimizer(BaseOptimizer):
                     continue
 
                 # Step 2. Score each candidate prompt
-                prompt_scores = []
+                prompt_scores: List[Tuple[chat_prompt.ChatPrompt, float]] = []
                 for candidate_count, prompt in enumerate(candidate_prompts):
-                    with reporting.display_prompt_candidate_scoring_report(candidate_count, prompt, verbose=self.verbose) as eval_report:
+                    with reporting.display_prompt_candidate_scoring_report(
+                        verbose=self.verbose
+                    ) as eval_report:
                         eval_report.set_generated_prompts(candidate_count, prompt)
 
+                        new_prompt = current_prompt.copy()
+                        new_prompt.set_messages(prompt.get_messages())
+
                         try:
-                            prompt_score = self.evaluate_prompt(
-                                prompt=chat_prompt.ChatPrompt(messages=prompt),
+                            prompt_score = self._evaluate_prompt(
+                                prompt=new_prompt,
                                 optimization_id=optimization_id,
                                 dataset=dataset,
                                 metric=metric,
@@ -485,34 +538,37 @@ class MetaPromptOptimizer(BaseOptimizer):
                             )
 
                             eval_report.set_final_score(best_score, prompt_score)
-                        except Exception as e:
-                            raise ValueError(f"Error evaluating candidate prompt: {e}")
+                        except Exception:
+                            print("Failed evaluating agent; continuing...")
+                            prompt_score = 0
 
                     prompt_scores.append((prompt, prompt_score))
-                
+
                 # Step 3. Identify potential improvements
                 if not prompt_scores:
-                    logger.warning("No prompts were successfully evaluated in this round")
+                    logger.warning(
+                        "No prompts were successfully evaluated in this round"
+                    )
                     break
 
                 prompt_scores.sort(key=lambda x: x[1], reverse=True)
-                best_candidate_this_round, best_cand_score_avg = (
-                    prompt_scores[0]
+                best_candidate_this_round, best_cand_score_avg = prompt_scores[0]
+                improvement = self._calculate_improvement(
+                    best_cand_score_avg, best_score
                 )
-                improvement = self._calculate_improvement(best_cand_score_avg, best_score)
-                round_reporter.round_end(round_num, best_cand_score_avg, best_score, best_prompt)
-                
+                round_reporter.round_end(round_num, best_cand_score_avg, best_score)
+
                 round_data = self._create_round_data(
                     round_num=round_num,
-                    current_best_prompt=chat_prompt.ChatPrompt(messages=best_candidate_this_round),
-                    current_best_score=best_cand_score_avg,
-                    best_prompt_overall=chat_prompt.ChatPrompt(messages=best_prompt),
+                    current_best_prompt=best_prompt,
+                    current_best_score=best_score,
+                    best_prompt_overall=best_prompt,
                     evaluated_candidates=prompt_scores,
                     previous_best_score=previous_best_score,
                     improvement_this_round=improvement,
                 )
                 rounds.append(round_data)
-                self._add_to_history(round_data.model_dump())
+                self._add_to_history(round_data)
 
                 if improvement > 0:
                     best_score = best_cand_score_avg
@@ -521,17 +577,21 @@ class MetaPromptOptimizer(BaseOptimizer):
         reporting.display_result(
             initial_score,
             best_score,
-            best_prompt,
-            verbose=self.verbose
+            best_prompt.get_messages() if best_prompt is not None else [],
+            verbose=self.verbose,
         )
 
         return self._create_result(
             metric,
-            prompt,
-            best_prompt,
-            best_score,
-            initial_score,
-            rounds,
+            initial_prompt=initial_prompt.get_messages()
+            if initial_prompt is not None
+            else [],
+            best_prompt=best_prompt.get_messages() if best_prompt is not None else [],
+            best_score=best_score,
+            initial_score=initial_score,
+            rounds=rounds,
+            dataset_id=dataset.id,
+            optimization_id=optimization_id,
         )
 
     def _calculate_improvement(
@@ -550,7 +610,7 @@ class MetaPromptOptimizer(BaseOptimizer):
         current_best_prompt: chat_prompt.ChatPrompt,
         current_best_score: float,
         best_prompt_overall: chat_prompt.ChatPrompt,
-        evaluated_candidates: List[tuple[str, float, List[float]]],
+        evaluated_candidates: List[Tuple[chat_prompt.ChatPrompt, float]],
         previous_best_score: float,
         improvement_this_round: float,
     ) -> OptimizationRound:
@@ -562,7 +622,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             )
             generated_prompts_log.append(
                 {
-                    "prompt": prompt,
+                    "prompt": prompt.get_messages(),
                     "score": score,
                     "improvement": improvement_vs_prev,
                 }
@@ -581,21 +641,21 @@ class MetaPromptOptimizer(BaseOptimizer):
     def _create_result(
         self,
         metric: Callable,
-        prompt: chat_prompt.ChatPrompt,
-        best_prompt: str,
+        initial_prompt: List[Dict[str, str]],
+        best_prompt: List[Dict[str, str]],
         best_score: float,
         initial_score: float,
         rounds: List[OptimizationRound],
+        dataset_id: Optional[str],
+        optimization_id: Optional[str],
     ) -> OptimizationResult:
         """Create the final OptimizationResult object."""
         details = {
-            "initial_prompt": prompt,
-            "initial_score": initial_score,
             "final_prompt": best_prompt,
             "final_score": best_score,
             "rounds": rounds,
             "total_rounds": len(rounds),
-            "metric_name": metric.__name__,
+            "metric_name": getattr(metric, "__name__", str(metric)),
             "model": self.model,
             "temperature": self.model_kwargs.get("temperature"),
         }
@@ -604,9 +664,13 @@ class MetaPromptOptimizer(BaseOptimizer):
             optimizer=self.__class__.__name__,
             prompt=best_prompt,
             score=best_score,
-            metric_name=metric.__name__,
+            initial_prompt=initial_prompt,
+            initial_score=initial_score,
+            metric_name=getattr(metric, "__name__", str(metric)),
             details=details,
-            llm_calls=self.llm_call_counter
+            llm_calls=self.llm_call_counter,
+            dataset_id=dataset_id,
+            optimization_id=optimization_id,
         )
 
     def _get_task_context(self, metric: Callable) -> str:
@@ -632,7 +696,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             context += f"Dataset fields (includes both input and optionally the expected output): {', '.join([x for x in sample.keys() if x != 'id'])}\n"
             context += f"Evaluation Metric:\n{metrics_str}\n"
             context += f"\nExample:\n{json.dumps(sample)}\n"
-        
+
         return context
 
     def _generate_candidate_prompts(
@@ -643,14 +707,14 @@ class MetaPromptOptimizer(BaseOptimizer):
         previous_rounds: List[OptimizationRound],
         metric: Callable,
         optimization_id: Optional[str] = None,
-    ) -> List[str]:
+        project_name: Optional[str] = None,
+    ) -> List[chat_prompt.ChatPrompt]:
         """Generate candidate prompts using meta-prompting."""
         with reporting.display_candidate_generation_report(
-            self.num_prompts_per_round,
-            verbose=self.verbose
-        ) as candidate_generation_report: 
+            self.num_prompts_per_round, verbose=self.verbose
+        ) as candidate_generation_report:
             logger.debug(f"\nGenerating candidate prompts for round {round_num + 1}")
-            logger.debug(f"Generating from prompt: {current_prompt}")
+            logger.debug(f"Generating from prompt: {current_prompt.get_messages()}")
             logger.debug(f"Current best score: {best_score:.4f}")
 
             history_context = self._build_history_context(previous_rounds)
@@ -662,16 +726,22 @@ class MetaPromptOptimizer(BaseOptimizer):
             if self.enable_context:
                 task_context_str = self._get_task_context(metric=metric)
                 analysis_instruction = "Analyze the example provided (if any), the metric description (if any), and the history of scores."
-                metric_focus_instruction = f"Focus on improving the score for the metric: {metric.__name__}."
+                metric_focus_instruction = (
+                    f"Focus on improving the score for the metric: {metric.__name__}."
+                )
                 improvement_point_1 = "1. Be more specific and clear about expectations based on the metric and task."
-                logger.debug("Task context and metric-specific instructions enabled for reasoning prompt.")
+                logger.debug(
+                    "Task context and metric-specific instructions enabled for reasoning prompt."
+                )
             else:
-                analysis_instruction = "Analyze the history of scores and the current prompt\'s performance."
+                analysis_instruction = "Analyze the history of scores and the current prompt's performance."
                 metric_focus_instruction = "Focus on generating diverse and effective prompt variations based on the history."
                 improvement_point_1 = "1. Be more specific and clear about expectations based on the task."
-                logger.debug("Task context and metric-specific instructions disabled for reasoning prompt.")
+                logger.debug(
+                    "Task context and metric-specific instructions disabled for reasoning prompt."
+                )
 
-            user_prompt = f"""Current prompt: {current_prompt}
+            user_prompt = f"""Current prompt: {current_prompt.get_messages()}
             Current score: {best_score}
             {history_context}
             {task_context_str}
@@ -691,9 +761,10 @@ class MetaPromptOptimizer(BaseOptimizer):
             try:
                 # Use _call_model which handles selecting reasoning_model
                 content = self._call_model(
+                    project_name,
                     messages=[
                         {"role": "system", "content": self._REASONING_SYSTEM_PROMPT},
-                        {"role": "user", "content": user_prompt}
+                        {"role": "user", "content": user_prompt},
                     ],
                     is_reasoning=True,
                     optimization_id=optimization_id,
@@ -713,9 +784,13 @@ class MetaPromptOptimizer(BaseOptimizer):
                         try:
                             json_result = json.loads(json_match.group())
                         except json.JSONDecodeError as e:
-                            raise ValueError(f"Could not parse JSON extracted via regex: {e} - received: {json_match.group()}")
+                            raise ValueError(
+                                f"Could not parse JSON extracted via regex: {e} - received: {json_match.group()}"
+                            )
                     else:
-                        raise ValueError(f"No JSON object found in response via regex. - received: {content}")
+                        raise ValueError(
+                            f"No JSON object found in response via regex. - received: {content}"
+                        )
 
                 # Validate the parsed JSON structure
                 if isinstance(json_result, list) and len(json_result) == 1:
@@ -723,27 +798,46 @@ class MetaPromptOptimizer(BaseOptimizer):
 
                 if not isinstance(json_result, dict) or "prompts" not in json_result:
                     logger.debug(f"Parsed JSON content: {json_result}")
-                    raise ValueError(f"Parsed JSON is not a dictionary or missing 'prompts' key. - received: {json_result}")
+                    raise ValueError(
+                        f"Parsed JSON is not a dictionary or missing 'prompts' key. - received: {json_result}"
+                    )
 
                 if not isinstance(json_result["prompts"], list):
                     logger.debug(f"Content of 'prompts': {json_result.get('prompts')}")
-                    raise ValueError(f"'prompts' key does not contain a list. - received: {json_result.get('prompts')}")
+                    raise ValueError(
+                        f"'prompts' key does not contain a list. - received: {json_result.get('prompts')}"
+                    )
 
                 # Extract and log valid prompts
-                valid_prompts = []
+                valid_prompts: List[chat_prompt.ChatPrompt] = []
                 for item in json_result["prompts"]:
                     if (
                         isinstance(item, dict)
                         and "prompt" in item
                         and isinstance(item["prompt"], list)
                     ):
-                        prompt_text = item["prompt"]
-                        valid_prompts.append(prompt_text)
-                        
+                        # NOTE: might be brittle
+                        if current_prompt.user:
+                            user_text = current_prompt.user
+                        else:
+                            if current_prompt.messages is not None:
+                                user_text = current_prompt.messages[-1]["content"]
+                            else:
+                                raise Exception(
+                                    "User content not found in chat-prompt!"
+                                )
+
+                        valid_prompts.append(
+                            chat_prompt.ChatPrompt(
+                                system=item["prompt"][0]["content"],
+                                user=user_text,
+                            )
+                        )
+
                         # Log details
                         focus = item.get("improvement_focus", "N/A")
                         reasoning = item.get("reasoning", "N/A")
-                        logger.debug(f"Generated prompt: {prompt_text}")
+                        logger.debug(f"Generated prompt: {item['prompt']}")
                         logger.debug(f"  Improvement focus: {focus}")
                         logger.debug(f"  Reasoning: {reasoning}")
                     else:
@@ -752,17 +846,19 @@ class MetaPromptOptimizer(BaseOptimizer):
                         )
 
                 if not valid_prompts:
-                    raise ValueError("No valid prompts found in the parsed JSON response after validation.")
-                
-                candidate_generation_report.set_generated_prompts(
-                    self.num_prompts_per_round
-                )
-                
+                    raise ValueError(
+                        "No valid prompts found in the parsed JSON response after validation."
+                    )
+
+                candidate_generation_report.set_generated_prompts()
+
                 return valid_prompts
                 # --- End Robust Parsing ---
 
             except Exception as e:
-                raise ValueError(f"Unexpected error during candidate prompt generation: {e}")
+                raise ValueError(
+                    f"Unexpected error during candidate prompt generation: {e}"
+                )
 
     def _build_history_context(self, previous_rounds: List[OptimizationRound]) -> str:
         """Build context from previous optimization rounds."""