opik-optimizer 0.8.1__py3-none-any.whl → 0.9.0rc0__py3-none-any.whl

opik_optimizer/optimization_result.py

@@ -1,47 +1,36 @@
 """Module containing the OptimizationResult class."""
 
-from typing import Dict, List, Any, Optional, Union, Literal
+from typing import Any, Dict, List, Literal, Optional
+
 import pydantic
-from opik.evaluation.metrics import BaseMetric
-from pydantic import BaseModel, Field
-from .base_optimizer import OptimizationRound  # Adjust import as necessary
 import rich
 
-class OptimizationStep(BaseModel):
-    """Represents a single step or trial in an optimization process."""
-    step: int
-    score: Optional[float] = None
-    prompt: Optional[Union[str, List[Dict[str, str]]]] = None
-    parameters: Optional[Dict[str, Any]] = None
-    timestamp: Optional[str] = None
-    # Add other relevant details per step if needed
+from .reporting_utils import get_console
 
 
 class OptimizationResult(pydantic.BaseModel):
-    """Result of an optimization run."""
+    """Result oan optimization run."""
 
-    prompt: Union[str, List[Dict[Literal["role", "content"], str]]]
+    optimizer: str = "Optimizer"
+    
+    prompt: List[Dict[Literal["role", "content"], str]]
     score: float
     metric_name: str
-    metadata: Dict[str, Any] = pydantic.Field(
-        default_factory=dict
-    )  # Default empty dict
-    details: Dict[str, Any] = pydantic.Field(default_factory=dict)  # Default empty dict
-    best_prompt: Optional[str] = None
-    best_score: Optional[float] = None
-    best_metric_name: Optional[str] = None
-    best_details: Optional[Dict[str, Any]] = None
-    all_results: Optional[List[Dict[str, Any]]] = None
+    
+    details: Dict[str, Any] = pydantic.Field(default_factory=dict)
     history: List[Dict[str, Any]] = []
-    metric: Optional[BaseMetric] = None
-    demonstrations: Optional[List[Dict[str, Any]]] = None
-    optimizer: str = "Optimizer"
-    tool_prompts: Optional[Dict[str, str]] = None
-    opik_metadata: Optional[Dict[str, Any]] = None
     llm_calls: Optional[int] = None
 
+    # MIPRO specific
+    demonstrations: Optional[List[Dict[str, Any]]] = None
+    mipro_prompt: Optional[str] = None
+    tool_prompts: Optional[Dict[str, str]] = None
+    
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
+    def model_dump(self, *kargs, **kwargs) -> Dict[str, Any]:
+        return super().model_dump(*kargs, **kwargs)
+
     def _calculate_improvement_str(self) -> str:
         """Helper to calculate improvement percentage string."""
         initial_s = self.details.get("initial_score")
@@ -91,24 +80,19 @@ class OptimizationResult(pydantic.BaseModel):
         temp = self.details.get("temperature")
         temp_str = f"{temp:.1f}" if isinstance(temp, (int, float)) else "N/A"
 
-        final_prompt_display = self.prompt
-        if self.details.get("prompt_type") == "chat" and self.details.get(
-            "chat_messages"
-        ):
-            try:
-                chat_display = "\n".join(
-                    [
-                        f"  {msg.get('role', 'unknown')}: {str(msg.get('content', ''))[:150]}..."
-                        for msg in self.details["chat_messages"]
-                    ]
-                )
-                final_prompt_display = f"Instruction:\n  {self.prompt}\nFew-Shot Examples (Chat Structure):\n{chat_display}"
-            except Exception:
-                pass
+        try:
+            final_prompt_display = "\n".join(
+                [
+                    f"  {msg.get('role', 'unknown')}: {str(msg.get('content', ''))[:150]}..."
+                    for msg in self.prompt
+                ]
+            )
+        except Exception:
+            final_prompt_display = str(self.prompt)
 
         output = [
             f"\n{separator}",
-            f"OPTIMIZATION COMPLETE",
+            "OPTIMIZATION COMPLETE",
             f"{separator}",
             f"Optimizer:        {self.optimizer}",
             f"Model Used:       {model_name} (Temp: {temp_str})",
@@ -118,10 +102,10 @@ class OptimizationResult(pydantic.BaseModel):
             f"Total Improvement:{improvement_str.rjust(max(0, 18 - len('Total Improvement:')))}",
             f"Rounds Completed: {rounds_ran}",
             f"Stopped Early:    {stopped_early}",
-            f"\nFINAL OPTIMIZED PROMPT / STRUCTURE:",
-            f"--------------------------------------------------------------------------------",
+            "\nFINAL OPTIMIZED PROMPT / STRUCTURE:",
+            "--------------------------------------------------------------------------------",
             f"{final_prompt_display}",
-            f"--------------------------------------------------------------------------------",
+            "--------------------------------------------------------------------------------",
             f"{separator}",
         ]
         return "\n".join(output)
