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

opik_optimizer/meta_prompt_optimizer/reporting.py

@@ -1,8 +1,9 @@
 from contextlib import contextmanager
+from typing import Any, Iterator
 
-import rich
 from rich.text import Text
 
+from ..optimization_config import chat_prompt
 from ..reporting_utils import (
     convert_tqdm_to_rich,
     display_configuration,  # noqa: F401
@@ -18,33 +19,66 @@ console = get_console()
 
 
 @contextmanager
-def display_round_progress(max_rounds: int, verbose: int = 1):
+def display_round_progress(max_rounds: int, verbose: int = 1) -> Any:
     """Context manager to display messages during an evaluation phase."""
-    
+
     # Create a simple object with a method to set the score
     class Reporter:
-        def failed_to_generate(self, num_prompts, error):
+        def failed_to_generate(self, num_prompts: int, error: str) -> None:
             if verbose >= 1:
-                console.print(Text(f"│    Failed to generate {num_prompts} candidate prompt{'' if num_prompts == 1 else 's'}: {error}", style="red"))
+                console.print(
+                    Text(
+                        f"│    Failed to generate {num_prompts} candidate prompt{'' if num_prompts == 1 else 's'}: {error}",
+                        style="red",
+                    )
+                )
                 console.print(Text("│"))
-            
-        def round_start(self, round_number):
+
+        def round_start(self, round_number: int) -> None:
             if verbose >= 1:
-                console.print(Text(f"│ - Starting optimization round {round_number + 1} of {max_rounds}"))
+                console.print(
+                    Text(
+                        f"│ - Starting optimization round {round_number + 1} of {max_rounds}"
+                    )
+                )
 
-        def round_end(self, round_number, score, best_score, best_prompt):
+        def round_end(self, round_number: int, score: float, best_score: float) -> None:
             if verbose >= 1:
-                console.print(Text(f"│    Completed optimization round {round_number + 1} of {max_rounds}"))
+                console.print(
+                    Text(
+                        f"│    Completed optimization round {round_number + 1} of {max_rounds}"
+                    )
+                )
                 if best_score == 0 and score == 0:
-                    console.print(Text("│    No improvement in this optimization round - score is 0", style="yellow"))
+                    console.print(
+                        Text(
+                            "│    No improvement in this optimization round - score is 0",
+                            style="yellow",
+                        )
+                    )
                 elif best_score == 0:
-                    console.print(Text(f"│    Found a new best performing prompt: {score:.4f}", style="green"))
+                    console.print(
+                        Text(
+                            f"│    Found a new best performing prompt: {score:.4f}",
+                            style="green",
+                        )
+                    )
                 elif score > best_score:
                     perc_change = (score - best_score) / best_score
-                    console.print(Text(f"│    Found a new best performing prompt: {score:.4f} ({perc_change:.2%})", style="green"))
+                    console.print(
+                        Text(
+                            f"│    Found a new best performing prompt: {score:.4f} ({perc_change:.2%})",
+                            style="green",
+                        )
+                    )
                 elif score <= best_score:
-                    console.print(Text("│    No improvement in this optimization round", style="red"))
-                
+                    console.print(
+                        Text(
+                            "│    No improvement in this optimization round",
+                            style="red",
+                        )
+                    )
+
                 console.print(Text("│"))
 
     # Use our log suppression context manager and yield the reporter
@@ -57,20 +91,22 @@ def display_round_progress(max_rounds: int, verbose: int = 1):
 
 
 @contextmanager
-def display_evaluation(message: str = "First we will establish the baseline performance:", verbose: int = 1):
+def display_evaluation(
+    message: str = "First we will establish the baseline performance:", verbose: int = 1
+) -> Any:
     """Context manager to display messages during an evaluation phase."""
-    score = None
-    
     # Entry point
     if verbose >= 1:
         console.print(Text(f"> {message}"))
-    
+
     # Create a simple object with a method to set the score
     class Reporter:
-        def set_score(self, s):
+        def set_score(self, s: float) -> None:
             if verbose >= 1:
-                console.print(Text(f"\r  Baseline score was: {s:.4f}.\n", style="green"))
-    
+                console.print(
+                    Text(f"\r  Baseline score was: {s:.4f}.\n", style="green")
+                )
+
     # Use our log suppression context manager and yield the reporter
     with suppress_opik_logs():
         with convert_tqdm_to_rich("  Evaluation", verbose=verbose):
@@ -79,62 +115,100 @@ def display_evaluation(message: str = "First we will establish the baseline perf
             finally:
                 pass
 
-def display_optimization_start_message(verbose: int = 1):
+
+def display_optimization_start_message(verbose: int = 1) -> None:
     if verbose >= 1:
         console.print(Text("> Starting the optimization run"))
         console.print(Text("│"))
 
 
+class CandidateGenerationReporter:
+    def __init__(self, num_prompts: int):
+        self.num_prompts = num_prompts
+
+    def set_generated_prompts(self) -> None:
+        console.print(
+            Text(
+                f"│      Successfully generated {self.num_prompts} candidate prompt{'' if self.num_prompts == 1 else 's'}",
+                style="dim",
+            )
+        )
+        console.print(Text("│"))
+
+
 @contextmanager
-def display_candidate_generation_report(num_prompts: int, verbose: int = 1):
-    """Context manager to display messages during an evaluation phase."""
-    # Entry point
+def display_candidate_generation_report(
+    num_prompts: int, verbose: int = 1
+) -> Iterator[CandidateGenerationReporter]:
     if verbose >= 1:
-        console.print(Text(f"│    Generating candidate prompt{'' if num_prompts == 1 else 's'}:"))
-    
-    # Create a simple object with a method to set the score
-    class Reporter:
-        def set_generated_prompts(self, prompts):
-            console.print(Text(f"│      Successfully generated {num_prompts} candidate prompt{'' if num_prompts == 1 else 's'}", style="dim"))
-            console.print(Text("│"))
+        console.print(
+            Text(f"│    Generating candidate prompt{'' if num_prompts == 1 else 's'}:")
+        )
 
     try:
-        yield Reporter()
+        yield CandidateGenerationReporter(num_prompts)
     finally:
         pass
 
 
 @contextmanager
-def display_prompt_candidate_scoring_report(candidate_count, prompt, verbose: int = 1):
+def display_prompt_candidate_scoring_report(verbose: int = 1) -> Any:
     """Context manager to display messages during an evaluation phase."""
+
     # Create a simple object with a method to set the score
     class Reporter:
-        def set_generated_prompts(self, candidate_count, prompt):
+        def set_generated_prompts(
+            self, candidate_count: int, prompt: chat_prompt.ChatPrompt
+        ) -> None:
             if verbose >= 1:
-                console.print(Text(f"│    Evaluating candidate prompt {candidate_count+1}:"))
-                display_messages(prompt, "│         ")
-        
-        def set_final_score(self, best_score, score):
+                console.print(
+                    Text(f"│    Evaluating candidate prompt {candidate_count+1}:")
+                )
+                display_messages(prompt.get_messages(), "│         ")
+
+        def set_final_score(self, best_score: float, score: float) -> None:
             if verbose >= 1:
                 if best_score == 0 and score > 0:
-                    console.print(Text(f"│          Evaluation score: {score:.4f}", style="green"))
+                    console.print(
+                        Text(f"│          Evaluation score: {score:.4f}", style="green")
+                    )
                 elif best_score == 0 and score == 0:
-                    console.print(Text(f"│         Evaluation score: {score:.4f}", style="dim yellow"))
+                    console.print(
+                        Text(
+                            f"│         Evaluation score: {score:.4f}",
+                            style="dim yellow",
+                        )
+                    )
                 elif score > best_score:
                     perc_change = (score - best_score) / best_score
-                    console.print(Text(f"│          Evaluation score: {score:.4f} ({perc_change:.2%})", style="green"))
+                    console.print(
+                        Text(
+                            f"│          Evaluation score: {score:.4f} ({perc_change:.2%})",
+                            style="green",
+                        )
+                    )
                 elif score < best_score:
                     perc_change = (score - best_score) / best_score
-                    console.print(Text(f"│          Evaluation score: {score:.4f} ({perc_change:.2%})", style="red"))
+                    console.print(
+                        Text(
+                            f"│          Evaluation score: {score:.4f} ({perc_change:.2%})",
+                            style="red",
+                        )
+                    )
                 else:
-                    console.print(Text(f"│         Evaluation score: {score:.4f}", style="dim yellow"))
-            
+                    console.print(
+                        Text(
+                            f"│         Evaluation score: {score:.4f}",
+                            style="dim yellow",
+                        )
+                    )
+
                 console.print(Text("│"))
                 console.print(Text("│"))
+
     try:
         with suppress_opik_logs():
             with convert_tqdm_to_rich("│         Evaluation", verbose=verbose):
                 yield Reporter()
     finally:
         pass
-

opik_optimizer/mipro_optimizer/__init__.py

@@ -1 +1,3 @@
 from .mipro_optimizer import MiproOptimizer, MIPROv2
+
+__all__ = ["MiproOptimizer", "MIPROv2"]

opik_optimizer/mipro_optimizer/_lm.py

@@ -22,12 +22,13 @@ from dspy.dsp.utils.settings import settings
 from dspy.utils.callback import BaseCallback, with_callbacks
 from dspy.clients.base_lm import BaseLM
 
-from .._throttle import RateLimiter, rate_limited, get_rate_limiter_for_current_opik_installation
+from .._throttle import rate_limited, get_rate_limiter_for_current_opik_installation
 
 logger = logging.getLogger(__name__)
 # Limit how fast an LLM can be called:
 limiter = get_rate_limiter_for_current_opik_installation()
 
+
 class LM(BaseLM):
     """
     A language model supporting chat or text completion requests for use with DSPy modules.
@@ -82,6 +83,7 @@ class LM(BaseLM):
         self.finetuning_model = finetuning_model
         self.launch_kwargs = launch_kwargs or {}
         self.train_kwargs = train_kwargs or {}
+        self.llm_call_counter = 0
 
         # Handle model-specific configuration for different model families
         model_family = model.split("/")[-1].lower() if "/" in model else model.lower()
@@ -94,7 +96,9 @@ class LM(BaseLM):
             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)
+            self.kwargs = dict(
+                temperature=temperature, max_completion_tokens=max_tokens, **kwargs
+            )
         else:
             self.kwargs = dict(temperature=temperature, max_tokens=max_tokens, **kwargs)
 
@@ -110,14 +114,22 @@ class LM(BaseLM):
 
         # Make the request and handle LRU & disk caching.
         if cache_in_memory:
-            completion = cached_litellm_completion if self.model_type == "chat" else cached_litellm_text_completion
+            completion = (
+                cached_litellm_completion
+                if self.model_type == "chat"
+                else cached_litellm_text_completion
+            )
 
             results = completion(
                 request=dict(model=self.model, messages=messages, **kwargs),
                 num_retries=self.num_retries,
             )
         else:
-            completion = litellm_completion if self.model_type == "chat" else litellm_text_completion
+            completion = (
+                litellm_completion
+                if self.model_type == "chat"
+                else litellm_text_completion
+            )
 
             results = completion(
                 request=dict(model=self.model, messages=messages, **kwargs),
@@ -126,9 +138,14 @@ class LM(BaseLM):
                 cache={"no-cache": not cache, "no-store": not cache},
             )
 
-        if not getattr(results, "cache_hit", False) and dspy.settings.usage_tracker and hasattr(results, "usage"):
+        if (
+            not getattr(results, "cache_hit", False)
+            and dspy.settings.usage_tracker
+            and hasattr(results, "usage")
+        ):
             settings.usage_tracker.add_usage(self.model, dict(results.usage))
 
+        self.llm_call_counter += 1
         return results
 
     def launch(self, launch_kwargs: Optional[Dict[str, Any]] = None):
@@ -237,7 +254,11 @@ def request_cache(maxsize: Optional[int] = None):
                 return value.model_json_schema()
             elif isinstance(value, pydantic.BaseModel):
                 return value.model_dump()
-            elif callable(value) and hasattr(value, "__code__") and hasattr(value.__code__, "co_code"):
+            elif (
+                callable(value)
+                and hasattr(value, "__code__")
+                and hasattr(value.__code__, "co_code")
+            ):
                 return value.__code__.co_code.decode("utf-8")
             else:
                 # Note: We don't attempt to compute a hash of the value, since the default
@@ -290,7 +311,11 @@ def cached_litellm_completion(request: Dict[str, Any], num_retries: int):
     )
 
 
-def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
+def litellm_completion(
+    request: Dict[str, Any],
+    num_retries: int,
+    cache={"no-cache": True, "no-store": True},
+):
     retry_kwargs = dict(
         retry_policy=_get_litellm_retry_policy(num_retries),
         retry_strategy="exponential_backoff_retry",
@@ -323,6 +348,7 @@ def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cac
             **retry_kwargs,
             **request,
         )
+
         chunks = []
         async for chunk in response:
             if caller_predict_id:
@@ -344,7 +370,11 @@ def cached_litellm_text_completion(request: Dict[str, Any], num_retries: int):
     )
 
 
-def litellm_text_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
+def litellm_text_completion(
+    request: Dict[str, Any],
+    num_retries: int,
+    cache={"no-cache": True, "no-store": True},
+):
     # Extract the provider and model from the model string.
     # TODO: Not all the models are in the format of "provider/model"
     model = request.pop("model").split("/", 1)
@@ -355,7 +385,9 @@ def litellm_text_completion(request: Dict[str, Any], num_retries: int, cache={"n
     api_base = request.pop("api_base", None) or os.getenv(f"{provider}_API_BASE")
 
     # Build the prompt from the messages.
-    prompt = "\n\n".join([x["content"] for x in request.pop("messages")] + ["BEGIN RESPONSE:"])
+    prompt = "\n\n".join(
+        [x["content"] for x in request.pop("messages")] + ["BEGIN RESPONSE:"]
+    )
 
     return litellm.text_completion(
         cache=cache,

opik_optimizer/mipro_optimizer/_mipro_optimizer_v2.py

@@ -23,12 +23,15 @@ from dspy.teleprompt.utils import (
 from optuna.distributions import CategoricalDistribution
 
 from ..optimization_config.configs import TaskConfig
+from opik_optimizer import task_evaluator
+from opik_optimizer.optimization_config import mappers
 
 
-class Logger():
+class Logger:
     def info(self, *args, **kwargs):
         print(*args)
 
+
 logger = Logger()
 
 # Constants
@@ -49,10 +52,6 @@ BLUE = "\033[94m"
 BOLD = "\033[1m"
 ENDC = "\033[0m"  # Resets the color to default
 
-import opik
-from opik_optimizer import task_evaluator
-from opik_optimizer.optimization_config.configs import TaskConfig
-from opik_optimizer.optimization_config import mappers
 
 def get_prompt(program):
     """
@@ -65,6 +64,7 @@ def get_prompt(program):
 
     return instructions
 
+
 class MIPROv2(Teleprompter):
     def __init__(
         self,
@@ -554,18 +554,19 @@ class MIPROv2(Teleprompter):
 
         examples = []
         for demo in demo_candidates.values():
-            for l in demo:
-                for example in l:
+            for l_ in demo:
+                for example in l_:
                     examples.append(example.toDict())
         prompt = get_prompt(program)
         experiment_config = {
             **self.experiment_config,
-            **{"configuration": {
-                "prompt": prompt,
-                "examples": examples,
+            **{
+                "configuration": {
+                    "prompt": prompt,
+                    "examples": examples,
+                },
+                "evaluation": "initial",
             },
-               "evaluation": "initial",
-            }
         }
 
         default_score = eval_candidate_program_with_opik(
@@ -579,7 +580,7 @@ class MIPROv2(Teleprompter):
             experiment_config=experiment_config,
             optimization_id=self.opik_optimization_id,
         )
-        
+
         logger.info(f"Default program score: {default_score}\n")
 
         trial_logs = {}
@@ -606,7 +607,13 @@ class MIPROv2(Teleprompter):
 
         # Define the objective function
         def objective(trial):
-            nonlocal program, best_program, best_score, trial_logs, total_eval_calls, score_data
+            nonlocal \
+                program, \
+                best_program, \
+                best_score, \
+                trial_logs, \
+                total_eval_calls, \
+                score_data
 
             trial_num = trial.number + 1
             if minibatch:
@@ -927,18 +934,19 @@ class MIPROv2(Teleprompter):
 
         examples = []
         for demo in demo_candidates.values():
-            for l in demo:
-                for example in l:
+            for l_ in demo:
+                for example in l_:
                     examples.append(example.toDict())
         prompt = get_prompt(highest_mean_program)
         experiment_config = {
             **self.experiment_config,
-            **{"configuration": {
-                "prompt": prompt,
-                "examples": examples,
-              },
-               "evaluation": "full",
-            }
+            **{
+                "configuration": {
+                    "prompt": prompt,
+                    "examples": examples,
+                },
+                "evaluation": "full",
+            },
         }
 
         full_eval_score = eval_candidate_program_with_opik(
@@ -988,7 +996,7 @@ class MIPROv2(Teleprompter):
         trial_logs[trial_num + 1]["full_eval_score"] = full_eval_score
 
         if full_eval_score == 1.0:
-            return self.early_stop(default_score, program)
+            return self.early_stop(default_score, program)  # noqa
 
         # Update best score and program if necessary
         if full_eval_score > best_score:
@@ -1042,9 +1050,12 @@ def eval_candidate_program_with_opik(
             candidate_program._assert_failures += dspy.settings.get("assert_failures")
         if hasattr(candidate_program, "_suggest_failures"):
             candidate_program._suggest_failures += dspy.settings.get("suggest_failures")
-        
-        return {mappers.from_llm_response_text(): prediction[prompt_task_config.output_dataset_field]}
 
+        return {
+            mappers.from_llm_response_text(): prediction[
+                prompt_task_config.output_dataset_field
+            ]
+        }
 
     score = task_evaluator.evaluate(
         dataset=opik_dataset,
@@ -1056,5 +1067,5 @@ def eval_candidate_program_with_opik(
         experiment_config=experiment_config,
         optimization_id=optimization_id,
     )
-    
+
     return score