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

opik_optimizer/meta_prompt_optimizer/reporting.py

@@ -0,0 +1,140 @@
+from contextlib import contextmanager
+
+import rich
+from rich.text import Text
+
+from ..reporting_utils import (
+    convert_tqdm_to_rich,
+    display_configuration,  # noqa: F401
+    display_header,  # noqa: F401
+    display_messages,
+    display_result,  # noqa: F401
+    get_console,
+    suppress_opik_logs,
+)
+
+PANEL_WIDTH = 70
+console = get_console()
+
+
+@contextmanager
+def display_round_progress(max_rounds: int, verbose: int = 1):
+    """Context manager to display messages during an evaluation phase."""
+    
+    # Create a simple object with a method to set the score
+    class Reporter:
+        def failed_to_generate(self, num_prompts, error):
+            if verbose >= 1:
+                console.print(Text(f"│    Failed to generate {num_prompts} candidate prompt{'' if num_prompts == 1 else 's'}: {error}", style="red"))
+                console.print(Text("│"))
+            
+        def round_start(self, round_number):
+            if verbose >= 1:
+                console.print(Text(f"│ - Starting optimization round {round_number + 1} of {max_rounds}"))
+
+        def round_end(self, round_number, score, best_score, best_prompt):
+            if verbose >= 1:
+                console.print(Text(f"│    Completed optimization round {round_number + 1} of {max_rounds}"))
+                if best_score == 0 and score == 0:
+                    console.print(Text("│    No improvement in this optimization round - score is 0", style="yellow"))
+                elif best_score == 0:
+                    console.print(Text(f"│    Found a new best performing prompt: {score:.4f}", style="green"))
+                elif score > best_score:
+                    perc_change = (score - best_score) / best_score
+                    console.print(Text(f"│    Found a new best performing prompt: {score:.4f} ({perc_change:.2%})", style="green"))
+                elif score <= best_score:
+                    console.print(Text("│    No improvement in this optimization round", style="red"))
+                
+                console.print(Text("│"))
+
+    # Use our log suppression context manager and yield the reporter
+    with suppress_opik_logs():
+        with convert_tqdm_to_rich(verbose=verbose):
+            try:
+                yield Reporter()
+            finally:
+                pass
+
+
+@contextmanager
+def display_evaluation(message: str = "First we will establish the baseline performance:", verbose: int = 1):
+    """Context manager to display messages during an evaluation phase."""
+    score = None
+    
+    # Entry point
+    if verbose >= 1:
+        console.print(Text(f"> {message}"))
+    
+    # Create a simple object with a method to set the score
+    class Reporter:
+        def set_score(self, s):
+            if verbose >= 1:
+                console.print(Text(f"\r  Baseline score was: {s:.4f}.\n", style="green"))
+    
+    # Use our log suppression context manager and yield the reporter
+    with suppress_opik_logs():
+        with convert_tqdm_to_rich("  Evaluation", verbose=verbose):
+            try:
+                yield Reporter()
+            finally:
+                pass
+
+def display_optimization_start_message(verbose: int = 1):
+    if verbose >= 1:
+        console.print(Text("> Starting the optimization run"))
+        console.print(Text("│"))
+
+
+@contextmanager
+def display_candidate_generation_report(num_prompts: int, verbose: int = 1):
+    """Context manager to display messages during an evaluation phase."""
+    # Entry point
+    if verbose >= 1:
+        console.print(Text(f"│    Generating candidate prompt{'' if num_prompts == 1 else 's'}:"))
+    
+    # Create a simple object with a method to set the score
+    class Reporter:
+        def set_generated_prompts(self, prompts):
+            console.print(Text(f"│      Successfully generated {num_prompts} candidate prompt{'' if num_prompts == 1 else 's'}", style="dim"))
+            console.print(Text("│"))
+
+    try:
+        yield Reporter()
+    finally:
+        pass
+
+
+@contextmanager
+def display_prompt_candidate_scoring_report(candidate_count, prompt, verbose: int = 1):
+    """Context manager to display messages during an evaluation phase."""
+    # Create a simple object with a method to set the score
+    class Reporter:
+        def set_generated_prompts(self, candidate_count, prompt):
+            if verbose >= 1:
+                console.print(Text(f"│    Evaluating candidate prompt {candidate_count+1}:"))
+                display_messages(prompt, "│         ")
+        
+        def set_final_score(self, best_score, score):
+            if verbose >= 1:
+                if best_score == 0 and score > 0:
+                    console.print(Text(f"│          Evaluation score: {score:.4f}", style="green"))
+                elif best_score == 0 and score == 0:
+                    console.print(Text(f"│         Evaluation score: {score:.4f}", style="dim yellow"))
+                elif score > best_score:
+                    perc_change = (score - best_score) / best_score
+                    console.print(Text(f"│          Evaluation score: {score:.4f} ({perc_change:.2%})", style="green"))
+                elif score < best_score:
+                    perc_change = (score - best_score) / best_score
+                    console.print(Text(f"│          Evaluation score: {score:.4f} ({perc_change:.2%})", style="red"))
+                else:
+                    console.print(Text(f"│         Evaluation score: {score:.4f}", style="dim yellow"))
+            
+                console.print(Text("│"))
+                console.print(Text("│"))
+    try:
+        with suppress_opik_logs():
+            with convert_tqdm_to_rich("│         Evaluation", verbose=verbose):
+                yield Reporter()
+    finally:
+        pass
+

