opik-optimizer 0.9.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

opik_optimizer/__init__.py

@@ -3,15 +3,18 @@ import logging
 
 from opik.evaluation.models.litellm import warning_filters
 
-from opik_optimizer.evolutionary_optimizer.evolutionary_optimizer import EvolutionaryOptimizer
+from opik_optimizer.evolutionary_optimizer.evolutionary_optimizer import (
+    EvolutionaryOptimizer,
+)
 
 from . import datasets
+from .optimizable_agent import OptimizableAgent
+from .optimization_config.chat_prompt import ChatPrompt
 from .base_optimizer import BaseOptimizer
 from .few_shot_bayesian_optimizer import FewShotBayesianOptimizer
 from .logging_config import setup_logging
 from .meta_prompt_optimizer import MetaPromptOptimizer
 from .mipro_optimizer import MiproOptimizer
-from .optimization_config.chat_prompt import ChatPrompt
 from .optimization_config.configs import TaskConfig
 from .optimization_result import OptimizationResult
 
@@ -30,7 +33,8 @@ __all__ = [
     "MiproOptimizer",
     "EvolutionaryOptimizer",
     "OptimizationResult",
+    "OptimizableAgent",
     "setup_logging",
     "datasets",
-    "TaskConfig"
+    "TaskConfig",
 ]

opik_optimizer/_throttle.py

@@ -10,34 +10,34 @@ class RateLimiter:
     """
     Rate limiter that enforces a maximum number of calls across all threads using pyrate_limiter.
     """
+
     def __init__(self, max_calls_per_second: int):
         self.max_calls_per_second = max_calls_per_second
         rate = pyrate_limiter.Rate(max_calls_per_second, pyrate_limiter.Duration.SECOND)
 
         self.limiter = pyrate_limiter.Limiter(rate, raise_when_fail=False)
         self.bucket_key = "global_rate_limit"
-        
+
     def acquire(self) -> None:
         while not self.limiter.try_acquire(self.bucket_key):
             time.sleep(0.01)
 
+
 def rate_limited(limiter: RateLimiter) -> Callable[[Callable], Callable]:
     """Decorator to rate limit a function using the provided limiter"""
 
     def decorator(func: Callable) -> Callable:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs) -> Any:
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
             limiter.acquire()
             return func(*args, **kwargs)
+
         return wrapper
+
     return decorator
 
 
 def get_rate_limiter_for_current_opik_installation() -> RateLimiter:
     opik_config = opik.config.OpikConfig()
-    max_calls_per_second = (
-        10
-        if opik_config.is_cloud_installation
-        else 50
-    )
-    return RateLimiter(max_calls_per_second=max_calls_per_second)
+    max_calls_per_second = 10 if opik_config.is_cloud_installation else 50
+    return RateLimiter(max_calls_per_second=max_calls_per_second)

opik_optimizer/base_optimizer.py

@@ -1,16 +1,23 @@
+from typing import Any, Callable, Dict, List, Optional, Type
+
 import logging
 import time
 from abc import abstractmethod
-from typing import Any, Callable, Dict, List, Optional
+import random
+
 
 import litellm
-import opik
 from opik.rest_api.core import ApiError
+from opik.api_objects import optimization
+from opik import Dataset
 from pydantic import BaseModel
 
 from . import _throttle, optimization_result
 from .cache_config import initialize_cache
-from .optimization_config import chat_prompt
+from .optimization_config import chat_prompt, mappers
+from .optimizable_agent import OptimizableAgent
+from .utils import create_litellm_agent_class
+from . import task_evaluator
 
 _limiter = _throttle.get_rate_limiter_for_current_opik_installation()
 
@@ -34,22 +41,25 @@ class OptimizationRound(BaseModel):
 
 
 class BaseOptimizer:
-    def __init__(self, model: str, project_name: Optional[str] = None, verbose: int = 1, **model_kwargs):
+    def __init__(
+        self,
+        model: str,
+        verbose: int = 1,
+        **model_kwargs: Any,
+    ) -> None:
         """
         Base class for optimizers.
 
         Args:
            model: LiteLLM model name
-           project_name: Opik project name
            verbose: Controls internal logging/progress bars (0=off, 1=on).
            model_kwargs: additional args for model (eg, temperature)
         """
         self.model = model
         self.reasoning_model = model
         self.model_kwargs = model_kwargs
-        self.project_name = project_name
         self.verbose = verbose
-        self._history = []
+        self._history: List[OptimizationRound] = []
         self.experiment_config = None
         self.llm_call_counter = 0
 
@@ -59,18 +69,18 @@ class BaseOptimizer:
     @abstractmethod
     def optimize_prompt(
         self,
-        prompt: chat_prompt.ChatPrompt,
-        dataset: opik.Dataset,
-        metrics: List[Callable],
+        prompt: "chat_prompt.ChatPrompt",
+        dataset: Dataset,
+        metric: Callable,
         experiment_config: Optional[Dict] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> optimization_result.OptimizationResult:
         """
         Optimize a prompt.
 
         Args:
            dataset: Opik dataset name, or Opik dataset
-           metrics: A list of metric functions, these functions should have two arguments:
+           metric: A metric function, this function should have two arguments:
                dataset_item and llm_output
            prompt: the prompt to optimize
            input_key: input field of dataset
@@ -80,36 +90,7 @@ class BaseOptimizer:
         """
         pass
 
-    @abstractmethod
-    def evaluate_prompt(
-        self,
-        prompt: chat_prompt.ChatPrompt,
-        dataset: opik.Dataset,
-        metrics: List[Callable],
-        n_samples: Optional[int] = None,
-        dataset_item_ids: Optional[List[str]] = None,
-        experiment_config: Optional[Dict] = None,
-        **kwargs,
-    ) -> float:
-        """
-        Evaluate a prompt.
-
-        Args:
-           prompt: the prompt to evaluate
-           dataset: Opik dataset name, or Opik dataset
-           metrics: A list of metric functions, these functions should have two arguments:
-               dataset_item and llm_output
-           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:
-            float: The evaluation score
-        """
-        pass
-
-    def get_history(self) -> List[Dict[str, Any]]:
+    def get_history(self) -> List[OptimizationRound]:
         """
         Get the optimization history.
 
@@ -118,7 +99,7 @@ class BaseOptimizer:
         """
         return self._history
 
-    def _add_to_history(self, round_data: Dict[str, Any]):
+    def _add_to_history(self, round_data: OptimizationRound) -> None:
         """
         Add a round to the optimization history.
 
@@ -127,8 +108,9 @@ class BaseOptimizer:
         """
         self._history.append(round_data)
 
-            
-    def update_optimization(self, optimization, status: str) -> None:
+    def update_optimization(
+        self, optimization: optimization.Optimization, status: str
+    ) -> None:
         """
         Update the optimization status
         """
@@ -143,3 +125,74 @@ class BaseOptimizer:
                 time.sleep(5)
         if count == 3:
             logger.warning("Unable to update optimization status; continuing...")
+
+    def evaluate_prompt(
+        self,
+        prompt: chat_prompt.ChatPrompt,
+        dataset: Dataset,
+        metric: Callable,
+        n_threads: int,
+        verbose: int = 1,
+        dataset_item_ids: Optional[List[str]] = None,
+        experiment_config: Optional[Dict] = None,
+        n_samples: Optional[int] = None,
+        seed: Optional[int] = None,
+        agent_class: Optional[Type[OptimizableAgent]] = None,
+    ) -> float:
+        random.seed(seed)
+
+        if prompt.model is None:
+            prompt.model = self.model
+        if prompt.model_kwargs is None:
+            prompt.model_kwargs = self.model_kwargs
+
+        self.agent_class: Type[OptimizableAgent]
+
+        if agent_class is None:
+            self.agent_class = create_litellm_agent_class(prompt)
+        else:
+            self.agent_class = agent_class
+
+        agent = self.agent_class(prompt)
+
+        def llm_task(dataset_item: Dict[str, Any]) -> Dict[str, str]:
+            messages = prompt.get_messages(dataset_item)
+            raw_model_output = agent.invoke(messages)
+            cleaned_model_output = raw_model_output.strip()
+            result = {
+                mappers.EVALUATED_LLM_TASK_OUTPUT: cleaned_model_output,
+            }
+            return result
+
+        experiment_config = experiment_config or {}
+        experiment_config["project_name"] = self.__class__.__name__
+        experiment_config = {
+            **experiment_config,
+            **{
+                "agent_class": self.agent_class.__name__,
+                "agent_config": prompt.to_dict(),
+                "metric": metric.__name__,
+                "dataset": dataset.name,
+                "configuration": {"prompt": (prompt.get_messages() if prompt else [])},
+            },
+        }
+
+        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)
+
+        score = task_evaluator.evaluate(
+            dataset=dataset,
+            dataset_item_ids=dataset_item_ids,
+            metric=metric,
+            evaluated_task=llm_task,
+            num_threads=n_threads,
+            project_name=self.agent_class.project_name,
+            experiment_config=experiment_config,
+            optimization_id=None,
+            verbose=verbose,
+        )
+        return score

opik_optimizer/cache_config.py

@@ -13,12 +13,14 @@ CACHE_CONFIG = {
     "disk_cache_dir": CACHE_DIR,
 }
 
-def initialize_cache():
+
+def initialize_cache() -> Cache:
     """Initialize the LiteLLM cache with custom configuration."""
     litellm.cache = Cache(**CACHE_CONFIG)
     return litellm.cache
 
-def clear_cache():
+
+def clear_cache() -> None:
     """Clear the LiteLLM cache."""
     if litellm.cache:
-        litellm.cache.clear() 
+        litellm.cache.clear()

opik_optimizer/datasets/ai2_arc.py

@@ -1,8 +1,7 @@
 import opik
 
-def ai2_arc(
-    test_mode: bool = False
-) -> opik.Dataset:
+
+def ai2_arc(test_mode: bool = False) -> opik.Dataset:
     """
     Dataset containing the first 300 samples of the AI2 ARC dataset.
     """
@@ -11,12 +10,14 @@ def ai2_arc(
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         import datasets as ds
 
@@ -24,19 +25,20 @@ def ai2_arc(
         download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
         ds.disable_progress_bar()
         hf_dataset = ds.load_dataset(
-            "ai2_arc", "ARC-Challenge",
-            streaming=True, download_config=download_config
+            "ai2_arc", "ARC-Challenge", streaming=True, download_config=download_config
         )
-        
+
         data = []
         for i, item in enumerate(hf_dataset["train"]):
             if i >= nb_items:
                 break
-            data.append({
-                "question": item["question"],
-                "answer": item["answerKey"],
-                "choices": item["choices"],
-            })
+            data.append(
+                {
+                    "question": item["question"],
+                    "answer": item["answerKey"],
+                    "choices": item["choices"],
+                }
+            )
         ds.enable_progress_bar()
 
         dataset.insert(data)

opik_optimizer/datasets/cnn_dailymail.py

@@ -1,8 +1,7 @@
 import opik
 
-def cnn_dailymail(
-    test_mode: bool = False
-) -> opik.Dataset:
+
+def cnn_dailymail(test_mode: bool = False) -> opik.Dataset:
     """
     Dataset containing the first 100 samples of the CNN Daily Mail dataset.
     """
@@ -11,30 +10,35 @@ def cnn_dailymail(
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         import datasets as ds
-        
+
         download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
         ds.disable_progress_bar()
-        hf_dataset = ds.load_dataset("cnn_dailymail", "3.0.0", streaming=True, download_config=download_config)
-        
+        hf_dataset = ds.load_dataset(
+            "cnn_dailymail", "3.0.0", streaming=True, download_config=download_config
+        )
+
         data = []
         for i, item in enumerate(hf_dataset["validation"]):
             if i >= nb_items:
                 break
-            data.append({
-                "article": item["article"],
-                "highlights": item["highlights"],
-            })
+            data.append(
+                {
+                    "article": item["article"],
+                    "highlights": item["highlights"],
+                }
+            )
         ds.enable_progress_bar()
-        
+
         dataset.insert(data)
-        
+
         return dataset
-    

opik_optimizer/datasets/election_questions.py

@@ -1,33 +1,32 @@
 import opik
 
 
-def election_questions(
-    test_mode: bool = False
-) -> opik.Dataset:
+def election_questions(test_mode: bool = False) -> opik.Dataset:
     dataset_name = "election_questions" if not test_mode else "election_questions_test"
     nb_items = 300 if not test_mode else 5
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         import datasets as ds
 
         # Load data from file and insert into the dataset
         download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
         ds.disable_progress_bar()
-        hf_dataset = ds.load_dataset("Anthropic/election_questions", download_config=download_config)
-        
+        hf_dataset = ds.load_dataset(
+            "Anthropic/election_questions", download_config=download_config
+        )
+
         data = [
-            {
-                "question": item["question"],
-                "label": item["label"]
-            }
+            {"question": item["question"], "label": item["label"]}
             for item in hf_dataset["test"].select(range(nb_items))
         ]
         ds.enable_progress_bar()

opik_optimizer/datasets/gsm8k.py

@@ -1,8 +1,7 @@
 import opik
 
-def gsm8k(
-    test_mode: bool = False
-) -> opik.Dataset:
+
+def gsm8k(test_mode: bool = False) -> opik.Dataset:
     """
     Dataset containing the first 300 samples of the GSM8K dataset.
     """
@@ -11,28 +10,34 @@ def gsm8k(
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         import datasets as ds
 
         # Load data from file and insert into the dataset
         download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
         ds.disable_progress_bar()
-        hf_dataset = ds.load_dataset("gsm8k", "main", streaming=True, download_config=download_config)
-        
+        hf_dataset = ds.load_dataset(
+            "gsm8k", "main", streaming=True, download_config=download_config
+        )
+
         data = []
         for i, item in enumerate(hf_dataset["train"]):
             if i >= nb_items:
                 break
-            data.append({
-                "question": item["question"],
-                "answer": item["answer"],
-            })
+            data.append(
+                {
+                    "question": item["question"],
+                    "answer": item["answer"],
+                }
+            )
         ds.enable_progress_bar()
 
         dataset.insert(data)

opik_optimizer/datasets/halu_eval.py

@@ -1,8 +1,7 @@
 import opik
 
-def halu_eval_300(
-    test_mode: bool = False
-) -> opik.Dataset:
+
+def halu_eval_300(test_mode: bool = False) -> opik.Dataset:
     """
     Dataset containing the first 300 samples of the HaluEval dataset.
     """
@@ -11,12 +10,14 @@ def halu_eval_300(
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         import pandas as pd
 

opik_optimizer/datasets/hotpot_qa.py

@@ -3,9 +3,7 @@ from importlib.resources import files
 import json
 
 
-def hotpot_300(
-    test_mode: bool = False
-) -> opik.Dataset:
+def hotpot_300(test_mode: bool = False) -> opik.Dataset:
     """
     Dataset containing the first 300 samples of the HotpotQA dataset.
     """
@@ -14,15 +12,19 @@ def hotpot_300(
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         # Load data from file and insert into the dataset
-        json_content = (files('opik_optimizer') / 'data' / 'hotpot-500.json').read_text(encoding='utf-8')
+        json_content = (files("opik_optimizer") / "data" / "hotpot-500.json").read_text(
+            encoding="utf-8"
+        )
         all_data = json.loads(json_content)
         trainset = all_data[:nb_items]
 
@@ -33,9 +35,8 @@ def hotpot_300(
         dataset.insert(data)
         return dataset
 
-def hotpot_500(
-    test_mode: bool = False
-) -> opik.Dataset:
+
+def hotpot_500(test_mode: bool = False) -> opik.Dataset:
     """
     Dataset containing the first 500 samples of the HotpotQA dataset.
     """
@@ -44,15 +45,19 @@ def hotpot_500(
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         # Load data from file and insert into the dataset
-        json_content = (files('opik_optimizer') / 'data' / 'hotpot-500.json').read_text(encoding='utf-8')
+        json_content = (files("opik_optimizer") / "data" / "hotpot-500.json").read_text(
+            encoding="utf-8"
+        )
         all_data = json.loads(json_content)
         trainset = all_data[:nb_items]
 
@@ -62,7 +67,3 @@ def hotpot_500(
 
         dataset.insert(data)
         return dataset
-
-    
-
-    

opik_optimizer/datasets/medhallu.py

@@ -1,27 +1,30 @@
 import opik
 
-def medhallu(
-    test_mode: bool = False
-) -> opik.Dataset:
+
+def medhallu(test_mode: bool = False) -> opik.Dataset:
     dataset_name = "medhallu" if not test_mode else "medhallu_test"
     nb_items = 300 if not test_mode else 5
 
     client = opik.Opik()
     dataset = client.get_or_create_dataset(dataset_name)
-    
+
     items = dataset.get_items()
     if len(items) == nb_items:
         return dataset
     elif len(items) != 0:
-        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+        raise ValueError(
+            f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it."
+        )
     elif len(items) == 0:
         import datasets as ds
 
         # Load data from file and insert into the dataset
         download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
         ds.disable_progress_bar()
-        hf_dataset = ds.load_dataset("UTAustin-AIHealth/MedHallu", "pqa_labeled", download_config=download_config)
-        
+        hf_dataset = ds.load_dataset(
+            "UTAustin-AIHealth/MedHallu", "pqa_labeled", download_config=download_config
+        )
+
         data = [
             {
                 "question": item["Question"],