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

opik_optimizer/few_shot_bayesian_optimizer/few_shot_bayesian_optimizer.py

@@ -5,6 +5,7 @@ import optuna
 import optuna.samplers
 import logging
 import json
+from datetime import datetime
 
 from opik import Dataset
 from opik_optimizer.optimization_config import mappers
@@ -26,20 +27,6 @@ _limiter = _throttle.get_rate_limiter_for_current_opik_installation()
 logger = logging.getLogger(__name__)
 
 
-@_throttle.rate_limited(_limiter)
-def _call_model(model, messages, seed, model_kwargs):
-    model_kwargs = opik_litellm_monitor.try_add_opik_monitoring_to_params(model_kwargs)
-
-    response = litellm.completion(
-        model=model,
-        messages=messages,
-        seed=seed,
-        **model_kwargs,
-    )
-
-    return response
-
-
 class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
     def __init__(
         self,
@@ -51,6 +38,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         n_threads: int = 8,
         n_initial_prompts: int = 5,
         n_iterations: int = 10,
+        verbose: int = 1,
         **model_kwargs,
     ) -> None:
         super().__init__(model, project_name, **model_kwargs)
@@ -60,9 +48,37 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         self.n_threads = n_threads
         self.n_initial_prompts = n_initial_prompts
         self.n_iterations = n_iterations
+        self.verbose = verbose
         self._opik_client = opik.Opik()
+        self.llm_call_counter = 0
         logger.debug(f"Initialized FewShotBayesianOptimizer with model: {model}")
 
+    @_throttle.rate_limited(_limiter)
+    def _call_model(self, model, messages, seed, model_kwargs):
+        self.llm_call_counter += 1
+
+        current_model_kwargs = self.model_kwargs.copy()
+        current_model_kwargs.update(model_kwargs)
+
+        filtered_call_kwargs = current_model_kwargs.copy()
+        filtered_call_kwargs.pop('n_trials', None)
+        filtered_call_kwargs.pop('n_samples', None)
+        filtered_call_kwargs.pop('n_iterations', None)
+        filtered_call_kwargs.pop('min_examples', None)
+        filtered_call_kwargs.pop('max_examples', None)
+        filtered_call_kwargs.pop('n_initial_prompts', None)
+
+        final_params_for_litellm = opik_litellm_monitor.try_add_opik_monitoring_to_params(filtered_call_kwargs)
+
+        response = litellm.completion(
+            model=self.model,
+            messages=messages,
+            seed=seed,
+            num_retries=6,
+            **final_params_for_litellm,
+        )
+        return response
+
     def _split_dataset(
         self, dataset: List[Dict[str, Any]], train_ratio: float
     ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
@@ -96,6 +112,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         n_samples: int = None,
     ) -> optimization_result.OptimizationResult:
         random.seed(self.seed)
+        self.llm_call_counter = 0
 
         if not task_config.use_chat_prompt:
             raise ValueError(
@@ -161,6 +178,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             project_name=self.project_name,
             experiment_config=initial_eval_config,
             optimization_id=optimization_id,
+            verbose=self.verbose,
         )
         logger.info(f"Initial (zero-shot) score: {initial_score:.4f}")
 
@@ -222,6 +240,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
                 project_name=self.project_name,
                 experiment_config=trial_config,
                 optimization_id=optimization_id,
+                verbose=self.verbose,
             )
             logger.debug(f"Trial {trial.number} score: {score:.4f}")
 
@@ -242,11 +261,59 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         except Exception as e:
             logger.warning(f"Could not configure Optuna logging within optimizer: {e}")
 
+        # Explicitly create and seed the sampler for Optuna
         sampler = optuna.samplers.TPESampler(seed=self.seed)
         study = optuna.create_study(direction="maximize", sampler=sampler)
-        study.optimize(optimization_objective, n_trials=n_trials)
+        
+        study.optimize(optimization_objective, n_trials=n_trials, show_progress_bar=(self.verbose >= 1))
         logger.info("Optuna study finished.")
 