opik_optimizer/mipro_optimizer/__init__.py

@@ -1 +1 @@
-from .mipro_optimizer import MiproOptimizer
+from .mipro_optimizer import MiproOptimizer, MIPROv2

opik_optimizer/mipro_optimizer/_mipro_optimizer_v2.py

@@ -1,15 +1,12 @@
-import logging
 import random
 import textwrap
 from collections import defaultdict
 from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
 
+import dspy
 import numpy as np
 import opik
 import optuna
-from optuna.distributions import CategoricalDistribution
-
-import dspy
 from dspy.evaluate.evaluate import Evaluate
 from dspy.propose import GroundedProposer
 from dspy.teleprompt.teleprompt import Teleprompter
@@ -23,6 +20,9 @@ from dspy.teleprompt.utils import (
     save_candidate_program,
     set_signature,
 )
+from optuna.distributions import CategoricalDistribution
+
+from ..optimization_config.configs import TaskConfig
 
 
 class Logger():
@@ -51,7 +51,7 @@ ENDC = "\033[0m"  # Resets the color to default
 
 import opik
 from opik_optimizer import task_evaluator
-from opik_optimizer.optimization_config.configs import MetricConfig, TaskConfig
+from opik_optimizer.optimization_config.configs import TaskConfig
 from opik_optimizer.optimization_config import mappers
 
 def get_prompt(program):
@@ -85,7 +85,7 @@ class MIPROv2(Teleprompter):
         log_dir: Optional[str] = None,
         metric_threshold: Optional[float] = None,
         opik_dataset: Optional[opik.Dataset] = None,
-        opik_metric_config: Optional[MetricConfig] = None,
+        opik_metric: Optional[Callable] = None,
         opik_prompt_task_config: Optional[TaskConfig] = None,
         opik_project_name: Optional[str] = None,
         opik_optimization_id: Optional[str] = None,
@@ -119,7 +119,7 @@ class MIPROv2(Teleprompter):
         self.rng = None
 
         self.opik_dataset = opik_dataset
-        self.opik_metric_config = opik_metric_config
+        self.opik_metric = opik_metric
         self.opik_prompt_task_config = opik_prompt_task_config
         self.opik_project_name = opik_project_name
         self.opik_optimization_id = opik_optimization_id
@@ -572,7 +572,7 @@ class MIPROv2(Teleprompter):
             opik_dataset=self.opik_dataset,
             trainset=valset,
             candidate_program=program,
-            metric_config=self.opik_metric_config,
+            metric=self.opik_metric,
             prompt_task_config=self.opik_prompt_task_config,
             project_name=self.opik_project_name,
             num_threads=self.num_threads,
@@ -643,15 +643,7 @@ class MIPROv2(Teleprompter):
             score = eval_candidate_program(
                 batch_size, valset, candidate_program, evaluate, self.rng
             )
