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

opik_optimizer/meta_prompt_optimizer/reporting.py

@@ -1,6 +1,8 @@
 from contextlib import contextmanager
-from typing import Any, Iterator
+from typing import Any
+from collections.abc import Iterator
 
+from rich.panel import Panel
 from rich.text import Text
 
 from ..optimization_config import chat_prompt
@@ -136,6 +138,18 @@ class CandidateGenerationReporter:
         console.print(Text("│"))
 
 
+def display_tool_description(description: str, label: str, color: str) -> None:
+    if not description.strip():
+        return
+    console.print(
+        Panel(
+            description.strip(),
+            title=label,
+            border_style=color,
+        )
+    )
+
+
 @contextmanager
 def display_candidate_generation_report(
     num_prompts: int, verbose: int = 1
@@ -162,7 +176,7 @@ def display_prompt_candidate_scoring_report(verbose: int = 1) -> Any:
         ) -> None:
             if verbose >= 1:
                 console.print(
-                    Text(f"│    Evaluating candidate prompt {candidate_count+1}:")
+                    Text(f"│    Evaluating candidate prompt {candidate_count + 1}:")
                 )
                 display_messages(prompt.get_messages(), "│         ")
 

opik_optimizer/mipro_optimizer/_lm.py

@@ -4,7 +4,7 @@ import os
 import re
 import threading
 from hashlib import sha256
-from typing import Any, Dict, List, Literal, Optional, cast
+from typing import Any, Literal, cast
 
 import litellm
 import pydantic
@@ -42,12 +42,12 @@ class LM(BaseLM):
         max_tokens: int = 1000,
         cache: bool = True,
         cache_in_memory: bool = True,
-        callbacks: Optional[List[BaseCallback]] = None,
+        callbacks: list[BaseCallback] | None = None,
         num_retries: int = 8,
         provider=None,
-        finetuning_model: Optional[str] = None,
-        launch_kwargs: Optional[dict[str, Any]] = None,
-        train_kwargs: Optional[dict[str, Any]] = None,
+        finetuning_model: str | None = None,
+        launch_kwargs: dict[str, Any] | None = None,
+        train_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ):
         """
@@ -93,9 +93,9 @@ class LM(BaseLM):
 
         if model_pattern:
             # Handle OpenAI reasoning models (o1, o3)
-            assert (
-                max_tokens >= 20_000 and temperature == 1.0
-            ), "OpenAI's reasoning models require passing temperature=1.0 and max_tokens >= 20_000 to `dspy.LM(...)`"
+            assert max_tokens >= 20_000 and temperature == 1.0, (
+                "OpenAI's reasoning models require passing temperature=1.0 and max_tokens >= 20_000 to `dspy.LM(...)`"
+            )
             self.kwargs = dict(
                 temperature=temperature, max_completion_tokens=max_tokens, **kwargs
             )
@@ -148,17 +148,17 @@ class LM(BaseLM):
         self.llm_call_counter += 1
         return results
 
-    def launch(self, launch_kwargs: Optional[Dict[str, Any]] = None):
+    def launch(self, launch_kwargs: dict[str, Any] | None = None):
         self.provider.launch(self, launch_kwargs)
 
-    def kill(self, launch_kwargs: Optional[Dict[str, Any]] = None):
+    def kill(self, launch_kwargs: dict[str, Any] | None = None):
         self.provider.kill(self, launch_kwargs)
 
     def finetune(
         self,
-        train_data: List[Dict[str, Any]],
-        train_data_format: Optional[TrainDataFormat],
-        train_kwargs: Optional[Dict[str, Any]] = None,
+        train_data: list[dict[str, Any]],
+        train_data_format: TrainDataFormat | None,
+        train_kwargs: dict[str, Any] | None = None,
     ) -> TrainingJob:
         from dspy import settings as settings
 
@@ -222,7 +222,7 @@ class LM(BaseLM):
         return {key: getattr(self, key) for key in state_keys} | self.kwargs
 
 
-def request_cache(maxsize: Optional[int] = None):
+def request_cache(maxsize: int | None = None):
     """
     A threadsafe decorator to create an in-memory LRU cache for LM inference functions that accept
     a dictionary-like LM request. An in-memory cache for LM calls is critical for ensuring
