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

opik_optimizer/gepa_optimizer/gepa_optimizer.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import logging
 from contextlib import nullcontext
 from typing import Any, ContextManager
@@ -12,15 +10,18 @@ from opik.evaluation.metrics.score_result import ScoreResult
 from ..base_optimizer import BaseOptimizer
 from ..optimization_config import chat_prompt, mappers
 from ..optimization_result import OptimizationResult
-from ..utils import optimization_context, create_litellm_agent_class
-from ..logging_config import setup_logging as _setup_logging
+from ..optimizable_agent import OptimizableAgent
+from ..utils import (
+    optimization_context,
+    create_litellm_agent_class,
+    disable_experiment_reporting,
+    enable_experiment_reporting,
+)
 from .. import task_evaluator
 from . import reporting as gepa_reporting
 from .adapter import OpikDataInst, OpikGEPAAdapter
 
-
-_setup_logging()
-LOGGER = logging.getLogger("opik_optimizer.gepa.optimizer")
+logger = logging.getLogger(__name__)
 
 
 class GepaOptimizer(BaseOptimizer):
@@ -32,14 +33,63 @@ class GepaOptimizer(BaseOptimizer):
         project_name: str | None = None,
         reflection_model: str | None = None,
         verbose: int = 1,
+        seed: int = 42,
         **model_kwargs: Any,
     ) -> None:
-        super().__init__(model=model, verbose=verbose, **model_kwargs)
+        # Validate required parameters
+        if model is None:
+            raise ValueError("model parameter is required and cannot be None")
+        if not isinstance(model, str):
+            raise ValueError(f"model must be a string, got {type(model).__name__}")
+        if not model.strip():
+            raise ValueError("model cannot be empty or whitespace-only")
+
+        # Validate optional parameters
+        if project_name is not None and not isinstance(project_name, str):
+            raise ValueError(
+                f"project_name must be a string or None, got {type(project_name).__name__}"
+            )
+
+        if reflection_model is not None and not isinstance(reflection_model, str):
+            raise ValueError(
+                f"reflection_model must be a string or None, got {type(reflection_model).__name__}"
+            )
+
+        if not isinstance(verbose, int):
+            raise ValueError(
+                f"verbose must be an integer, got {type(verbose).__name__}"
+            )
+        if verbose < 0:
+            raise ValueError("verbose must be non-negative")
+
+        if not isinstance(seed, int):
+            raise ValueError(f"seed must be an integer, got {type(seed).__name__}")
+
+        super().__init__(model=model, verbose=verbose, seed=seed, **model_kwargs)
         self.project_name = project_name
         self.reflection_model = reflection_model or model
         self.num_threads = self.model_kwargs.pop("num_threads", 6)
-        self.seed = self.model_kwargs.pop("seed", 42)
         self._gepa_live_metric_calls = 0
+        self._adapter = None  # Will be set during optimization
+
+    def get_optimizer_metadata(self) -> dict[str, Any]:
+        return {
+            "project_name": self.project_name,
+            "reflection_model": self.reflection_model,
+        }
+
+    def cleanup(self) -> None:
+        """
+        Clean up GEPA-specific resources.
+        """
+        # Call parent cleanup
+        super().cleanup()
+
+        # Clear GEPA-specific resources
+        self._adapter = None
+        self._gepa_live_metric_calls = 0
+
+        logger.debug("Cleaned up GEPA-specific resources")
 
     # ------------------------------------------------------------------
     # Helpers
@@ -105,21 +155,62 @@ class GepaOptimizer(BaseOptimizer):
     def optimize_prompt(
         self,
         prompt: chat_prompt.ChatPrompt,
-        dataset: str | Dataset,
-        metric: Callable[[dict[str, Any], str], ScoreResult],
-        experiment_config: dict[str, Any] | None = None,
+        dataset: Dataset,
+        metric: Callable,
+        experiment_config: dict | None = None,
+        n_samples: int | None = None,
+        auto_continue: bool = False,
+        agent_class: type[OptimizableAgent] | None = None,
         **kwargs: Any,
     ) -> OptimizationResult:
