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

opik_optimizer/meta_prompt_optimizer/meta_prompt_optimizer.py

@@ -3,6 +3,7 @@ import json
 import logging
 import os
 import textwrap
+import warnings
 from typing import Any, cast
 from collections.abc import Callable
 
@@ -11,12 +12,10 @@ import opik
 from litellm.caching import Cache
 from litellm.types.caching import LiteLLMCacheType
 from opik import Dataset
-from opik.api_objects import opik_client
 from opik.environment import get_tqdm_for_current_environment
 from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
 
 from opik_optimizer import task_evaluator
-from ..utils.core import create_litellm_agent_class
 
 from .. import _throttle
 from ..base_optimizer import BaseOptimizer, OptimizationRound
@@ -143,6 +142,7 @@ class MetaPromptOptimizer(BaseOptimizer):
         verbose: int = 1,
         enable_context: bool = True,
         n_threads: int = 12,
+        seed: int = 42,
         **model_kwargs: Any,
     ) -> None:
         """
@@ -157,22 +157,28 @@ class MetaPromptOptimizer(BaseOptimizer):
             **model_kwargs: Additional model parameters
         """
         if "project_name" in model_kwargs:
-            print(
-                "Removing `project_name` from constructor; it now belongs in the ChatPrompt()"
+            warnings.warn(
+                "The 'project_name' parameter in optimizer constructor is deprecated. "
+                "Set project_name in the ChatPrompt instead.",
+                DeprecationWarning,
+                stacklevel=2,
             )
             del model_kwargs["project_name"]
 
-        super().__init__(model=model, verbose=verbose, **model_kwargs)
+        super().__init__(model=model, verbose=verbose, seed=seed, **model_kwargs)
         self.reasoning_model = reasoning_model if reasoning_model is not None else model
         self.rounds = rounds
         self.num_prompts_per_round = num_prompts_per_round
         if num_threads is not None:
-            print("num_threads is deprecated; use n_threads instead")
+            warnings.warn(
+                "The 'num_threads' parameter is deprecated and will be removed in a future version. "
+                "Use 'n_threads' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
             n_threads = num_threads
         self.num_threads = n_threads
         self.dataset: Dataset | None = None
-        self._opik_client = opik_client.get_client_cached()
-        self.llm_call_counter = 0
         self.enable_context = enable_context
         logger.debug(
             f"Initialized MetaPromptOptimizer with model={model}, reasoning_model={self.reasoning_model}"
@@ -181,6 +187,14 @@ class MetaPromptOptimizer(BaseOptimizer):
             f"Optimization rounds: {rounds}, Prompts/round: {num_prompts_per_round}"
         )
 
+    def get_optimizer_metadata(self) -> dict[str, Any]:
+        return {
+            "rounds": self.rounds,
+            "num_prompts_per_round": self.num_prompts_per_round,
+            "reasoning_model": self.reasoning_model,
+            "enable_context": self.enable_context,
+        }
+
     @_throttle.rate_limited(_rate_limiter)
     def _call_model(
         self,
@@ -190,7 +204,7 @@ class MetaPromptOptimizer(BaseOptimizer):
         optimization_id: str | None = None,
     ) -> str:
         """Call the model with the given prompt and return the response."""
-        self.llm_call_counter += 1
+        self.increment_llm_counter()
         # Note: Basic retry logic could be added here using tenacity
         try:
             # Basic LLM parameters (e.g., temperature, max_tokens)
@@ -321,25 +335,28 @@ class MetaPromptOptimizer(BaseOptimizer):
             subset_size = None  # Use all items for final checks
             logger.debug("Using full dataset for evaluation")
 
-        experiment_config = experiment_config or {}
-        experiment_config = {
-            **experiment_config,
-            **{
-                "optimizer": self.__class__.__name__,
-                "agent_class": self.agent_class.__name__,
-                "agent_config": prompt.to_dict(),
-                "metric": getattr(metric, "__name__", str(metric)),
-                "dataset": dataset.name,
-                "configuration": {
-                    "prompt": prompt.get_messages(),
-                    "tools": getattr(prompt, "tools", None),
-                    "n_samples": subset_size,
-                    "use_full_dataset": use_full_dataset,
-                },
-            },
-        }
-        if optimization_id:
-            experiment_config["optimization_id"] = optimization_id
+        configuration_updates = self._drop_none(
+            {
+                "n_samples": subset_size,
+                "use_full_dataset": use_full_dataset,
+            }
+        )
+        meta_metadata = self._drop_none(
+            {
+                "optimization_id": optimization_id,
+                "stage": "trial_evaluation" if not use_full_dataset else "final_eval",
+            }
+        )
+        experiment_config = self._prepare_experiment_config(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            configuration_updates=configuration_updates,
+            additional_metadata={"meta_prompt": meta_metadata}
+            if meta_metadata
+            else None,
+        )
 
         def llm_task(dataset_item: dict[str, Any]) -> dict[str, str]:
             new_prompt = prompt.copy()
@@ -357,7 +374,7 @@ class MetaPromptOptimizer(BaseOptimizer):
                     )
                     raw_model_output = agent.llm_invoke(
                         messages=messages,
-                        seed=None,
+                        seed=self.seed,
                         allow_tool_use=True,
                     )
                 except Exception as exc:
@@ -391,7 +408,7 @@ class MetaPromptOptimizer(BaseOptimizer):
                     )
                     final_response = agent.llm_invoke(
                         messages=second_pass_messages,
-                        seed=None,
+                        seed=self.seed,
                         allow_tool_use=mcp_config.allow_tool_use_on_second_pass,
                     )
                 else:
@@ -459,36 +476,25 @@ class MetaPromptOptimizer(BaseOptimizer):
         Optimize a prompt using meta-reasoning.
 
         Args:
+            prompt: The prompt to optimize
             dataset: The dataset to evaluate against
             metric: The metric to use for evaluation
             experiment_config: A dictionary to log with the experiments
             n_samples: The number of dataset items to use for evaluation
             auto_continue: If True, the algorithm may continue if goal not met
-            **kwargs: Additional arguments for evaluation
+            agent_class: Optional agent class to use
+            **kwargs: Additional arguments for evaluation, including:
+                mcp_config (MCPExecutionConfig | None): MCP tool calling configuration (default: None)
+                candidate_generator: Optional candidate generator
+                candidate_generator_kwargs: Optional kwargs for candidate generator
 
         Returns:
             OptimizationResult: Structured result containing optimization details
         """
-        if not isinstance(prompt, chat_prompt.ChatPrompt):
-            raise ValueError("Prompt must be a ChatPrompt object")
-
-        if not isinstance(dataset, Dataset):
-            raise ValueError("Dataset must be a Dataset object")
-
-        if not callable(metric):
-            raise ValueError(
-                "Metric must be a function that takes `dataset_item` and `llm_output` as arguments."
-            )
-
-        if prompt.model is None:
-            prompt.model = self.model
-        if prompt.model_kwargs is None:
-            prompt.model_kwargs = self.model_kwargs
-
-        if agent_class is None:
-            self.agent_class = create_litellm_agent_class(prompt)
-        else:
-            self.agent_class = agent_class
+        # Use base class validation and setup methods
+        self.validate_optimization_inputs(prompt, dataset, metric)
+        self.configure_prompt_model(prompt)
+        self.agent_class = self.setup_agent_class(prompt, agent_class)
 
         total_items = len(dataset.get_items())
         if n_samples is not None and n_samples > total_items:
@@ -499,7 +505,7 @@ class MetaPromptOptimizer(BaseOptimizer):
 
         optimization = None
         try:
-            optimization = self._opik_client.create_optimization(
+            optimization = self.opik_client.create_optimization(
                 dataset_name=dataset.name,
                 objective_name=getattr(metric, "__name__", str(metric)),
                 metadata={"optimizer": self.__class__.__name__},
@@ -633,26 +639,25 @@ class MetaPromptOptimizer(BaseOptimizer):
         self.auto_continue = auto_continue
         self.dataset = dataset
         self.prompt = prompt
-        self.llm_call_counter = 0  # Reset counter for run
+        self.reset_counters()  # Reset counters for run
         initial_prompt = prompt
 
         current_prompt = prompt
-        experiment_config = experiment_config or {}
-        experiment_config = {
-            **experiment_config,
-            **{
-                "optimizer": self.__class__.__name__,
-                "agent_class": self.agent_class.__name__,
-                "agent_config": prompt.to_dict(),
-                "metric": getattr(metric, "__name__", str(metric)),
-                "dataset": dataset.name,
-                "configuration": {
-                    "prompt": prompt.get_messages(),
-                    "rounds": self.rounds,
-                    "num_prompts_per_round": self.num_prompts_per_round,
-                },
-            },
-        }
+        configuration_updates = self._drop_none(
+            {
+                "rounds": self.rounds,
+                "num_prompts_per_round": self.num_prompts_per_round,
+            }
+        )
+        meta_metadata = {"stage": "initial"}
+        experiment_config = self._prepare_experiment_config(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            configuration_updates=configuration_updates,
+            additional_metadata={"meta_prompt": meta_metadata},
+        )
 
         with reporting.display_evaluation(verbose=self.verbose) as baseline_reporter:
             initial_score = self._evaluate_prompt(
@@ -887,6 +892,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             metric_name=getattr(metric, "__name__", str(metric)),
             details=details,
             llm_calls=self.llm_call_counter,
+            tool_calls=self.tool_call_counter,
             dataset_id=dataset_id,
             optimization_id=optimization_id,
             tool_prompts=tool_prompts,

opik_optimizer/mipro_optimizer/_lm.py

@@ -145,9 +145,16 @@ class LM(BaseLM):
         ):
             settings.usage_tracker.add_usage(self.model, dict(results.usage))
 
-        self.llm_call_counter += 1
+        self.increment_llm_counter()
         return results
 
+    def increment_llm_counter(self) -> None:
+        """Increment the LLM call counter."""
+        self.llm_call_counter += 1
+        parent = getattr(self, "parent_optimizer", None)
+        if parent is not None and hasattr(parent, "increment_llm_counter"):
+            parent.increment_llm_counter()
+
     def launch(self, launch_kwargs: dict[str, Any] | None = None):
         self.provider.launch(self, launch_kwargs)
 
@@ -302,7 +309,7 @@ def request_cache(maxsize: int | None = None):
     return decorator
 
 
-@request_cache(maxsize=None)
+@request_cache(maxsize=2000)
 def cached_litellm_completion(request: dict[str, Any], num_retries: int):
     return litellm_completion(
         request,
@@ -361,7 +368,7 @@ def litellm_completion(
     return stream_completion()
 
 
-@request_cache(maxsize=None)
+@request_cache(maxsize=2000)
 def cached_litellm_text_completion(request: dict[str, Any], num_retries: int):
     return litellm_text_completion(
         request,

opik_optimizer/mipro_optimizer/_mipro_optimizer_v2.py

@@ -79,7 +79,7 @@ class MIPROv2(Teleprompter):
         num_candidates: int = 10,
         num_threads: int | None = None,
         max_errors: int = 10,
-        seed: int = 9,
+        seed: int = 42,
         init_temperature: float = 0.5,
         verbose: bool = False,
         track_stats: bool = True,

opik_optimizer/mipro_optimizer/mipro_optimizer.py

@@ -1,7 +1,7 @@
 import os
 import random
 from datetime import datetime
-from typing import Literal
+from typing import Any, Literal
 from collections.abc import Callable
 import logging
 
@@ -15,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 (
@@ -45,14 +45,26 @@ class MiproOptimizer(BaseOptimizer):
         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: str | Dataset,
@@ -84,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
 
@@ -239,23 +251,57 @@ class MiproOptimizer(BaseOptimizer):
 
     def optimize_prompt(
         self,
+        prompt: chat_prompt.ChatPrompt,
         dataset: str | Dataset,
         metric: Callable,
-        task_config: TaskConfig,
-        num_candidates: int = 10,
         experiment_config: dict | None = None,
-        num_trials: int | None = 3,
         n_samples: int | None = 10,
-        auto: Literal["light", "medium", "heavy"] | None = "light",
+        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,
@@ -315,19 +361,18 @@ class MiproOptimizer(BaseOptimizer):
         **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):
@@ -396,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.
@@ -516,7 +574,8 @@ class MiproOptimizer(BaseOptimizer):
                 ),
                 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"]
@@ -548,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:
@@ -556,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=[
@@ -574,7 +642,8 @@ class MiproOptimizer(BaseOptimizer):
                 ),
                 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"]
@@ -592,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",
@@ -601,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/optimizable_agent.py

@@ -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

opik_optimizer/optimization_result.py

@@ -27,6 +27,7 @@ class OptimizationResult(pydantic.BaseModel):
     details: dict[str, Any] = pydantic.Field(default_factory=dict)
     history: list[dict[str, Any]] = []
     llm_calls: int | None = None
+    tool_calls: int | None = None
 
     # MIPRO specific
     demonstrations: list[dict[str, Any]] | None = None