@@ -235,7 +235,7 @@ def request_cache(maxsize: Optional[int] = None):
         A decorator that wraps the target function with caching.
     """
 
-    def cache_key(request: Dict[str, Any]) -> str:
+    def cache_key(request: dict[str, Any]) -> str:
         """
         Obtain a unique cache key for the given request dictionary by hashing its JSON
         representation. For request fields having types that are known to be JSON-incompatible,
@@ -278,7 +278,7 @@ def request_cache(maxsize: Optional[int] = None):
             # concurrently, e.g. during optimization and evaluation
             lock=threading.RLock(),
         )
-        def func_cached(key: str, request: Dict[str, Any], *args, **kwargs):
+        def func_cached(key: str, request: dict[str, Any], *args, **kwargs):
             return func(request, *args, **kwargs)
 
         @functools.wraps(func)
@@ -303,7 +303,7 @@ def request_cache(maxsize: Optional[int] = None):
 
 
 @request_cache(maxsize=None)
-def cached_litellm_completion(request: Dict[str, Any], num_retries: int):
+def cached_litellm_completion(request: dict[str, Any], num_retries: int):
     return litellm_completion(
         request,
         cache={"no-cache": False, "no-store": False},
@@ -312,7 +312,7 @@ def cached_litellm_completion(request: Dict[str, Any], num_retries: int):
 
 
 def litellm_completion(
-    request: Dict[str, Any],
+    request: dict[str, Any],
     num_retries: int,
     cache={"no-cache": True, "no-store": True},
 ):
@@ -362,7 +362,7 @@ def litellm_completion(
 
 
 @request_cache(maxsize=None)
-def cached_litellm_text_completion(request: Dict[str, Any], num_retries: int):
+def cached_litellm_text_completion(request: dict[str, Any], num_retries: int):
     return litellm_text_completion(
         request,
         num_retries=num_retries,
@@ -371,7 +371,7 @@ def cached_litellm_text_completion(request: Dict[str, Any], num_retries: int):
 
 
 def litellm_text_completion(
-    request: Dict[str, Any],
+    request: dict[str, Any],
     num_retries: int,
     cache={"no-cache": True, "no-store": True},
 ):

opik_optimizer/mipro_optimizer/_mipro_optimizer_v2.py

@@ -1,7 +1,8 @@
 import random
 import textwrap
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
+from typing import Any, Literal
+from collections.abc import Callable
 
 import dspy
 import numpy as np
@@ -69,27 +70,27 @@ class MIPROv2(Teleprompter):
     def __init__(
         self,
         metric: Callable,
-        prompt_model: Optional[Any] = None,
-        task_model: Optional[Any] = None,
-        teacher_settings: Dict = {},
+        prompt_model: Any | None = None,
+        task_model: Any | None = None,
+        teacher_settings: dict = {},
         max_bootstrapped_demos: int = 4,
         max_labeled_demos: int = 4,
-        auto: Optional[Literal["light", "medium", "heavy"]] = "medium",
+        auto: Literal["light", "medium", "heavy"] | None = "medium",
         num_candidates: int = 10,
-        num_threads: Optional[int] = None,
+        num_threads: int | None = None,
         max_errors: int = 10,
         seed: int = 9,
         init_temperature: float = 0.5,
         verbose: bool = False,
         track_stats: bool = True,
-        log_dir: Optional[str] = None,
-        metric_threshold: Optional[float] = None,
-        opik_dataset: Optional[opik.Dataset] = None,
-        opik_metric: Optional[Callable] = None,
-        opik_prompt_task_config: Optional[TaskConfig] = None,
-        opik_project_name: Optional[str] = None,
-        opik_optimization_id: Optional[str] = None,
-        experiment_config: Optional[Dict[str, Any]] = None,
+        log_dir: str | None = None,
+        metric_threshold: float | None = None,
+        opik_dataset: opik.Dataset | None = None,
+        opik_metric: Callable | None = None,
+        opik_prompt_task_config: TaskConfig | None = None,
+        opik_project_name: str | None = None,
+        opik_optimization_id: str | None = None,
+        experiment_config: dict[str, Any] | None = None,
     ):
         # Validate 'auto' parameter
         allowed_modes = {None, "light", "medium", "heavy"}
@@ -129,13 +130,13 @@ class MIPROv2(Teleprompter):
         self,
         student: Any,
         *,
-        trainset: List,
+        trainset: list,
         teacher: Any = None,
-        valset: Optional[List] = None,
+        valset: list | None = None,
         num_trials: int = 30,
-        max_bootstrapped_demos: Optional[int] = None,
-        max_labeled_demos: Optional[int] = None,
-        seed: Optional[int] = None,
+        max_bootstrapped_demos: int | None = None,
+        max_labeled_demos: int | None = None,
+        seed: int | None = None,
         minibatch: bool = True,
         minibatch_size: int = 35,
         minibatch_full_eval_steps: int = 5,
@@ -145,7 +146,7 @@ class MIPROv2(Teleprompter):
         tip_aware_proposer: bool = True,
         fewshot_aware_proposer: bool = True,
         requires_permission_to_run: bool = True,
-        provide_traceback: Optional[bool] = None,
+        provide_traceback: bool | None = None,
     ) -> Any:
         # Set random seeds
         seed = seed or self.seed
@@ -252,8 +253,8 @@ class MIPROv2(Teleprompter):
         num_trials: int,
         minibatch: bool,
         zeroshot_opt: bool,
-        valset: List,
-    ) -> Tuple[int, List, bool]:
+        valset: list,
+    ) -> tuple[int, list, bool]:
         if self.auto is None:
             return num_trials, valset, minibatch
 
