opik-optimizer 0.7.0__py3-none-any.whl

opik_optimizer/__init__.py

@@ -0,0 +1,65 @@
+import importlib.metadata
+import logging
+from .logging_config import setup_logging
+
+__version__ = importlib.metadata.version("opik_optimizer")
+
+# Using WARNING as a sensible default to avoid flooding users with INFO/DEBUG
+setup_logging(level=logging.WARNING)
+
+
+# Lazy imports to avoid circular dependencies
+def __getattr__(name):
+    if name == "MiproOptimizer":
+        from .mipro_optimizer import MiproOptimizer
+
+        return MiproOptimizer
+    elif name == "BaseOptimizer":
+        from .base_optimizer import BaseOptimizer
+
+        return BaseOptimizer
+    elif name == "MetaPromptOptimizer":
+        from .meta_prompt_optimizer import MetaPromptOptimizer
+
+        return MetaPromptOptimizer
+    elif name == "FewShotBayesianOptimizer":
+        from .few_shot_bayesian_optimizer import FewShotBayesianOptimizer
+
+        return FewShotBayesianOptimizer
+    elif name in ["MetricConfig", "OptimizationConfig", "TaskConfig"]:
+        from .optimization_config.configs import (
+            MetricConfig,
+            OptimizationConfig,
+            TaskConfig,
+        )
+
+        return locals()[name]
+    elif name in ["from_dataset_field", "from_llm_response_text"]:
+        from .optimization_config.mappers import (
+            from_dataset_field,
+            from_llm_response_text,
+        )
+
+        return locals()[name]
+    raise AttributeError(f"module 'opik_optimizer' has no attribute '{name}'")
+
+
+from opik.evaluation.models.litellm import warning_filters
+
+warning_filters.add_warning_filters()
+
+from .optimization_result import OptimizationResult
+
+__all__ = [
+    "BaseOptimizer",
+    "FewShotBayesianOptimizer",
+    "MetaPromptOptimizer",
+    "MiproOptimizer",
+    "MetricConfig",
+    "OptimizationConfig",
+    "TaskConfig",
+    "from_dataset_field",
+    "from_llm_response_text",
+    "OptimizationResult",
+    "setup_logging",
+]

opik_optimizer/_throttle.py

@@ -0,0 +1,43 @@
+import threading
+import time
+import queue
+from functools import wraps
+
+class RateLimiter:
+    """
+    Rate limiter that enforces a maximum number of calls across all threads.
+    """
+    def __init__(self, max_calls_per_second):
+        self.max_calls_per_second = max_calls_per_second
+        self.interval = 1.0 / max_calls_per_second  # Time between allowed calls
+        self.last_call_time = 0
+        self.lock = threading.Lock()
+        
+    def acquire(self):
+        """
+        Wait until a call is allowed according to the global rate limit.
+        Returns immediately if the call is allowed, otherwise blocks until it's time.
+        """
+        with self.lock:
+            current_time = time.time()
+            time_since_last = current_time - self.last_call_time
+            
+            # If we haven't waited long enough since the last call
+            if time_since_last < self.interval:
+                # Calculate how much longer we need to wait
+                sleep_time = self.interval - time_since_last
+                time.sleep(sleep_time)
+                
+            # Update the last call time (after potential sleep)
+            self.last_call_time = time.time()
+
+def rate_limited(limiter):
+    """Decorator to rate limit a function using the provided limiter"""
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            limiter.acquire()
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator
+

opik_optimizer/base_optimizer.py

