Inline STT and tray

src/leld.py

@@ -9,13 +9,18 @@ import threading
 import time
 from pathlib import Path
 
+import gi
+from faster_whisper import WhisperModel
+
 from config import Config, load, redacted_dict
 from recorder import start_recording, stop_recording
-from stt import FasterWhisperSTT, STTConfig
 from aiprocess import AIConfig, build_processor
 from inject import inject
 from x11_hotkey import listen
-from tray import run_tray
+
+gi.require_version("Gtk", "3.0")
+
+from gi.repository import GLib, Gtk  # type: ignore[import-not-found]
 
 
 class State:
@@ -26,6 +31,13 @@ class State:
     OUTPUTTING = "outputting"
 
 
+def _compute_type(device: str) -> str:
+    dev = (device or "cpu").lower()
+    if dev == "cuda":
+        return "float16"
+    return "int8"
+
+
 class Daemon:
     def __init__(self, cfg: Config):
         self.cfg = cfg
@@ -34,15 +46,12 @@ class Daemon:
         self.proc = None
         self.record = None
         self.timer = None
-        self.stt = FasterWhisperSTT(
-            STTConfig(
-                model=cfg.stt.get("model", "base"),
-                language=None,
-                device=cfg.stt.get("device", "cpu"),
-                vad_filter=True,
-            )
+        self.model = WhisperModel(
+            cfg.stt.get("model", "base"),
+            device=cfg.stt.get("device", "cpu"),
+            compute_type=_compute_type(cfg.stt.get("device", "cpu")),
         )
-        self.ai = None
+        self.tray = _Tray(self.get_state, self._quit)
 
     def set_state(self, state: str):
         with self.lock:
@@ -55,6 +64,9 @@ class Daemon:
         with self.lock:
             return self.state
 
+    def _quit(self):
+        os._exit(0)
+
     def toggle(self):
         with self.lock:
             if self.state == State.IDLE:
@@ -118,7 +130,7 @@ class Daemon:
         try:
             self.set_state(State.STT)
             logging.info("stt started")
-            text = self.stt.transcribe(record.wav_path, language="en")
+            text = self._transcribe(record.wav_path)
         except Exception as exc:
             logging.error("stt failed: %s", exc)
             self.set_state(State.IDLE)
@@ -131,7 +143,6 @@ class Daemon:
 
         logging.info("stt: %s", text)
 
-        ai_prompt_file = ""
         ai_model = (self.cfg.ai_cleanup.get("model") or "").strip()
         ai_base_url = (self.cfg.ai_cleanup.get("base_url") or "").strip()
         if ai_model and ai_base_url:
@@ -141,8 +152,6 @@ class Daemon:
                 processor = build_processor(
                     AIConfig(
                         model=ai_model,
-                        temperature=self.cfg.ai_cleanup.get("temperature", 0.0),
-                        system_prompt_file=ai_prompt_file,
                         base_url=ai_base_url,
                         api_key=self.cfg.ai_cleanup.get("api_key", ""),
                         timeout_sec=25,
@@ -174,6 +183,69 @@ class Daemon:
             self.state = State.STT
         threading.Thread(target=self._stop_and_process, daemon=True).start()
 
+    def _transcribe(self, wav_path: str) -> str:
+        segments, _info = self.model.transcribe(
+            wav_path,
+            language=None,
+            vad_filter=True,
+        )
+        parts = []
+        for seg in segments:
+            text = (seg.text or "").strip()
+            if text:
+                parts.append(text)
+        return " ".join(parts).strip()
+
+    def run_tray(self):
+        self.tray.run()
+
+
+class _Tray:
+    def __init__(self, state_getter, on_quit):
+        self.state_getter = state_getter
+        self.on_quit = on_quit
+        self.base = Path(__file__).parent / "assets"
+        self.icon = Gtk.StatusIcon()
+        self.icon.set_visible(True)
+        self.icon.connect("popup-menu", self._on_menu)
+        self.menu = Gtk.Menu()
+        quit_item = Gtk.MenuItem(label="Quit")
+        quit_item.connect("activate", lambda *_: self.on_quit())
+        self.menu.append(quit_item)
+        self.menu.show_all()
+
+    def _on_menu(self, _icon, _button, _time):
+        self.menu.popup(None, None, None, None, 0, _time)
+
+    def _icon_path(self, state: str) -> str:
+        if state == State.RECORDING:
+            return str(self.base / "recording.png")
+        if state == State.STT:
+            return str(self.base / "transcribing.png")
+        if state == State.PROCESSING:
+            return str(self.base / "processing.png")
+        return str(self.base / "idle.png")
+
+    def _title(self, state: str) -> str:
+        if state == State.RECORDING:
+            return "Recording"
+        if state == State.STT:
+            return "STT"
+        if state == State.PROCESSING:
+            return "AI Processing"
+        return "Idle"
+
+    def update(self):
+        state = self.state_getter()
+        self.icon.set_from_file(self._icon_path(state))
+        self.icon.set_tooltip_text(self._title(state))
+        return True
+
+    def run(self):
+        self.update()
+        GLib.timeout_add(250, self.update)
+        Gtk.main()
+
 
 def _lock_single_instance():
     runtime_dir = Path(os.getenv("XDG_RUNTIME_DIR", "/tmp")) / "lel"
@@ -199,8 +271,6 @@ def main():
 
     logging.basicConfig(stream=sys.stderr, level=logging.INFO, format="leld: %(asctime)s %(message)s")
     cfg = load(args.config)
-    config_path = Path(args.config) if args.config else Path.home() / ".config" / "lel" / "config.json"
-
     _lock_single_instance()
 
     logging.info("ready (hotkey: %s)", cfg.daemon.get("hotkey", ""))
@@ -208,9 +278,6 @@ def main():
 
     daemon = Daemon(cfg)
 
-    def on_quit():
-        os._exit(0)
-
     def handle_signal(_sig, _frame):
         logging.info("signal received, shutting down")
         daemon.stop_recording()
@@ -236,7 +303,7 @@ def main():
         ),
         daemon=True,
     ).start()
-    run_tray(daemon.get_state, on_quit, None)
+    daemon.run_tray()
 
 
 if __name__ == "__main__":