opik-optimizer 0.7.0__py3-none-any.whl

opik_optimizer/mipro_optimizer/mipro_optimizer.py

@@ -0,0 +1,395 @@
+from typing import Any, Dict, List, Tuple, Union, Optional
+import os
+import random
+
+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
+
+import dspy
+
+import litellm
+from litellm.caching import Cache
+
+from ..optimization_result import OptimizationResult
+from ..base_optimizer import BaseOptimizer
+from ._mipro_optimizer_v2 import MIPROv2
+from ._lm import LM
+from ..optimization_config.configs import MetricConfig, TaskConfig
+from .utils import (
+    create_dspy_signature,
+    opik_metric_to_dspy,
+    create_dspy_training_set,
+    get_tool_prompts,
+)
+
+# Using disk cache for LLM calls
+disk_cache_dir = os.path.expanduser("~/.litellm_cache")
+litellm.cache = Cache(type="disk", disk_cache_dir=disk_cache_dir)
+
+# Set up logging
+import logging
+
+logger = logging.getLogger(__name__)  # Inherits config from setup_logging
+
+
+class MiproOptimizer(BaseOptimizer):
+    def __init__(self, model, project_name: Optional[str] = None, **model_kwargs):
+        super().__init__(model, project_name, **model_kwargs)
+        self.tools = []
+        self.num_threads = self.model_kwargs.pop("num_threads", 6)
+        self.model_kwargs["model"] = self.model
+        lm = LM(**self.model_kwargs)
+        opik_callback = OpikCallback(project_name=self.project_name, log_graph=True)
+        dspy.configure(lm=lm, callbacks=[opik_callback])
+        logger.debug(f"Initialized MiproOptimizer with model: {model}")
+
+    def evaluate_prompt(
+        self,
+        dataset: Union[str, Dataset],
+        metric_config: MetricConfig,
+        task_config: TaskConfig,
+        prompt: Union[str, dspy.Module, OptimizationResult] = None,
+        n_samples: int = 10,
+        dataset_item_ids: Optional[List[str]] = None,
+        experiment_config: Optional[Dict] = None,
+        **kwargs,
+    ) -> float:
+        """
+        Compute the score of a prompt on dataset (or part thereof)
+
+        Args:
+            dataset: Opik dataset name or dataset
+            metric_config: A MetricConfig instance
+            task_config: A TaskConfig instance
+            prompt: The prompt to evaluate
+            n_samples: number of items to test in the dataset
+            dataset_item_ids: Optional list of dataset item IDs to evaluate
+            experiment_config: Optional configuration for the experiment
+            **kwargs: Additional arguments for evaluation
+
+        Returns:
+            Evaluation score
+        """
+        # FIMXE: call super when it is ready
+        # FIXME: Intermediate values:
+        metric = metric_config.metric
+        input_key = task_config.input_dataset_fields[0]  # FIXME: allow all inputs
+        output_key = task_config.output_dataset_field
+
+        if isinstance(dataset, str):
+            opik_client = opik.Opik(project_name=self.project_name)
+            dataset = opik_client.get_dataset(dataset)
+
+        def LLM(input: str) -> str:
+            if isinstance(prompt, str):
+                response = litellm.completion(
+                    messages=[
+                        {"role": "system", "content": prompt},
+                        {"role": "user", "content": input},
+                    ],
+                    metadata={
+                        "opik": {
+                            "current_span_data": get_current_span_data(),
+                            "tags": ["optimizer"],
+                        },
+                    },
+                    **self.model_kwargs,
+                )
+                return response.choices[0].message.content
+            elif isinstance(prompt, OptimizationResult):
+                if prompt.optimizer == "MiproOptimizer" and getattr(prompt, "details"):
+                    program = prompt.details["program"]
+                    result = program(**{input_key: input})
+                    return getattr(result, output_key)
+                else:
+                    response = litellm.completion(
+                        messages=[
+                            {"role": "system", "content": prompt.prompt},
+                            # FIXME: insert demonstrations here
+                            {"role": "user", "content": input},
+                        ],
+                        metadata={
+                            "opik": {
+                                "current_span_data": get_current_span_data(),
+                                "tags": ["optimizer"],
+                            },
+                        },
+                        **self.model_kwargs,
+                    )
+                    return response.choices[0].message.content
+            elif isinstance(prompt, dspy.Module):
+                result = prompt(**{input_key: input})
+                return getattr(result, output_key)
+            else:
+                raise Exception("I don't know how to evaluate this prompt: %r" % prompt)
+
+        def evaluation_task(dataset_item):
+            # Get the model output
+            model_output = LLM(dataset_item[input_key])
+
+            # Prepare the result with all required fields
+            result = {
+                "input": dataset_item[input_key],
+                "output": model_output,
+                "expected_output": dataset_item[output_key],
+                "reference": dataset_item[output_key],
+            }
+
+            # Add context if available, otherwise use input as context
+            result["context"] = dataset_item.get("context", dataset_item[input_key])
+
+            return result
+
+        if n_samples is not None:
+            if dataset_item_ids is not None:
+                raise Exception("Can't use n_samples and dataset_item_ids")
+
+            all_ids = [dataset_item["id"] for dataset_item in dataset.get_items()]
+            dataset_item_ids = random.sample(all_ids, n_samples)
+
+        experiment_config = experiment_config or {}
+        experiment_config = {
+            **experiment_config,
+            **{
+                "optimizer": self.__class__.__name__,
+                "tools": (
+                    [f.__name__ for f in task_config.tools] if task_config.tools else []
+                ),
+                "metric": metric_config.metric.name,
+                "dataset": dataset.name,
+            },
+        }
+        # Run evaluation with all metrics at once
+        evaluation = evaluate(
+            dataset=dataset,
+            task=evaluation_task,
+            scoring_metrics=[metric],
+            # "reference" needs to match metric
+            scoring_key_mapping={"reference": output_key},
+            task_threads=self.num_threads,
+            dataset_item_ids=dataset_item_ids,
+            project_name=self.project_name,
+            experiment_config=experiment_config,
+        )
+
+        # Calculate average score across all metrics
+        total_score = 0
+        count = len(evaluation.test_results)
+        for i in range(count):
+            total_score += evaluation.test_results[i].score_results[0].value
+        score = total_score / count if count > 0 else 0.0
+
+        logger.debug(
+            f"Starting Mipro evaluation for prompt type: {type(prompt).__name__}"
+        )
+        logger.debug(f"Evaluation score: {score:.4f}")
+        return score
+
+    def optimize_prompt(
+        self,
+        dataset: Union[str, Dataset],
+        metric_config: MetricConfig,
+        task_config: TaskConfig,
+        num_candidates: int = 10,
+        experiment_config: Optional[Dict] = None,
+        **kwargs,
+    ) -> OptimizationResult:
+        self._opik_client = opik.Opik()
+        optimization = None
+        try:
+            optimization = self._opik_client.create_optimization(
+                dataset_name=dataset.name,
+                objective_name=metric_config.metric.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:
+            result = self._optimize_prompt(
+                dataset=dataset,
+                metric_config=metric_config,
+                task_config=task_config,
+                num_candidates=num_candidates,
+                experiment_config=experiment_config,
+                optimization_id=optimization.id if optimization is not None else None,
+                **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,
+        task_config: TaskConfig,
+        num_candidates: int = 10,
+        experiment_config: Optional[Dict] = None,
+        optimization_id: Optional[str] = None,
+        **kwargs,
+    ) -> OptimizationResult:
+        logger.info("Preparing MIPRO optimization...")
+        self.prepare_optimize_prompt(
+            dataset=dataset,
+            metric_config=metric_config,
+            task_config=task_config,
+            num_candidates=num_candidates,
+            experiment_config=experiment_config,
+            optimization_id=optimization_id,
+            **kwargs,
+        )
+        logger.info("Starting MIPRO compilation...")
+        result = self.continue_optimize_prompt()
+        logger.info("MIPRO optimization complete.")
+        return result
+
+    def prepare_optimize_prompt(
+        self,
+        dataset,
+        metric_config,
+        task_config,
+        num_candidates: int = 10,
+        experiment_config: Optional[Dict] = None,
+        optimization_id: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        # FIXME: Intermediate values:
+        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
+        self.tools = task_config.tools
+        self.num_candidates = num_candidates
+        self.seed = 9
+        self.input_key = input_key
+        self.output_key = output_key
+        self.prompt = prompt
+
+        # Convert to values for MIPRO:
+        if isinstance(dataset, str):
+            opik_client = opik.Opik(project_name=self.project_name)
+            self.dataset = opik_client.get_dataset(dataset).get_items()
+        else:
+            self.dataset = dataset.get_items()
+
+        # Validate dataset:
+        for row in self.dataset:
+            if self.input_key not in row:
+                raise Exception("row does not contain input_key: %r" % self.input_key)
+            if self.output_key not in row:
+                raise Exception("row does not contain output_key: %r" % self.output_key)
+
+        self.trainset = create_dspy_training_set(self.dataset, self.input_key)
+        self.data_signature = create_dspy_signature(
+            self.input_key, self.output_key, self.prompt
+        )
+
+        if self.tools:
+            self.module = dspy.ReAct(self.data_signature, tools=self.tools)
+        else:
+            self.module = dspy.Predict(self.data_signature)
+
+        # Convert the metric to a DSPy-compatible function
+        self.metric_function = opik_metric_to_dspy(metric, self.output_key)
+        self.opik_metric = metric
+        log_dir = os.path.expanduser("~/.opik-optimizer-checkpoints")
+        os.makedirs(log_dir, exist_ok=True)
+
+        experiment_config = experiment_config or {}
+        experiment_config = {
+            **experiment_config,
+            **{
+                "optimizer": self.__class__.__name__,
+                "tools": [f.__name__ for f in self.tools],
+                "metric": metric.name,
+                "num_threads": self.num_threads,
+                "num_candidates": self.num_candidates,
+                "dataset": dataset.name,
+            },
+        }
+
+        # Initialize the optimizer:
+        self.optimizer = MIPROv2(
+            metric=self.metric_function,
+            auto="light",
+            num_threads=self.num_threads,
+            verbose=False,
+            num_candidates=self.num_candidates,
+            seed=self.seed,
+            opik_prompt_task_config=task_config,
+            opik_dataset=dataset,
+            opik_project_name=self.project_name,
+            opik_metric_config=metric_config,
+            opik_optimization_id=optimization_id,
+            log_dir=log_dir,
+            experiment_config=experiment_config,
+        )
+
+        logger.debug("Created DSPy training set.")
+        logger.debug(f"Using DSPy module: {type(self.module).__name__}")
+        logger.debug(f"Using metric function: {self.metric_function.__name__}")
+
+    def load_from_checkpoint(self, filename):
+        """
+        Load the module from a checkpoint.
+        """
+        self.module.load(os.path.expanduser(filename))
+
+    def continue_optimize_prompt(self):
+        """
+        Continue to look for optimizations
+        """
+        self.results = self.optimizer.compile(
+            student=self.module,
+            trainset=self.trainset,
+            provide_traceback=True,
+            requires_permission_to_run=False,
+            num_trials=3,
+        )
+        self.best_programs = sorted(
+            self.results.candidate_programs,
+            key=lambda item: item["score"],
+            reverse=True,
+        )
+        self.module = self.get_best().details["program"]
+        return self.get_best()
+
+    def get_best(self, position: int = 0) -> OptimizationResult:
+        score = self.best_programs[position]["score"]
+        state = self.best_programs[position]["program"].dump_state()
+        if self.tools:
+            tool_names = [tool.__name__ for tool in self.tools]
+            tool_prompts = get_tool_prompts(
+                tool_names, state["react"]["signature"]["instructions"]
+            )
+            best_prompt = state["react"]["signature"]["instructions"]
+            demos = [x.toDict() for x in state["react"]["demos"]]
+        else:
+            tool_prompts = None
+            best_prompt = state["signature"]["instructions"]
+            demos = [x.toDict() for x in state["demos"]]
+
+        return OptimizationResult(
+            optimizer="MiproOptimizer",
+            prompt=best_prompt,
+            tool_prompts=tool_prompts,
+            score=score,
+            metric_name=self.opik_metric.name,
+            demonstrations=demos,
+            details={"program": self.best_programs[position]["program"]},
+        )

opik_optimizer/mipro_optimizer/utils.py

@@ -0,0 +1,107 @@
+from typing import Any, Dict, List, Tuple, Union, Optional
+
+import uuid
+import dspy
+import re
+
+from dspy.signatures.signature import make_signature
+
+
+class State(dict):
+    def __getattr__(self, key):
+        try:
+            return self[key]
+        except KeyError as e:
+            raise AttributeError(e)
+
+    def __setattr__(self, key, value):
+        self[key] = value
+
+    def __delattr__(self, key):
+        try:
+            del self[key]
+        except KeyError as e:
+            raise AttributeError(e)
+
+
+def create_dspy_signature(
+    input: str,
+    output: str,
+    prompt: str = None,
+):
+    """
+    Create a dspy Signature given inputs, outputs, prompt
+    """
+    # FIXME: allow multiple inputs, input/ouput descriptions
+    return make_signature(
+        signature={input: (str, dspy.InputField()), output: (str, dspy.OutputField())},
+        instructions=prompt,
+    )
+
+
+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)
+            return (
+                score_result.value if hasattr(score_result, "value") else score_result
+            )
+        except Exception as e:
+            print(f"Error calculating metric score: {e}")
+            return 0.0
+
+    return opik_metric_score_wrapper
+
+
+def create_dspy_training_set(data: list[dict], input: str) -> list[dspy.Example]:
+    """
+    Turn a list of dicts into a list of dspy Examples
+    """
+    output = []
+    for example in data:
+        example_obj = dspy.Example(
+            **example, dspy_uuid=str(uuid.uuid4()), dspy_split="train"
+        )
+        example_obj = example_obj.with_inputs(input)
+        output.append(example_obj)
+    return output
+
+
+def get_tool_prompts(tool_names, text: str) -> Dict[str, str]:
+    """
+    Extract the embedded tool prompts from a text.
+    """
+    tool_prompts = {}
+    for count, tool_name in enumerate(tool_names):
+        pattern = rf"\b{tool_name}\b[, \.]*([^{count + 2}]*)"
+        match = re.search(pattern, text)
+        if match:
+            description = match.groups()[0]
+            if description:
+                tool_prompts[tool_name] = description.strip()
+    return tool_prompts

opik_optimizer/optimization_config/configs.py

@@ -0,0 +1,35 @@
+"""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]]]
+
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+
+class TaskConfig(pydantic.BaseModel):
+    """Configuration for a prompt task."""
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+    
+    instruction_prompt: Union[str, List[Dict[Literal["role", "content"], str]]]
+    use_chat_prompt: bool = False
+    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

@@ -0,0 +1,49 @@
+from typing import Dict, Callable, Optional, Any, Union
+
+EVALUATED_LLM_TASK_OUTPUT = "_llm_task_output"
+
+class Mapper:
+    """Base class for mapping functions that transform data between different formats."""
+    
+    def __init__(self, name: Optional[str] = None, transform: Optional[Callable[[Any], Any]] = None):
+        if name is not None and transform is not None:
+            raise ValueError("Only one of name or transform can be provided")
+        
+        self.name = name
+        self.transform = transform
+    
+    def __call__(self, data: Any) -> Any:
+        if self.transform is not None:
+            return self.transform(data)
+        if self.name is not None:
+            return data[self.name]
+        return data
+
+def from_dataset_field(*, name: str = None, transform: Optional[Callable[[Dict[str, Any]], Any]] = None) -> Union[str, Callable[[Dict[str, Any]], Any]]:
+    if name is not None and transform is not None:
+        raise ValueError("Only one of name or transform can be provided")
+
+    if name is not None:
+        return name
+
+    if transform is not None:
+        return transform
+
+    raise ValueError("At least one of name or transform must be provided")
+
+
+def from_llm_response_text() -> str:
+    return EVALUATED_LLM_TASK_OUTPUT
+
+
+def from_agent_output(*, name: str = None, transform: Optional[Callable[[Any], Any]] = None) -> Union[str, Callable[[Any], Any]]:
+    if name is not None and transform is not None:
+        raise ValueError("Only one of name or transform can be provided")
+
+    if name is not None:
+        return lambda agent_output: agent_output[name]
+
+    if transform is not None:
+        return transform
+
+    return EVALUATED_LLM_TASK_OUTPUT