@@ -273,7 +274,7 @@ class MIPROv2(Teleprompter):
 
         return num_trials, valset, minibatch
 
-    def _set_and_validate_datasets(self, trainset: List, valset: Optional[List]):
+    def _set_and_validate_datasets(self, trainset: list, valset: list | None):
         if not trainset:
             raise ValueError("Trainset cannot be empty.")
 
@@ -292,7 +293,7 @@ class MIPROv2(Teleprompter):
 
         return trainset, valset
 
-    def _print_auto_run_settings(self, num_trials: int, minibatch: bool, valset: List):
+    def _print_auto_run_settings(self, num_trials: int, minibatch: bool, valset: list):
         logger.info(
             f"\nRUNNING WITH THE FOLLOWING {self.auto.upper()} AUTO RUN SETTINGS:"
             f"\nnum_trials: {num_trials}"
@@ -308,9 +309,9 @@ class MIPROv2(Teleprompter):
         minibatch: bool,
         minibatch_size: int,
         minibatch_full_eval_steps: int,
-        valset: List,
+        valset: list,
         program_aware_proposer: bool,
-    ) -> Tuple[str, str]:
+    ) -> tuple[str, str]:
         num_predictors = len(program.predictors())
 
         # Estimate prompt model calls
@@ -359,7 +360,7 @@ class MIPROv2(Teleprompter):
         minibatch: bool,
         minibatch_size: int,
         minibatch_full_eval_steps: int,
-        valset: List,
+        valset: list,
         program_aware_proposer: bool,
     ) -> bool:
         prompt_model_line, task_model_line = self._estimate_lm_calls(
@@ -414,8 +415,8 @@ class MIPROv2(Teleprompter):
         return user_input == "y"
 
     def _bootstrap_fewshot_examples(
-        self, program: Any, trainset: List, seed: int, teacher: Any
-    ) -> Optional[List]:
+        self, program: Any, trainset: list, seed: int, teacher: Any
+    ) -> list | None:
         logger.info("\n==> STEP 1: BOOTSTRAP FEWSHOT EXAMPLES <==")
         if self.max_bootstrapped_demos > 0:
             logger.info(
@@ -461,14 +462,14 @@ class MIPROv2(Teleprompter):
     def _propose_instructions(
         self,
         program: Any,
-        trainset: List,
-        demo_candidates: Optional[List],
+        trainset: list,
+        demo_candidates: list | None,
         view_data_batch_size: int,
         program_aware_proposer: bool,
         data_aware_proposer: bool,
         tip_aware_proposer: bool,
         fewshot_aware_proposer: bool,
-    ) -> Dict[int, List[str]]:
+    ) -> dict[int, list[str]]:
         logger.info("\n==> STEP 2: PROPOSE INSTRUCTION CANDIDATES <==")
         logger.info(
             "We will use the few-shot examples from the previous step, a generated dataset summary, a summary of the program code, and a randomly selected prompting tip to propose instructions."
@@ -513,16 +514,16 @@ class MIPROv2(Teleprompter):
     def _optimize_prompt_parameters(
         self,
         program: Any,
-        instruction_candidates: Dict[int, List[str]],
-        demo_candidates: Optional[List],
+        instruction_candidates: dict[int, list[str]],
+        demo_candidates: list | None,
         evaluate: Evaluate,
-        valset: List,
+        valset: list,
         num_trials: int,
         minibatch: bool,
         minibatch_size: int,
         minibatch_full_eval_steps: int,
         seed: int,
-    ) -> Optional[Any]:
+    ) -> Any | None:
         # Run optimization
         optuna.logging.set_verbosity(optuna.logging.WARNING)
         logger.info("==> STEP 3: FINDING OPTIMAL PROMPT PARAMETERS <==")