-        if isinstance(dataset, str):
-            client = opik.Opik(project_name=self.project_name)
-            dataset = client.get_dataset(dataset)
-
-        max_metric_calls: int = int(kwargs.get("max_metric_calls", 30))
+        """
+        Optimize a prompt using GEPA (Genetic-Pareto) algorithm.
+
+        Args:
+            prompt: The prompt to optimize
+            dataset: Opik Dataset to optimize on
+            metric: Metric function to evaluate on
+            experiment_config: Optional configuration for the experiment
+            n_samples: Optional number of items to test in the dataset
+            auto_continue: Whether to auto-continue optimization
+            agent_class: Optional agent class to use
+            **kwargs: GEPA-specific parameters:
+                max_metric_calls (int | None): Maximum number of metric evaluations (default: 30)
+                reflection_minibatch_size (int): Size of reflection minibatches (default: 3)
+                candidate_selection_strategy (str): Strategy for candidate selection (default: "pareto")
+                skip_perfect_score (bool): Skip candidates with perfect scores (default: True)
+                perfect_score (float): Score considered perfect (default: 1.0)
+                use_merge (bool): Enable merge operations (default: False)
+                max_merge_invocations (int): Maximum merge invocations (default: 5)
+                run_dir (str | None): Directory for run outputs (default: None)
+                track_best_outputs (bool): Track best outputs during optimization (default: False)
+                display_progress_bar (bool): Display progress bar (default: False)
+                seed (int): Random seed for reproducibility (default: 42)
+                raise_on_exception (bool): Raise exceptions instead of continuing (default: True)
+                mcp_config (MCPExecutionConfig | None): MCP tool calling configuration (default: None)
+
+        Returns:
+            OptimizationResult: Result of the optimization
+        """
+        # Use base class validation and setup methods
+        self.validate_optimization_inputs(prompt, dataset, metric)
+
+        # Extract GEPA-specific parameters from kwargs
+        max_metric_calls: int | None = kwargs.get("max_metric_calls", 30)
         reflection_minibatch_size: int = int(kwargs.get("reflection_minibatch_size", 3))
         candidate_selection_strategy: str = str(
             kwargs.get("candidate_selection_strategy", "pareto")
         )
-        n_samples: int | None = kwargs.get("n_samples")
+        skip_perfect_score: bool = kwargs.get("skip_perfect_score", True)
+        perfect_score: float = float(kwargs.get("perfect_score", 1.0))
+        use_merge: bool = kwargs.get("use_merge", False)
+        max_merge_invocations: int = int(kwargs.get("max_merge_invocations", 5))
+        run_dir: str | None = kwargs.get("run_dir", None)
+        track_best_outputs: bool = kwargs.get("track_best_outputs", False)
+        display_progress_bar: bool = kwargs.get("display_progress_bar", False)
+        seed: int = int(kwargs.get("seed", 42))
+        raise_on_exception: bool = kwargs.get("raise_on_exception", True)
+        kwargs.pop("mcp_config", None)  # Added for MCP support (for future use)
 
         prompt = prompt.copy()
         if self.project_name:
@@ -147,16 +238,19 @@ class GepaOptimizer(BaseOptimizer):
 
         opik_client = opik.Opik(project_name=self.project_name)
 
-        with optimization_context(
-            client=opik_client,
-            dataset_name=dataset.name,
-            objective_name=metric.__name__,
-            metadata={"optimizer": self.__class__.__name__},
-        ) as optimization:
-            try:
-                opt_id = optimization.id if optimization is not None else None
-            except Exception:
-                opt_id = None
+        disable_experiment_reporting()
+
+        try:
+            with optimization_context(
+                client=opik_client,
+                dataset_name=dataset.name,
+                objective_name=metric.__name__,
+                metadata={"optimizer": self.__class__.__name__},
+            ) as optimization:
+                try:
+                    opt_id = optimization.id if optimization is not None else None
+                except Exception:
+                    opt_id = None
 
             gepa_reporting.display_header(
                 algorithm="GEPA",
@@ -210,7 +304,7 @@ class GepaOptimizer(BaseOptimizer):
                         )
                     baseline.set_score(initial_score)
                 except Exception:
