opik-optimizer 1.0.6__py3-none-any.whl → 2.0.0__py3-none-any.whl

opik_optimizer/mipro_optimizer/mipro_optimizer.py

@@ -1,7 +1,8 @@
 import os
 import random
 from datetime import datetime
-from typing import Callable, Dict, List, Literal, Optional, Union
+from typing import Any, Literal
+from collections.abc import Callable
 import logging
 
 import dspy
@@ -14,9 +15,9 @@ 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 ..optimization_config.configs import TaskConfig
+from ..optimization_config import chat_prompt
 from ._lm import LM
 from ._mipro_optimizer_v2 import MIPROv2
 from .utils import (
@@ -37,30 +38,42 @@ class MiproOptimizer(BaseOptimizer):
     def __init__(
         self,
         model,
-        project_name: Optional[str] = None,
+        project_name: str | None = None,
         verbose: int = 1,
         **model_kwargs,
     ):
         super().__init__(model=model, verbose=verbose, **model_kwargs)
         self.tools = []
         self.project_name = project_name
+        if "n_threads" in self.model_kwargs:
+            # To allow compatibility with other optimizers:
+            self.model_kwargs["num_threads"] = self.model_kwargs["n_threads"]
         self.num_threads = self.model_kwargs.pop("num_threads", 6)
         self.model_kwargs["model"] = self.model
         # FIXME: add mipro_optimizer=True - It does not count the LLM calls made internally by DSPy during MiproOptimizer.optimizer.compile().
         self.lm = LM(**self.model_kwargs)
+        setattr(self.lm, "parent_optimizer", self)
         opik_callback = OpikCallback(project_name=self.project_name, log_graph=True)
         dspy.configure(lm=self.lm, callbacks=[opik_callback])
         logger.debug(f"Initialized MiproOptimizer with model: {model}")
 
+    def get_optimizer_metadata(self) -> dict[str, Any]:
+        return self._drop_none(
+            {
+                "project_name": self.project_name,
+                "num_threads": self.num_threads,
+            }
+        )
+
     def evaluate_prompt(
         self,
-        dataset: Union[str, Dataset],
+        dataset: str | Dataset,
         metric: Callable,
         task_config: TaskConfig,
-        prompt: Optional[Union[str, dspy.Module, OptimizationResult]] = None,
+        prompt: str | dspy.Module | OptimizationResult | None = None,
         n_samples: int = 10,
-        dataset_item_ids: Optional[List[str]] = None,
-        experiment_config: Optional[Dict] = None,
+        dataset_item_ids: list[str] | None = None,
+        experiment_config: dict | None = None,
         verbose: int = 1,
         **kwargs,
     ) -> float:
@@ -83,7 +96,7 @@ class MiproOptimizer(BaseOptimizer):
         """
         # FIMXE: call super when it is ready
         # FIXME: Intermediate values:
-        self.llm_call_counter += 1
+        self.increment_llm_counter()
         input_key = task_config.input_dataset_fields[0]  # FIXME: allow all inputs
         output_key = task_config.output_dataset_field
 
@@ -238,23 +251,57 @@ class MiproOptimizer(BaseOptimizer):
 
     def optimize_prompt(
         self,
-        dataset: Union[str, Dataset],
+        prompt: chat_prompt.ChatPrompt,
+        dataset: str | Dataset,
         metric: Callable,
-        task_config: TaskConfig,
-        num_candidates: int = 10,
-        experiment_config: Optional[Dict] = None,
-        num_trials: Optional[int] = 3,
-        n_samples: Optional[int] = 10,
-        auto: Optional[Literal["light", "medium", "heavy"]] = "light",
+        experiment_config: dict | None = None,
+        n_samples: int | None = 10,
+        auto_continue: bool = False,
+        agent_class: str | None = None,
         **kwargs,
     ) -> OptimizationResult:
-        self._opik_client = opik.Opik()
-        with optimization_context(
-            client=self._opik_client,
-            dataset_name=dataset.name,
-            objective_name=metric.__name__,
-            metadata={"optimizer": self.__class__.__name__},
-        ) as optimization:
+        """
+        Optimize a prompt using MIPRO (Multi-Input Prompt Optimization).
+
+        Args:
+            prompt: The chat prompt to optimize
+            dataset: Opik dataset (or dataset name) containing evaluation data
+            metric: Evaluation function that takes (dataset_item, llm_output) and returns a score
+            experiment_config: Optional configuration for the experiment
+            n_samples: Number of samples to use for optimization (default: 10)
+            auto_continue: Whether to auto-continue optimization (default: False)
+            agent_class: Custom agent class to use (default: None)
+            **kwargs: Additional arguments including:
+                task_config: TaskConfig instance (required)
+                num_candidates: Number of candidates to generate (default: 10)
+                num_trials: Number of trials to run (default: 3)
+                auto: Optimization mode - "light", "medium", or "heavy" (default: "light")
+
+        Returns:
+            OptimizationResult: The optimization result containing the optimized prompt and metrics
+
+        Raises:
+            ValueError: If task_config is not provided
+        """
+        # Resolve dataset names to Dataset objects for validation compatibility
+        if isinstance(dataset, str):
+            dataset_name = dataset
+            client = opik.Opik(project_name=self.project_name)
+            dataset = client.get_dataset(dataset_name)
+
+        # Use base class validation and setup methods
+        self.validate_optimization_inputs(prompt, dataset, metric)
+
+        # Extract MIPRO-specific parameters from kwargs
+        task_config = kwargs.pop("task_config", None)
+        if task_config is None:
+            raise ValueError("task_config is required for MiproOptimizer")
+
+        num_candidates = kwargs.pop("num_candidates", 10)
+        num_trials = kwargs.pop("num_trials", 3)
+        auto = kwargs.pop("auto", "light")
+
+        with self.create_optimization_context(dataset, metric) as optimization:
             result = self._optimize_prompt(
                 dataset=dataset,
                 metric=metric,
@@ -271,15 +318,15 @@ class MiproOptimizer(BaseOptimizer):
 
     def _optimize_prompt(
         self,
-        dataset: Union[str, Dataset],
+        dataset: str | Dataset,
         metric: Callable,
         task_config: TaskConfig,
         num_candidates: int = 10,
-        experiment_config: Optional[Dict] = None,
-        optimization_id: Optional[str] = None,
-        num_trials: Optional[int] = 3,
-        n_samples: Optional[int] = 10,
-        auto: Optional[Literal["light", "medium", "heavy"]] = "light",
+        experiment_config: dict | None = None,
+        optimization_id: str | None = None,
+        num_trials: int | None = 3,
+        n_samples: int | None = 10,
+        auto: Literal["light", "medium", "heavy"] | None = "light",
         **kwargs,
     ) -> OptimizationResult:
         logger.info("Preparing MIPRO optimization...")
@@ -306,27 +353,26 @@ class MiproOptimizer(BaseOptimizer):
         metric,
         task_config,
         num_candidates: int = 10,
-        experiment_config: Optional[Dict] = None,
-        optimization_id: Optional[str] = None,
-        num_trials: Optional[int] = 3,
-        n_samples: Optional[int] = 10,
-        auto: Optional[Literal["light", "medium", "heavy"]] = "light",
+        experiment_config: dict | None = None,
+        optimization_id: str | None = None,
+        num_trials: int | None = 3,
+        n_samples: int | None = 10,
+        auto: Literal["light", "medium", "heavy"] | None = "light",
         **kwargs,
     ) -> None:
         # FIXME: Intermediate values:
-        self.llm_call_counter = 0
+        self.reset_counters()  # Reset counters for run
         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 = 42
+        self.auto = auto
         self.input_key = input_key
         self.output_key = output_key
         self.prompt = prompt
         self.num_trials = num_trials
         self.n_samples = n_samples
-        self.auto = auto
 
         # Convert to values for MIPRO:
         if isinstance(dataset, str):
@@ -395,6 +441,19 @@ class MiproOptimizer(BaseOptimizer):
         logger.debug(f"Using DSPy module: {type(self.module).__name__}")
         logger.debug(f"Using metric function: {self.metric_function.__name__}")
 
+    def cleanup(self) -> None:
+        """
+        Clean up MIPRO-specific resources.
+        """
+        # Call parent cleanup
+        super().cleanup()
+
+        # Clear MIPRO-specific resources
+        self.tools = None
+        self.prompt = None
+
+        logger.debug("Cleaned up MIPRO-specific resources")
+
     def load_from_checkpoint(self, filename):
         """
         Load the module from a checkpoint.
@@ -508,12 +567,15 @@ class MiproOptimizer(BaseOptimizer):
                     }
                 ],
                 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.lm.llm_call_counter,