@@ -160,43 +144,33 @@ class OptimizationResult(pydantic.BaseModel):
         table.add_row("Stopped Early:", str(stopped_early))
 
         # Display Chat Structure if available
-        prompt_renderable: Any = rich.text.Text(
-            self.prompt or "", overflow="fold"
-        )  # Default to text
-        panel_title = "[bold]Final Optimized Prompt (Instruction)[/bold]"
-
-        if self.details.get("prompt_type") == "chat" and self.details.get(
-            "chat_messages"
-        ):
-            panel_title = "[bold]Final Optimized Prompt (Chat Structure)[/bold]"
-            try:
-                chat_group_items = [
-                    f"[dim]Instruction:[/dim] [i]{self.prompt}[/i]\n---"
-                ]
-                for msg in self.details["chat_messages"]:
-                    role = msg.get("role", "unknown")
-                    content = str(msg.get("content", ""))
-                    role_style = (
-                        "bold green"
-                        if role == "user"
-                        else (
-                            "bold blue"
-                            if role == "assistant"
-                            else ("bold magenta" if role == "system" else "")
-                        )
-                    )
-                    chat_group_items.append(
-                        f"[{role_style}]{role.capitalize()}:[/] {content}"
+        panel_title = "[bold]Final Optimized Prompt[/bold]"
+        try:
+            chat_group_items = []
+            for msg in self.prompt:
+                role = msg.get("role", "unknown")
+                content = str(msg.get("content", ""))
+                role_style = (
+                    "bold green"
+                    if role == "user"
+                    else (
+                        "bold blue"
+                        if role == "assistant"
+                        else ("bold magenta" if role == "system" else "")
                     )
-                    chat_group_items.append("---")  # Separator
-                prompt_renderable = rich.console.Group(*chat_group_items)
-
-            except Exception:
-                # Fallback to simple text prompt
-                prompt_renderable = rich.text.Text(self.prompt or "", overflow="fold")
-                panel_title = (
-                    "[bold]Final Optimized Prompt (Instruction - fallback)[/bold]"
                 )
+                chat_group_items.append(
+                    f"[{role_style}]{role.capitalize()}:[/] {content}"
+                )
+                chat_group_items.append("---")  # Separator
+            prompt_renderable = rich.console.Group(*chat_group_items)
+
+        except Exception:
+            # Fallback to simple text prompt
+            prompt_renderable = rich.text.Text(str(self.prompt or ""), overflow="fold")
+            panel_title = (
+                "[bold]Final Optimized Prompt (Instruction - fallback)[/bold]"
+            )
 
         prompt_panel = rich.panel.Panel(
             prompt_renderable, title=panel_title, border_style="blue", padding=(1, 2)
@@ -212,11 +186,9 @@ class OptimizationResult(pydantic.BaseModel):
             padding=1,
         )
 
-    def model_dump(self) -> Dict[str, Any]:
-        return super().model_dump()
-
     def display(self) -> None:
         """
         Displays the OptimizationResult using rich formatting
         """
-        rich.print(self)
+        console = get_console()
+        console.print(self)

opik_optimizer/reporting_utils.py

@@ -0,0 +1,180 @@
+import logging
+from contextlib import contextmanager
+from typing import Dict, List, Optional
+
+import rich
+from rich import box
+from rich.console import Console, Group
+from rich.panel import Panel
+from rich.progress import track
+from rich.text import Text
+
+PANEL_WIDTH = 70
+
+def get_console(*args, **kwargs):
+    console = Console(*args, **kwargs)
+    console.is_jupyter = False
+    return console
+
+@contextmanager
+def convert_tqdm_to_rich(description: Optional[str] = None, verbose: int = 1):
+    """Context manager to convert tqdm to rich."""
+    import opik.evaluation.engine.evaluation_tasks_executor
+    
+    optimizer_logger = logging.getLogger('opik_optimizer')
+
+    def _tqdm_to_track(iterable, desc, disable, total):
+        disable = verbose == 0 or optimizer_logger.level > logging.INFO
+        return track(
+            iterable,
+            description=description or desc,
+            disable=disable,
+            total=total
+        )
+
+    original__tqdm = opik.evaluation.engine.evaluation_tasks_executor._tqdm
+    opik.evaluation.engine.evaluation_tasks_executor._tqdm = _tqdm_to_track
+
+
+    from opik.evaluation import report
+    report.display_experiment_results = lambda *args, **kwargs: None
+    report.display_experiment_link = lambda *args, **kwargs: None
+
+    try:
+        yield
+    finally:
+        opik.evaluation.engine.evaluation_tasks_executor._tqdm = original__tqdm
+
+
+
+@contextmanager
+def suppress_opik_logs():
+    """Suppress Opik startup logs by temporarily increasing the log level."""
+    # Optimizer log level
+    optimizer_logger = logging.getLogger('opik_optimizer')
+    
+    # Get the Opik logger
+    opik_logger = logging.getLogger("opik.api_objects.opik_client")
+    
+    # Store original log level
+    original_level = opik_logger.level
+    
+    # Set log level to ERROR to suppress INFO messages
+    opik_logger.setLevel(optimizer_logger.level)
+    
+    try:
+        yield
+    finally:
+        # Restore original log level
+        opik_logger.setLevel(original_level)
+
+def display_messages(messages: List[Dict[str, str]], prefix: str = ""):
+    for i, msg in enumerate(messages):
+        panel = Panel(
+            Text(msg.get('content', ''), overflow="fold"),
+            title=f"{msg.get('role', 'message')}",
+            title_align="left",
+            border_style="dim",
+            width=PANEL_WIDTH,
+            padding=(1, 2),
+        )
+
+        # Capture the panel as rendered text with ANSI styles
+        console = get_console()
+        with console.capture() as capture:
+            console.print(panel)
+
+        # Retrieve the rendered string (with ANSI)
+        rendered_panel = capture.get()
+
+        # Prefix each line with '| ', preserving ANSI styles
+        for line in rendered_panel.splitlines():
+            console.print(Text(prefix) + Text.from_ansi(line))
+
+def display_header(algorithm: str, verbose: int = 1):
+    if verbose < 1:
+        return
+
+    content = Text.assemble(
+        ("● ", "green"),  
+        "Running Opik Evaluation - ",
+        (algorithm, "blue")
+    )
+
+    panel = Panel(
+        content,
+        box=box.ROUNDED,
+        width=PANEL_WIDTH
+    )
+
+    console = get_console()
+    console.print(panel)
+    console.print("\n")
+
+
+def display_result(initial_score, best_score, best_prompt, verbose: int = 1):
+    if verbose < 1:
+        return
+
+    console = get_console()
+    console.print(Text("\n> Optimization complete\n"))
+    
+    if best_score > initial_score:
+        if initial_score == 0:
+            content = [Text(f"Prompt was optimized and improved from {initial_score:.4f} to {best_score:.4f}", style="bold green")]
+        else:
+            perc_change = (best_score - initial_score) / initial_score
+            content = [Text(f"Prompt was optimized and improved from {initial_score:.4f} to {best_score:.4f} ({perc_change:.2%})", style="bold green")]
+    else:
+        content = [Text("Optimization trial did not find a better prompt than the initial one.", style="bold red")]
+    
+    content.append(Text("\nOptimized prompt:"))
+    for i, msg in enumerate(best_prompt):
+        content.append(
+            Panel(
+                Text(msg.get('content', ''), overflow="fold"),
+                title=f"{msg.get('role', 'message')}",
+                title_align="left",
+                border_style="dim",
+                width=PANEL_WIDTH,
+                padding=(1, 2),
+            )
+        )
+
+    console.print(
+        Panel(
+            Group(*content),
+            title="Optimization results",
+            title_align="left",
+            border_style="green",
+            width=PANEL_WIDTH,
+            padding=(1, 2)
+        )
+    )
+
+
+def display_configuration(messages: List[Dict[str, str]], optimizer_config: Dict[str, str], verbose: int = 1):
+    """Displays the LLM messages and optimizer configuration using Rich panels."""
+
+    if verbose < 1:
+        return
+
+    # Panel for Optimizer configuration
+    console = get_console()
+    console.print(Text("> Let's optimize the prompt:\n"))
+
+    display_messages(messages)
+
+    # Panel for configuration
+    console.print(Text(f"\nUsing {optimizer_config['optimizer']} with the parameters: "))
+    
+    for key, value in optimizer_config.items():
+        if key == "optimizer":  # Already displayed in the introductory text
+            continue
+        parameter_text = Text.assemble(
+            Text(f"  - {key}: ", style="dim"), 
+            Text(str(value), style="cyan")      
+        )
+        console.print(parameter_text)
+    
+    console.print("\n")

opik_optimizer/task_evaluator.py

@@ -1,17 +1,40 @@
-import opik
 import logging
 from typing import Any, Callable, Dict, List, Optional
-from opik_optimizer.optimization_config.configs import MetricConfig
-from opik.evaluation.metrics import score_result
 
+import opik
 from opik.evaluation import evaluator as opik_evaluator
+from opik.evaluation.metrics import base_metric, score_result
 
 logger = logging.getLogger(__name__)
 
+def _create_metric_class(metric: Callable):
+    class MetricClass(base_metric.BaseMetric):
+        def __init__(self):
+            self.name = metric.__name__
+        
+        def score(self, llm_output, **kwargs) -> score_result.ScoreResult:
+            try:
+                metric_val = metric(dataset_item=kwargs, llm_output=llm_output)
+                if isinstance(metric_val , score_result.ScoreResult):
+                    return metric_val
+                else:
+                    return score_result.ScoreResult(
+                        name = self.name,
+                        value = metric_val
+                    )
+            except Exception:
+                return score_result.ScoreResult(
+                    name = self.name,
+                    value = 0,
+                    scoring_failed=True
+                )
+
+    return MetricClass()
+
 def evaluate(
     dataset: opik.Dataset,
     evaluated_task: Callable[[Dict[str, Any]], Dict[str, Any]],
-    metric_config: MetricConfig,
+    metric: Callable,
     num_threads: int,
     optimization_id: Optional[str] = None,
     dataset_item_ids: Optional[List[str]] = None,
@@ -25,7 +48,8 @@ def evaluate(
 
     Args:
         dataset: A list of dictionaries representing the dataset.
-        metric_config: The metric configuration to use for evaluation.
+        metric: A metric function, this function should have two arguments:
+            dataset_item and llm_output
         evaluated_task: A function that takes a dataset item dict as input and returns a dictionary with output(s).
         dataset_item_ids: Optional list of dataset item IDs to evaluate.
         project_name: Optional project name for evaluation.
@@ -38,7 +62,7 @@ def evaluate(
     Returns:
         float: The average score of the evaluated task.
     """
-    items = dataset.get_items(dataset_item_ids)
+    items = dataset.get_items(n_samples)
     if not items:
         print("[DEBUG] Empty dataset, returning 0.0")
         return 0.0
@@ -46,21 +70,7 @@ def evaluate(
     if dataset_item_ids:
         items = [item for item in items if item.get("id") in dataset_item_ids]
 
-    if n_samples:
-        items = items[:n_samples]
-
-    # TODO: move to debug logger
-    # print(f"[DEBUG] Starting evaluation with task: {evaluated_task}")
-    # print(f"[DEBUG] Items to evaluate: {items}")
-    # print(f"[DEBUG] Metric config inputs: {metric_config.inputs}")
-    # print(f"[DEBUG] Number of threads: {num_threads}")
-    # print(f"[DEBUG] Project name: {project_name}")
-
-    scoring_key_mapping = {
-        key: value if isinstance(value, str) else value.__name__
-        for key, value in metric_config.inputs.items()
-    }
-    scoring_key_mapping["output"] = "_llm_task_output"
+    eval_metrics = [_create_metric_class(metric)]
 
     if optimization_id is not None:
         result = opik_evaluator.evaluate_optimization_trial(
@@ -68,9 +78,8 @@ def evaluate(
             dataset=dataset,
             task=evaluated_task,
             project_name=project_name,
-            scoring_key_mapping=scoring_key_mapping,
             dataset_item_ids=dataset_item_ids,
-            scoring_metrics=[metric_config.metric],
+            scoring_metrics=eval_metrics,
             task_threads=num_threads,
             nb_samples=n_samples,
             experiment_config=experiment_config,
@@ -81,9 +90,8 @@ def evaluate(
             dataset=dataset,
             task=evaluated_task,
             project_name=project_name,
-            scoring_key_mapping=scoring_key_mapping,
             dataset_item_ids=dataset_item_ids,
-            scoring_metrics=[metric_config.metric],
+            scoring_metrics=eval_metrics,
             task_threads=num_threads,
             nb_samples=n_samples,
             experiment_config=experiment_config,