-                    LOGGER.exception("Baseline evaluation failed")
+                    logger.exception("Baseline evaluation failed")
 
             adapter_prompt = self._apply_system_text(base_prompt, seed_prompt_text)
             adapter_prompt.project_name = self.project_name
@@ -244,10 +338,17 @@ class GepaOptimizer(BaseOptimizer):
                 "task_lm": None,
                 "reflection_lm": self.reflection_model,
                 "candidate_selection_strategy": candidate_selection_strategy,
+                "skip_perfect_score": skip_perfect_score,
                 "reflection_minibatch_size": reflection_minibatch_size,
+                "perfect_score": perfect_score,
+                "use_merge": use_merge,
+                "max_merge_invocations": max_merge_invocations,
                 "max_metric_calls": max_metric_calls,
-                "display_progress_bar": False,
-                "track_best_outputs": False,
+                "run_dir": run_dir,
+                "track_best_outputs": track_best_outputs,
+                "display_progress_bar": display_progress_bar,
+                "seed": seed,
+                "raise_on_exception": raise_on_exception,
                 "logger": gepa_reporting.RichGEPAOptimizerLogger(
                     self, verbose=self.verbose
                 ),
@@ -265,10 +366,13 @@ class GepaOptimizer(BaseOptimizer):
             with gepa_reporting.start_gepa_optimization(verbose=self.verbose):
                 gepa_result = gepa.optimize(**kwargs_gepa)
 
-            try:
-                opt_id = optimization.id if optimization is not None else None
-            except Exception:
-                opt_id = None
+                try:
+                    opt_id = optimization.id if optimization is not None else None
+                except Exception:
+                    opt_id = None
+
+        finally:
+            enable_experiment_reporting()
 
         # ------------------------------------------------------------------
         # Rescoring & result assembly
@@ -308,7 +412,7 @@ class GepaOptimizer(BaseOptimizer):
             try:
                 score = float(self._evaluate_prompt_logged(**eval_kwargs))
             except Exception:
-                LOGGER.debug("Rescoring failed for candidate %s", idx, exc_info=True)
+                logger.debug("Rescoring failed for candidate %s", idx, exc_info=True)
                 score = 0.0
 
             rescored.append(score)
@@ -382,12 +486,12 @@ class GepaOptimizer(BaseOptimizer):
             try:
                 self._evaluate_prompt_logged(**final_eval_kwargs)
             except Exception:
-                LOGGER.debug("Final evaluation failed", exc_info=True)
+                logger.debug("Final evaluation failed", exc_info=True)
 
         per_item_scores: list[dict[str, Any]] = []
         try:
             analysis_prompt = final_prompt.copy()
-            agent_cls = create_litellm_agent_class(analysis_prompt)
+            agent_cls = create_litellm_agent_class(analysis_prompt, optimizer_ref=self)
             agent = agent_cls(analysis_prompt)
             for item in items:
                 messages = analysis_prompt.get_messages(item)
@@ -408,7 +512,7 @@ class GepaOptimizer(BaseOptimizer):
                     }
                 )
         except Exception:
-            LOGGER.debug("Per-item diagnostics failed", exc_info=True)
+            logger.debug("Per-item diagnostics failed", exc_info=True)
 
         details: dict[str, Any] = {
             "model": self.model,
@@ -420,13 +524,13 @@ class GepaOptimizer(BaseOptimizer):
             "val_scores": val_scores,
             "opik_rescored_scores": rescored,
             "candidate_summary": candidate_rows,
-            "best_candidate_iteration": candidate_rows[best_idx]["iteration"]
-            if candidate_rows
-            else 0,
+            "best_candidate_iteration": (
+                candidate_rows[best_idx]["iteration"] if candidate_rows else 0
+            ),
             "selected_candidate_index": best_idx,
-            "selected_candidate_gepa_score": val_scores[best_idx]
-            if best_idx < len(val_scores)
-            else None,
+            "selected_candidate_gepa_score": (
+                val_scores[best_idx] if best_idx < len(val_scores) else None
+            ),
             "selected_candidate_opik_score": best_score,
             "gepa_live_metric_used": True,
             "gepa_live_metric_call_count": self._gepa_live_metric_calls,
@@ -446,16 +550,16 @@ class GepaOptimizer(BaseOptimizer):
                 best_prompt_text, best_score, verbose=self.verbose
             )
 
-        if LOGGER.isEnabledFor(logging.DEBUG):
+        if logger.isEnabledFor(logging.DEBUG):
             for idx, row in enumerate(candidate_rows):
