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

opik_optimizer/few_shot_bayesian_optimizer/few_shot_bayesian_optimizer.py

@@ -1,5 +1,8 @@
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+from typing import Any
+from collections.abc import Callable
+import warnings
 
+import copy
 import json
 import logging
 import random
@@ -15,7 +18,6 @@ from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
 from pydantic import BaseModel
 
 from opik_optimizer import base_optimizer
-from ..utils import create_litellm_agent_class
 from ..optimization_config import chat_prompt, mappers
 from ..optimizable_agent import OptimizableAgent
 from .. import _throttle, optimization_result, task_evaluator, utils
@@ -56,7 +58,7 @@ Respond only with the JSON object. Do not include any explanation or extra text.
 
 
 class FewShotPromptTemplate(BaseModel):
-    message_list_with_placeholder: List[Dict[str, str]]
+    message_list_with_placeholder: list[dict[str, str]]
     example_template: str
 
 
@@ -94,8 +96,11 @@ class FewShotBayesianOptimizer(base_optimizer.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"]
 
@@ -111,18 +116,22 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         elif self.verbose == 2:
             logger.setLevel(logging.DEBUG)
 
-        self._opik_client = opik.Opik()
-        self.llm_call_counter = 0
         logger.debug(f"Initialized FewShotBayesianOptimizer with model: {model}")
 
+    def get_optimizer_metadata(self) -> dict[str, Any]:
+        return {
+            "min_examples": self.min_examples,
+            "max_examples": self.max_examples,
+        }
+
     @_throttle.rate_limited(_limiter)
     def _call_model(
         self,
         model: str,
-        messages: List[Dict[str, str]],
+        messages: list[dict[str, str]],
         seed: int,
-        model_kwargs: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        model_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
         """
         Args:
             model: The model to use for the call
@@ -133,7 +142,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         Returns:
             Dict containing the model's response
         """
-        self.llm_call_counter += 1
+        self.increment_llm_counter()
 
         current_model_kwargs = self.model_kwargs.copy()
         current_model_kwargs.update(model_kwargs)
@@ -159,8 +168,8 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         return response
 
     def _split_dataset(
-        self, dataset: List[Dict[str, Any]], train_ratio: float
-    ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+        self, dataset: list[dict[str, Any]], train_ratio: float
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
         """
         Split the dataset into training and validation sets.
 
@@ -194,7 +203,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         self,
         model: str,
         prompt: chat_prompt.ChatPrompt,
-        few_shot_examples: List[Dict[str, Any]],
+        few_shot_examples: list[dict[str, Any]],
     ) -> FewShotPromptTemplate:
         """
         Generate a few-shot prompt template that can be used to insert examples into the prompt.
@@ -215,7 +224,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             "examples": few_shot_examples,
         }
 
-        messages: List[Dict[str, str]] = [
+        messages: list[dict[str, str]] = [
             {"role": "system", "content": SYSTEM_PROMPT_TEMPLATE},
             {"role": "user", "content": json.dumps(user_message)},
         ]
@@ -244,9 +253,9 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         metric: Callable,
         baseline_score: float,
         n_trials: int = 10,
-        optimization_id: Optional[str] = None,
-        experiment_config: Optional[Dict] = None,
-        n_samples: Optional[int] = None,
+        optimization_id: str | None = None,
+        experiment_config: dict | None = None,
+        n_samples: int | None = None,
     ) -> optimization_result.OptimizationResult:
         reporting.start_optimization_run(verbose=self.verbose)
 
@@ -259,19 +268,20 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         if n_samples is not None and n_samples < len(dataset_items):
             eval_dataset_item_ids = random.sample(all_dataset_item_ids, n_samples)
 
-        # Define the experiment configuration
-        experiment_config = experiment_config or {}
-        base_experiment_config = {  # Base config for reuse
-            **experiment_config,
-            **{
-                "optimizer": self.__class__.__name__,
-                "agent_class": self.agent_class.__name__,
-                "agent_config": prompt.to_dict(),
-                "metric": metric.__name__,
-                "dataset": dataset.name,
-                "configuration": {},
-            },
-        }
+        configuration_updates = self._drop_none(
+            {
+                "n_trials": n_trials,
+                "n_samples": n_samples,
+                "baseline_score": baseline_score,
+            }
+        )
+        base_experiment_config = self._prepare_experiment_config(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            configuration_updates=configuration_updates,
+        )
 
         # Start Optuna Study
         def optimization_objective(trial: optuna.Trial) -> float:
@@ -326,7 +336,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             ]
 
             # Log trial config
-            trial_config = base_experiment_config.copy()
+            trial_config = copy.deepcopy(base_experiment_config)
             trial_config["configuration"]["prompt"] = (
                 messages_for_reporting  # Base instruction
             )
@@ -450,6 +460,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             best_score=best_score,
             best_prompt=best_prompt,
             verbose=self.verbose,
+            tools=getattr(prompt, "tools", None),
         )
 
         return optimization_result.OptimizationResult(
@@ -479,6 +490,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             },
             history=optuna_history_processed,
             llm_calls=self.llm_call_counter,
+            tool_calls=self.tool_call_counter,
             dataset_id=dataset.id,
             optimization_id=optimization_id,
         )
@@ -488,47 +500,39 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         prompt: chat_prompt.ChatPrompt,
         dataset: Dataset,
         metric: Callable,
-        n_trials: int = 10,
-        agent_class: Optional[Type[OptimizableAgent]] = None,
-        experiment_config: Optional[Dict] = None,
-        n_samples: Optional[int] = None,
+        experiment_config: dict | None = None,
+        n_samples: int | None = None,
+        auto_continue: bool = False,
+        agent_class: type[OptimizableAgent] | None = None,
+        **kwargs: Any,
     ) -> optimization_result.OptimizationResult:
         """
         Args:
-            prompt:
+            prompt: The prompt to optimize
             dataset: Opik Dataset to optimize on
             metric: Metric function to evaluate on
-            n_trials: Number of trials for Bayesian Optimization
             experiment_config: Optional configuration for the experiment, useful to log additional metadata
             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: Additional parameters including:
+                n_trials (int): Number of trials for Bayesian Optimization (default: 10)
+                mcp_config (MCPExecutionConfig | None): MCP tool calling configuration (default: None)
 
         Returns:
             OptimizationResult: Result of the optimization
         """
-        if not isinstance(prompt, chat_prompt.ChatPrompt):
-            raise ValueError("Prompt must be a ChatPrompt object")
+        # 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)
 
-        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
+        # Extract n_trials from kwargs for backward compatibility
+        n_trials = kwargs.get("n_trials", 10)
 
         optimization = None
         try:
-            optimization = self._opik_client.create_optimization(
+            optimization = self.opik_client.create_optimization(
                 dataset_name=dataset.name,
                 objective_name=metric.__name__,
                 metadata={"optimizer": self.__class__.__name__},
@@ -557,6 +561,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
                 "n_samples": n_samples,
             },
             verbose=self.verbose,
+            tools=getattr(prompt, "tools", None),
         )
 
         utils.disable_experiment_reporting()
@@ -614,10 +619,10 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         prompt: chat_prompt.ChatPrompt,
         dataset: opik.Dataset,
         metric: Callable,
-        n_samples: Optional[int] = None,
-        dataset_item_ids: Optional[List[str]] = None,
-        experiment_config: Optional[Dict] = None,
-        optimization_id: Optional[str] = None,
+        n_samples: int | None = None,
+        dataset_item_ids: list[str] | None = None,
+        experiment_config: dict | None = None,
+        optimization_id: str | None = None,
         **kwargs: Any,
     ) -> float:
         """
@@ -633,20 +638,6 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
         """
         llm_task = self._build_task_from_messages(prompt, prompt.get_messages())
 
-        experiment_config = experiment_config or {}
-        experiment_config["project_name"] = self.agent_class.__name__
-        experiment_config = {
-            **experiment_config,
-            **{
-                "optimizer": self.__class__.__name__,
-                "agent_class": self.agent_class.__name__,
-                "agent_config": prompt.to_dict(),
-                "metric": metric.__name__,
-                "dataset": dataset.name,
-                "configuration": {"prompt": prompt.get_messages()},
-            },
-        }
-
         if n_samples is not None:
             if dataset_item_ids is not None:
                 raise Exception("Can't use n_samples and dataset_item_ids")
@@ -654,6 +645,24 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             all_ids = [dataset_item["id"] for dataset_item in dataset.get_items()]
             dataset_item_ids = random.sample(all_ids, n_samples)
 
+        configuration_updates = self._drop_none(
+            {
+                "n_samples": n_samples,
+                "dataset_item_ids": dataset_item_ids,
+            }
+        )
+        additional_metadata = (
+            {"optimization_id": optimization_id} if optimization_id else None
+        )
+        experiment_config = self._prepare_experiment_config(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            configuration_updates=configuration_updates,
+            additional_metadata=additional_metadata,
+        )
+
         logger.debug("Starting FewShotBayesian evaluation...")
         score = task_evaluator.evaluate(
             dataset=dataset,
@@ -661,7 +670,7 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
             metric=metric,
             evaluated_task=llm_task,
             num_threads=self.n_threads,
-            project_name=self.agent_class.project_name,
+            project_name=experiment_config.get("project_name"),
             experiment_config=experiment_config,
             optimization_id=optimization_id,
             verbose=self.verbose,
@@ -673,14 +682,14 @@ class FewShotBayesianOptimizer(base_optimizer.BaseOptimizer):
     def _build_task_from_messages(
         self,
         prompt: chat_prompt.ChatPrompt,
-        messages: List[Dict[str, str]],
-        few_shot_examples: Optional[str] = None,
-    ) -> Callable[[Dict[str, Any]], Dict[str, Any]]:
+        messages: list[dict[str, str]],
+        few_shot_examples: str | None = None,
+    ) -> Callable[[dict[str, Any]], dict[str, Any]]:
         new_prompt = prompt.copy()
         new_prompt.set_messages(messages)
         agent = self.agent_class(new_prompt)
 
-        def llm_task(dataset_item: Dict[str, Any]) -> Dict[str, Any]:
+        def llm_task(dataset_item: dict[str, Any]) -> dict[str, Any]:
             """
             Process a single dataset item through the LLM task.
 

opik_optimizer/few_shot_bayesian_optimizer/reporting.py

@@ -1,6 +1,6 @@
 from contextlib import contextmanager
 from io import StringIO
-from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from typing import Any, Optional, TYPE_CHECKING
 
 from rich.panel import Panel
 from rich.text import Text
@@ -46,9 +46,16 @@ def display_evaluation(
                 yield Reporter()
             finally:
                 if verbose >= 1:
-                    console.print(
-                        Text(f"\r  Baseline score was: {score:.4f}.\n", style="green")
-                    )
+                    if score is not None:
+                        console.print(
+                            Text(
+                                f"\r  Baseline score was: {score:.4f}.\n", style="green"
+                            )
+                        )
+                    else:
+                        console.print(
+                            Text("\r  Baseline score was: None\n", style="red")
+                        )
 
 
 @contextmanager
@@ -121,7 +128,7 @@ def start_optimization_trial(
 
     # Create a simple object with a method to set the score
     class Reporter:
-        def start_trial(self, messages: List[Dict[str, str]]) -> None:
+        def start_trial(self, messages: list[dict[str, str]]) -> None:
             if verbose >= 1:
                 console.print(
                     Text(

opik_optimizer/gepa_optimizer/__init__.py

@@ -0,0 +1,3 @@
+from .gepa_optimizer import GepaOptimizer
+
+__all__ = ["GepaOptimizer"]

opik_optimizer/gepa_optimizer/adapter.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+from collections.abc import Callable, Iterable
+
+import logging
+
+from gepa.core.adapter import EvaluationBatch, GEPAAdapter
+
+from ..optimization_config import chat_prompt
+from ..utils import create_litellm_agent_class
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OpikDataInst:
+    """Data instance handed to GEPA.
+
+    We keep the original Opik dataset item so metrics and prompt formatting can use it
+    directly without duplicated bookkeeping.
+    """
+
+    input_text: str
+    answer: str
+    additional_context: dict[str, str]
+    opik_item: dict[str, Any]
+
+
+def _extract_system_text(candidate: dict[str, str], fallback: str) -> str:
+    for key in ("system_prompt", "system", "prompt"):
+        value = candidate.get(key)
+        if isinstance(value, str) and value.strip():
+            return value
+    return fallback
+
+
+def _apply_system_text(
+    prompt_obj: chat_prompt.ChatPrompt, system_text: str
+) -> chat_prompt.ChatPrompt:
+    updated = prompt_obj.copy()
+    if updated.messages is not None:
+        messages = updated.get_messages()
+        if messages and messages[0].get("role") == "system":
+            messages[0]["content"] = system_text
+        else:
+            messages.insert(0, {"role": "system", "content": system_text})
+        updated.set_messages(messages)
+    else:
+        updated.system = system_text
+    return updated
+
+
+class OpikGEPAAdapter(GEPAAdapter[OpikDataInst, dict[str, Any], dict[str, Any]]):
+    """Minimal GEPA adapter that routes evaluation through Opik's metric."""
+
+    def __init__(
+        self,
+        base_prompt: chat_prompt.ChatPrompt,
+        optimizer: Any,
+        metric: Callable[[dict[str, Any], str], Any],
+        system_fallback: str,
+    ) -> None:
+        self._base_prompt = base_prompt
+        self._optimizer = optimizer
+        self._metric = metric
+        self._system_fallback = system_fallback
+
+    def evaluate(
+        self,
+        batch: list[OpikDataInst],
+        candidate: dict[str, str],
+        capture_traces: bool = False,
+    ) -> EvaluationBatch[dict[str, Any], dict[str, Any]]:
+        system_text = _extract_system_text(candidate, self._system_fallback)
+        prompt_variant = _apply_system_text(self._base_prompt, system_text)
+
+        agent_class = create_litellm_agent_class(
+            prompt_variant, optimizer_ref=self._optimizer
+        )
+        agent = agent_class(prompt_variant)
+
+        outputs: list[dict[str, Any]] = []
+        scores: list[float] = []
+        trajectories: list[dict[str, Any]] | None = [] if capture_traces else None
+
+        for inst in batch:
+            dataset_item = inst.opik_item
+            messages = prompt_variant.get_messages(dataset_item)
+            raw_output = agent.invoke(messages).strip()
+
+            metric_result = self._metric(dataset_item, raw_output)
+            if hasattr(metric_result, "value"):
+                score = float(metric_result.value)
+            elif hasattr(metric_result, "score"):
+                score = float(metric_result.score)
+            else:
+                score = float(metric_result)
+
+            outputs.append({"output": raw_output})
+            scores.append(score)
+            try:
+                self._optimizer._gepa_live_metric_calls += 1
+            except Exception:
+                pass
+
+            if trajectories is not None:
+                trajectories.append(
+                    {
+                        "input": dataset_item,
+                        "output": raw_output,
+                        "score": score,
+                    }
+                )
+
+        return EvaluationBatch(
+            outputs=outputs, scores=scores, trajectories=trajectories
+        )
+
+    def make_reflective_dataset(
+        self,
+        candidate: dict[str, str],
+        eval_batch: EvaluationBatch[dict[str, Any], dict[str, Any]],
+        components_to_update: list[str],
+    ) -> dict[str, list[dict[str, Any]]]:
+        components = components_to_update or ["system_prompt"]
+        trajectories = eval_batch.trajectories or []
+
+        def _records() -> Iterable[dict[str, Any]]:
+            for traj in trajectories:
+                dataset_item = traj.get("input", {})
+                output_text = traj.get("output", "")
+                score = traj.get("score", 0.0)
+                feedback = f"Observed score={score:.4f}. Expected answer: {dataset_item.get('answer', '')}"
+                yield {
+                    "Inputs": {
+                        "text": dataset_item.get("input")
+                        or dataset_item.get("question")
+                        or "",
+                    },
+                    "Generated Outputs": output_text,
+                    "Feedback": feedback,
+                }
+
+        reflective_records = list(_records())
+        if not reflective_records:
+            logger.debug(
+                "No trajectories captured for candidate; returning empty reflective dataset"
+            )
+            reflective_records = []
+
+        return {component: reflective_records for component in components}