opik-optimizer 2.1.1__tar.gz → 2.1.3__tar.gz

{opik_optimizer-2.1.1 → opik_optimizer-2.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opik_optimizer
-Version: 2.1.1
+Version: 2.1.3
 Summary: Agent optimization with Opik
 Home-page: https://github.com/comet-ml/opik
 Author: Comet ML

{opik_optimizer-2.1.1 → opik_optimizer-2.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "opik_optimizer"
-version = "2.1.1"
+version = "2.1.3"
 description = "Agent optimization with Opik"
 authors = [
     {name = "Comet ML", email = "support@comet.com"}

{opik_optimizer-2.1.1 → opik_optimizer-2.1.3}/src/opik_optimizer/__init__.py

@@ -19,6 +19,7 @@ from .mipro_optimizer import MiproOptimizer
 from .hierarchical_reflective_optimizer import HierarchicalReflectiveOptimizer
 from .optimization_config.configs import TaskConfig
 from .optimization_result import OptimizationResult
+from .multi_metric_objective import MultiMetricObjective
 from .parameter_optimizer import (
     ParameterOptimizer,
     ParameterSearchSpace,
@@ -48,6 +49,7 @@ __all__ = [
     "setup_logging",
     "datasets",
     "TaskConfig",
+    "MultiMetricObjective",
     "ParameterSearchSpace",
     "ParameterSpec",
     "ParameterType",

{opik_optimizer-2.1.1 → opik_optimizer-2.1.3}/src/opik_optimizer/base_optimizer.py

@@ -280,6 +280,7 @@ class BaseOptimizer(ABC):
         agent_config["project_name"] = getattr(prompt, "project_name", None)
         agent_config["model"] = getattr(prompt, "model", None) or self.model
         agent_config["tools"] = self._serialize_tools(prompt)
+        agent_config["optimizer"] = self.__class__.__name__
         return self._drop_none(agent_config)
 
     def get_optimizer_metadata(self) -> dict[str, Any]:
@@ -341,6 +342,7 @@ class BaseOptimizer(ABC):
             "metric": getattr(metric, "__name__", str(metric)),
             "dataset": getattr(dataset, "name", None),
             "dataset_id": dataset_id,
+            "optimizer": self.__class__.__name__,
             "optimizer_metadata": self._build_optimizer_metadata(),
             "tool_signatures": self._summarize_tool_signatures(prompt),
             "configuration": {

{opik_optimizer-2.1.1 → opik_optimizer-2.1.3}/src/opik_optimizer/gepa_optimizer/gepa_optimizer.py

@@ -253,7 +253,7 @@ class GepaOptimizer(BaseOptimizer):
                     opt_id = None
 
             gepa_reporting.display_header(
-                algorithm="GEPA",
+                algorithm=self.__class__.__name__,
                 optimization_id=opt_id,
                 dataset_id=getattr(dataset, "id", None),
                 verbose=self.verbose,
@@ -264,7 +264,7 @@ class GepaOptimizer(BaseOptimizer):
             _display_config(
                 messages=prompt.get_messages(),
                 optimizer_config={
-                    "optimizer": "GEPA",
+                    "optimizer": self.__class__.__name__,
                     "model": self.model,
                     "reflection_model": self.reflection_model,
                     "max_metric_calls": max_metric_calls,
@@ -422,7 +422,7 @@ class GepaOptimizer(BaseOptimizer):
                     "system_prompt": candidate_prompt,
                     "gepa_score": val_scores[idx] if idx < len(val_scores) else None,
                     "opik_score": score,
-                    "source": "GEPA",
+                    "source": self.__class__.__name__,
                 }
             )
             history.append(

{opik_optimizer-2.1.1 → opik_optimizer-2.1.3}/src/opik_optimizer/hierarchical_reflective_optimizer/hierarchical_reflective_optimizer.py

@@ -1,4 +1,3 @@
-from opik.environment import get_tqdm_for_current_environment
 import os
 import logging
 
@@ -29,8 +28,6 @@ from .types import (
 )
 from .prompts import IMPROVE_PROMPT_TEMPLATE
 
-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)
@@ -61,10 +58,14 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
         max_parallel_batches: Maximum number of batches to process concurrently during
             hierarchical root cause analysis (default: 5)
         batch_size: Number of test cases per batch for root cause analysis (default: 25)
+        max_iterations: Maximum number of optimization iterations (default: 5)
+        convergence_threshold: Stop if relative improvement is below this threshold (default: 0.01)
         **model_kwargs: Additional arguments passed to the LLM model
     """
 
     DEFAULT_ROUNDS = 10
+    DEFAULT_MAX_ITERATIONS = 5
+    DEFAULT_CONVERGENCE_THRESHOLD = 0.01  # Stop if improvement is less than 1%
 
     def __init__(
         self,
@@ -74,6 +75,8 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
         seed: int = 42,
         max_parallel_batches: int = 5,
         batch_size: int = 25,
+        max_iterations: int = DEFAULT_MAX_ITERATIONS,
+        convergence_threshold: float = DEFAULT_CONVERGENCE_THRESHOLD,
         **model_kwargs: Any,
     ):
         super().__init__(
@@ -83,6 +86,8 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
         self.num_threads = num_threads
         self.max_parallel_batches = max_parallel_batches
         self.batch_size = batch_size
+        self.max_iterations = max_iterations
+        self.convergence_threshold = convergence_threshold
 
         # Initialize hierarchical analyzer
         self._hierarchical_analyzer = HierarchicalRootCauseAnalyzer(
@@ -247,6 +252,8 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
             "reasoning_model": self.reasoning_model,
             "num_threads": self.num_threads,
             "max_parallel_batches": self.max_parallel_batches,
+            "max_iterations": self.max_iterations,
+            "convergence_threshold": self.convergence_threshold,
             "seed": self.seed,
             "verbose": self.verbose,
         }
@@ -417,7 +424,7 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
         n_samples: int | None,
         attempt: int,
         max_attempts: int,
-    ) -> tuple[chat_prompt.ChatPrompt, float]:
+    ) -> tuple[chat_prompt.ChatPrompt, float, EvaluationResult]:
         """
         Generate and evaluate a single improvement attempt for a failure mode.
 
@@ -434,7 +441,7 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
             max_attempts: Total number of attempts
 
         Returns:
-            Tuple of (improved_prompt, improved_score)
+            Tuple of (improved_prompt, improved_score, improved_experiment_result)
         """
         # Generate improvement with progress indication
         with reporting.display_prompt_improvement(
@@ -485,7 +492,7 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
             ) / len(improved_experiment_result.test_results)
             improved_reporter.set_score(improved_score)
 
-        return improved_chat_prompt, improved_score
+        return improved_chat_prompt, improved_score, improved_experiment_result
 
     def optimize_prompt(
         self,
@@ -528,6 +535,8 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
                 "n_samples": n_samples,
                 "auto_continue": auto_continue,
                 "max_retries": max_retries,
+                "max_iterations": self.max_iterations,
+                "convergence_threshold": self.convergence_threshold,
             },
             verbose=self.verbose,
             tools=getattr(prompt, "tools", None),
@@ -557,53 +566,62 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
             prompt.get_messages()
         )  # Store copy of initial messages for diff
 
-        # Iteration 1: Analyze and improve (structure ready for future multi-iteration support)
-        with reporting.display_optimization_iteration(
-            iteration=1, verbose=self.verbose
-        ) as iteration_reporter:
-            # Perform hierarchical root cause analysis
-            with reporting.display_root_cause_analysis(
-                verbose=self.verbose
-            ) as analysis_reporter:
-                hierarchical_analysis = self._hierarchical_root_cause_analysis(
-                    experiment_result
-                )
-                analysis_reporter.set_completed(
-                    total_test_cases=hierarchical_analysis.total_test_cases,
-                    num_batches=hierarchical_analysis.num_batches,
-                )
-
-            # Display hierarchical synthesis and failure modes
-            if self.verbose:
-                reporting.display_hierarchical_synthesis(
-                    total_test_cases=hierarchical_analysis.total_test_cases,
-                    num_batches=hierarchical_analysis.num_batches,
-                    synthesis_notes=hierarchical_analysis.synthesis_notes,
-                    verbose=self.verbose,
-                )
+        # Multi-iteration optimization loop
+        iteration = 0
+        previous_iteration_score = initial_score
+
+        for iteration in range(1, self.max_iterations + 1):
+            logger.info(f"Starting iteration {iteration}/{self.max_iterations}")
+
+            with reporting.display_optimization_iteration(
+                iteration=iteration, verbose=self.verbose
+            ) as iteration_reporter:
+                # Perform hierarchical root cause analysis
+                with reporting.display_root_cause_analysis(
+                    verbose=self.verbose
+                ) as analysis_reporter:
+                    hierarchical_analysis = self._hierarchical_root_cause_analysis(
+                        experiment_result
+                    )
+                    analysis_reporter.set_completed(
+                        total_test_cases=hierarchical_analysis.total_test_cases,
+                        num_batches=hierarchical_analysis.num_batches,
+                    )
 
-            reporting.display_failure_modes(
-                failure_modes=hierarchical_analysis.unified_failure_modes,
-                verbose=self.verbose,
-            )
+                # Display hierarchical synthesis and failure modes
+                if self.verbose:
+                    reporting.display_hierarchical_synthesis(
+                        total_test_cases=hierarchical_analysis.total_test_cases,
+                        num_batches=hierarchical_analysis.num_batches,
+                        synthesis_notes=hierarchical_analysis.synthesis_notes,
+                        verbose=self.verbose,
+                    )
 
-            # Generate improved prompt for each failure mode
-            for idx, root_cause in enumerate(
-                hierarchical_analysis.unified_failure_modes, 1
-            ):
-                logger.debug(
-                    f"Addressing failure mode {idx}/{len(hierarchical_analysis.unified_failure_modes)}: {root_cause.name}"
+                reporting.display_failure_modes(
+                    failure_modes=hierarchical_analysis.unified_failure_modes,
+                    verbose=self.verbose,
                 )
 
-                # Try multiple attempts if needed
-                max_attempts = max_retries + 1
-                improved_chat_prompt = None
-                improved_score = None
+                # Generate improved prompt for each failure mode
+                for idx, root_cause in enumerate(
+                    hierarchical_analysis.unified_failure_modes, 1
+                ):
+                    logger.debug(
+                        f"Addressing failure mode {idx}/{len(hierarchical_analysis.unified_failure_modes)}: {root_cause.name}"
+                    )
 
-                for attempt in range(1, max_attempts + 1):
-                    # Generate and evaluate improvement
-                    improved_chat_prompt, improved_score = (
-                        self._generate_and_evaluate_improvement(
+                    # Try multiple attempts if needed
+                    max_attempts = max_retries + 1
+                    improved_chat_prompt = None
+                    improved_score = None
+
+                    for attempt in range(1, max_attempts + 1):
+                        # Generate and evaluate improvement
+                        (
+                            improved_chat_prompt,
+                            improved_score,
+                            improved_experiment_result,
+                        ) = self._generate_and_evaluate_improvement(
                             root_cause=root_cause,
                             best_prompt=best_prompt,
                             best_score=best_score,
@@ -615,64 +633,86 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
                             attempt=attempt,
                             max_attempts=max_attempts,
                         )
-                    )
 
-                    # Check if we got improvement
-                    if improved_score > best_score:
-                        logger.info(
-                            f"Improvement found for '{root_cause.name}' on attempt {attempt}"
+                        # Check if we got improvement
+                        if improved_score > best_score:
+                            logger.info(
+                                f"Improvement found for '{root_cause.name}' on attempt {attempt}"
+                            )
+                            break
+
+                        # No improvement - should we retry?
+                        if attempt < max_attempts:
+                            reporting.display_retry_attempt(
+                                attempt=attempt,
+                                max_attempts=max_attempts,
+                                failure_mode_name=root_cause.name,
+                                verbose=self.verbose,
+                            )
+                        else:
+                            logger.debug(
+                                f"No improvement after {attempt} attempts for '{root_cause.name}'"
+                            )
+
+                    # Check if final result is an improvement
+                    if (
+                        improved_score is not None
+                        and improved_chat_prompt is not None
+                        and improved_score > best_score
+                    ):
+                        improvement = self._calculate_improvement(
+                            improved_score, best_score
                         )
-                        break
 
-                    # No improvement - should we retry?
-                    if attempt < max_attempts:
-                        reporting.display_retry_attempt(
-                            attempt=attempt,
-                            max_attempts=max_attempts,
-                            failure_mode_name=root_cause.name,
+                        # Display improvement for this iteration
+                        reporting.display_iteration_improvement(
+                            improvement=improvement,
+                            current_score=improved_score,
+                            best_score=best_score,
                             verbose=self.verbose,
                         )
+
+                        # Update best
+                        best_score = improved_score
+                        best_prompt = improved_chat_prompt
+                        best_messages = improved_chat_prompt.get_messages()
+                        experiment_result = improved_experiment_result
+                        logger.info(
+                            f"Updated best prompt after addressing '{root_cause.name}'"
+                        )
                     else:
                         logger.debug(
-                            f"No improvement after {attempt} attempts for '{root_cause.name}'"
+                            f"Keeping previous best prompt, no improvement from '{root_cause.name}'"
                         )
 
-                # Check if final result is an improvement
-                if (
-                    improved_score is not None
-                    and improved_chat_prompt is not None
-                    and improved_score > best_score
-                ):
-                    improvement = self._calculate_improvement(
-                        improved_score, best_score
-                    )
-
-                    # Display improvement for this iteration
-                    reporting.display_iteration_improvement(
-                        improvement=improvement,
-                        current_score=improved_score,
-                        best_score=best_score,
-                        verbose=self.verbose,
-                    )
+                # Mark iteration complete
+                improved_since_start = best_score > initial_score
+                iteration_reporter.iteration_complete(
+                    best_score=best_score, improved=improved_since_start
+                )
 
-                    # Update best
-                    best_score = improved_score
-                    best_prompt = improved_chat_prompt
-                    best_messages = improved_chat_prompt.get_messages()
-                    logger.info(
-                        f"Updated best prompt after addressing '{root_cause.name}'"
-                    )
-                else:
-                    logger.debug(
-                        f"Keeping previous best prompt, no improvement from '{root_cause.name}'"
-                    )
+            # Check for convergence after iteration
+            iteration_improvement = self._calculate_improvement(
+                best_score, previous_iteration_score
+            )
 
-            # Mark iteration complete
-            improved_since_start = best_score > initial_score
-            iteration_reporter.iteration_complete(
-                best_score=best_score, improved=improved_since_start
+            logger.info(
+                f"Iteration {iteration} complete. Score: {best_score:.4f}, "
+                f"Improvement: {iteration_improvement:.2%}"
             )
 
+            # Stop if improvement is below convergence threshold
+            if abs(iteration_improvement) < self.convergence_threshold:
+                logger.info(
+                    f"Convergence achieved: improvement ({iteration_improvement:.2%}) "
+                    f"below threshold ({self.convergence_threshold:.2%}). "
+                    f"Stopping after {iteration} iterations."
+                )
+                break
+
+            # Update previous score for next iteration
+            previous_iteration_score = best_score
+
         # Display final optimization result with diff
         reporting.display_optimized_prompt_diff(
             initial_messages=initial_messages,
@@ -682,14 +722,27 @@ class HierarchicalReflectiveOptimizer(BaseOptimizer):
             verbose=self.verbose,
         )
 
+        # Update optimization status to completed
+        try:
+            optimization.update(status="completed")
+            logger.info(f"Optimization {optimization.id} status updated to completed.")
+        except Exception as e:
+            logger.warning(f"Failed to update optimization status: {e}")
+
         # Prepare details for the result
         details = {
+            "model": best_prompt.model or self.model,
+            "temperature": (best_prompt.model_kwargs or {}).get("temperature")
+            or self.model_kwargs.get("temperature"),
             "reasoning_model": self.reasoning_model,
             "num_threads": self.num_threads,
             "max_parallel_batches": self.max_parallel_batches,
             "max_retries": max_retries,
             "n_samples": n_samples,
             "auto_continue": auto_continue,
+            "max_iterations": self.max_iterations,
+            "convergence_threshold": self.convergence_threshold,
+            "iterations_completed": iteration,
         }
 
         # Extract tool prompts if tools exist

{opik_optimizer-2.1.1 → opik_optimizer-2.1.3}/src/opik_optimizer/hierarchical_reflective_optimizer/hierarchical_root_cause_analyzer.py

@@ -1,8 +1,8 @@
 import logging
 import asyncio
 from typing import Any
-from tqdm import tqdm
 
+from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn
 from opik.evaluation.evaluation_result import EvaluationResult
 from .types import (
     RootCauseAnalysis,
@@ -11,6 +11,7 @@ from .types import (
 )
 from . import reporting
 from .prompts import BATCH_ANALYSIS_PROMPT, SYNTHESIS_PROMPT
+from ..reporting_utils import get_console
 
 logger = logging.getLogger(__name__)
 
@@ -285,13 +286,11 @@ Scores:
 
         semaphore = asyncio.Semaphore(self.max_parallel_batches)
 
-        # Create progress bar for batch processing
-        pbar = tqdm(
-            total=len(batch_tasks), desc="Processing batches", unit="batch", leave=False
-        )
+        # Create progress bar for batch processing using Rich
+        console = get_console()
 
         async def run_with_semaphore(
-            batch_num: int, task: Any
+            batch_num: int, task: Any, progress: Progress | None, task_id: Any | None
         ) -> tuple[int, BatchAnalysis]:
             async with semaphore:
                 try:
@@ -300,19 +299,44 @@ Scores:
                         f"Completed batch {batch_num}: "
                         f"identified {len(result.failure_modes)} failure modes"
                     )
-                    pbar.update(1)  # Update progress bar
+                    if progress and task_id is not None:
+                        progress.update(task_id, advance=1)  # Update progress bar
                     return batch_num, result
                 except Exception as exc:
                     logger.error(f"Batch {batch_num} failed: {exc}")
-                    pbar.update(1)  # Update progress bar even on error
+                    if progress and task_id is not None:
+                        progress.update(
+                            task_id, advance=1
+                        )  # Update progress bar even on error
                     raise
 
-        # Run all tasks with semaphore control
-        results = await asyncio.gather(
-            *[run_with_semaphore(num, task) for num, task in batch_tasks]
-        )
-
-        pbar.close()  # Close progress bar
+        # Run all tasks with semaphore control and rich progress bar
+        if self.verbose >= 1:
+            with Progress(
+                TextColumn("│      "),
+                TextColumn("[progress.description]{task.description}"),
+                BarColumn(),
+                TaskProgressColumn(),
+                console=console,
+                transient=True,
+            ) as progress:
+                task_id = progress.add_task(
+                    "Processing batches", total=len(batch_tasks)
+                )
+                results = await asyncio.gather(
+                    *[
+                        run_with_semaphore(num, task, progress, task_id)
+                        for num, task in batch_tasks
+                    ]
+                )
+        else:
+            # No progress bar in non-verbose mode
+            results = await asyncio.gather(
+                *[
+                    run_with_semaphore(num, task, None, None)
+                    for num, task in batch_tasks
+                ]
+            )
 
         # Sort by batch number to maintain order
         batch_analyses = [result for _, result in sorted(results)]