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

opik_optimizer/few_shot_bayesian_optimizer/reporting.py

@@ -1,11 +1,13 @@
 from contextlib import contextmanager
 from io import StringIO
-from typing import Dict, List
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
-import rich
 from rich.panel import Panel
 from rich.text import Text
 
+if TYPE_CHECKING:
+    from .few_shot_bayesian_optimizer import FewShotPromptTemplate
+
 from ..reporting_utils import (
     convert_tqdm_to_rich,
     display_configuration,  # noqa: F401
@@ -21,20 +23,22 @@ console = get_console()
 
 
 @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:
             nonlocal score
             score = s
-    
+
     # Use our log suppression context manager and yield the reporter
     with suppress_opik_logs():
         with convert_tqdm_to_rich(verbose=verbose):
@@ -42,31 +46,50 @@ def display_evaluation(message: str = "First we will establish the baseline perf
                 yield Reporter()
             finally:
                 if verbose >= 1:
-                    console.print(Text(f"\r  Baseline score was: {score:.4f}.\n", style="green"))
+                    console.print(
+                        Text(f"\r  Baseline score was: {score:.4f}.\n", style="green")
+                    )
+
 
 @contextmanager
-def creation_few_shot_prompt_template(verbose: int = 1):
+def creation_few_shot_prompt_template(verbose: int = 1) -> Any:
     """Context manager to display messages during the creation of a few-shot prompt template."""
-    console.print(Text("> Let's add a placeholder for few-shot examples in the messages:"))
+    console.print(
+        Text("> Let's add a placeholder for few-shot examples in the messages:")
+    )
+
+    fewshot_template: Optional["FewShotPromptTemplate"] = None
 
-    fewshot_template = None
-    
     # Create a simple object with a method to set the prompt template
     class Reporter:
-        def set_fewshot_template(self, s):
+        def set_fewshot_template(self, s: "FewShotPromptTemplate") -> None:
             nonlocal fewshot_template
             fewshot_template = s
-    
+
     # Use our log suppression context manager and yield the reporter
     try:
         yield Reporter()
     finally:
         if verbose >= 1:
-            console.print(Text("│    Created the prompt template:\n│", style="dim yellow"))
-            display_messages(fewshot_template.message_list_with_placeholder, prefix="│    ")
-            console.print(Text("│\n│   With the FEW_SHOT_EXAMPLE_PLACEHOLDER following the format:"))
-            
-            panel = Panel(Text(fewshot_template.example_template), width=PANEL_WIDTH, border_style="dim")
+            console.print(
+                Text("│    Created the prompt template:\n│", style="dim yellow")
+            )
+            if fewshot_template is not None:
+                display_messages(
+                    fewshot_template.message_list_with_placeholder, prefix="│    "
+                )
+            console.print(
+                Text(
+                    "│\n│   With the FEW_SHOT_EXAMPLE_PLACEHOLDER following the format:"
+                )
+            )
+
+            if fewshot_template is not None:
+                panel = Panel(
+                    Text(fewshot_template.example_template),
+                    width=PANEL_WIDTH,
+                    border_style="dim",
+                )
             # Use a temporary buffer to render the panel
             buffer = StringIO()
             temp_console = get_console(file=buffer, width=console.width)
@@ -80,7 +103,8 @@ def creation_few_shot_prompt_template(verbose: int = 1):
             console.print(prefixed)
             console.print()
 
-def start_optimization_run(verbose: int = 1):
+
+def start_optimization_run(verbose: int = 1) -> None:
     """Start the optimization run"""
     if verbose >= 1:
         console.print(Text("\n> Starting the optimization run"))
@@ -88,27 +112,54 @@ def start_optimization_run(verbose: int = 1):
 
 
 @contextmanager
-def start_optimization_trial(trial_number: int, total_trials: int, verbose: int = 1):
+def start_optimization_trial(
+    trial_number: int, total_trials: 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 start_trial(self, messages: List[Dict[str, str]]):
+        def start_trial(self, messages: List[Dict[str, str]]) -> None:
             if verbose >= 1:
-                console.print(Text(f"│ - Starting optimization round {trial_number + 1} of {total_trials}"))
+                console.print(
+                    Text(
+                        f"│ - Starting optimization round {trial_number + 1} of {total_trials}"
+                    )
+                )
                 console.print(Text("│"))
                 display_messages(messages, prefix="│    ")
                 console.print("│")
 
-        def set_score(self, baseline_score, score):
+        def set_score(self, baseline_score: float, score: float) -> None:
             if verbose >= 1:
                 if baseline_score == 0:
-                    console.print(Text(f"│    Trial {trial_number + 1} - score was: {score:.4f}\n│", style="green"))
+                    console.print(
+                        Text(
+                            f"│    Trial {trial_number + 1} - score was: {score:.4f}\n│",
+                            style="green",
+                        )
+                    )
                 elif score is not None and score > baseline_score:
-                    console.print(Text(f"│    Trial {trial_number + 1} - score was: {score:.4f} ({(score - baseline_score) / baseline_score * 100:.2f}%).\n│", style="green"))
+                    console.print(
+                        Text(
+                            f"│    Trial {trial_number + 1} - score was: {score:.4f} ({(score - baseline_score) / baseline_score * 100:.2f}%).\n│",
+                            style="green",
+                        )
+                    )
                 elif score is not None and score <= baseline_score:
-                    console.print(Text(f"│    Trial {trial_number + 1} - score was: {score:.4f} ({(score - baseline_score) / baseline_score * 100:.2f}%).\n│", style="red"))
+                    console.print(
+                        Text(
+                            f"│    Trial {trial_number + 1} - score was: {score:.4f} ({(score - baseline_score) / baseline_score * 100:.2f}%).\n│",
+                            style="red",
+                        )
+                    )
                 else:
-                    console.print(Text(f"│    Trial {trial_number + 1} - score was not set.\n│", style="dim yellow"))
+                    console.print(
+                        Text(
+                            f"│    Trial {trial_number + 1} - score was not set.\n│",
+                            style="dim yellow",
+                        )
+                    )
 
     # Use our log suppression context manager and yield the reporter
     with suppress_opik_logs():

opik_optimizer/logging_config.py

@@ -1,18 +1,19 @@
 import logging
 from rich.logging import RichHandler
 
-DEFAULT_LOG_FORMAT = '%(message)s'
-DEFAULT_DATE_FORMAT = '%Y-%m-%d %H:%M:%S'
+DEFAULT_LOG_FORMAT = "%(message)s"
+DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
 
 # Store configured state to prevent reconfiguration
 _logging_configured = False
 
+
 def setup_logging(
-    level=logging.WARNING,
-    format_string=DEFAULT_LOG_FORMAT,
-    date_format=DEFAULT_DATE_FORMAT,
-    force=False,
-):
+    level: int = logging.WARNING,
+    format_string: str = DEFAULT_LOG_FORMAT,
+    date_format: str = DEFAULT_DATE_FORMAT,
+    force: bool = False,
+) -> None:
     """
     Configures logging for the opik_optimizer package using rich.
 
@@ -28,7 +29,7 @@ def setup_logging(
         return
 
     # Configure opik_optimizer package logger
-    package_logger = logging.getLogger('opik_optimizer')
+    package_logger = logging.getLogger("opik_optimizer")
 
     # Avoid adding handlers repeatedly if force=True replaces them
     if not package_logger.handlers or force:
@@ -36,19 +37,19 @@ def setup_logging(
         if force and package_logger.handlers:
             for handler in package_logger.handlers[:]:
                 package_logger.removeHandler(handler)
-                
+
         console_handler = RichHandler(
             rich_tracebacks=True,
-            markup=True, # Enable rich markup in log messages
-            log_time_format=f"[{date_format}]" # Apply date format
+            markup=True,  # Enable rich markup in log messages
+            log_time_format=f"[{date_format}]",  # Apply date format
         )
         # RichHandler manages formatting, so we don't need a separate formatter
         # formatter = logging.Formatter(format_string, datefmt=date_format)
         # console_handler.setFormatter(formatter)
         package_logger.addHandler(console_handler)
-        
+
     package_logger.setLevel(level)
-    package_logger.propagate = False # Don't duplicate messages in root logger
+    package_logger.propagate = False  # Don't duplicate messages in root logger
 
     # Set levels for noisy libraries like LiteLLM and httpx
     logging.getLogger("LiteLLM").setLevel(logging.WARNING)
@@ -61,9 +62,12 @@ def setup_logging(
     logging.getLogger("filelock").setLevel(logging.WARNING)
 
     _logging_configured = True
-    
+
     # Use level name provided by rich handler by default
-    package_logger.info(f"Opik Agent Optimizer logging configured to level: [bold cyan]{logging.getLevelName(level)}[/bold cyan]")
+    package_logger.info(
+        f"Opik Agent Optimizer logging configured to level: [bold cyan]{logging.getLevelName(level)}[/bold cyan]"
+    )
+
 
 # Ensure logger obtained after setup can be used immediately if needed
 logger = logging.getLogger(__name__)