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

opik_optimizer/meta_prompt_optimizer/meta_prompt_optimizer.py

@@ -0,0 +1,816 @@
+import json
+import logging
+import os
+from typing import Any, Callable, Dict, List, Optional, overload
+
+import litellm
+import opik
+from litellm.caching import Cache
+from litellm.types.caching import LiteLLMCacheType
+from opik import Dataset
+from opik.api_objects import opik_client
+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 .. import _throttle
+from ..base_optimizer import BaseOptimizer, OptimizationRound
+from ..optimization_config import chat_prompt, mappers
+from ..optimization_result import OptimizationResult
+from . import reporting
+
+tqdm = get_tqdm_for_current_environment()
+
+# Using disk cache for LLM calls
+disk_cache_dir = os.path.expanduser("~/.litellm_cache")
+litellm.cache = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=disk_cache_dir)
+
+# Set up logging
+logger = logging.getLogger(__name__)  # Gets logger configured by setup_logging
+
+_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
+        3. Guiding the model to produce the desired output format
+        4. Removing ambiguity and unnecessary elements
+        5. Maintaining conciseness while being complete
+
+        Instructions:
+        1. If there is a system prompt, prioritize adding instructions there if and only if it makes sense.
+        2. DO NOT add any variables or parameters to the prompt you are editing.
+        3. You can reuse variables that already exist in the prompt.
+
+        Return a JSON array of prompts with the following structure. Make sure to return a valid
+        JSON object with correct use of double quotes and single quotes. JSON keys should be
+        double-quoted:
+        {
+            "prompts": [
+                {
+                    "prompt": [{"role": "<role>", "content": "<content>"}],
+                    "improvement_focus": "what aspect this prompt improves",
+                    "reasoning": "why this improvement should help"
+                },
+                {
+                    "prompt": [{"role": "<role>", "content": "<content>"}],
+                    "improvement_focus": "what aspect this prompt improves",
+                    "reasoning": "why this improvement should help"
+                }
+            ]
+        }"""
+
+    def __init__(
+        self,
+        model: str,
+        reasoning_model: Optional[str] = None,
+        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,
+    ):
+        """
+        Args:
+            model: The model to use for evaluation
+            reasoning_model: The model to use for reasoning and prompt generation
+            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)
+        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._opik_client = opik_client.get_client_cached()
+        self.llm_call_counter = 0
+        self.enable_context = enable_context
+        logger.debug(
+            f"Initialized MetaPromptOptimizer with model={model}, reasoning_model={self.reasoning_model}"
+        )
+        logger.debug(
+            f"Optimization rounds: {rounds}, Prompts/round: {num_prompts_per_round}"
+        )
+
+    @_throttle.rate_limited(_rate_limiter)
+    def _call_model(
+        self,
+        messages: List[Dict[str, str]],
+        is_reasoning: bool = False,
+        optimization_id: Optional[str] = None,
+    ) -> str:
+        """Call the model with the given prompt and return the response."""
+        self.llm_call_counter += 1
+        # Note: Basic retry logic could be added here using tenacity
+        try:
+            # Basic LLM parameters (e.g., temperature, max_tokens)
+            base_temperature = getattr(self, "temperature", 0.3)
+            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
+            # 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 
+
+            llm_config_params = {
+                "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),
+                "presence_penalty": getattr(self, "presence_penalty", 0.0),
+            }
+
+            # Prepare metadata that we want to be part of the LLM call context.
+            metadata_for_opik = {}
+            if self.project_name:
+                metadata_for_opik["project_name"] = (
+                    self.project_name
+                )  # Top-level for general use
+                metadata_for_opik["opik"] = {"project_name": self.project_name}
+
+            if optimization_id:
+                # Also add to opik-specific structure if project_name was added
+                if "opik" in metadata_for_opik:
+                    metadata_for_opik["opik"]["optimization_id"] = optimization_id
+
+            metadata_for_opik["optimizer_name"] = self.__class__.__name__
+            metadata_for_opik["opik_call_type"] = (
+                "reasoning" if is_reasoning else "evaluation_llm_task_direct"
+            )
+
+            if metadata_for_opik:
+                llm_config_params["metadata"] = metadata_for_opik
+
+            model_to_use = self.reasoning_model if is_reasoning else self.model
+
+            # Pass llm_config_params (which now includes our metadata) to the Opik monitor.
+            # The monitor is expected to return a dictionary suitable for spreading into litellm.completion,
+            # having handled our metadata and added any Opik-specific configurations.
+            final_call_params = opik_litellm_monitor.try_add_opik_monitoring_to_params(
+                llm_config_params.copy()
+            )
+
+            logger.debug(
+                f"Calling model '{model_to_use}' with messages: {messages}, "
+                f"final params for litellm (from monitor): {final_call_params}"
+            )
+
+            response = litellm.completion(
+                model=model_to_use, 
+                messages=messages, 
+                num_retries=6,
+                **final_call_params
+            )
+            return response.choices[0].message.content
+        except litellm.exceptions.RateLimitError as e:
+            logger.error(f"LiteLLM Rate Limit Error: {e}")
+            raise
+        except litellm.exceptions.APIConnectionError as e:
+            logger.error(f"LiteLLM API Connection Error: {e}")
+            raise
+        except litellm.exceptions.ContextWindowExceededError as e:
+            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}"
+            )
+            raise
+
+     # type: ignore
+    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,
+        optimization_id: Optional[str] = None,
+        verbose: int = 1,
+    ) -> 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
+            experiment_config: Optional configuration for the experiment, useful to log additional metadata
+            n_samples: Optional number of items to test in the dataset
+            optimization_id: Optional ID of the optimization
+            verbose: Controls internal logging/progress bars (0=off, 1=on).
+        Returns:
+            float: The evaluation score
+        """
+        # Calculate subset size for trials
+        if not use_full_dataset:
+            total_items = len(dataset.get_items())
+            if n_samples is not None:
+                if n_samples > total_items:
+                    logger.warning(
+                        f"Requested n_samples ({n_samples}) is larger than dataset size ({total_items}). Using full dataset."
+                    )
+                    subset_size = None
+                else:
+                    subset_size = n_samples
+                    logger.debug(f"Using specified n_samples: {subset_size} items")
+            else:
+                # Calculate 20% of total, but no more than 20 items and no more than total items
+                subset_size = min(total_items, min(20, max(10, int(total_items * 0.2))))
+                logger.debug(
+                    f"Using automatic subset size calculation: {subset_size} items (20% of {total_items} total items)"
+                )
+        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__,
+                "dataset": dataset.name,
+                "configuration": {
+                    "prompt": prompt.formatted_messages,
+                    "n_samples": subset_size,
+                    "use_full_dataset": use_full_dataset,
+                },
+            },
+        }
+        if optimization_id:
+            experiment_config["optimization_id"] = optimization_id
+
+        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]
+
+            # --- 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"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)}")
+                raise
+
+            # --- Step 3: Clean the model's output before metric evaluation ---
+            cleaned_model_output = raw_model_output.strip()
+
+            result = {
+                mappers.EVALUATED_LLM_TASK_OUTPUT: cleaned_model_output,
+            }
+            return result
+
+        # 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__}"
+        )
+        score = task_evaluator.evaluate(
+            dataset=dataset,
+            metric=metric,
+            evaluated_task=llm_task,
+            num_threads=self.num_threads,
+            project_name=self.project_name,
+            n_samples=subset_size,  # Use subset_size for trials, None for full dataset
+            experiment_config=experiment_config,
+            optimization_id=optimization_id,
+            verbose=self.verbose,
+        )
+        logger.debug(f"Evaluation score: {score:.4f}")
+        return score
+
+    def optimize_prompt( # type: ignore[override]
+        self,
+        prompt: chat_prompt.ChatPrompt,
+        dataset: Dataset,
+        metric: Callable,
+        experiment_config: Optional[Dict] = None,
+        n_samples: Optional[int] = None,
+        auto_continue: bool = False,
+        **kwargs,
+    ) -> 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
+            n_samples: The number of dataset items to use for evaluation
+            auto_continue: If True, the algorithm may continue if goal not met
+            **kwargs: Additional arguments for evaluation
+
+        Returns:
+            OptimizationResult: Structured result containing optimization details
+        """
+        reporting.display_header(self.__class__.__name__, verbose=self.verbose)
+        
+        total_items = len(dataset.get_items())
+        if n_samples is not None and n_samples > total_items:
+            logger.warning(
+                f"Requested n_samples ({n_samples}) is larger than dataset size ({total_items}). Using full dataset."
+            )
+            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__,
+                metadata={"optimizer": self.__class__.__name__},
+            )
+            logger.debug(f"Created optimization with ID: {optimization.id}")
+        except Exception as e:
+            logger.warning(
+                f"Opik server does not support optimizations: {e}. Please upgrade opik."
+            )
+            optimization = None
+
+        try:
+            result = self._optimize_prompt(
+                optimization_id=optimization.id if optimization is not None else None,
+                prompt=prompt,
+                dataset=dataset,
+                metric=metric,
+                experiment_config=experiment_config,
+                n_samples=n_samples,
+                auto_continue=auto_continue,
+                **kwargs,
+            )
+            if optimization:
+                self.update_optimization(optimization, status="completed")
+                logger.debug("Optimization completed successfully")
+            return result
+        except Exception as e:
+            logger.error(f"Optimization failed: {e}")
+            if optimization:
+                self.update_optimization(optimization, status="cancelled")
+                logger.debug("Optimization marked as cancelled")
+            raise e
+
+    def _optimize_prompt(
+        self,
+        optimization_id: str,
+        prompt: chat_prompt.ChatPrompt,
+        dataset: Dataset,
+        metric: Callable,
+        experiment_config: Optional[Dict],
+        n_samples: int,
+        auto_continue: bool,
+        **kwargs,
+    ) -> OptimizationResult:
+        self.auto_continue = auto_continue
+        self.dataset = dataset
+        self.prompt = prompt
+        self.llm_call_counter = 0 # Reset counter for run
+
+        current_prompt = prompt.formatted_messages
+        experiment_config = experiment_config or {}
+        experiment_config = {
+            **experiment_config,
+            **{
+                "optimizer": self.__class__.__name__,
+                "metric": metric.__name__,
+                "dataset": self.dataset.name,
+                "configuration": {
+                    "prompt": current_prompt,
+                    "rounds": self.rounds,
+                    "num_prompts_per_round": self.num_prompts_per_round,
+                },
+            },
+        }
+
+        with reporting.display_evaluation(verbose=self.verbose) as baseline_reporter:
+            initial_score = self.evaluate_prompt(
+                prompt=prompt,
+                optimization_id=optimization_id,
+                dataset=dataset,
+                metric=metric,
+                n_samples=n_samples,
+                experiment_config=experiment_config,
+                use_full_dataset=n_samples is None,
+                verbose=self.verbose,
+            )
+            best_score = initial_score
+            best_prompt = current_prompt
+            rounds = []
+
+            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:
+            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(
+                        current_prompt=best_prompt,
+                        best_score=best_score,
+                        round_num=round_num,
+                        previous_rounds=rounds,
+                        metric=metric,
+                        optimization_id=optimization_id,
+                    )
+                except Exception as e:
+                    round_reporter.failed_to_generate(self.num_prompts_per_round, e)
+                    continue
+
+                # Step 2. Score each candidate prompt
+                prompt_scores = []
+                for candidate_count, prompt in enumerate(candidate_prompts):
+                    with reporting.display_prompt_candidate_scoring_report(candidate_count, prompt, verbose=self.verbose) as eval_report:
+                        eval_report.set_generated_prompts(candidate_count, prompt)
+
+                        try:
+                            prompt_score = self.evaluate_prompt(
+                                prompt=chat_prompt.ChatPrompt(messages=prompt),
+                                optimization_id=optimization_id,
+                                dataset=dataset,
+                                metric=metric,
+                                n_samples=n_samples,
+                                use_full_dataset=False,
+                                experiment_config=experiment_config,
+                                verbose=self.verbose,
+                            )
+
+                            eval_report.set_final_score(best_score, prompt_score)
+                        except Exception as e:
+                            raise ValueError(f"Error evaluating candidate prompt: {e}")
+
+                    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")
+                    break
+
+                prompt_scores.sort(key=lambda x: x[1], reverse=True)
+                best_candidate_this_round, best_cand_score_avg = (
+                    prompt_scores[0]
+                )
+                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_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),
+                    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())
+
+                if improvement > 0:
+                    best_score = best_cand_score_avg
+                    best_prompt = best_candidate_this_round
+
+        reporting.display_result(
+            initial_score,
+            best_score,
+            best_prompt,
+            verbose=self.verbose
+        )
+
+        return self._create_result(
+            metric,
+            prompt,
+            best_prompt,
+            best_score,
+            initial_score,
+            rounds,
+        )
+
+    def _calculate_improvement(
+        self, current_score: float, previous_score: float
+    ) -> float:
+        """Calculate the improvement percentage between scores."""
+        return (
+            (current_score - previous_score) / previous_score
+            if previous_score > 0
+            else 0
+        )
+
+    def _create_round_data(
+        self,
+        round_num: int,
+        current_best_prompt: chat_prompt.ChatPrompt,
+        current_best_score: float,
+        best_prompt_overall: chat_prompt.ChatPrompt,
+        evaluated_candidates: List[tuple[str, float, List[float]]],
+        previous_best_score: float,
+        improvement_this_round: float,
+    ) -> OptimizationRound:
+        """Create an OptimizationRound object with the current round's data."""
+        generated_prompts_log = []
+        for prompt, score in evaluated_candidates:
+            improvement_vs_prev = self._calculate_improvement(
+                score, previous_best_score
+            )
+            generated_prompts_log.append(
+                {
+                    "prompt": prompt,
+                    "score": score,
+                    "improvement": improvement_vs_prev,
+                }
+            )
+
+        return OptimizationRound(
+            round_number=round_num + 1,
+            current_prompt=current_best_prompt,
+            current_score=current_best_score,
+            generated_prompts=generated_prompts_log,
+            best_prompt=best_prompt_overall,
+            best_score=current_best_score,
+            improvement=improvement_this_round,
+        )
+
+    def _create_result(
+        self,
+        metric: Callable,
+        prompt: chat_prompt.ChatPrompt,
+        best_prompt: str,
+        best_score: float,
+        initial_score: float,
+        rounds: List[OptimizationRound],
+    ) -> 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__,
+            "model": self.model,
+            "temperature": self.model_kwargs.get("temperature"),
+        }
+
+        return OptimizationResult(
+            optimizer=self.__class__.__name__,
+            prompt=best_prompt,
+            score=best_score,
+            metric_name=metric.__name__,
+            details=details,
+            llm_calls=self.llm_call_counter
+        )
+
+    def _get_task_context(self, metric: Callable) -> str:
+        """Get task-specific context from the dataset and metric configuration."""
+        if self.dataset is None:
+            return ""
+
+        try:
+            # Try get_items() first as it's the preferred method
+            items = self.dataset.get_items()
+            sample = items[0]  # Get first sample
+        except Exception as e:
+            logger.warning(f"Could not get sample from dataset: {e}")
+
+        # Describe Single Metric
+        if sample is not None:
+            metric_name = metric.__name__
+            description = metric.__doc__ or "No description available."
+
+            metrics_str = f"- {metric_name}: {description}"
+
+            context = "\nTask Context:\n"
+            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(
+        self,
+        current_prompt: chat_prompt.ChatPrompt,
+        best_score: float,
+        round_num: int,
+        previous_rounds: List[OptimizationRound],
+        metric: Callable,
+        optimization_id: Optional[str] = None,
+    ) -> List[str]:
+        """Generate candidate prompts using meta-prompting."""
+        with reporting.display_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"Current best score: {best_score:.4f}")
+
+            history_context = self._build_history_context(previous_rounds)
+            task_context_str = ""
+            analysis_instruction = ""
+            metric_focus_instruction = ""
+            improvement_point_1 = ""
+
+            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__}."
+                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.")
+            else:
+                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.")
+
+            user_prompt = f"""Current prompt: {current_prompt}
+            Current score: {best_score}
+            {history_context}
+            {task_context_str}
+
+            {analysis_instruction}
+            Generate {self.num_prompts_per_round} improved versions of this prompt.
+            {metric_focus_instruction}
+            Each version should aim to:
+            {improvement_point_1}
+            2. Provide necessary context and constraints (if applicable, without relying on disabled external context).
+            3. Guide the model to produce the desired output format suitable for the task.
+            4. Remove ambiguity and unnecessary elements.
+            5. Maintain conciseness while being complete.
+
+            Return a valid JSON array as specified."""
+
+            try:
+                # Use _call_model which handles selecting reasoning_model
+                content = self._call_model(
+                    messages=[
+                        {"role": "system", "content": self._REASONING_SYSTEM_PROMPT},
+                        {"role": "user", "content": user_prompt}
+                    ],
+                    is_reasoning=True,
+                    optimization_id=optimization_id,
+                )
+                logger.debug(f"Raw response from reasoning model: {content}")
+
+                # --- Robust JSON Parsing and Validation ---
+                json_result = None
+                try:
+                    # Try direct JSON parsing
+                    json_result = json.loads(content)
+                except json.JSONDecodeError:
+                    import re
+
+                    json_match = re.search(r"\{.*\}", content, re.DOTALL)
+                    if json_match:
+                        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()}")
+                    else:
+                        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:
+                    json_result = json_result[0]
+
+                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}")
+
+                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')}")
+
+                # Extract and log valid prompts
+                valid_prompts = []
+                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)
+                        
+                        # 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"  Improvement focus: {focus}")
+                        logger.debug(f"  Reasoning: {reasoning}")
+                    else:
+                        logger.warning(
+                            f"Skipping invalid prompt item structure in JSON response: {item}"
+                        )
+
+                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
+                )
+                
+                return valid_prompts
+                # --- End Robust Parsing ---
+
+            except Exception as 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."""
+        if not previous_rounds:
+            return ""
+
+        context = "\nPrevious rounds (latest first):\n"
+        for round_data in reversed(previous_rounds[-3:]):
+            context += f"\nRound {round_data.round_number}:\n"
+            context += f"Best score this round: {round_data.best_score:.4f}\n"
+            context += "Generated prompts this round (best first):\n"
+
+            sorted_generated = sorted(
+                round_data.generated_prompts,
+                key=lambda p: p.get("score", -float("inf")),
+                reverse=True,
+            )
+
+            for p in sorted_generated[:3]:
+                prompt_text = p.get("prompt", "N/A")
+                score = p.get("score", float("nan"))
+                context += f"- Prompt: {prompt_text[:150]}...\n"
+                context += f"  Avg Score: {score:.4f}\n"
+        return context
+
+    def _get_evaluation_subset(
+        self, dataset: opik.Dataset, min_size: int = 20, max_size: int = 100
+    ) -> List[Dict[str, Any]]:
+        """Get a random subset of the dataset for evaluation.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dataset items to evaluate against
+        """
+        try:
+            # Get all items from the dataset
+            all_items = dataset.get_items()
+            if not all_items:
+                return all_items
+
+            # Calculate subset size
+            total_size = len(all_items)
+            subset_size = min(max(min_size, int(total_size * 0.2)), max_size)
+
+            # Get random subset of items
+            import random
+
+            return random.sample(all_items, subset_size)
+
+        except Exception as e:
+            logger.warning(f"Could not create evaluation subset: {e}")
+            return all_items