local-deep-research 0.3.12__py3-none-any.whl → 0.4.1__py3-none-any.whl

local_deep_research/benchmarks/optimization/optuna_optimizer.py

@@ -0,0 +1,1163 @@
+"""
+Optuna-based parameter optimizer for Local Deep Research.
+
+This module provides the core optimization functionality using Optuna
+to find optimal parameters for the research system, balancing quality
+and performance metrics.
+"""
+
+import json
+import logging
+import os
+import time
+from datetime import datetime
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import joblib
+import numpy as np
+import optuna
+from optuna.visualization import (
+    plot_contour,
+    plot_optimization_history,
+    plot_param_importances,
+    plot_slice,
+)
+
+from local_deep_research.benchmarks.efficiency.speed_profiler import SpeedProfiler
+from local_deep_research.benchmarks.evaluators import CompositeBenchmarkEvaluator
+
+# Import benchmark evaluator components
+
+logger = logging.getLogger(__name__)
+
+# Try to import visualization libraries, but don't fail if not available
+try:
+    import matplotlib.pyplot as plt
+    from matplotlib.lines import Line2D
+
+    # We'll use matplotlib for plotting visualization results
+
+    PLOTTING_AVAILABLE = True
+except ImportError:
+    PLOTTING_AVAILABLE = False
+    logger.warning("Matplotlib not available, visualization will be limited")
+
+
+class OptunaOptimizer:
+    """
+    Optimize parameters for Local Deep Research using Optuna.
+
+    This class provides functionality to:
+    1. Define search spaces for parameter optimization
+    2. Evaluate parameter combinations using objective functions
+    3. Find optimal parameters via Optuna
+    4. Visualize and analyze optimization results
+    """
+
+    def __init__(
+        self,
+        base_query: str,
+        output_dir: str = "optimization_results",
+        model_name: Optional[str] = None,
+        provider: Optional[str] = None,
+        search_tool: Optional[str] = None,
+        temperature: float = 0.7,
+        n_trials: int = 30,
+        timeout: Optional[int] = None,
+        n_jobs: int = 1,
+        study_name: Optional[str] = None,
+        optimization_metrics: Optional[List[str]] = None,
+        metric_weights: Optional[Dict[str, float]] = None,
+        progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
+        benchmark_weights: Optional[Dict[str, float]] = None,
+    ):
+        """
+        Initialize the optimizer.
+
+        Args:
+            base_query: The research query to use for all experiments
+            output_dir: Directory to save optimization results
+            model_name: Name of the LLM model to use
+            provider: LLM provider
+            search_tool: Search engine to use
+            temperature: LLM temperature
+            n_trials: Number of parameter combinations to try
+            timeout: Maximum seconds to run optimization (None for no limit)
+            n_jobs: Number of parallel jobs for optimization
+            study_name: Name of the Optuna study
+            optimization_metrics: List of metrics to optimize (default: ["quality", "speed"])
+            metric_weights: Dictionary of weights for each metric (e.g., {"quality": 0.6, "speed": 0.4})
+            progress_callback: Optional callback for progress updates
+            benchmark_weights: Dictionary mapping benchmark types to weights
+                (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
+                If None, only SimpleQA is used with weight 1.0
+        """
+        self.base_query = base_query
+        self.output_dir = output_dir
+        self.model_name = model_name
+        self.provider = provider
+        self.search_tool = search_tool
+        self.temperature = temperature
+        self.n_trials = n_trials
+        self.timeout = timeout
+        self.n_jobs = n_jobs
+        self.optimization_metrics = optimization_metrics or ["quality", "speed"]
+        self.metric_weights = metric_weights or {"quality": 0.6, "speed": 0.4}
+        self.progress_callback = progress_callback
+
+        # Initialize benchmark evaluator with weights
+        self.benchmark_weights = benchmark_weights or {"simpleqa": 1.0}
+        self.benchmark_evaluator = CompositeBenchmarkEvaluator(self.benchmark_weights)
+
+        # Normalize weights to sum to 1.0
+        total_weight = sum(self.metric_weights.values())
+        if total_weight > 0:
+            self.metric_weights = {
+                k: v / total_weight for k, v in self.metric_weights.items()
+            }
+
+        # Generate a unique study name if not provided
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.study_name = study_name or f"ldr_opt_{timestamp}"
+
+        # Create output directory
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Store the trial history for analysis
+        self.trials_history = []
+
+        # Storage for the best parameters and study
+        self.best_params = None
+        self.study = None
+
+    def optimize(
+        self, param_space: Optional[Dict[str, Any]] = None
+    ) -> Tuple[Dict[str, Any], float]:
+        """
+        Run the optimization process using Optuna.
+
+        Args:
+            param_space: Dictionary defining parameter search spaces
+                         (if None, use default spaces)
+
+        Returns:
+            Tuple containing (best_parameters, best_score)
+        """
+        param_space = param_space or self._get_default_param_space()
+
+        # Create a study object
+        storage_name = f"sqlite:///{self.output_dir}/{self.study_name}.db"
+        self.study = optuna.create_study(
+            study_name=self.study_name,
+            storage=storage_name,
+            load_if_exists=True,
+            direction="maximize",
+            sampler=optuna.samplers.TPESampler(seed=42),
+        )
+
+        # Create partial function with param_space
+        objective = partial(self._objective, param_space=param_space)
+
+        # Log optimization start
+        logger.info(
+            f"Starting optimization with {self.n_trials} trials, {self.n_jobs} parallel jobs"
+        )
+        logger.info(f"Parameter space: {param_space}")
+        logger.info(f"Metric weights: {self.metric_weights}")
+        logger.info(f"Benchmark weights: {self.benchmark_weights}")
+
+        # Initialize progress tracking
+        if self.progress_callback:
+            self.progress_callback(
+                0,
+                self.n_trials,
+                {
+                    "status": "starting",
+                    "stage": "initialization",
+                    "trials_completed": 0,
+                    "total_trials": self.n_trials,
+                },
+            )
+
+        try:
+            # Run optimization
+            self.study.optimize(
+                objective,
+                n_trials=self.n_trials,
+                timeout=self.timeout,
+                n_jobs=self.n_jobs,
+                callbacks=[self._optimization_callback],
+                show_progress_bar=True,
+            )
+
+            # Store best parameters
+            self.best_params = self.study.best_params
+
+            # Save the results
+            self._save_results()
+
+            # Create visualizations
+            self._create_visualizations()
+
+            logger.info(f"Optimization complete. Best parameters: {self.best_params}")
+            logger.info(f"Best value: {self.study.best_value}")
+
+            # Report completion
+            if self.progress_callback:
+                self.progress_callback(
+                    self.n_trials,
+                    self.n_trials,
+                    {
+                        "status": "completed",
+                        "stage": "finished",
+                        "trials_completed": len(self.study.trials),
+                        "total_trials": self.n_trials,
+                        "best_params": self.best_params,
+                        "best_value": self.study.best_value,
+                    },
+                )
+
+            return self.best_params, self.study.best_value
+
+        except KeyboardInterrupt:
+            logger.info("Optimization interrupted by user")
+            # Still save what we have
+            self._save_results()
+            self._create_visualizations()
+
+            # Report interruption
+            if self.progress_callback:
+                self.progress_callback(
+                    len(self.study.trials),
+                    self.n_trials,
+                    {
+                        "status": "interrupted",
+                        "stage": "interrupted",
+                        "trials_completed": len(self.study.trials),
+                        "total_trials": self.n_trials,
+                        "best_params": self.study.best_params,
+                        "best_value": self.study.best_value,
+                    },
+                )
+
+            return self.study.best_params, self.study.best_value
+
+    def _get_default_param_space(self) -> Dict[str, Any]:
+        """
+        Get default parameter search space.
+
+        Returns:
+            Dictionary defining the default parameter search spaces
+        """
+        return {
+            "iterations": {
+                "type": "int",
+                "low": 1,
+                "high": 5,
+                "step": 1,
+            },
+            "questions_per_iteration": {
+                "type": "int",
+                "low": 1,
+                "high": 5,
+                "step": 1,
+            },
+            "search_strategy": {
+                "type": "categorical",
+                "choices": [
+                    "iterdrag",
+                    "standard",
+                    "rapid",
+                    "parallel",
+                    "source_based",
+                ],
+            },
+            "max_results": {
+                "type": "int",
+                "low": 10,
+                "high": 100,
+                "step": 10,
+            },
+        }
+
+    def _objective(self, trial: optuna.Trial, param_space: Dict[str, Any]) -> float:
+        """
+        Objective function for Optuna optimization.
+
+        Args:
+            trial: Optuna trial object
+            param_space: Dictionary defining parameter search spaces
+
+        Returns:
+            Score to maximize
+        """
+        # Generate parameters for this trial
+        params = {}
+        for param_name, param_config in param_space.items():
+            param_type = param_config["type"]
+
+            if param_type == "int":
+                params[param_name] = trial.suggest_int(
+                    param_name,
+                    param_config["low"],
+                    param_config["high"],
+                    step=param_config.get("step", 1),
+                )
+            elif param_type == "float":
+                params[param_name] = trial.suggest_float(
+                    param_name,
+                    param_config["low"],
+                    param_config["high"],
+                    step=param_config.get("step"),
+                    log=param_config.get("log", False),
+                )
+            elif param_type == "categorical":
+                params[param_name] = trial.suggest_categorical(
+                    param_name, param_config["choices"]
+                )
+
+        # Log the trial parameters
+        logger.info(f"Trial {trial.number}: {params}")
+
+        # Update progress callback if available
+        if self.progress_callback:
+            self.progress_callback(
+                trial.number,
+                self.n_trials,
+                {
+                    "status": "running",
+                    "stage": "trial_started",
+                    "trial_number": trial.number,
+                    "params": params,
+                    "trials_completed": trial.number,
+                    "total_trials": self.n_trials,
+                },
+            )
+
+        # Run an experiment with these parameters
+        try:
+            start_time = time.time()
+            result = self._run_experiment(params)
+            duration = time.time() - start_time
+
+            # Store details about the trial
+            trial_info = {
+                "trial_number": trial.number,
+                "params": params,
+                "result": result,
+                "score": result.get("score", 0),
+                "duration": duration,
+                "timestamp": datetime.now().isoformat(),
+            }
+            self.trials_history.append(trial_info)
+
+            # Update callback with results
+            if self.progress_callback:
+                self.progress_callback(
+                    trial.number,
+                    self.n_trials,
+                    {
+                        "status": "completed",
+                        "stage": "trial_completed",
+                        "trial_number": trial.number,
+                        "params": params,
+                        "score": result.get("score", 0),
+                        "trials_completed": trial.number + 1,
+                        "total_trials": self.n_trials,
+                    },
+                )
+
+            logger.info(
+                f"Trial {trial.number} completed: {params}, score: {result['score']:.4f}"
+            )
+
+            return result["score"]
+        except Exception as e:
+            logger.error(f"Error in trial {trial.number}: {str(e)}")
+
+            # Update callback with error
+            if self.progress_callback:
+                self.progress_callback(
+                    trial.number,
+                    self.n_trials,
+                    {
+                        "status": "error",
+                        "stage": "trial_error",
+                        "trial_number": trial.number,
+                        "params": params,
+                        "error": str(e),
+                        "trials_completed": trial.number,
+                        "total_trials": self.n_trials,
+                    },
+                )
+
+            return float("-inf")  # Return a very low score for failed trials
+
+    def _run_experiment(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Run a single experiment with the given parameters.
+
+        Args:
+            params: Dictionary of parameters to test
+
+        Returns:
+            Results dictionary with metrics and score
+        """
+        # Extract parameters
+        iterations = params.get("iterations", 2)
+        questions_per_iteration = params.get("questions_per_iteration", 2)
+        search_strategy = params.get("search_strategy", "iterdrag")
+        max_results = params.get("max_results", 50)
+
+        # Initialize profiling tools
+        speed_profiler = SpeedProfiler()
+
+        # Start profiling
+        speed_profiler.start()
+
+        try:
+            # Create system configuration
+            system_config = {
+                "iterations": iterations,
+                "questions_per_iteration": questions_per_iteration,
+                "search_strategy": search_strategy,
+                "search_tool": self.search_tool,
+                "max_results": max_results,
+                "model_name": self.model_name,
+                "provider": self.provider,
+            }
+
+            # Evaluate quality using composite benchmark evaluator
+            # Use a small number of examples for efficiency
+            benchmark_dir = os.path.join(self.output_dir, "benchmark_temp")
+            quality_results = self.benchmark_evaluator.evaluate(
+                system_config=system_config,
+                num_examples=5,  # Small number for optimization efficiency
+                output_dir=benchmark_dir,
+            )
+
+            # Stop timing
+            speed_profiler.stop()
+            timing_results = speed_profiler.get_summary()
+
+            # Extract key metrics
+            quality_score = quality_results.get("quality_score", 0.0)
+            benchmark_results = quality_results.get("benchmark_results", {})
+
+            # Speed score: convert duration to a 0-1 score where faster is better
+            # Using a reasonable threshold (e.g., 180 seconds for 5 examples)
+            # Below this threshold: high score, above it: declining score
+            total_duration = timing_results.get("total_duration", 180)
+            speed_score = max(0.0, min(1.0, 1.0 - (total_duration - 60) / 180))
+
+            # Calculate combined score based on weights
+            combined_score = (
+                self.metric_weights.get("quality", 0.6) * quality_score
+                + self.metric_weights.get("speed", 0.4) * speed_score
+            )
+
+            # Return streamlined results
+            return {
+                "quality_score": quality_score,
+                "benchmark_results": benchmark_results,
+                "speed_score": speed_score,
+                "total_duration": total_duration,
+                "score": combined_score,
+                "success": True,
+            }
+
+        except Exception as e:
+            # Stop profiling on error
+            speed_profiler.stop()
+
+            # Log error
+            logger.error(f"Error in experiment: {str(e)}")
+
+            # Return error information
+            return {"error": str(e), "score": 0.0, "success": False}
+
+    def _optimization_callback(self, study: optuna.Study, trial: optuna.Trial):
+        """
+        Callback for the Optuna optimization process.
+
+        Args:
+            study: Optuna study object
+            trial: Current trial
+        """
+        # Save intermediate results periodically
+        if trial.number % 10 == 0 and trial.number > 0:
+            self._save_results()
+            self._create_quick_visualizations()
+
+    def _save_results(self):
+        """Save the optimization results to disk."""
+        # Create a timestamp for filenames
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+        # Save trial history
+        history_file = os.path.join(self.output_dir, f"{self.study_name}_history.json")
+        with open(history_file, "w") as f:
+            # Convert numpy values to native Python types for JSON serialization
+            clean_history = []
+            for trial in self.trials_history:
+                clean_trial = {}
+                for k, v in trial.items():
+                    if isinstance(v, dict):
+                        clean_trial[k] = {
+                            dk: (float(dv) if isinstance(dv, np.number) else dv)
+                            for dk, dv in v.items()
+                        }
+                    elif isinstance(v, np.number):
+                        clean_trial[k] = float(v)
+                    else:
+                        clean_trial[k] = v
+                clean_history.append(clean_trial)
+
+            json.dump(clean_history, f, indent=2)
+
+        # Save current best parameters
+        if self.study and hasattr(self.study, "best_params") and self.study.best_params:
+            best_params_file = os.path.join(
+                self.output_dir, f"{self.study_name}_best_params.json"
+            )
+            with open(best_params_file, "w") as f:
+                json.dump(
+                    {
+                        "best_params": self.study.best_params,
+                        "best_value": float(self.study.best_value),
+                        "n_trials": len(self.study.trials),
+                        "timestamp": timestamp,
+                        "base_query": self.base_query,
+                        "model_name": self.model_name,
+                        "provider": self.provider,
+                        "search_tool": self.search_tool,
+                        "metric_weights": self.metric_weights,
+                        "benchmark_weights": self.benchmark_weights,
+                    },
+                    f,
+                    indent=2,
+                )
+
+        # Save the Optuna study
+        if self.study:
+            study_file = os.path.join(self.output_dir, f"{self.study_name}_study.pkl")
+            joblib.dump(self.study, study_file)
+
+        logger.info(f"Results saved to {self.output_dir}")
+
+    def _create_visualizations(self):
+        """Create and save comprehensive visualizations of the optimization results."""
+        if not PLOTTING_AVAILABLE:
+            logger.warning("Matplotlib not available, skipping visualization creation")
+            return
+
+        if not self.study or len(self.study.trials) < 2:
+            logger.warning("Not enough trials to create visualizations")
+            return
+
+        # Create directory for visualizations
+        viz_dir = os.path.join(self.output_dir, "visualizations")
+        os.makedirs(viz_dir, exist_ok=True)
+
+        # Create Optuna visualizations
+        self._create_optuna_visualizations(viz_dir)
+
+        # Create custom visualizations
+        self._create_custom_visualizations(viz_dir)
+
+        logger.info(f"Visualizations saved to {viz_dir}")
+
+    def _create_quick_visualizations(self):
+        """Create a smaller set of visualizations for intermediate progress."""
+        if not PLOTTING_AVAILABLE or not self.study or len(self.study.trials) < 2:
+            return
+
+        # Create directory for visualizations
+        viz_dir = os.path.join(self.output_dir, "visualizations")
+        os.makedirs(viz_dir, exist_ok=True)
+
+        # Create optimization history only (faster than full visualization)
+        try:
+            fig = plot_optimization_history(self.study)
+            fig.write_image(
+                os.path.join(
+                    viz_dir, f"{self.study_name}_optimization_history_current.png"
+                )
+            )
+        except Exception as e:
+            logger.error(f"Error creating optimization history plot: {str(e)}")
+
+    def _create_optuna_visualizations(self, viz_dir: str):
+        """
+        Create and save Optuna's built-in visualizations.
+
+        Args:
+            viz_dir: Directory to save visualizations
+        """
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+        # 1. Optimization history
+        try:
+            fig = plot_optimization_history(self.study)
+            fig.write_image(
+                os.path.join(
+                    viz_dir, f"{self.study_name}_optimization_history_{timestamp}.png"
+                )
+            )
+        except Exception as e:
+            logger.error(f"Error creating optimization history plot: {str(e)}")
+
+        # 2. Parameter importances
+        try:
+            fig = plot_param_importances(self.study)
+            fig.write_image(
+                os.path.join(
+                    viz_dir, f"{self.study_name}_param_importances_{timestamp}.png"
+                )
+            )
+        except Exception as e:
+            logger.error(f"Error creating parameter importances plot: {str(e)}")
+
+        # 3. Slice plot for each parameter
+        try:
+            for param_name in self.study.best_params.keys():
+                fig = plot_slice(self.study, [param_name])
+                fig.write_image(
+                    os.path.join(
+                        viz_dir, f"{self.study_name}_slice_{param_name}_{timestamp}.png"
+                    )
+                )
+        except Exception as e:
+            logger.error(f"Error creating slice plots: {str(e)}")
+
+        # 4. Contour plots for important parameter pairs
+        try:
+            # Get all parameter names
+            param_names = list(self.study.best_params.keys())
+
+            # Create contour plots for each pair
+            for i in range(len(param_names)):
+                for j in range(i + 1, len(param_names)):
+                    try:
+                        fig = plot_contour(
+                            self.study, params=[param_names[i], param_names[j]]
+                        )
+                        fig.write_image(
+                            os.path.join(
+                                viz_dir,
+                                f"{self.study_name}_contour_{param_names[i]}_{param_names[j]}_{timestamp}.png",
+                            )
+                        )
+                    except Exception as e:
+                        logger.warning(
+                            f"Error creating contour plot for {param_names[i]} vs {param_names[j]}: {str(e)}"
+                        )
+        except Exception as e:
+            logger.error(f"Error creating contour plots: {str(e)}")
+
+    def _create_custom_visualizations(self, viz_dir: str):
+        """
+        Create custom visualizations based on trial history.
+
+        Args:
+            viz_dir: Directory to save visualizations
+        """
+        if not self.trials_history:
+            return
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+        # Create quality vs speed plot
+        self._create_quality_vs_speed_plot(viz_dir, timestamp)
+
+        # Create parameter evolution plots
+        self._create_parameter_evolution_plots(viz_dir, timestamp)
+
+        # Create trial duration vs score plot
+        self._create_duration_vs_score_plot(viz_dir, timestamp)
+
+    def _create_quality_vs_speed_plot(self, viz_dir: str, timestamp: str):
+        """Create a plot showing quality vs. speed trade-off."""
+        if not self.trials_history:
+            return
+
+        # Extract data from successful trials
+        successful_trials = [
+            t for t in self.trials_history if t.get("result", {}).get("success", False)
+        ]
+
+        if not successful_trials:
+            logger.warning("No successful trials for visualization")
+            return
+
+        try:
+            plt.figure(figsize=(10, 8))
+
+            # Extract metrics
+            quality_scores = []
+            speed_scores = []
+            labels = []
+            iterations_values = []
+            questions_values = []
+
+            for trial in successful_trials:
+                result = trial["result"]
+                quality = result.get("quality_score", 0)
+                speed = result.get("speed_score", 0)
+                iterations = trial["params"].get("iterations", 0)
+                questions = trial["params"].get("questions_per_iteration", 0)
+
+                quality_scores.append(quality)
+                speed_scores.append(speed)
+                labels.append(f"Trial {trial['trial_number']}")
+                iterations_values.append(iterations)
+                questions_values.append(questions)
+
+            # Create scatter plot with size based on iterations*questions
+            sizes = [i * q * 5 for i, q in zip(iterations_values, questions_values)]
+            scatter = plt.scatter(
+                quality_scores,
+                speed_scores,
+                s=sizes,
+                alpha=0.7,
+                c=range(len(quality_scores)),
+                cmap="viridis",
+            )
+
+            # Highlight best trial
+            best_trial = max(
+                successful_trials, key=lambda x: x.get("result", {}).get("score", 0)
+            )
+            best_quality = best_trial["result"].get("quality_score", 0)
+            best_speed = best_trial["result"].get("speed_score", 0)
+            best_iter = best_trial["params"].get("iterations", 0)
+            best_questions = best_trial["params"].get("questions_per_iteration", 0)
+
+            plt.scatter(
+                [best_quality],
+                [best_speed],
+                s=200,
+                facecolors="none",
+                edgecolors="red",
+                linewidth=2,
+                label=f"Best: {best_iter}×{best_questions}",
+            )
+
+            # Add annotations for key points
+            for i, (q, s, l) in enumerate(zip(quality_scores, speed_scores, labels)):
+                if i % max(1, len(quality_scores) // 5) == 0:  # Label ~5 points
+                    plt.annotate(
+                        f"{iterations_values[i]}×{questions_values[i]}",
+                        (q, s),
+                        xytext=(5, 5),
+                        textcoords="offset points",
+                    )
+
+            # Add colorbar and labels
+            cbar = plt.colorbar(scatter)
+            cbar.set_label("Trial Progression")
+
+            # Add benchmark weight information
+            weights_str = ", ".join(
+                [f"{k}:{v:.1f}" for k, v in self.benchmark_weights.items()]
+            )
+            plt.title(f"Quality vs. Speed Trade-off\nBenchmark Weights: {weights_str}")
+            plt.xlabel("Quality Score (Benchmark Accuracy)")
+            plt.ylabel("Speed Score")
+            plt.grid(True, linestyle="--", alpha=0.7)
+
+            # Add legend explaining size
+            legend_elements = [
+                Line2D(
+                    [0],
+                    [0],
+                    marker="o",
+                    color="w",
+                    markerfacecolor="gray",
+                    markersize=np.sqrt(n * 5 / np.pi),
+                    label=f"{n} Total Questions",
+                )
+                for n in [5, 10, 15, 20, 25]
+            ]
+            plt.legend(handles=legend_elements, title="Workload")
+
+            # Save the figure
+            plt.tight_layout()
+            plt.savefig(
+                os.path.join(
+                    viz_dir, f"{self.study_name}_quality_vs_speed_{timestamp}.png"
+                )
+            )
+            plt.close()
+        except Exception as e:
+            logger.error(f"Error creating quality vs speed plot: {str(e)}")
+
+    def _create_parameter_evolution_plots(self, viz_dir: str, timestamp: str):
+        """Create plots showing how parameter values evolve over trials."""
+        try:
+            successful_trials = [
+                t
+                for t in self.trials_history
+                if t.get("result", {}).get("success", False)
+            ]
+
+            if not successful_trials or len(successful_trials) < 5:
+                return
+
+            # Get key parameters
+            main_params = list(successful_trials[0]["params"].keys())
+
+            # For each parameter, plot its values over trials
+            for param_name in main_params:
+                plt.figure(figsize=(12, 6))
+
+                trial_numbers = []
+                param_values = []
+                scores = []
+
+                for trial in self.trials_history:
+                    if "params" in trial and param_name in trial["params"]:
+                        trial_numbers.append(trial["trial_number"])
+                        param_values.append(trial["params"][param_name])
+                        scores.append(trial.get("score", 0))
+
+                # Create evolution plot
+                scatter = plt.scatter(
+                    trial_numbers,
+                    param_values,
+                    c=scores,
+                    cmap="plasma",
+                    alpha=0.8,
+                    s=80,
+                )
+
+                # Add best trial marker
+                best_trial_idx = scores.index(max(scores))
+                plt.scatter(
+                    [trial_numbers[best_trial_idx]],
+                    [param_values[best_trial_idx]],
+                    s=150,
+                    facecolors="none",
+                    edgecolors="red",
+                    linewidth=2,
+                    label=f"Best Value: {param_values[best_trial_idx]}",
+                )
+
+                # Add colorbar
+                cbar = plt.colorbar(scatter)
+                cbar.set_label("Score")
+
+                # Set chart properties
+                plt.title(f"Evolution of {param_name} Values")
+                plt.xlabel("Trial Number")
+                plt.ylabel(param_name)
+                plt.grid(True, linestyle="--", alpha=0.7)
+                plt.legend()
+
+                # For categorical parameters, adjust y-axis
+                if isinstance(param_values[0], str):
+                    unique_values = sorted(set(param_values))
+                    plt.yticks(range(len(unique_values)), unique_values)
+
+                # Save the figure
+                plt.tight_layout()
+                plt.savefig(
+                    os.path.join(
+                        viz_dir,
+                        f"{self.study_name}_param_evolution_{param_name}_{timestamp}.png",
+                    )
+                )
+                plt.close()
+        except Exception as e:
+            logger.error(f"Error creating parameter evolution plots: {str(e)}")
+
+    def _create_duration_vs_score_plot(self, viz_dir: str, timestamp: str):
+        """Create a plot showing trial duration vs score."""
+        try:
+            plt.figure(figsize=(10, 6))
+
+            successful_trials = [
+                t
+                for t in self.trials_history
+                if t.get("result", {}).get("success", False)
+            ]
+
+            if not successful_trials:
+                return
+
+            trial_durations = []
+            trial_scores = []
+            trial_iterations = []
+            trial_questions = []
+
+            for trial in successful_trials:
+                duration = trial.get("duration", 0)
+                score = trial.get("score", 0)
+                iterations = trial.get("params", {}).get("iterations", 1)
+                questions = trial.get("params", {}).get("questions_per_iteration", 1)
+
+                trial_durations.append(duration)
+                trial_scores.append(score)
+                trial_iterations.append(iterations)
+                trial_questions.append(questions)
+
+            # Total questions per trial
+            total_questions = [i * q for i, q in zip(trial_iterations, trial_questions)]
+
+            # Create scatter plot with size based on total questions
+            plt.scatter(
+                trial_durations,
+                trial_scores,
+                s=[q * 5 for q in total_questions],  # Size based on total questions
+                alpha=0.7,
+                c=range(len(trial_durations)),
+                cmap="viridis",
+            )
+
+            # Add labels
+            plt.xlabel("Trial Duration (seconds)")
+            plt.ylabel("Score")
+            plt.title("Trial Duration vs. Score")
+            plt.grid(True, linestyle="--", alpha=0.7)
+
+            # Add trial number annotations for selected points
+            for i, (d, s) in enumerate(zip(trial_durations, trial_scores)):
+                if i % max(1, len(trial_durations) // 5) == 0:  # Annotate ~5 points
+                    plt.annotate(
+                        f"{trial_iterations[i]}×{trial_questions[i]}",
+                        (d, s),
+                        xytext=(5, 5),
+                        textcoords="offset points",
+                    )
+
+            # Save the figure
+            plt.tight_layout()
+            plt.savefig(
+                os.path.join(
+                    viz_dir, f"{self.study_name}_duration_vs_score_{timestamp}.png"
+                )
+            )
+            plt.close()
+        except Exception as e:
+            logger.error(f"Error creating duration vs score plot: {str(e)}")
+
+
+def optimize_parameters(
+    query: str,
+    param_space: Optional[Dict[str, Any]] = None,
+    output_dir: str = os.path.join("data", "optimization_results"),
+    model_name: Optional[str] = None,
+    provider: Optional[str] = None,
+    search_tool: Optional[str] = None,
+    temperature: float = 0.7,
+    n_trials: int = 30,
+    timeout: Optional[int] = None,
+    n_jobs: int = 1,
+    study_name: Optional[str] = None,
+    optimization_metrics: Optional[List[str]] = None,
+    metric_weights: Optional[Dict[str, float]] = None,
+    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
+    benchmark_weights: Optional[Dict[str, float]] = None,
+) -> Tuple[Dict[str, Any], float]:
+    """
+    Optimize parameters for Local Deep Research.
+
+    Args:
+        query: The research query to use for all experiments
+        param_space: Dictionary defining parameter search spaces (optional)
+        output_dir: Directory to save optimization results
+        model_name: Name of the LLM model to use
+        provider: LLM provider
+        search_tool: Search engine to use
+        temperature: LLM temperature
+        n_trials: Number of parameter combinations to try
+        timeout: Maximum seconds to run optimization (None for no limit)
+        n_jobs: Number of parallel jobs for optimization
+        study_name: Name of the Optuna study
+        optimization_metrics: List of metrics to optimize (default: ["quality", "speed"])
+        metric_weights: Dictionary of weights for each metric (e.g., {"quality": 0.6, "speed": 0.4})
+        progress_callback: Optional callback for progress updates
+        benchmark_weights: Dictionary mapping benchmark types to weights
+            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
+            If None, only SimpleQA is used with weight 1.0
+
+    Returns:
+        Tuple of (best_parameters, best_score)
+    """
+    # Create optimizer
+    optimizer = OptunaOptimizer(
+        base_query=query,
+        output_dir=output_dir,
+        model_name=model_name,
+        provider=provider,
+        search_tool=search_tool,
+        temperature=temperature,
+        n_trials=n_trials,
+        timeout=timeout,
+        n_jobs=n_jobs,
+        study_name=study_name,
+        optimization_metrics=optimization_metrics,
+        metric_weights=metric_weights,
+        progress_callback=progress_callback,
+        benchmark_weights=benchmark_weights,
+    )
+
+    # Run optimization
+    return optimizer.optimize(param_space)
+
+
+def optimize_for_speed(
+    query: str,
+    n_trials: int = 20,
+    output_dir: str = os.path.join("data", "optimization_results"),
+    model_name: Optional[str] = None,
+    provider: Optional[str] = None,
+    search_tool: Optional[str] = None,
+    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
+    benchmark_weights: Optional[Dict[str, float]] = None,
+) -> Tuple[Dict[str, Any], float]:
+    """
+    Optimize parameters with a focus on speed performance.
+
+    Args:
+        query: The research query to use for all experiments
+        n_trials: Number of parameter combinations to try
+        output_dir: Directory to save optimization results
+        model_name: Name of the LLM model to use
+        provider: LLM provider
+        search_tool: Search engine to use
+        progress_callback: Optional callback for progress updates
+        benchmark_weights: Dictionary mapping benchmark types to weights
+            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
+            If None, only SimpleQA is used with weight 1.0
+
+    Returns:
+        Tuple of (best_parameters, best_score)
+    """
+    # Focus on speed with reduced parameter space
+    param_space = {
+        "iterations": {
+            "type": "int",
+            "low": 1,
+            "high": 3,
+            "step": 1,
+        },
+        "questions_per_iteration": {
+            "type": "int",
+            "low": 1,
+            "high": 3,
+            "step": 1,
+        },
+        "search_strategy": {
+            "type": "categorical",
+            "choices": ["rapid", "parallel", "source_based"],
+        },
+    }
+
+    # Speed-focused weights
+    metric_weights = {"speed": 0.8, "quality": 0.2}
+
+    return optimize_parameters(
+        query=query,
+        param_space=param_space,
+        output_dir=output_dir,
+        model_name=model_name,
+        provider=provider,
+        search_tool=search_tool,
+        n_trials=n_trials,
+        metric_weights=metric_weights,
+        optimization_metrics=["speed", "quality"],
+        progress_callback=progress_callback,
+        benchmark_weights=benchmark_weights,
+    )
+
+
+def optimize_for_quality(
+    query: str,
+    n_trials: int = 30,
+    output_dir: str = os.path.join("data", "optimization_results"),
+    model_name: Optional[str] = None,
+    provider: Optional[str] = None,
+    search_tool: Optional[str] = None,
+    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
+    benchmark_weights: Optional[Dict[str, float]] = None,
+) -> Tuple[Dict[str, Any], float]:
+    """
+    Optimize parameters with a focus on result quality.
+
+    Args:
+        query: The research query to use for all experiments
+        n_trials: Number of parameter combinations to try
+        output_dir: Directory to save optimization results
+        model_name: Name of the LLM model to use
+        provider: LLM provider
+        search_tool: Search engine to use
+        progress_callback: Optional callback for progress updates
+        benchmark_weights: Dictionary mapping benchmark types to weights
+            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
+            If None, only SimpleQA is used with weight 1.0
+
+    Returns:
+        Tuple of (best_parameters, best_score)
+    """
+    # Quality-focused weights
+    metric_weights = {"quality": 0.9, "speed": 0.1}
+
+    return optimize_parameters(
+        query=query,
+        output_dir=output_dir,
+        model_name=model_name,
+        provider=provider,
+        search_tool=search_tool,
+        n_trials=n_trials,
+        metric_weights=metric_weights,
+        optimization_metrics=["quality", "speed"],
+        progress_callback=progress_callback,
+        benchmark_weights=benchmark_weights,
+    )
+
+
+def optimize_for_efficiency(
+    query: str,
+    n_trials: int = 25,
+    output_dir: str = os.path.join("data", "optimization_results"),
+    model_name: Optional[str] = None,
+    provider: Optional[str] = None,
+    search_tool: Optional[str] = None,
+    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
+    benchmark_weights: Optional[Dict[str, float]] = None,
+) -> Tuple[Dict[str, Any], float]:
+    """
+    Optimize parameters with a focus on resource efficiency.
+
+    Args:
+        query: The research query to use for all experiments
+        n_trials: Number of parameter combinations to try
+        output_dir: Directory to save optimization results
+        model_name: Name of the LLM model to use
+        provider: LLM provider
+        search_tool: Search engine to use
+        progress_callback: Optional callback for progress updates
+        benchmark_weights: Dictionary mapping benchmark types to weights
+            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
+            If None, only SimpleQA is used with weight 1.0
+
+    Returns:
+        Tuple of (best_parameters, best_score)
+    """
+    # Balance of quality, speed and resource usage
+    metric_weights = {"quality": 0.4, "speed": 0.3, "resource": 0.3}
+
+    return optimize_parameters(
+        query=query,
+        output_dir=output_dir,
+        model_name=model_name,
+        provider=provider,
+        search_tool=search_tool,
+        n_trials=n_trials,
+        metric_weights=metric_weights,
+        optimization_metrics=["quality", "speed", "resource"],
+        progress_callback=progress_callback,
+        benchmark_weights=benchmark_weights,
+    )