+                llm_calls=self.llm_call_counter,
+                tool_calls=self.tool_call_counter,
             )
 
         self.module = self.get_best().details["program"]
@@ -545,7 +607,8 @@ class MiproOptimizer(BaseOptimizer):
             demonstrations=best_program_details.demonstrations,
             details=best_program_details.details,
             history=mipro_history_processed,
-            llm_calls=self.lm.llm_call_counter,
+            llm_calls=self.llm_call_counter,
+            tool_calls=self.tool_call_counter,
         )
 
     def get_best(self, position: int = 0) -> OptimizationResult:
@@ -553,6 +616,14 @@ class MiproOptimizer(BaseOptimizer):
             logger.error(
                 "get_best() called but no best_programs found. MIPRO compile might have failed or yielded no results."
             )
+            # Get LLM call count from the optimizer if available
+            dspy_llm_calls = (
+                getattr(self.optimizer, "total_calls", 0)
+                if hasattr(self, "optimizer") and self.optimizer
+                else 0
+            )
+            actual_llm_calls = max(self.llm_call_counter, dspy_llm_calls)
+
             return OptimizationResult(
                 optimizer="MiproOptimizer",
                 prompt=[
@@ -564,12 +635,15 @@ class MiproOptimizer(BaseOptimizer):
                     }
                 ],
                 score=0.0,