+        optuna_history_processed = []
+        for trial_idx, trial in enumerate(study.trials):
+            if trial.state == optuna.trial.TrialState.COMPLETE:
+                param_obj: Optional[prompt_parameter.ChatPromptParameter] = trial.user_attrs.get("param")
+                prompt_cand_display = None # Default to None
+                if param_obj and hasattr(param_obj, 'as_template') and callable(param_obj.as_template):
+                    try:
+                        # .format() on ChatPromptTemplate returns the list of messages
+                        chat_messages_for_history = param_obj.as_template().format()
+                        prompt_cand_display = json.dumps(chat_messages_for_history) 
+                    except Exception as e_param_format:
+                        logger.warning(f"Trial {trial.number}: Error formatting prompt from param_obj: {e_param_format}")
+                        prompt_cand_display = "Error: Could not format prompt content."
+                elif not param_obj:
+                    logger.warning(f"Trial {trial.number}: 'param' object not found in user_attrs.")
+                    prompt_cand_display = "Error: Prompt data missing in trial."
+                else:
+                    logger.warning(f"Trial {trial.number}: 'param' object is not of expected type or lacks methods.")
+                    prompt_cand_display = "Error: Invalid prompt data structure in trial."
+
+                score_val = trial.value # This can be None if trial failed to produce a score
+                duration_val = None
+                if trial.datetime_complete and trial.datetime_start:
+                    duration_val = (trial.datetime_complete - trial.datetime_start).total_seconds()
+
+                iter_detail = {
+                    "iteration": trial.number + 1, 
+                    "timestamp": trial.datetime_start.isoformat() if trial.datetime_start else datetime.now().isoformat(),
+                    "prompt_candidate": prompt_cand_display,
+                    "parameters_used": { 
+                        "optuna_params": trial.params, 
+                        "example_indices": trial.user_attrs.get("example_indices", []) # Default to empty list
+                    },
+                    "scores": [{
+                        "metric_name": metric_config.metric.name, 
+                        "score": score_val, # Can be None
+                        "opik_evaluation_id": None # TODO
+                    }],
+                    "tokens_used": None, # TODO
+                    "cost": None, # TODO
+                    "duration_seconds": duration_val,
+                }
+                optuna_history_processed.append(iter_detail)
+            else:
+                logger.warning(f"Skipping trial {trial.number} from history due to state: {trial.state}. Value: {trial.value}")
+
         best_trial = study.best_trial
         best_score = best_trial.value
         best_n_examples = best_trial.params["n_examples"]
@@ -279,6 +346,8 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
                 "model": self.model,
                 "temperature": self.model_kwargs.get("temperature"),
             },
+            history=optuna_history_processed,
+            llm_calls=self.llm_call_counter
         )
 
     def optimize_prompt(
@@ -295,6 +364,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             optimization = self._opik_client.create_optimization(
                 dataset_name=dataset.name,
                 objective_name=metric_config.metric.name,
+                metadata={"optimizer": self.__class__.__name__},
             )
         except Exception:
             logger.warning(
@@ -389,6 +459,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             num_threads=self.n_threads,
             project_name=self.project_name,
             experiment_config=experiment_config,
+            verbose=self.verbose,
         )
         logger.debug(f"Evaluation score: {score:.4f}")
 
@@ -400,11 +471,11 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         def llm_task(dataset_item: Dict[str, Any]) -> Dict[str, Any]:
             prompt_ = template.format(**dataset_item)
 
-            response = _call_model(
+            response = self._call_model(
                 model=self.model,
                 messages=prompt_,
                 seed=self.seed,
-                model_kwargs=self.model_kwargs,
+                model_kwargs=self.model_kwargs
             )
 
             return {

opik_optimizer/logging_config.py

@@ -63,7 +63,7 @@ def setup_logging(
     _logging_configured = True
     
     # Use level name provided by rich handler by default
-    package_logger.info(f"Opik Optimizer logging configured to level: [bold cyan]{logging.getLevelName(level)}[/bold cyan]")
+    package_logger.info(f"Opik Agent Optimizer logging configured to level: [bold cyan]{logging.getLevelName(level)}[/bold cyan]")
 
 # Ensure logger obtained after setup can be used immediately if needed
 logger = logging.getLogger(__name__)

opik_optimizer/meta_prompt_optimizer.py

@@ -101,6 +101,8 @@ class MetaPromptOptimizer(BaseOptimizer):
         adaptive_trial_threshold: Optional[float] = DEFAULT_ADAPTIVE_THRESHOLD,
         num_threads: int = 12,
         project_name: Optional[str] = None,
+        verbose: int = 1,
+        enable_context: bool = True,
         **model_kwargs,
     ):
         """
@@ -117,6 +119,8 @@ class MetaPromptOptimizer(BaseOptimizer):
             adaptive_trial_threshold: If not None, prompts scoring below `best_score * adaptive_trial_threshold` after initial trials won't get max trials.
             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)
@@ -128,9 +132,12 @@ class MetaPromptOptimizer(BaseOptimizer):
         self.max_trials = max_trials_per_candidate
         self.adaptive_threshold = adaptive_trial_threshold
         self.num_threads = num_threads
+        self.verbose = verbose
         self.dataset = None
         self.task_config = 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}"
         )
@@ -151,6 +158,7 @@ class MetaPromptOptimizer(BaseOptimizer):
         experiment_config: Optional[Dict] = None,
         n_samples: Optional[int] = None,
         optimization_id: Optional[str] = None,
+        verbose: int = 1,
     ) -> float:
         """
         Evaluate a prompt using the given dataset and metric configuration.
@@ -177,6 +185,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             experiment_config=experiment_config,
             n_samples=n_samples,
             optimization_id=optimization_id,
+            verbose=self.verbose,
         )
 
     @_throttle.rate_limited(_rate_limiter)
@@ -188,12 +197,21 @@ class MetaPromptOptimizer(BaseOptimizer):
         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": getattr(self, "temperature", 0.3),
-                "max_tokens": getattr(self, "max_tokens", 1000),
+                "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),
@@ -242,7 +260,10 @@ class MetaPromptOptimizer(BaseOptimizer):
             )
 
             response = litellm.completion(
-                model=model_to_use, messages=messages, **final_call_params
+                model=model_to_use, 
+                messages=messages, 
+                num_retries=6,
+                **final_call_params
             )
             return response.choices[0].message.content
         except litellm.exceptions.RateLimitError as e:
@@ -271,6 +292,7 @@ class MetaPromptOptimizer(BaseOptimizer):
         experiment_config: Optional[Dict],
         n_samples: Optional[int],
         optimization_id: Optional[str] = None,
+        verbose: int = 1,
     ) -> float:
         # Calculate subset size for trials
         if not use_full_dataset:
@@ -429,6 +451,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             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
@@ -474,7 +497,9 @@ class MetaPromptOptimizer(BaseOptimizer):
         optimization = None
         try:
             optimization = self._opik_client.create_optimization(
-                dataset_name=dataset.name, objective_name=metric_config.metric.name
+                dataset_name=dataset.name,
+                objective_name=metric_config.metric.name,
+                metadata={"optimizer": self.__class__.__name__},
             )
             logger.info(f"Created optimization with ID: {optimization.id}")
         except Exception as e:
@@ -519,6 +544,7 @@ class MetaPromptOptimizer(BaseOptimizer):
         self.auto_continue = auto_continue
         self.dataset = dataset
         self.task_config = task_config
+        self.llm_call_counter = 0 # Reset counter for run
 
         current_prompt = task_config.instruction_prompt
         experiment_config = experiment_config or {}
@@ -550,6 +576,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             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
@@ -617,6 +644,7 @@ class MetaPromptOptimizer(BaseOptimizer):
                             n_samples=n_samples,
                             use_full_dataset=False,
                             experiment_config=experiment_config,
+                            verbose=self.verbose,
                         )
                         scores.append(score)
                         logger.debug(f"Trial {trial+1} score: {score:.4f}")
@@ -659,6 +687,7 @@ class MetaPromptOptimizer(BaseOptimizer):
                                 n_samples=n_samples,
                                 use_full_dataset=False,
                                 experiment_config=experiment_config,
+                                verbose=self.verbose,
                             )
                             scores.append(score)
                             logger.debug(
@@ -710,6 +739,7 @@ class MetaPromptOptimizer(BaseOptimizer):
                     experiment_config=experiment_config,
                     n_samples=n_samples,
                     use_full_dataset=n_samples is None,
+                    verbose=self.verbose,
                 )
                 logger.info(
                     f"Final evaluation score for best candidate: {final_score_best_cand:.4f}"
@@ -879,6 +909,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             score=best_score,
             metric_name=metric_config.metric.name,
             details=details,
+            llm_calls=self.llm_call_counter
         )
 
     def _get_task_context(self, metric_config: MetricConfig) -> str:
@@ -952,20 +983,35 @@ class MetaPromptOptimizer(BaseOptimizer):
 
         # Pass single metric_config
         history_context = self._build_history_context(previous_rounds)
-        task_context = self._get_task_context(metric_config=metric_config)
+        task_context_str = ""
+        analysis_instruction = ""
+        metric_focus_instruction = ""
+        improvement_point_1 = ""
+
+        if self.enable_context:
+            task_context_str = self._get_task_context(metric_config=metric_config)
+            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_config.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}
+        {task_context_str}
 
-        Analyze the example provided, the metric description, and the history of scores.
+        {analysis_instruction}
         Generate {self.num_prompts_per_round} improved versions of this prompt.
-        Focus on improving the score for the metric: {metric_config.metric.name}.
+        {metric_focus_instruction}
         Each version should aim to:
-        1. Be more specific and clear about expectations based on the metric and task.
-        2. Provide necessary context and constraints.
-        3. Guide the model to produce the desired output format suitable for the metric.
+        {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.