@@ -0,0 +1,240 @@
+from typing import Optional, Union, List, Dict, Any
+import opik
+import logging
+import time
+
+import litellm
+from opik.evaluation import metrics
+from opik.opik_context import get_current_span_data
+from opik.rest_api.core import ApiError
+
+from pydantic import BaseModel
+from ._throttle import RateLimiter, rate_limited
+from .cache_config import initialize_cache
+from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
+from .optimization_config.configs import TaskConfig, MetricConfig
+
+limiter = RateLimiter(max_calls_per_second=15)
+
+# Don't use unsupported params:
+litellm.drop_params = True
+
+# Set up logging:
+logger = logging.getLogger(__name__)
+
+
+class OptimizationRound(BaseModel):
+    round_number: int
+    current_prompt: str
+    current_score: float
+    generated_prompts: List[Dict[str, Any]]
+    best_prompt: str
+    best_score: float
+    improvement: float
+
+
+class BaseOptimizer:
+    def __init__(self, model: str, project_name: Optional[str] = None, **model_kwargs):
+        """
+        Base class for optimizers.
+
+        Args:
+           model: LiteLLM model name
+           project_name: Opik project name
+           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._history = []
+        self.experiment_config = None
+        self.llm_call_counter = 0
+
+        # Initialize shared cache
+        initialize_cache()
+
+    def optimize_prompt(
+        self,
+        dataset: Union[str, opik.Dataset],
+        metric_config: MetricConfig,
+        task_config: TaskConfig,
+        prompt: str,
+        input_key: str,
+        output_key: str,
+        experiment_config: Optional[Dict] = None,
+        **kwargs,
+    ):
+        """
+        Optimize a prompt.
+
+        Args:
+           dataset: Opik dataset name, or Opik dataset
+           metric_config: instance of a MetricConfig
+           task_config: instance of a TaskConfig
+           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
+        """
+        self.dataset = dataset
+        self.metric = metric
+        self.prompt = prompt
+        self.input_key = input_key
+        self.output_key = output_key
+        self.experiment_config = experiment_config
+
+    def evaluate_prompt(
+        self,
+        dataset: Union[str, opik.Dataset],
+        metric_config: MetricConfig,
+        prompt: str,
+        input_key: str,
+        output_key: str,
+        n_samples: int = 10,
+        task_config: Optional[TaskConfig] = None,
+        dataset_item_ids: Optional[List[str]] = None,
+        experiment_config: Optional[Dict] = None,
+        **kwargs,
+    ) -> float:
+        """
+        Evaluate a prompt.
+
+        Args:
+           dataset: Opik dataset name, or Opik dataset
+           metric_config: instance of a MetricConfig
+           task_config: instance of a TaskConfig
+           prompt: the prompt to evaluate
+           input_key: input field of dataset
+           output_key: output field of dataset
+           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
+        """
+        self.dataset = dataset
+        self.metric_config = metric_config
+        self.task_config = task_config
+        self.prompt = prompt
+        self.input_key = input_key
+        self.output_key = output_key
+        self.experiment_config = experiment_config
+        return 0.0  # Base implementation returns 0
+
+    def get_history(self) -> List[Dict[str, Any]]:
+        """
+        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: Dict[str, Any]):
+        """
+        Add a round to the optimization history.
+
+        Args:
+            round_data: Dictionary containing round details
+        """
+        self._history.append(round_data)
+
+    @rate_limited(limiter)
+    def _call_model(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        is_reasoning: bool = False,
+    ) -> str:
+        """Call the model to get suggestions based on the meta-prompt."""
+        model = self.reasoning_model if is_reasoning else self.model
+        messages = []
+
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+            logger.debug(f"Using custom system prompt: {system_prompt[:100]}...")
+        else:
+            messages.append(
+                {"role": "system", "content": "You are a helpful assistant."}
+            )
+
+        messages.append({"role": "user", "content": prompt})
+        logger.debug(f"Calling model {model} with prompt: {prompt[:100]}...")
+
+        api_params = self.model_kwargs.copy()
+        api_params.update(
+            {
+                "model": model,
+                "messages": messages,
+                # Ensure required params like 'temperature', 'max_tokens' are present
+                # Defaults added here for safety, though usually set in __init__ kwargs
+                "temperature": api_params.get("temperature", 0.3),
+                "max_tokens": api_params.get("max_tokens", 1000),
+            }
+        )
+
+        # Attempt to add Opik monitoring if available
+        try:
+            # Assuming opik_litellm_monitor is imported and configured elsewhere
+            api_params = opik_litellm_monitor.try_add_opik_monitoring_to_params(
+                api_params
+            )
+            logger.debug("Opik monitoring hooks added to LiteLLM params.")
+        except Exception as e:
+            logger.warning(f"Could not add Opik monitoring to LiteLLM params: {e}")
+
+        logger.debug(
+            f"Final API params (excluding messages): { {k:v for k,v in api_params.items() if k != 'messages'} }"
+        )
+
+        # Increment Counter
+        self.llm_call_counter += 1
+        logger.debug(f"LLM Call Count: {self.llm_call_counter}")
+
+        try:
+            response = litellm.completion(**api_params)
+            model_output = response.choices[0].message.content.strip()
+            logger.debug(f"Model response from {model_to_use}: {model_output[:100]}...")
+            return model_output
+        except litellm.exceptions.RateLimitError as e:
+            logger.error(f"LiteLLM Rate Limit Error for model {model_to_use}: {e}")
+            # Consider adding retry logic here with tenacity
+            raise
+        except litellm.exceptions.APIConnectionError as e:
+            logger.error(f"LiteLLM API Connection Error for model {model_to_use}: {e}")
+            # Consider adding retry logic here
+            raise
+        except litellm.exceptions.ContextWindowExceededError as e:
+            logger.error(
+                f"LiteLLM Context Window Exceeded Error for model {model_to_use}. Prompt length: {len(prompt)}. Details: {e}"
+            )
+            raise
+        except litellm.exceptions.APIError as e:  # Catch broader API errors
+            logger.error(f"LiteLLM API Error for model {model_to_use}: {e}")
+            raise
+        except Exception as e:
+            # Catch any other unexpected errors
+            logger.error(
+                f"Unexpected error during model call to {model_to_use}: {type(e).__name__} - {e}"
+            )
+            raise
+
+    def update_optimization(self, 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...")

opik_optimizer/cache_config.py

@@ -0,0 +1,24 @@
+import os
+from pathlib import Path
+import litellm
+from litellm.caching import Cache
+
+# Configure cache directory
+CACHE_DIR = os.path.expanduser("~/.litellm_cache")
+Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)
+
+# Configure cache settings
+CACHE_CONFIG = {
+    "type": "disk",
+    "disk_cache_dir": CACHE_DIR,
+}
+
+def initialize_cache():
+    """Initialize the LiteLLM cache with custom configuration."""
+    litellm.cache = Cache(**CACHE_CONFIG)
+    return litellm.cache
+
+def clear_cache():
+    """Clear the LiteLLM cache."""
+    if litellm.cache:
+        litellm.cache.clear() 