-                metric_name=getattr(self, "opik_metric", None).name
-                if hasattr(self, "opik_metric") and self.opik_metric
-                else "unknown_metric",
+                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"},
                 history=[],
-                llm_calls=self.lm.llm_call_counter,
+                llm_calls=actual_llm_calls,
+                tool_calls=self.tool_call_counter,
             )
 
         score = self.best_programs[position]["score"]
@@ -587,6 +661,11 @@ class MiproOptimizer(BaseOptimizer):
             best_prompt = state["signature"]["instructions"]
             demos = [x.toDict() for x in state["demos"]]
 
+        # Get LLM call count from the DSPy program module
+        dspy_llm_calls = getattr(program_module, "total_calls", 0)
+        # Use the higher of our counter or DSPy's counter
+        actual_llm_calls = max(self.llm_call_counter, dspy_llm_calls)
+
         print(best_prompt)
         return OptimizationResult(
             optimizer="MiproOptimizer",
@@ -596,5 +675,6 @@ class MiproOptimizer(BaseOptimizer):
             metric_name=self.opik_metric.__name__,
             demonstrations=demos,
             details={"program": program_module},
-            llm_calls=self.lm.llm_call_counter,
+            llm_calls=actual_llm_calls,
+            tool_calls=self.tool_call_counter,
         )

opik_optimizer/mipro_optimizer/utils.py

@@ -1,5 +1,3 @@
-from typing import Dict, Optional
-
 import uuid
 import dspy
 import re
@@ -61,7 +59,7 @@ def opik_metric_to_dspy(metric, output):
 
 
 def create_dspy_training_set(
-    data: list[dict], input: str, n_samples: Optional[int] = None
+    data: list[dict], input: str, n_samples: int | None = None
 ) -> list[dspy.Example]:
     """
     Turn a list of dicts into a list of dspy Examples
@@ -80,7 +78,7 @@ def create_dspy_training_set(
     return output
 
 
-def get_tool_prompts(tool_names, text: str) -> Dict[str, str]:
+def get_tool_prompts(tool_names, text: str) -> dict[str, str]:
     """
     Extract the embedded tool prompts from a text.
     """

opik_optimizer/optimizable_agent.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, List, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 import json
 import os
 
@@ -16,7 +16,7 @@ if TYPE_CHECKING:
     from .optimization_config.chat_prompt import ChatPrompt
 
 
-def tools_to_dict(tools: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
+def tools_to_dict(tools: dict[str, dict[str, Any]]) -> dict[str, Any]:
     retval = {}
     for name in tools:
         parts = {}
@@ -38,11 +38,11 @@ class OptimizableAgent:
         project_name (Optional[str]): The project name for tracking
     """
 
-    model: Optional[str] = None
-    model_kwargs: Dict[str, Any] = {}
-    project_name: Optional[str] = "Default Project"
-    input_dataset_field: Optional[str] = None
-    prompts: Dict[str, "ChatPrompt"]
+    model: str | None = None
+    model_kwargs: dict[str, Any] = {}
+    project_name: str | None = "Default Project"
+    input_dataset_field: str | None = None
+    prompts: dict[str, "ChatPrompt"]
     prompt: "ChatPrompt"
 
     def __init__(self, prompt: "ChatPrompt") -> None:
@@ -71,8 +71,8 @@ class OptimizableAgent:
     @_throttle.rate_limited(_limiter)
     def _llm_complete(
         self,
-        messages: List[Dict[str, str]],
-        tools: Optional[List[Dict[str, str]]],
+        messages: list[dict[str, str]],
+        tools: list[dict[str, str]] | None,
         seed: int,
     ) -> Any:
         response = litellm.completion(
@@ -91,10 +91,10 @@ class OptimizableAgent:
 
     def llm_invoke(
         self,
-        query: Optional[str] = None,
-        messages: Optional[List[Dict[str, str]]] = None,
-        seed: Optional[int] = None,
-        allow_tool_use: Optional[bool] = False,
+        query: str | None = None,
+        messages: list[dict[str, str]] | None = None,
+        seed: int | None = None,
+        allow_tool_use: bool | None = False,
     ) -> str:
         """
         NOTE: this is the default LiteLLM API. It is used
@@ -147,6 +147,11 @@ class OptimizableAgent:
                                 "content": str(tool_result),
                             }
                         )
+                        # Increment tool call counter if we have access to the optimizer
+                        if hasattr(self, "optimizer") and hasattr(
+                            self.optimizer, "increment_tool_counter"
+                        ):
+                            self.optimizer.increment_tool_counter()
                 else:
                     final_response = msg["content"]
                     break
@@ -156,14 +161,14 @@ class OptimizableAgent:
             result = response.choices[0].message.content
         return result
 
-    def invoke_dataset_item(self, dataset_item: Dict[str, str]) -> str:
+    def invoke_dataset_item(self, dataset_item: dict[str, str]) -> str:
         messages = self.prompt.get_messages(dataset_item)
         return self.invoke(messages)
 
     def invoke(
         self,
-        messages: List[Dict[str, str]],
-        seed: Optional[int] = None,
+        messages: list[dict[str, str]],
+        seed: int | None = None,
     ) -> str:
         """
         Invoke the agent with a dataset item.

opik_optimizer/optimization_config/chat_prompt.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any
+from collections.abc import Callable
 
 import copy
 
@@ -10,7 +11,7 @@ from opik import track
 class Tool(BaseModel):
     name: str = Field(..., description="Name of the tool")
     description: str = Field(..., description="Description of the tool")
-    parameters: Dict[str, Any] = Field(
+    parameters: dict[str, Any] = Field(
         ..., description="JSON Schema defining the input parameters for the tool"
     )
 
@@ -33,14 +34,14 @@ class ChatPrompt:
     def __init__(
         self,
         name: str = "chat-prompt",
-        system: Optional[str] = None,
-        user: Optional[str] = None,
-        messages: Optional[List[Dict[str, str]]] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        function_map: Optional[Dict[str, Callable]] = None,
-        model: Optional[str] = None,
-        invoke: Optional[Callable] = None,
-        project_name: Optional[str] = "Default Project",
+        system: str | None = None,
+        user: str | None = None,
+        messages: list[dict[str, str]] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        function_map: dict[str, Callable] | None = None,
+        model: str | None = None,
+        invoke: Callable | None = None,
+        project_name: str | None = "Default Project",
         **model_kwargs: Any,
     ) -> None:
         if system is None and user is None and messages is None:
@@ -97,8 +98,8 @@ class ChatPrompt:
 
     def get_messages(
         self,
-        dataset_item: Optional[Dict[str, str]] = None,
-    ) -> List[Dict[str, str]]:
+        dataset_item: dict[str, str] | None = None,
+    ) -> list[dict[str, str]]:
         # This is a copy, so we can alter the messages:
         messages = self._standardize_prompts()
 
@@ -113,8 +114,8 @@ class ChatPrompt:
                         )
         return messages
 
-    def _standardize_prompts(self, **kwargs: Any) -> List[Dict[str, str]]:
-        standardize_messages: List[Dict[str, str]] = []
+    def _standardize_prompts(self, **kwargs: Any) -> list[dict[str, str]]:
+        standardize_messages: list[dict[str, str]] = []
 
         if self.system is not None:
             standardize_messages.append({"role": "system", "content": self.system})
