opik-optimizer 0.9.1__tar.gz → 1.0.0__tar.gz

{opik_optimizer-0.9.1/src/opik_optimizer.egg-info → opik_optimizer-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opik_optimizer
-Version: 0.9.1
+Version: 1.0.0
 Summary: Agent optimization with Opik
 Home-page: https://github.com/comet-ml/opik
 Author: Comet ML
@@ -12,17 +12,17 @@ Classifier: Programming Language :: Python :: 3.10
 Requires-Python: >=3.9,<3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: opik>=1.7.17
-Requires-Dist: dspy<=2.6.24,>=2.6.18
-Requires-Dist: litellm
-Requires-Dist: tqdm
 Requires-Dist: datasets
+Requires-Dist: deap>=1.4.3
+Requires-Dist: diskcache
+Requires-Dist: hf_xet
+Requires-Dist: litellm
+Requires-Dist: opik>=1.7.17
 Requires-Dist: optuna
-Requires-Dist: pydantic
 Requires-Dist: pandas
-Requires-Dist: hf_xet
+Requires-Dist: pydantic
 Requires-Dist: pyrate-limiter
-Requires-Dist: deap>=1.4.3
+Requires-Dist: tqdm
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-conv; extra == "dev"

opik_optimizer-1.0.0/pyproject.toml

@@ -0,0 +1,11 @@
+[tool.mypy]
+follow_imports = "skip"
+ignore_missing_imports = true
+disallow_untyped_defs = true
+disallow_untyped_calls = true
+check_untyped_defs = true
+exclude = "src/opik_optimizer/mipro_optimizer/"
+
+
+[tool.uv]
+managed = false

{opik_optimizer-0.9.1 → opik_optimizer-1.0.0}/setup.py

@@ -2,38 +2,35 @@ from setuptools import find_packages, setup
 
 setup(
     name="opik_optimizer",
-    version="0.9.1",
+    version="1.0.0",
     description="Agent optimization with Opik",
     author="Comet ML",
     author_email="support@comet.com",
     long_description=open("README.md", encoding="utf-8").read(),
-    long_description_content_type='text/markdown',
+    long_description_content_type="text/markdown",
     url="https://github.com/comet-ml/opik",
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     package_data={
-        'opik_optimizer': ['data/*.json'],
+        "opik_optimizer": ["data/*.json"],
     },
     python_requires=">=3.9,<3.13",
     install_requires=[
-        "opik>=1.7.17",
-        "dspy>=2.6.18,<=2.6.24",
-        "litellm",
-        "tqdm",
         "datasets",
+        "deap>=1.4.3",
+        "diskcache",
+        "hf_xet",
+        "litellm",
+        "opik>=1.7.17",
         "optuna",
-        "pydantic",
         "pandas",
-        "hf_xet",
+        "pydantic",
         "pyrate-limiter",
-        "deap>=1.4.3",
+        "tqdm",
     ],
     # dev requirements
     extras_require={
-        "dev": [
-            "pytest",
-            "pytest-conv"
-        ],
+        "dev": ["pytest", "pytest-conv"],
     },
     classifiers=[
         "Development Status :: 3 - Alpha",

{opik_optimizer-0.9.1 → opik_optimizer-1.0.0}/src/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-0.9.1 → opik_optimizer-1.0.0}/src/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-1.0.0/src/opik_optimizer/base_optimizer.py

@@ -0,0 +1,198 @@
+from typing import Any, Callable, Dict, List, Optional, Type
+
+import logging
+import time
+from abc import abstractmethod
+import random
+
+
+import litellm
+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, 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()
+
+# Don't use unsupported params:
+litellm.drop_params = True
+
+# Set up logging:
+logger = logging.getLogger(__name__)
+
+
+class OptimizationRound(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+
+    round_number: int
+    current_prompt: "chat_prompt.ChatPrompt"
+    current_score: float
+    generated_prompts: Any
+    best_prompt: "chat_prompt.ChatPrompt"
+    best_score: float
+    improvement: float
+
+
+class BaseOptimizer:
+    def __init__(
+        self,
+        model: str,
+        verbose: int = 1,
+        **model_kwargs: Any,
+    ) -> None:
+        """
+        Base class for optimizers.
+
+        Args:
+           model: LiteLLM model 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.verbose = verbose
+        self._history: List[OptimizationRound] = []
+        self.experiment_config = None
+        self.llm_call_counter = 0
+
+        # Initialize shared cache
+        initialize_cache()
+
+    @abstractmethod
+    def optimize_prompt(
+        self,
+        prompt: "chat_prompt.ChatPrompt",
+        dataset: Dataset,
+        metric: Callable,
+        experiment_config: Optional[Dict] = None,
+        **kwargs: Any,
+    ) -> optimization_result.OptimizationResult:
+        """
+        Optimize a prompt.
+
+        Args:
+           dataset: Opik dataset name, or Opik dataset
+           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
+           output_key: output field of dataset
+           experiment_config: Optional configuration for the experiment
+           **kwargs: Additional arguments for optimization
+        """
+        pass
+
+    def get_history(self) -> List[OptimizationRound]:
+        """
+        Get the optimization history.
+
+        Returns:
+            List[Dict[str, Any]]: List of optimization rounds with their details
+        """
+        return self._history
+
+    def _add_to_history(self, round_data: OptimizationRound) -> None:
+        """
+        Add a round to the optimization history.
+
+        Args:
+            round_data: Dictionary containing round details
+        """
+        self._history.append(round_data)
+
+    def update_optimization(
+        self, optimization: optimization.Optimization, status: str
+    ) -> None:
+        """
+        Update the optimization status
+        """
+        # FIXME: remove when a solution is added to opik's optimization.update method
+        count = 0
+        while count < 3:
+            try:
+                optimization.update(status="completed")
+                break
+            except ApiError:
+                count += 1
+                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-0.9.1 → opik_optimizer-1.0.0}/src/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-0.9.1 → opik_optimizer-1.0.0}/src/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-0.9.1 → opik_optimizer-1.0.0}/src/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-0.9.1 → opik_optimizer-1.0.0}/src/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-0.9.1 → opik_optimizer-1.0.0}/src/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-0.9.1 → opik_optimizer-1.0.0}/src/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