@@ -848,12 +849,12 @@ class MIPROv2(Teleprompter):
     def _select_and_insert_instructions_and_demos(
         self,
         candidate_program: Any,
-        instruction_candidates: Dict[int, List[str]],
-        demo_candidates: Optional[List],
+        instruction_candidates: dict[int, list[str]],
+        demo_candidates: list | None,
         trial: optuna.trial.Trial,
-        trial_logs: Dict,
+        trial_logs: dict,
         trial_num: int,
-    ) -> List[str]:
+    ) -> list[str]:
         chosen_params = []
         raw_chosen_params = {}
 
@@ -902,18 +903,18 @@ class MIPROv2(Teleprompter):
         self,
         trial_num: int,
         adjusted_num_trials: int,
-        param_score_dict: Dict,
-        fully_evaled_param_combos: Dict,
+        param_score_dict: dict,
+        fully_evaled_param_combos: dict,
         evaluate: Evaluate,
-        valset: List,
-        trial_logs: Dict,
+        valset: list,
+        trial_logs: dict,
         total_eval_calls: int,
         score_data,
         best_score: float,
         best_program: Any,
         study: optuna.Study,
-        instruction_candidates: List,
-        demo_candidates: List,
+        instruction_candidates: list,
+        demo_candidates: list,
     ):
         logger.info(
             f"===== Trial {trial_num + 1} / {adjusted_num_trials} - Full Evaluation ====="
@@ -1026,19 +1027,19 @@ class MIPROv2(Teleprompter):
 
 def eval_candidate_program_with_opik(
     opik_dataset: opik.Dataset,
-    trainset: List,
+    trainset: list,
     candidate_program: Any,
     project_name: str,
     metric: Callable,
     prompt_task_config: TaskConfig,
     num_threads: int,
-    experiment_config: Optional[Dict[str, Any]] = None,
-    optimization_id: Optional[str] = None,
+    experiment_config: dict[str, Any] | None = None,
+    optimization_id: str | None = None,
 ):
     """Evaluate a candidate program on the trainset, using the specified batch size."""
     dataset_item_ids = [example["id"] for example in trainset]
 
-    def program_task(dataset_item: Dict[str, Any]) -> Dict[str, Any]:
+    def program_task(dataset_item: dict[str, Any]) -> dict[str, Any]:
         program_inputs = {
             input_key: dataset_item[input_key]
             for input_key in prompt_task_config.input_dataset_fields

opik_optimizer/mipro_optimizer/mipro_optimizer.py

@@ -1,7 +1,8 @@
 import os
 import random
 from datetime import datetime
-from typing import Callable, Dict, List, Literal, Optional, Union
+from typing import Literal
+from collections.abc import Callable
 import logging
 
 import dspy
@@ -37,7 +38,7 @@ class MiproOptimizer(BaseOptimizer):
     def __init__(
         self,
         model,
-        project_name: Optional[str] = None,
+        project_name: str | None = None,
         verbose: int = 1,
         **model_kwargs,
     ):
@@ -54,13 +55,13 @@ class MiproOptimizer(BaseOptimizer):
 
     def evaluate_prompt(
         self,
-        dataset: Union[str, Dataset],
+        dataset: str | Dataset,
         metric: Callable,
         task_config: TaskConfig,
-        prompt: Optional[Union[str, dspy.Module, OptimizationResult]] = None,
+        prompt: str | dspy.Module | OptimizationResult | None = None,
         n_samples: int = 10,
-        dataset_item_ids: Optional[List[str]] = None,
-        experiment_config: Optional[Dict] = None,
+        dataset_item_ids: list[str] | None = None,
+        experiment_config: dict | None = None,
         verbose: int = 1,
         **kwargs,
     ) -> float:
@@ -238,14 +239,14 @@ class MiproOptimizer(BaseOptimizer):
 
     def optimize_prompt(
         self,
-        dataset: Union[str, Dataset],
+        dataset: str | Dataset,
         metric: Callable,
         task_config: TaskConfig,
         num_candidates: int = 10,
-        experiment_config: Optional[Dict] = None,
-        num_trials: Optional[int] = 3,
-        n_samples: Optional[int] = 10,
-        auto: Optional[Literal["light", "medium", "heavy"]] = "light",
+        experiment_config: dict | None = None,
+        num_trials: int | None = 3,
+        n_samples: int | None = 10,
+        auto: Literal["light", "medium", "heavy"] | None = "light",
         **kwargs,
     ) -> OptimizationResult:
         self._opik_client = opik.Opik()
@@ -271,15 +272,15 @@ class MiproOptimizer(BaseOptimizer):
 
     def _optimize_prompt(
         self,
-        dataset: Union[str, Dataset],
+        dataset: str | Dataset,
         metric: Callable,
         task_config: TaskConfig,
         num_candidates: int = 10,
-        experiment_config: Optional[Dict] = None,
-        optimization_id: Optional[str] = None,
-        num_trials: Optional[int] = 3,
-        n_samples: Optional[int] = 10,
-        auto: Optional[Literal["light", "medium", "heavy"]] = "light",
+        experiment_config: dict | None = None,
+        optimization_id: str | None = None,
+        num_trials: int | None = 3,
+        n_samples: int | None = 10,
+        auto: Literal["light", "medium", "heavy"] | None = "light",
         **kwargs,
     ) -> OptimizationResult:
         logger.info("Preparing MIPRO optimization...")
@@ -306,11 +307,11 @@ class MiproOptimizer(BaseOptimizer):
         metric,
         task_config,
         num_candidates: int = 10,
-        experiment_config: Optional[Dict] = None,
-        optimization_id: Optional[str] = None,
-        num_trials: Optional[int] = 3,
-        n_samples: Optional[int] = 10,
-        auto: Optional[Literal["light", "medium", "heavy"]] = "light",
+        experiment_config: dict | None = None,
+        optimization_id: str | None = None,
+        num_trials: int | None = 3,
+        n_samples: int | None = 10,
+        auto: Literal["light", "medium", "heavy"] | None = "light",
         **kwargs,
     ) -> None:
         # FIXME: Intermediate values:
@@ -508,9 +509,11 @@ class MiproOptimizer(BaseOptimizer):
                     }
                 ],
                 score=0.0,