@@ -128,13 +129,13 @@ class ChatPrompt:
 
         return copy.deepcopy(standardize_messages)
 
-    def to_dict(self) -> Dict[str, Union[str, List[Dict[str, str]]]]:
+    def to_dict(self) -> dict[str, str | list[dict[str, str]]]:
         """Convert ChatPrompt to a dictionary for JSON serialization.
 
         Returns:
             Dict containing the serializable representation of this ChatPrompt
         """
-        retval: Dict[str, Union[str, List[Dict[str, str]]]] = {}
+        retval: dict[str, str | list[dict[str, str]]] = {}
         if self.system is not None:
             retval["system"] = self.system
         if self.user is not None:
@@ -144,29 +145,49 @@ class ChatPrompt:
         return retval
 
     def copy(self) -> "ChatPrompt":
+        """Shallow clone preserving model configuration and tools."""
+
+        # TODO(opik-mcp): once we introduce a dedicated MCP prompt subclass,
+        # migrate callers away from generic copies so optimizer metadata stays typed.
+        model_kwargs = (
+            copy.deepcopy(self.model_kwargs) if self.model_kwargs is not None else {}
+        )
         return ChatPrompt(
+            name=self.name,
             system=self.system,
             user=self.user,
             messages=copy.deepcopy(self.messages),
-            tools=self.tools,
+            tools=copy.deepcopy(self.tools),
             function_map=self.function_map,
+            model=self.model,
+            invoke=self.invoke,
+            project_name=self.project_name,
+            **model_kwargs,
         )
 
-    def set_messages(self, messages: List[Dict[str, Any]]) -> None:
+    def set_messages(self, messages: list[dict[str, Any]]) -> None:
         self.system = None
         self.user = None
         self.messages = copy.deepcopy(messages)
 
+    # TODO(opik): remove this stop-gap once MetaPromptOptimizer supports MCP.
+    # Provides a second-pass flow so tool results can be appended before
+    # rerunning the model.
+    def with_messages(self, messages: list[dict[str, Any]]) -> "ChatPrompt":
+        cloned = self.copy()
+        cloned.set_messages(messages)
+        return cloned
+
     @classmethod
     def model_validate(
         cls,
         obj: Any,
         *,
-        strict: Optional[bool] = None,
-        from_attributes: Optional[bool] = None,
-        context: Optional[Any] = None,
-        by_alias: Optional[bool] = None,
-        by_name: Optional[bool] = None,
+        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(

opik_optimizer/optimization_config/configs.py

@@ -1,6 +1,6 @@
 """Module containing configuration classes for optimization."""
 
-from typing import Any, List
+from typing import Any
 
 import pydantic
 
@@ -12,6 +12,6 @@ class TaskConfig(pydantic.BaseModel):
 
     instruction_prompt: str
     use_chat_prompt: bool = False
-    input_dataset_fields: List[str]
+    input_dataset_fields: list[str]
     output_dataset_field: str
-    tools: List[Any] = []
+    tools: list[Any] = []

opik_optimizer/optimization_config/mappers.py

@@ -1,4 +1,5 @@
-from typing import Dict, Callable, Optional, Any, Union
+from typing import Any
+from collections.abc import Callable
 
 EVALUATED_LLM_TASK_OUTPUT = "llm_output"
 
@@ -8,8 +9,8 @@ class Mapper:
 
     def __init__(
         self,
-        name: Optional[str] = None,
-        transform: Optional[Callable[[Any], Any]] = None,
+        name: str | None = None,
+        transform: Callable[[Any], Any] | None = None,
     ):
         if name is not None and transform is not None:
             raise ValueError("Only one of name or transform can be provided")
@@ -27,9 +28,9 @@ class Mapper:
 
 def from_dataset_field(
     *,
-    name: Optional[str] = None,
-    transform: Optional[Callable[[Dict[str, Any]], Any]] = None,
-) -> Union[str, Callable[[Dict[str, Any]], Any]]:
+    name: str | None = None,
+    transform: Callable[[dict[str, Any]], Any] | None = None,
+) -> 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")
 
@@ -47,8 +48,8 @@ def from_llm_response_text() -> str:
 
 
 def from_agent_output(
-    *, name: Optional[str] = None, transform: Optional[Callable[[Any], Any]] = None
-) -> Union[str, Callable[[Any], Any]]:
+    *, name: str | None = None, transform: Callable[[Any], Any] | None = None
+) -> 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")