implement nested benchmarks with snapshots (#280)

BREAKING: added benchmark obj to Skill's execute_tool function

Author: Shackless

api/commands.py

@@ -1,7 +1,14 @@
 from typing import Literal, Optional
 from pydantic import BaseModel
-from api.enums import CommandTag, KeyboardRecordingType, LogSource, LogType, RecordingDevice, ToastType
-from api.interface import CommandActionConfig
+from api.enums import (
+    CommandTag,
+    KeyboardRecordingType,
+    LogSource,
+    LogType,
+    RecordingDevice,
+    ToastType,
+)
+from api.interface import CommandActionConfig, BenchmarkResult
 
 
 # We use this Marker base class for reflection to "iterate all commands"
@@ -28,10 +35,11 @@ class RecordKeyboardActionsCommand(WebSocketCommandModel):
     command: Literal["record_keyboard_actions"] = "record_keyboard_actions"
     recording_type: KeyboardRecordingType
 
+
 class RecordMouseActionsCommand(WebSocketCommandModel):
     command: Literal["record_mouse_actions"] = "record_mouse_actions"
 
-    
+
 class RecordJoystickActionsCommand(WebSocketCommandModel):
     command: Literal["record_joystick_actions"] = "record_joystick_actions"
 
@@ -55,6 +63,7 @@ class LogCommand(WebSocketCommandModel):
     tag: Optional[CommandTag] = None
     skill_name: Optional[str] = None
     additional_data: Optional[dict] = None
+    benchmark_result: Optional[BenchmarkResult] = None
 
 
 class PromptSecretCommand(WebSocketCommandModel):

api/interface.py

@@ -742,3 +742,13 @@ class SettingsConfig(BaseModel):
     wingman_pro: WingmanProSettings
     xvasynth: XVASynthSettings
     debug_mode: bool = False
+
+
+class BenchmarkResult(BaseModel):
+    label: str
+    execution_time_ms: float
+    formatted_execution_time: str
+    snapshots: Optional[list["BenchmarkResult"]] = None
+
+
+BenchmarkResult.model_rebuild()

main.py

@@ -14,6 +14,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.openapi.utils import get_openapi
 from api.commands import WebSocketCommandModel
+from api.interface import BenchmarkResult
 from api.enums import ENUM_TYPES, LogType, WingmanInitializationErrorType
 import keyboard.keyboard as keyboard
 from services.command_handler import CommandHandler
@@ -256,6 +257,23 @@ async def ping():
     return "Ok" if core.is_started else "Starting"
 
 
+# required to generate API specs for class BenchmarkResult that is only used internally
+@app.get("/dummy-benchmark", tags=["main"], response_model=BenchmarkResult)
+async def get_dummy_benchmark():
+    return BenchmarkResult(
+        label="Sample Benchmark",
+        execution_time_ms=150.0,
+        formatted_execution_time=0.15,
+        snapshots=[
+            BenchmarkResult(
+                label="Sub Benchmark",
+                execution_time_ms=75.0,
+                formatted_execution_time=0.075,
+            )
+        ],
+    )
+
+
 async def async_main(host: str, port: int, sidecar: bool):
     await core.config_service.migrate_configs(system_manager)
     await core.config_service.load_config()

services/benchmark.py

@@ -0,0 +1,63 @@
+import time
+from api.enums import LogType
+from api.interface import BenchmarkResult
+from services.printr import Printr
+
+
+class Benchmark:
+    def __init__(self, label: str):
+        self.label = label
+        self.snapshot_label: str = None
+        self.start_time = time.perf_counter()
+        self.snapshot_start_time: float = None
+        self.snapshots: list[BenchmarkResult] = []
+        self.printr = Printr()
+
+    def finish(self):
+        if self.snapshot_label or self.snapshot_start_time:
+            self.finish_snapshot()
+            self.printr.print(
+                f"Snapshot benchmark '{self.snapshot_label}' was still running when finishing '{self.label}'.",
+                color=LogType.WARNING,
+                server_only=True,
+            )
+        result = self._create_benchmark_result(self.label, self.start_time)
+        if len(self.snapshots) > 0:
+            result.snapshots = self.snapshots
+        return result
+
+    def start_snapshot(self, label: str):
+        if self.snapshot_label or self.snapshot_start_time:
+            self.finish_snapshot()
+            self.printr.print(
+                f"Snapshot benchmark '{self.snapshot_label}' was still running when starting '{label}'.",
+                color=LogType.WARNING,
+                server_only=True,
+            )
+        self.snapshot_label = label
+        self.snapshot_start_time = time.perf_counter()
+
+    def finish_snapshot(self):
+        try:
+            result = self._create_benchmark_result(
+                label=self.snapshot_label, start_time=self.snapshot_start_time
+            )
+            self.snapshots.append(result)
+        except Exception:
+            pass
+        self.snapshot_label = None
+        self.snapshot_start_time = None
+
+    def _create_benchmark_result(self, label: str, start_time: float):
+        end_time = time.perf_counter()
+        execution_time = (end_time - start_time) * 1000  # Convert to milliseconds
+        if execution_time >= 1000:
+            formatted_execution_time = f"{execution_time/1000:.1f}s"
+        else:
+            formatted_execution_time = f"{int(execution_time)}ms"
+
+        return BenchmarkResult(
+            label=label,
+            execution_time_ms=execution_time,
+            formatted_execution_time=formatted_execution_time,
+        )

services/printr.py

@@ -7,27 +7,30 @@
 from os import path
 from api.commands import LogCommand, ToastCommand
 from api.enums import CommandTag, LogSource, LogType, ToastType
+from api.interface import BenchmarkResult
 from services.file import get_writable_dir
 from services.websocket_user import WebSocketUser
 
+
 class StreamToLogger:
     def __init__(self, logger, log_level=logging.INFO, stream=sys.stdout):
         self.logger = logger
         self.log_level = log_level
-        self.linebuf = ''
+        self.linebuf = ""
         self.stream = stream
 
     def write(self, buf):
         for line in buf.rstrip().splitlines():
             self.logger.log(self.log_level, line.rstrip())
-            self.stream.write(line + '\n')
+            self.stream.write(line + "\n")
 
     def flush(self):
         self.stream.flush()
 
     def isatty(self):
         return False
 
+
 class Printr(WebSocketUser):
     """Singleton"""
 
@@ -50,7 +53,7 @@ class Printr(WebSocketUser):
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super(Printr, cls).__new__(cls)
-            cls._instance.logger = logging.getLogger('file_logger')
+            cls._instance.logger = logging.getLogger("file_logger")
             cls._instance.logger.setLevel(logging.INFO)
 
             timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
@@ -60,16 +63,18 @@ def __new__(cls):
                 path.join(get_writable_dir("logs"), f"wingman-core.{timestamp}.log")
             )
             fh.setLevel(logging.DEBUG)
-            file_formatter = Formatter('%(asctime)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
+            file_formatter = Formatter(
+                "%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+            )
             fh.setFormatter(file_formatter)
             cls._instance.logger.addHandler(fh)
 
             # Console logger with color
-            cls._instance.console_logger = logging.getLogger('console_logger')
+            cls._instance.console_logger = logging.getLogger("console_logger")
             cls._instance.console_logger.setLevel(logging.INFO)
             ch = logging.StreamHandler()
             ch.setLevel(logging.INFO)
-            console_formatter = Formatter('%(message)s')
+            console_formatter = Formatter("%(message)s")
             ch.setFormatter(console_formatter)
             cls._instance.console_logger.addHandler(ch)
 
@@ -88,6 +93,7 @@ async def __send_to_gui(
         command_tag: CommandTag = None,
         skill_name: str = "",
         additional_data: dict = None,
+        benchmark_result: BenchmarkResult = None,
     ):
         if self._connection_manager is None:
             raise ValueError("connection_manager has not been set.")
@@ -102,10 +108,13 @@ async def __send_to_gui(
             if current_frame is not None:
                 while current_frame:
                     # Check if the caller is a method of a class
-                    if 'self' in current_frame.f_locals:
-                        caller_instance = current_frame.f_locals['self']
+                    if "self" in current_frame.f_locals:
+                        caller_instance = current_frame.f_locals["self"]
                         caller_instance_name = caller_instance.__class__.__name__
-                        if caller_instance_name == "Wingman" or caller_instance_name == "OpenAiWingman":
+                        if (
+                            caller_instance_name == "Wingman"
+                            or caller_instance_name == "OpenAiWingman"
+                        ):
                             wingman_name = caller_instance.name
                             break
                     # Move to the previous frame in the call stack
@@ -121,6 +130,7 @@ async def __send_to_gui(
                     skill_name=skill_name,
                     additional_data=additional_data,
                     wingman_name=wingman_name,
+                    benchmark_result=benchmark_result,
                 )
             )
 
@@ -163,9 +173,23 @@ async def print_async(
         command_tag: CommandTag = None,
         skill_name: str = "",
         additional_data: dict = None,
+        benchmark_result: BenchmarkResult = None,
     ):
         # print to server (terminal)
-        self.print_colored(text, color=self.get_terminal_color(color))
+        self.print_colored(
+            (
+                text
+                if not benchmark_result
+                else f"{text} ({benchmark_result.formatted_execution_time})"
+            ),
+            color=self.get_terminal_color(color),
+        )
+        if benchmark_result and benchmark_result.snapshots:
+            for snapshot in benchmark_result.snapshots:
+                self.print_colored(
+                    f"  - {snapshot.label}: {snapshot.formatted_execution_time}",
+                    color=self.get_terminal_color(color),
+                )
 
         if not server_only and self._connection_manager is not None:
             await self.__send_to_gui(
@@ -177,6 +201,7 @@ async def print_async(
                 command_tag=command_tag,
                 skill_name=skill_name,
                 additional_data=additional_data,
+                benchmark_result=benchmark_result,
             )
 
     def toast(self, text: str):