-                metric_name=self.opik_metric.__name__
-                if hasattr(self, "opik_metric")
-                else "unknown_metric",
+                metric_name=(
+                    self.opik_metric.__name__
+                    if hasattr(self, "opik_metric")
+                    else "unknown_metric"
+                ),
                 details={"error": "No candidate programs generated by MIPRO"},
                 history=mipro_history_processed,
                 llm_calls=self.lm.llm_call_counter,
@@ -564,9 +567,11 @@ class MiproOptimizer(BaseOptimizer):
                     }
                 ],
                 score=0.0,
-                metric_name=getattr(self, "opik_metric", None).name
-                if hasattr(self, "opik_metric") and self.opik_metric
-                else "unknown_metric",
+                metric_name=(
+                    getattr(self, "opik_metric", None).name
+                    if hasattr(self, "opik_metric") and self.opik_metric
+                    else "unknown_metric"
+                ),
                 details={"error": "No programs generated or compile failed"},
                 history=[],
                 llm_calls=self.lm.llm_call_counter,

opik_optimizer/mipro_optimizer/utils.py

@@ -1,5 +1,3 @@
-from typing import Dict, Optional
-
 import uuid
 import dspy
 import re
@@ -61,7 +59,7 @@ def opik_metric_to_dspy(metric, output):
 
 
 def create_dspy_training_set(
-    data: list[dict], input: str, n_samples: Optional[int] = None
+    data: list[dict], input: str, n_samples: int | None = None
 ) -> list[dspy.Example]:
     """
     Turn a list of dicts into a list of dspy Examples
@@ -80,7 +78,7 @@ def create_dspy_training_set(
     return output
 
 
-def get_tool_prompts(tool_names, text: str) -> Dict[str, str]:
+def get_tool_prompts(tool_names, text: str) -> dict[str, str]:
     """
     Extract the embedded tool prompts from a text.
     """

opik_optimizer/optimizable_agent.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, List, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 import json
 import os
 
@@ -16,7 +16,7 @@ if TYPE_CHECKING:
     from .optimization_config.chat_prompt import ChatPrompt
 
 
-def tools_to_dict(tools: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
+def tools_to_dict(tools: dict[str, dict[str, Any]]) -> dict[str, Any]:
     retval = {}
     for name in tools:
         parts = {}
@@ -38,11 +38,11 @@ class OptimizableAgent:
         project_name (Optional[str]): The project name for tracking
     """
 
