opik-optimizer 0.9.2__py3-none-any.whl → 1.0.1__py3-none-any.whl

opik_optimizer/evolutionary_optimizer/reporting.py

@@ -1,6 +1,6 @@
 from contextlib import contextmanager
 from io import StringIO
-from typing import List
+from typing import Any, List
 
 from rich.panel import Panel
 from rich.text import Text
@@ -21,19 +21,30 @@ console = get_console()
 
 
 @contextmanager
-def infer_output_style(verbose: int = 1):
+def infer_output_style(verbose: int = 1) -> Any:
     class Reporter:
-        def start_style_inference(self, output_style_prompt):
+        def start_style_inference(self) -> None:
             if verbose >= 1:
                 console.print("> Infering the output style using the prompt:")
                 console.print("│")
-            
-        def error(self, error_message):
+
+        def error(self, error_message: str) -> None:
             if verbose >= 1:
-                console.print(Text("│    ").append(Text(f"Failed to infer output style: {error_message}", style="red")))
-                console.print(Text("│    ").append(Text("Continuing with default style", style="dim")))
-        
-        def display_style_inference_prompt(self, output_style_prompt):
+                console.print(
+                    Text("│    ").append(
+                        Text(
+                            f"Failed to infer output style: {error_message}",
+                            style="red",
+                        )
+                    )
+                )
+                console.print(
+                    Text("│    ").append(
+                        Text("Continuing with default style", style="dim")
+                    )
+                )
+
+        def display_style_inference_prompt(self, output_style_prompt: str) -> None:
             if verbose >= 1:
                 panel = Panel(
                     Text(output_style_prompt),
@@ -57,7 +68,7 @@ def infer_output_style(verbose: int = 1):
                 console.print(prefixed)
                 console.print(Text("│"))
 
-        def success(self, output_style_prompt):
+        def success(self, output_style_prompt: str) -> None:
             if verbose >= 1:
                 panel = Panel(
                     Text(output_style_prompt),
@@ -67,7 +78,7 @@ def infer_output_style(verbose: int = 1):
                     width=PANEL_WIDTH,
                     padding=(1, 2),
                 )
-                
+
                 # Capture the panel as rendered text with ANSI styles
                 with console.capture() as capture:
                     console.print(panel)
@@ -76,77 +87,102 @@ def infer_output_style(verbose: int = 1):
                 rendered_panel = capture.get()
 
                 # Prefix each line with '│ ', preserving ANSI styles
-                prefixed_output = "\n".join(f"│ {line}" for line in rendered_panel.splitlines())
+                prefixed_output = "\n".join(
+                    f"│ {line}" for line in rendered_panel.splitlines()
+                )
 
                 # Print the prefixed output (will include colors)
                 console.print(prefixed_output, highlight=False)
                 console.print(Text(""))
-    
+
     try:
         yield Reporter()
     finally:
         pass
 
+
 @contextmanager
-def initializing_population(verbose: int = 1):
+def initializing_population(verbose: int = 1) -> Any:
     class Reporter:
-        def start(self, population_size):
+        def start(self, population_size: int) -> None:
             if verbose >= 1:
-                console.print(f"> Creating {population_size - 1} variations of the initial prompt")
+                console.print(
+                    f"> Creating {population_size - 1} variations of the initial prompt"
+                )
                 console.print("│")
-            
-        def start_fresh_prompts(self, num_fresh_starts):
-            if verbose >= 1:
-                console.print(f"│    Generating {num_fresh_starts} fresh prompts based on the task description.")
-        
-        def success_fresh_prompts(self, num_fresh_starts):
+
+        def start_fresh_prompts(self, num_fresh_starts: int) -> None:
             if verbose >= 1:
-                console.print(Text("│       ").append(Text(f"Successfully generated {num_fresh_starts} fresh prompts based on the task description.", style="dim green")))
-                console.print("│")
-            
-        def failed_fresh_prompts(self, num_fresh_starts, error):
+                console.print(
+                    f"│    Generating {num_fresh_starts} fresh prompts based on the task description."
+                )
+
+        def success_fresh_prompts(self, num_fresh_starts: int) -> None:
             if verbose >= 1:
-                console.print(Text("│       ").append(Text(f"Failed to generate fresh prompts from LLM: {error}", style="dim red")))
+                console.print(
+                    Text("│       ").append(
+                        Text(
+                            f"Successfully generated {num_fresh_starts} fresh prompts based on the task description.",
+                            style="dim green",
+                        )
+                    )
+                )
                 console.print("│")
 
-        def start_variations(self, num_variations):
+        def start_variations(self, num_variations: int) -> None:
             if verbose >= 1:
-                console.print(f"│    Generating {num_variations} variations of the initial prompt.")
-        
-        def success_variations(self, num_variations):
+                console.print(
+                    f"│    Generating {num_variations} variations of the initial prompt."
+                )
+
+        def success_variations(self, num_variations: int) -> None:
             if verbose >= 1:
-                console.print(Text(f"│       Successfully generated {num_variations - 1} variations of the initial prompt).", style="dim green"))
+                console.print(
+                    Text(
+                        f"│       Successfully generated {num_variations - 1} variations of the initial prompt).",
+                        style="dim green",
+                    )
+                )
                 console.print("│")
-            
-        def failed_variations(self, num_variations, error):
+
+        def failed_variations(self, num_variations: int, error: str) -> None:
             if verbose >= 1:
-                console.print(Text(f"│       Failed to generate {num_variations - 1} variations of the initial prompt: {error}", style="dim red"))
+                console.print(
+                    Text(
+                        f"│       Failed to generate {num_variations - 1} variations of the initial prompt: {error}",
+                        style="dim red",
+                    )
+                )
                 console.print("│")
-            
-        def end(self, population_prompts: List[chat_prompt.ChatPrompt]):
+
+        def end(self, population_prompts: List[chat_prompt.ChatPrompt]) -> None:
             if verbose >= 1:
-                console.print(f"│ Successfully initialized population with {len(population_prompts)} prompts.")
+                console.print(
+                    f"│ Successfully initialized population with {len(population_prompts)} prompts."
+                )
                 console.print("")
-                
-    
+
     try:
         yield Reporter()
     finally:
         pass
 
+
 @contextmanager
-def baseline_performance(verbose: int = 1):
+def baseline_performance(verbose: int = 1) -> Any:
     """Context manager to display messages during an evaluation phase."""
     # Entry point
     if verbose >= 1:
         console.print(Text("> First we will establish the baseline performance."))
-    
+
     # 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):
@@ -155,22 +191,27 @@ def baseline_performance(verbose: int = 1):
             finally:
                 pass
 
+
 @contextmanager
-def evaluate_initial_population(verbose: int = 1):
+def evaluate_initial_population(verbose: int = 1) -> Any:
     """Context manager to display messages during an evaluation phase."""
     # Entry point
     if verbose >= 1:
         console.print(Text("> Let's now evaluate the initial population"))
-    
+
     # Create a simple object with a method to set the score
     class Reporter:
-        def set_score(self, index, score, baseline_score):
+        def set_score(self, index: int, score: float, baseline_score: float) -> None:
             if verbose >= 1:
                 if score >= baseline_score:
-                    console.print(Text(f"\r  Prompt {index+1} score was: {score}.", style="green"))
+                    console.print(
+                        Text(f"\r  Prompt {index+1} score was: {score}.", style="green")
+                    )
                 else:
-                    console.print(Text(f"\r  Prompt {index+1} score was: {score}.", style="dim"))
-        
+                    console.print(
+                        Text(f"\r  Prompt {index+1} score was: {score}.", style="dim")
+                    )
+
     # Use our log suppression context manager and yield the reporter
     with suppress_opik_logs():
         with convert_tqdm_to_rich("│   Evaluation", verbose=verbose):
@@ -180,38 +221,60 @@ def evaluate_initial_population(verbose: int = 1):
                 if verbose >= 1:
                     console.print("")
 
+
 @contextmanager
-def start_evolutionary_algo(verbose: int = 1):
+def start_evolutionary_algo(verbose: int = 1) -> Any:
     """Context manager to display messages during an evolutionary algorithm phase."""
     # Entry point
     if verbose >= 1:
         console.print(Text("> Starting evolutionary algorithm optimization"))
-    
+
     # Create a simple object with a method to set the score
     class Reporter:
-        def start_gen(self, gen, num_gens):
+        def start_gen(self, gen: int, num_gens: int) -> None:
             if verbose >= 1:
                 console.print(Text(f"│   Starting generation {gen} of {num_gens}"))
 
-        def restart_population(self, restart_generation_nb):
+        def restart_population(self, restart_generation_nb: int) -> None:
             if verbose >= 1:
-                console.print(Text(f"│      Re-creating the population as we have not made progress in {restart_generation_nb} generations."))
-        
-        def performing_crossover(self):
+                console.print(
+                    Text(
+                        f"│      Re-creating the population as we have not made progress in {restart_generation_nb} generations."
+                    )
+                )
+
+        def performing_crossover(self) -> None:
             if verbose >= 1:
-                console.print(Text("│      Performing crossover - Combining multiple prompts into a new one."))
-        
-        def performing_mutation(self):
+                console.print(
+                    Text(
+                        "│      Performing crossover - Combining multiple prompts into a new one."
+                    )
+                )
+
+        def performing_mutation(self) -> None:
             if verbose >= 1:
-                console.print(Text("│      Performing mutation - Altering prompts to improve their performance."))
-        
-        def performing_evaluation(self, num_prompts: int):
+                console.print(
+                    Text(
+                        "│      Performing mutation - Altering prompts to improve their performance."
+                    )
+                )
+
+        def performing_evaluation(self, num_prompts: int) -> None:
             if verbose >= 1:
-                console.print(Text(f"│      Performing evaluation - Assessing {num_prompts} prompts' performance."))
-        
-        def performed_evaluation(self, prompt_idx: int, score: float):
+                console.print(
+                    Text(
+                        f"│      Performing evaluation - Assessing {num_prompts} prompts' performance."
+                    )
+                )
+
+        def performed_evaluation(self, prompt_idx: int, score: float) -> None:
             if verbose >= 1:
-                console.print(Text(f"│      Performed evaluation for prompt {prompt_idx} - Score: {score:.4f}.", style="dim"))
+                console.print(
+                    Text(
+                        f"│      Performed evaluation for prompt {prompt_idx} - Score: {score:.4f}.",
+                        style="dim",
+                    )
+                )
 
     # Use our log suppression context manager and yield the reporter
     with suppress_opik_logs():
@@ -222,23 +285,41 @@ def start_evolutionary_algo(verbose: int = 1):
                 if verbose >= 1:
                     console.print("")
 
-def display_error(error_message, verbose: int = 1):
+
+def display_error(error_message: str, verbose: int = 1) -> None:
     if verbose >= 1:
         console.print(Text("│   ").append(Text(error_message, style="dim red")))
 
-def display_success(message, verbose: int = 1):
+
+def display_success(message: str, verbose: int = 1) -> None:
     if verbose >= 1:
         console.print(Text("│   ").append(Text(message, style="dim green")))
 
-def display_message(message, verbose: int = 1):
+
+def display_message(message: str, verbose: int = 1) -> None:
     if verbose >= 1:
         console.print(Text("│   ").append(Text(message, style="dim")))
 
-def end_gen(generation_idx, best_gen_score, initial_primary_score, verbose: int = 1):
+
+def end_gen(
+    generation_idx: int,
+    best_gen_score: float,
+    initial_primary_score: float,
+    verbose: int = 1,
+) -> None:
     if verbose >= 1:
         if best_gen_score >= initial_primary_score:
-            console.print(Text(f"│   Generation {generation_idx} completed. Found a new prompt with a score of {best_gen_score:.4f}.", style="green"))
+            console.print(
+                Text(
+                    f"│   Generation {generation_idx} completed. Found a new prompt with a score of {best_gen_score:.4f}.",
+                    style="green",
+                )
+            )
         else:
-            console.print(Text(f"│   Generation {generation_idx} completed. No improvement in this generation."))
+            console.print(
+                Text(
+                    f"│   Generation {generation_idx} completed. No improvement in this generation."
+                )
+            )
 
         console.print("│")