opik_optimizer/demo/__init__.py

@@ -0,0 +1,7 @@
+from .datasets import get_or_create_dataset
+from .cache import get_litellm_cache
+
+__all__ = [
+    "get_or_create_dataset",
+    "get_litellm_cache",
+]

opik_optimizer/demo/cache.py

@@ -0,0 +1,112 @@
+from urllib.parse import urlparse, parse_qs
+import sqlite3
+import shutil
+import os
+import litellm
+from litellm.caching import Cache
+import requests
+
+NAMED_CACHES = {
+    "test": "https://drive.google.com/file/d/1RifNtpN-pl0DW49daRaAMJwW7MCsOh6y/view?usp=sharing",
+    "test2": "https://drive.google.com/uc?id=1RifNtpN-pl0DW49daRaAMJwW7MCsOh6y&export=download",
+}
+CACHE_DIR = os.path.expanduser("~/.litellm_cache")
+
+
+def get_litellm_cache(name: str):
+    """
+    Get a LiteLLM cache from a remote location, and add it to the
+    local cache
+    """
+    # Try to close an existing one, if there is one:
+    try:
+        litellm.cache.cache.disk_cache.close()
+    except Exception:
+        pass
+
+    if not os.path.exists(CACHE_DIR):
+        os.makedirs(CACHE_DIR)
+
+    if name.lower() in NAMED_CACHES:
+        return get_litellm_cache(NAMED_CACHES[name.lower()])
+    elif name.startswith("https://drive.google.com/file/d/"):
+        file_id = name.split("/d/")[1].split("/view")[0]
+        download_url = f"https://drive.google.com/uc?id={file_id}&export=download"
+        file_path = _get_google_drive_file(download_url)
+    elif name.startswith("https://drive.google.com/uc"):
+        file_path = _get_google_drive_file(name)
+    else:
+        raise Exception("Unknown cache type: %r" % name)
+
+    dest_path = os.path.join(CACHE_DIR, "cache.db")
+
+    if os.path.exists(dest_path):
+        # Copy contents from source to dest:
+        _copy_cache(file_path, dest_path)
+    else:
+        # Just copy the file:
+        shutil.copy(file_path, dest_path)
+
+    # Update the cache to use the new database:
+    litellm.cache = Cache(type="disk", disk_cache_dir=CACHE_DIR)
+
+
+def _copy_cache(source_path, dest_path):
+    """
+    Copy cached items from a source to a destination cache.
+    """
+    source_conn = sqlite3.connect(source_path)
+    source_conn.row_factory = sqlite3.Row
+    source_cursor = source_conn.cursor()
+
+    dest_conn = sqlite3.connect(dest_path)
+    dest_cursor = dest_conn.cursor()
+
+    source_cursor.execute(f"PRAGMA table_info(Cache)")
+    columns_info = source_cursor.fetchall()
+    column_names = [info[1] for info in columns_info[1:]]  # Skip rowid
+    placeholders = ", ".join(["?"] * len(column_names))
+    columns_str = ", ".join(column_names)
+
+    inserted_count = 0
+    source_cursor.execute("SELECT * FROM Cache")
+    records = source_cursor.fetchall()
+    for record in records:
+        record = dict(record)
+        del record["rowid"]
+        key_value = record["key"]
+
+        dest_cursor.execute("SELECT 1 FROM Cache WHERE key = ?", (key_value,))
+        existing_record = dest_cursor.fetchone()
+
+        if not existing_record:
+            dest_cursor.execute(
+                f"INSERT INTO Cache ({columns_str}) VALUES ({placeholders})",
+                list(record.values()),
+            )
+            inserted_count += 1
+
+    print(f"Inserted {inserted_count} record(s) in litellm cache")
+    dest_conn.commit()
+
+
+def _get_google_drive_file(file_url):
+    """
+    Given a common google drive URL with id=ID
+    get it, or use cache.
+    """
+    parsed_url = urlparse(file_url)
+    query_params = parse_qs(parsed_url.query)
+    id_value = query_params.get("id")[0]
+
+    cache_file_path = os.path.join(CACHE_DIR, id_value)
+
+    if not os.path.exists(cache_file_path):
+        response = requests.get(file_url)
+        response.raise_for_status()
+
+        with open(cache_file_path, "wb") as tmp_file:
+            for chunk in response.iter_content(chunk_size=8192):
+                tmp_file.write(chunk)
+
+    return cache_file_path