-                LOGGER.debug(
+                logger.debug(
                     "candidate=%s source=%s gepa=%s opik=%s",
                     idx,
                     row.get("source"),
                     row.get("gepa_score"),
                     row.get("opik_score"),
                 )
-            LOGGER.debug(
+            logger.debug(
                 "selected candidate idx=%s gepa=%s opik=%.4f",
                 best_idx,
                 details.get("selected_candidate_gepa_score"),
@@ -516,7 +620,8 @@ class GepaOptimizer(BaseOptimizer):
         if prompt.model_kwargs is None:
             prompt.model_kwargs = self.model_kwargs
 
-        agent_class = create_litellm_agent_class(prompt)
+        agent_class = create_litellm_agent_class(prompt, optimizer_ref=self)
+        self.agent_class = agent_class
         agent = agent_class(prompt)
 
         def llm_task(dataset_item: dict[str, Any]) -> dict[str, str]:
@@ -524,22 +629,14 @@ class GepaOptimizer(BaseOptimizer):
             raw = agent.invoke(messages)
             return {mappers.EVALUATED_LLM_TASK_OUTPUT: raw.strip()}
 
-        experiment_config = experiment_config or {}
-        experiment_config["project_name"] = agent_class.__name__
-        experiment_config = {
-            **experiment_config,
-            **{
-                "optimizer": self.__class__.__name__,
-                "agent_class": agent_class.__name__,
-                "agent_config": prompt.to_dict(),
-                "metric": metric.__name__,
-                "dataset": dataset.name,
-                "configuration": {
-                    "prompt": prompt.get_messages(),
-                    "gepa": (extra_metadata or {}),
-                },
-            },
-        }
+        configuration_updates = self._drop_none({"gepa": extra_metadata})
+        experiment_config = self._prepare_experiment_config(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            configuration_updates=configuration_updates,
+        )
 
         score = task_evaluator.evaluate(
             dataset=dataset,
@@ -547,7 +644,7 @@ class GepaOptimizer(BaseOptimizer):
             metric=metric,
             evaluated_task=llm_task,
             num_threads=self.num_threads,
-            project_name=agent_class.project_name,
+            project_name=experiment_config.get("project_name"),
             experiment_config=experiment_config,
             optimization_id=optimization_id,
             n_samples=n_samples,

opik_optimizer/mcp_utils/mcp_workflow.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 import contextlib
 import copy
 import io
+import json
 import logging
 import os
 import textwrap
@@ -346,12 +347,19 @@ class MCPToolInvocation:
     preview_label: str | None = None
     preview_chars: int = 160
     rate_limit_sleep: float = DEFAULT_MCP_RATELIMIT_SLEEP
+    cache_enabled: bool = True
     _logger: logging.Logger = field(default_factory=lambda: logger)
+    _cache: dict[str, str] = field(default_factory=dict, init=False)
 
     def __call__(self, **arguments: Any) -> str:
         return self.invoke(arguments)
 
-    def invoke(self, arguments: Mapping[str, Any]) -> str:
+    def clear_cache(self) -> None:
+        self._cache.clear()
+
+    def invoke(
+        self, arguments: Mapping[str, Any], *, use_cache: bool | None = None
+    ) -> str:
         def call_tool(name: str, payload: dict[str, Any]) -> Any:
             if self.rate_limit_sleep > 0:
                 time.sleep(self.rate_limit_sleep)
@@ -367,6 +375,19 @@ class MCPToolInvocation:
         if self.argument_adapter:
             prepared = self.argument_adapter(prepared, call_tool)
 
+        effective_cache = self.cache_enabled if use_cache is None else use_cache
+        cache_key: str | None = None
+        if effective_cache:
+            cache_key = self._make_cache_key(prepared)
+            cached_summary = self._cache.get(cache_key)
+            if cached_summary is not None:
+                if self.summary_handler:
+                    self.summary_handler.record_summary(cached_summary)
+                self._logger.debug(
+                    "MCP tool %s cache hit arguments=%s", self.tool_name, prepared
+                )
+                return cached_summary
+
         # TODO(opik-mcp): reuse a persistent MCP client so we avoid spawning a
         # new stdio subprocess for each call. This currently mirrors the
         # original blocking behaviour for stability.
@@ -391,11 +412,41 @@ class MCPToolInvocation:
         if self.summary_handler:
             self.summary_handler.record_summary(summary)
 
+        if effective_cache and cache_key is not None:
+            self._cache[cache_key] = summary
+
         if os.getenv("OPIK_DEBUG_MCP"):
             self._logger.info("MCP %s raw response:\n%s", label, text)
 
         return summary
 
+    def _make_cache_key(self, payload: Mapping[str, Any]) -> str:
+        try:
+            return json.dumps(payload, sort_keys=True, default=str)
+        except TypeError:
+            normalised = self._normalise_cache_payload(payload)
+            return json.dumps(normalised, sort_keys=True, default=str)
+
+    @staticmethod
+    def _normalise_cache_payload(value: Any) -> Any:
+        if isinstance(value, Mapping):
+            return {
+                key: MCPToolInvocation._normalise_cache_payload(val)
+                for key, val in sorted(value.items(), key=lambda item: str(item[0]))
+            }
+        if isinstance(value, list):
+            return [MCPToolInvocation._normalise_cache_payload(item) for item in value]
+        if isinstance(value, tuple):
+            return [MCPToolInvocation._normalise_cache_payload(item) for item in value]
+        if isinstance(value, set):
+            return [
+                MCPToolInvocation._normalise_cache_payload(item)
+                for item in sorted(value, key=repr)
+            ]
+        if isinstance(value, (str, int, float, bool)) or value is None:
+            return value
+        return str(value)
+
 
 def summarise_with_template(template: str) -> SummaryBuilder:
     """Return a summary builder that fills the provided template."""
@@ -465,6 +516,7 @@ def preview_second_pass(
     dataset_item: dict[str, Any],
     coordinator: MCPSecondPassCoordinator,
     agent_factory: Callable[[Any], Any],
+    seed: int = 42,
 ) -> None:
     """Debug helper mirroring the old inline scripts."""
 
@@ -472,7 +524,9 @@ def preview_second_pass(
     agent = agent_factory(prompt)
     base_messages = prompt.get_messages(dataset_item)
 
-    raw_output = agent.llm_invoke(messages=base_messages, seed=42, allow_tool_use=True)
+    raw_output = agent.llm_invoke(
+        messages=base_messages, seed=seed, allow_tool_use=True
+    )
     logger.debug("Raw model output: %s", raw_output)
 
     second_pass_messages = coordinator.build_second_pass_messages(
@@ -484,7 +538,7 @@ def preview_second_pass(
         logger.debug("Second-pass messages: %s", second_pass_messages)
         final_output = agent.llm_invoke(
             messages=second_pass_messages,
-            seed=101,
+            seed=seed,
             allow_tool_use=True,
         )
     else: