opik-optimizer 0.8.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

opik_optimizer/evolutionary_optimizer/reporting.py

@@ -0,0 +1,246 @@
+from contextlib import contextmanager
+from io import StringIO
+from typing import List
+
+import rich
+from rich.console import Console
+from rich.panel import Panel
+from rich.text import Text
+
+from ..optimization_config import chat_prompt
+from ..reporting_utils import (
+    convert_tqdm_to_rich,
+    display_configuration,  # noqa: F401
+    display_header,  # noqa: F401
+    display_messages,  # noqa: F401
+    display_result,  # noqa: F401
+    get_console,
+    suppress_opik_logs,
+)
+
+PANEL_WIDTH = 70
+console = get_console()
+
+
+@contextmanager
+def infer_output_style(verbose: int = 1):
+    class Reporter:
+        def start_style_inference(self, output_style_prompt):
+            if verbose >= 1:
+                console.print("> Infering the output style using the prompt:")
+                console.print("│")
+            
+        def error(self, error_message):
+            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):
+            if verbose >= 1:
+                panel = Panel(
+                    Text(output_style_prompt),
+                    title="Output style inference prompt",
+                    title_align="left",
+                    border_style="dim",
+                    width=PANEL_WIDTH,
+                    padding=(1, 2),
+                )
+
+                # Use a temporary buffer to render the panel
+                buffer = StringIO()
+                temp_console = get_console(file=buffer, width=console.width)
+                temp_console.print(panel)
+
+                # Add prefix to each line
+                panel_output = buffer.getvalue()
+                prefixed = "\n".join(f"│ {line}" for line in panel_output.splitlines())
+
+                # Print the final result
+                console.print(prefixed)
+                console.print(Text("│"))
+
+        def success(self, output_style_prompt):
+            if verbose >= 1:
+                panel = Panel(
+                    Text(output_style_prompt),
+                    title="Successfully inferred output style",
+                    title_align="left",
+                    border_style="green",
+                    width=PANEL_WIDTH,
+                    padding=(1, 2),
+                )
+                
+                # Capture the panel as rendered text with ANSI styles
+                with console.capture() as capture:
+                    console.print(panel)
+
+                # Retrieve the rendered string (with ANSI)
+                rendered_panel = capture.get()
+
+                # Prefix each line with '│ ', preserving ANSI styles
+                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):
+    class Reporter:
+        def start(self, population_size):
+            if verbose >= 1:
+                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):
+            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):
+            if verbose >= 1:
+                console.print(Text("│       ").append(Text(f"Failed to generate fresh prompts from LLM: {error}", style="dim red")))
+                console.print("│")
+
+        def start_variations(self, num_variations):
+            if verbose >= 1:
+                console.print(f"│    Generating {num_variations} variations of the initial prompt.")
+        
+        def success_variations(self, num_variations):
+            if verbose >= 1:
+                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):
+            if verbose >= 1:
+                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]):
+            if verbose >= 1:
+                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):
+    """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):
+            if verbose >= 1:
+                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):
+            try:
+                yield Reporter()
+            finally:
+                pass
+
+@contextmanager
+def evaluate_initial_population(verbose: int = 1):
+    """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):
+            if verbose >= 1:
+                if score >= baseline_score:
+                    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"))
+        
+    # Use our log suppression context manager and yield the reporter
+    with suppress_opik_logs():
+        with convert_tqdm_to_rich("│   Evaluation", verbose=verbose):
+            try:
+                yield Reporter()
+            finally:
+                if verbose >= 1:
+                    console.print("")
+
+@contextmanager
+def start_evolutionary_algo(verbose: int = 1):
+    """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):
+            if verbose >= 1:
+                console.print(Text(f"│   Starting generation {gen} of {num_gens}"))
+
+        def restart_population(self, restart_generation_nb):
+            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):
+            if verbose >= 1:
+                console.print(Text("│      Performing crossover - Combining multiple prompts into a new one."))
+        
+        def performing_mutation(self):
+            if verbose >= 1:
+                console.print(Text("│      Performing mutation - Altering prompts to improve their performance."))
+        
+        def performing_evaluation(self, num_prompts: int):
+            if verbose >= 1:
+                console.print(Text(f"│      Performing evaluation - Assessing {num_prompts} prompts' performance."))
+        
+        def performed_evaluation(self, prompt_idx: int, score: float):
+            if verbose >= 1:
+                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():
+        with convert_tqdm_to_rich("│         Evaluation", verbose=verbose):
+            try:
+                yield Reporter()
+            finally:
+                if verbose >= 1:
+                    console.print("")
+
+def display_error(error_message, verbose: int = 1):
+    if verbose >= 1:
+        console.print(Text("│   ").append(Text(error_message, style="dim red")))
+
+def display_success(message, verbose: int = 1):
+    if verbose >= 1:
+        console.print(Text("│   ").append(Text(message, style="dim green")))
+
+def display_message(message, verbose: int = 1):
+    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):
+    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"))
+        else:
+            console.print(Text(f"│   Generation {generation_idx} completed. No improvement in this generation."))
+
+        console.print("│")