-    model: Optional[str] = None
-    model_kwargs: Dict[str, Any] = {}
-    project_name: Optional[str] = "Default Project"
-    input_dataset_field: Optional[str] = None
-    prompts: Dict[str, "ChatPrompt"]
+    model: str | None = None
+    model_kwargs: dict[str, Any] = {}
+    project_name: str | None = "Default Project"
+    input_dataset_field: str | None = None
+    prompts: dict[str, "ChatPrompt"]
     prompt: "ChatPrompt"
 
     def __init__(self, prompt: "ChatPrompt") -> None:
@@ -58,7 +58,8 @@ class OptimizableAgent:
     def init_llm(self) -> None:
         """Initialize the LLM with the appropriate callbacks."""
         # Litellm bug requires this (maybe problematic if multi-threaded)
-        os.environ["OPIK_PROJECT_NAME"] = str(self.project_name)
+        if "OPIK_PROJECT_NAME" not in os.environ:
+            os.environ["OPIK_PROJECT_NAME"] = str(self.project_name)
         self.opik_logger = OpikLogger()
         litellm.callbacks = [self.opik_logger]
 
@@ -70,8 +71,8 @@ class OptimizableAgent:
     @_throttle.rate_limited(_limiter)
     def _llm_complete(
         self,
-        messages: List[Dict[str, str]],
-        tools: Optional[List[Dict[str, str]]],
+        messages: list[dict[str, str]],
+        tools: list[dict[str, str]] | None,
         seed: int,
     ) -> Any:
         response = litellm.completion(
@@ -90,10 +91,10 @@ class OptimizableAgent:
 
     def llm_invoke(
         self,
-        query: Optional[str] = None,
-        messages: Optional[List[Dict[str, str]]] = None,
-        seed: Optional[int] = None,
-        allow_tool_use: Optional[bool] = False,
+        query: str | None = None,
+        messages: list[dict[str, str]] | None = None,
+        seed: int | None = None,
+        allow_tool_use: bool | None = False,
     ) -> str:
         """
         NOTE: this is the default LiteLLM API. It is used
@@ -155,14 +156,14 @@ class OptimizableAgent:
             result = response.choices[0].message.content
         return result
 
-    def invoke_dataset_item(self, dataset_item: Dict[str, str]) -> str:
+    def invoke_dataset_item(self, dataset_item: dict[str, str]) -> str:
         messages = self.prompt.get_messages(dataset_item)
         return self.invoke(messages)
 
     def invoke(
         self,
-        messages: List[Dict[str, str]],
-        seed: Optional[int] = None,
+        messages: list[dict[str, str]],
+        seed: int | None = None,
     ) -> str:
         """
         Invoke the agent with a dataset item.