-            # score = eval_candidate_program_with_opik(
-            #     opik_dataset=self.opik_dataset,
-            #     trainset=valset,
-            #     candidate_program=candidate_program,
-            #     metric_config=self.opik_metric_config,
-            #     prompt_task_config=self.opik_prompt_task_config,
-            #     project_name=self.opik_project_name,
-            #     experiment_config=experiment_config,
-            # )
+
             total_eval_calls += batch_size
 
             # Update best score and program
@@ -953,7 +945,7 @@ class MIPROv2(Teleprompter):
             opik_dataset=self.opik_dataset,
             trainset=valset,
             candidate_program=highest_mean_program,
-            metric_config=self.opik_metric_config,
+            metric=self.opik_metric,
             prompt_task_config=self.opik_prompt_task_config,
             project_name=self.opik_project_name,
             num_threads=self.num_threads,
@@ -1029,7 +1021,7 @@ def eval_candidate_program_with_opik(
     trainset: List,
     candidate_program: Any,
     project_name: str,
-    metric_config: MetricConfig,
+    metric: Callable,
     prompt_task_config: TaskConfig,
     num_threads: int,
     experiment_config: Optional[Dict[str, Any]] = None,
@@ -1057,7 +1049,7 @@ def eval_candidate_program_with_opik(
     score = task_evaluator.evaluate(
         dataset=opik_dataset,
         evaluated_task=program_task,
-        metric_config=metric_config,
+        metric=metric,
         dataset_item_ids=dataset_item_ids,
         project_name=project_name,
         num_threads=num_threads,

opik_optimizer/mipro_optimizer/mipro_optimizer.py

@@ -1,30 +1,29 @@
-from typing import Any, Dict, List, Tuple, Union, Optional, Literal
 import os
 import random
 from datetime import datetime
-
-import opik
-
-from opik.integrations.dspy.callback import OpikCallback
-from opik.opik_context import get_current_span_data
-from opik.evaluation import evaluate
-from opik import Dataset
+from typing import Callable, Dict, List, Literal, Optional, Union
 
 import dspy
-
 import litellm
+import opik
 from litellm.caching import Cache
+from opik import Dataset
+from opik.evaluation import evaluate
+from opik.integrations.dspy.callback import OpikCallback
+from opik.opik_context import get_current_span_data
 
 from ..optimization_result import OptimizationResult
+from ..utils import optimization_context
 from ..base_optimizer import BaseOptimizer
-from ._mipro_optimizer_v2 import MIPROv2
+from ..optimization_config.configs import TaskConfig
+from ..optimization_result import OptimizationResult
 from ._lm import LM
-from ..optimization_config.configs import MetricConfig, TaskConfig
+from ._mipro_optimizer_v2 import MIPROv2
 from .utils import (
     create_dspy_signature,
-    opik_metric_to_dspy,
     create_dspy_training_set,
     get_tool_prompts,
+    opik_metric_to_dspy,
 )
 
 # Using disk cache for LLM calls
@@ -53,7 +52,7 @@ class MiproOptimizer(BaseOptimizer):
     def evaluate_prompt(
         self,
         dataset: Union[str, Dataset],
-        metric_config: MetricConfig,
+        metric: Callable,
         task_config: TaskConfig,
         prompt: Union[str, dspy.Module, OptimizationResult] = None,
         n_samples: int = 10,
@@ -67,7 +66,7 @@ class MiproOptimizer(BaseOptimizer):
 
         Args:
             dataset: Opik dataset name or dataset
-            metric_config: A MetricConfig instance
+            metric: Metric function to optimize
             task_config: A TaskConfig instance
             prompt: The prompt to evaluate
             n_samples: number of items to test in the dataset
@@ -82,7 +81,6 @@ class MiproOptimizer(BaseOptimizer):
         # FIMXE: call super when it is ready
         # FIXME: Intermediate values:
         self.llm_call_counter += 1
-        metric = metric_config.metric
         input_key = task_config.input_dataset_fields[0]  # FIXME: allow all inputs
         output_key = task_config.output_dataset_field
 
@@ -188,7 +186,7 @@ class MiproOptimizer(BaseOptimizer):
                 "tools": (
                     [f.__name__ for f in task_config.tools] if task_config.tools else []
                 ),
-                "metric": metric_config.metric.name,
+                "metric": metric.__name__,
                 "dataset": dataset.name,
             },
         }
@@ -222,7 +220,7 @@ class MiproOptimizer(BaseOptimizer):
     def optimize_prompt(
         self,
         dataset: Union[str, Dataset],
-        metric_config: MetricConfig,
+        metric: Callable,
         task_config: TaskConfig,
         num_candidates: int = 10,
         experiment_config: Optional[Dict] = None,
@@ -232,26 +230,15 @@ class MiproOptimizer(BaseOptimizer):
         **kwargs,
     ) -> OptimizationResult:
         self._opik_client = opik.Opik()
-        optimization = None
-        try:
-            optimization = self._opik_client.create_optimization(
+        with optimization_context(
+                client=self._opik_client,
                 dataset_name=dataset.name,
-                objective_name=metric_config.metric.name,
+                objective_name=metric.__name__,
                 metadata={"optimizer": self.__class__.__name__},
-            )
-        except Exception:
-            logger.warning(
-                "Opik server does not support optimizations. Please upgrade opik."
-            )
-            optimization = None
-
-        if not optimization:
-            logger.warning("Continuing without Opik optimization tracking.")
-
-        try:
+        ) as optimization:
             result = self._optimize_prompt(
                 dataset=dataset,
-                metric_config=metric_config,
+                metric=metric,
                 task_config=task_config,
                 num_candidates=num_candidates,
                 experiment_config=experiment_config,
@@ -261,19 +248,12 @@ class MiproOptimizer(BaseOptimizer):
                 auto=auto,
                 **kwargs,
             )
-            if optimization:
-                self.update_optimization(optimization, status="completed")
             return result
-        except Exception as e:
-            logger.error(f"Mipro optimization failed: {e}", exc_info=True)
-            if optimization:
-                self.update_optimization(optimization, status="cancelled")
-            raise e
 
     def _optimize_prompt(
         self,
         dataset: Union[str, Dataset],
-        metric_config: MetricConfig,
+        metric: Callable,
         task_config: TaskConfig,
         num_candidates: int = 10,
         experiment_config: Optional[Dict] = None,
@@ -286,7 +266,7 @@ class MiproOptimizer(BaseOptimizer):
         logger.info("Preparing MIPRO optimization...")
         self.prepare_optimize_prompt(
             dataset=dataset,
-            metric_config=metric_config,
+            metric=metric,
             task_config=task_config,
             num_candidates=num_candidates,
             experiment_config=experiment_config,
@@ -304,7 +284,7 @@ class MiproOptimizer(BaseOptimizer):
     def prepare_optimize_prompt(
         self,
         dataset,
-        metric_config,
+        metric,
         task_config,
         num_candidates: int = 10,
         experiment_config: Optional[Dict] = None,
@@ -316,7 +296,6 @@ class MiproOptimizer(BaseOptimizer):
     ) -> None:
         # FIXME: Intermediate values:
         self.llm_call_counter = 0
-        metric = metric_config.metric
         prompt = task_config.instruction_prompt
         input_key = task_config.input_dataset_fields[0]  # FIXME: allow all
         output_key = task_config.output_dataset_field
@@ -366,7 +345,7 @@ class MiproOptimizer(BaseOptimizer):
             **{
                 "optimizer": self.__class__.__name__,
                 "tools": [f.__name__ for f in self.tools],
-                "metric": metric.name,
+                "metric": metric.__name__,
                 "num_threads": self.num_threads,
                 "num_candidates": self.num_candidates,
                 "num_trials": self.num_trials,
@@ -385,7 +364,7 @@ class MiproOptimizer(BaseOptimizer):
             opik_prompt_task_config=task_config,
             opik_dataset=dataset,
             opik_project_name=self.project_name,
-            opik_metric_config=metric_config,
+            opik_metric=metric,
             opik_optimization_id=optimization_id,
             log_dir=log_dir,
             experiment_config=experiment_config,
@@ -457,7 +436,7 @@ class MiproOptimizer(BaseOptimizer):
             }
 
             current_score = candidate_data.get("score")
-            metric_name_for_history = self.opik_metric.name if hasattr(self, 'opik_metric') and self.opik_metric else "unknown_metric"
+            metric_name_for_history = self.opik_metric.__name__
 
             # Unscale if it's a known 0-1 metric that MIPRO might scale to 0-100
             # For now, specifically targeting Levenshtein-like metrics
@@ -479,9 +458,9 @@ class MiproOptimizer(BaseOptimizer):
             logger.warning("MIPRO compile returned no candidate programs.")
             return OptimizationResult(
                 optimizer="MiproOptimizer",
-                prompt=self.prompt,
+                prompt=[{"role": "user", "content": getattr(self, 'prompt', "Error: Initial prompt not found")}], 
                 score=0.0,
-                metric_name=self.opik_metric.name if hasattr(self, 'opik_metric') else "unknown_metric",
+                metric_name=self.opik_metric.__name__ if hasattr(self, 'opik_metric') else "unknown_metric",
                 details={"error": "No candidate programs generated by MIPRO"},
                 history=mipro_history_processed,
                 llm_calls=self.llm_call_counter
@@ -517,7 +496,7 @@ class MiproOptimizer(BaseOptimizer):
             logger.error("get_best() called but no best_programs found. MIPRO compile might have failed or yielded no results.")
             return OptimizationResult(
                 optimizer="MiproOptimizer",
-                prompt=getattr(self, 'prompt', "Error: Initial prompt not found"), 
+                prompt=[{"role": "user", "content": getattr(self, 'prompt', "Error: Initial prompt not found")}], 
                 score=0.0, 
                 metric_name=getattr(self, 'opik_metric', None).name if hasattr(self, 'opik_metric') and self.opik_metric else "unknown_metric",
                 details={"error": "No programs generated or compile failed"},
@@ -540,12 +519,13 @@ class MiproOptimizer(BaseOptimizer):
             best_prompt = state["signature"]["instructions"]
             demos = [x.toDict() for x in state["demos"]]
 
+        print(best_prompt)
         return OptimizationResult(
             optimizer="MiproOptimizer",
-            prompt=best_prompt,
+            prompt=[{"role": "user", "content": best_prompt}],
             tool_prompts=tool_prompts,
             score=score,
-            metric_name=self.opik_metric.name,
+            metric_name=self.opik_metric.__name__,
             demonstrations=demos,
             details={"program": program_module},
             llm_calls=self.llm_call_counter

opik_optimizer/mipro_optimizer/utils.py

@@ -44,31 +44,9 @@ def opik_metric_to_dspy(metric, output):
     answer_field = output
 
     def opik_metric_score_wrapper(example, prediction, trace=None):
-        # Extract the input from the example
-        input_text = getattr(example, "input", "")
-        if isinstance(input_text, list):
-            input_text = input_text[0] if input_text else ""
-
-        # Extract the expected output
-        expected_output = getattr(example, answer_field, "")
-        if isinstance(expected_output, list):
-            expected_output = expected_output[0] if expected_output else ""
-
-        # Get the model output
-        model_output = getattr(prediction, answer_field, "")
-
-        # Create a result dictionary with all required fields
-        result = {
-            "input": input_text,
-            "output": model_output,
-            "expected_output": expected_output,
-            "reference": expected_output,
-            "context": getattr(example, "context", input_text),
-        }
-
         try:
             # Calculate the score using the metric
-            score_result = metric.score(**result)
+            score_result = metric(dataset_item=example.toDict(), llm_output=getattr(prediction, answer_field, ""))
             return (
                 score_result.value if hasattr(score_result, "value") else score_result
             )

opik_optimizer/optimization_config/chat_prompt.py

@@ -0,0 +1,106 @@
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, Field
+
+
+class Tool(BaseModel):
+    name: str =Field(
+        ...,
+        description="Name of the tool"
+    )
+    description: str = Field(
+        ...,
+        description="Description of the tool"
+    )
+    parameters: Dict[str, Any] = Field(
+        ...,
+        description="JSON Schema defining the input parameters for the tool"
+    )
+
+class ChatPrompt:
+    system: str
+    prompt: str
+    messages: List[Dict[Literal["role", "content"], str]]
+
+    def __init__(
+        self,
+        system: Optional[str] = None,
+        prompt: Optional[str] = None,
+        messages: Optional[List[Dict[Literal["role", "content"], str]]] = None,
+        tools: Optional[List[Tool]] = None
+    ):
+        self.system = system
+        self.prompt = prompt
+        self.messages = messages
+
+        self.formatted_messages = self._standardize_prompts()
+
+    def _standardize_prompts(
+        self, **kwargs: Any
+    ) -> List[Dict[Literal["role", "content"], str]]:
+        if (self.system is None and self.prompt is None and self.messages is None):
+            raise ValueError(
+                "At least one of `system`, `prompt` or `messages` must be provided"
+            )
+
+        if (self.prompt is not None and self.messages is not None):
+            raise ValueError(
+                "`prompt` and `messages` cannot be provided together"
+            )
+        
+        if (self.system is not None and not isinstance(self.system, str)):
+            raise ValueError(
+                "`system` must be a string"
+            )
+        
+        if (self.prompt is not None and not isinstance(self.prompt, str)):
+            raise ValueError(
+                "`prompt` must be a string"
+            )
+
+        if (self.messages is not None and not isinstance(self.messages, list)):
+            raise ValueError(
+                "`messages` must be a list"
+            )
+
+        standardize_messages = []
+            
+        if (self.system is not None):
+            standardize_messages.append({"role": "system", "content": self.system})
+        
+        if (self.prompt is not None):
+            standardize_messages.append({"role": "user", "content": self.prompt})
+        
+        if (self.messages is not None):
+            for message in self.messages:
+                standardize_messages.append(message)
+        
+        return standardize_messages
+
+    def format(self, **kwargs: Any) -> str:
+        return self.prompt.format(**kwargs)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert ChatPrompt to a dictionary for JSON serialization.
+        
+        Returns:
+            Dict containing the serializable representation of this ChatPrompt
+        """
+        return {
+            "system": self.system,
+            "prompt": self.prompt,
+            "messages": self.messages,
+            "formatted_messages": self.formatted_messages
+        }
+
+    @classmethod
+    def model_validate(cls, obj: Any, *, strict: bool | None = None, from_attributes: bool | None = None, 
+                      context: Any | None = None, by_alias: bool | None = None, by_name: bool | None = None) -> 'ChatPrompt':
+        """Custom validation method to handle nested objects during deserialization."""
+        return ChatPrompt(
+            system=obj.get('system', None),
+            prompt=obj.get('prompt', None),
+            messages=obj.get('messages', None),
+            
+        )
+        

opik_optimizer/optimization_config/configs.py

@@ -1,17 +1,8 @@
 """Module containing configuration classes for optimization."""
 
-import pydantic
-import opik
-from typing import Dict, Callable, Union, List, Literal, Any, Optional
-from opik.evaluation.metrics import BaseMetric
-
-
-class MetricConfig(pydantic.BaseModel):
-    """Configuration for a metric used in optimization."""
-    metric: BaseMetric
-    inputs: Dict[str, Union[str, Callable[[Any], Any]]]
+from typing import Any, Dict, List, Literal, Union
 
-    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+import pydantic
 
 
 class TaskConfig(pydantic.BaseModel):
@@ -23,13 +14,3 @@ class TaskConfig(pydantic.BaseModel):
     input_dataset_fields: List[str]
     output_dataset_field: str
     tools: List[Any] = []
-
-
-class OptimizationConfig(pydantic.BaseModel):
-    """Configuration for optimization."""
-    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
-
-    dataset: opik.Dataset
-    objective: MetricConfig
-    optimization_direction: Literal["maximize", "minimize"] = "maximize"
-    task: TaskConfig

opik_optimizer/optimization_config/mappers.py

@@ -1,6 +1,6 @@
 from typing import Dict, Callable, Optional, Any, Union
 
-EVALUATED_LLM_TASK_OUTPUT = "_llm_task_output"
+EVALUATED_LLM_TASK_OUTPUT = "llm_output"
 
 class Mapper:
     """Base class for mapping functions that transform data between different formats."""