lemonade-sdk 9.1.1__py3-none-any.whl

lemonade/tools/server/tray.py

@@ -0,0 +1,664 @@
+import os
+import sys
+import threading
+from concurrent.futures import ThreadPoolExecutor
+import subprocess
+import webbrowser
+from pathlib import Path
+import logging
+import tempfile
+import platform
+
+import requests
+from packaging.version import parse as parse_version
+
+from lemonade_server.pydantic_models import DEFAULT_CTX_SIZE
+
+from lemonade.version import __version__
+
+# Import the appropriate tray implementation based on platform
+if platform.system() == "Darwin":  # macOS
+    from lemonade.tools.server.utils.macos_tray import (
+        MacOSSystemTray as SystemTray,
+        Menu,
+        MenuItem,
+    )
+else:  # Windows/Linux
+    from lemonade.tools.server.utils.windows_tray import (
+        SystemTray,
+        Menu,
+        MenuItem,
+    )
+
+
+class OutputDuplicator:
+    """
+    Output duplicator that writes to both a file and a stream.
+    """
+
+    def __init__(self, file_path, stream):
+        self.file = open(file_path, "a", encoding="utf-8")
+        self.stream = stream
+
+    def write(self, data):
+        self.file.write(data)
+        self.stream.write(data)
+        self.file.flush()
+        self.stream.flush()
+
+    def flush(self):
+        self.file.flush()
+        self.stream.flush()
+
+    def isatty(self):
+        # Delegate the isatty check to the original stream
+        return self.stream.isatty()
+
+
+class LemonadeTray(SystemTray):
+    """
+    Lemonade-specific system tray implementation.
+    """
+
+    def __init__(self, log_file, port, server_factory, log_level="info"):
+        # Find the icon path
+        icon_path = Path(__file__).resolve().parents[0] / "static" / "favicon.ico"
+
+        # Initialize the base class
+        super().__init__("Lemonade Server", str(icon_path))
+
+        # Lemonade-specific attributes
+        self.loaded_llm = None
+        self.server = None
+        self.server_thread = None
+        self.executor = ThreadPoolExecutor(max_workers=1)
+        self.log_file = log_file
+        self.port = port
+        self.ctx_size = DEFAULT_CTX_SIZE
+        self.server_factory = server_factory
+        self.debug_logs_enabled = log_level == "debug"
+
+        # Get current and latest version
+        self.current_version = __version__
+        self.latest_version = __version__
+        self.latest_version_url = None
+
+        # We will get the models list with update_downloaded_models_background()
+        # so that getting the list isn't on the critical path of startup
+        self.downloaded_models = {}
+
+        # Set up logger
+        self.logger = logging.getLogger(__name__)
+
+        # Initialize console handler (will be set up in run method)
+        self.console_handler = None
+
+        # Background thread for updating model mapping
+        self.model_update_thread = None
+        self.stop_model_update = threading.Event()
+
+        # Background thread for version checking
+        self.version_check_thread = None
+        self.stop_version_check = threading.Event()
+
+        # Hook function for platform-specific initialization callback
+        self.on_ready = None
+
+    def get_latest_version(self):
+        """
+        Update the latest version information.
+        """
+        try:
+            # Prepare headers for GitHub API request
+            headers = {}
+            github_token = os.environ.get("GITHUB_TOKEN")
+            if github_token:
+                headers["Authorization"] = f"token {github_token}"
+
+            response = requests.get(
+                "https://api.github.com/repos/lemonade-sdk/lemonade/releases/latest",
+                headers=headers,
+                timeout=10,  # Add timeout to prevent hanging
+            )
+            response.raise_for_status()
+            release_data = response.json()
+
+            # Get version from tag name (typically in format "vX.Y.Z")
+            self.latest_version = release_data.get("tag_name", "").lstrip("v")
+
+            # Find the installer asset
+            self.latest_version_url = None
+            for asset in release_data.get("assets", []):
+                if asset.get("name", "").endswith(
+                    ".exe"
+                ) and "Lemonade_Server_Installer" in asset.get("name", ""):
+                    self.latest_version_url = asset.get("browser_download_url")
+                    break
+
+            self.logger.debug(
+                f"Updated version info: latest version {self.latest_version}"
+            )
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            self.logger.error(f"Error fetching latest version: {str(e)}")
+        return self.latest_version, self.latest_version_url
+
+    def update_version_background(self):
+        """
+        Background thread function to update version information periodically.
+        """
+        self.get_latest_version()
+        while not self.stop_version_check.wait(900):  # 900 seconds = 15 minutes
+            try:
+                self.get_latest_version()
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                self.logger.error(f"Error updating version in background: {str(e)}")
+
+    def update_downloaded_models_background(self):
+        """
+        Background thread function to update model mapping every 1 second until a valid
+        response is received, then every 10 seconds after that.
+        This is used to avoid a ~0.5s delay when opening the tray menu.
+        """
+        poll_interval = 1
+        while not self.stop_model_update.wait(poll_interval):
+            try:
+                response = requests.get(
+                    f"http://localhost:{self.port}/api/v0/models",
+                    timeout=0.1,  # Add timeout
+                )
+                response.raise_for_status()
+
+                # Update the model mapping
+                models_endpoint_response = response.json().get("data")
+
+                # Convert to dict
+                self.downloaded_models = {
+                    model["id"]: {"checkpoint": model["checkpoint"]}
+                    for model in models_endpoint_response
+                }
+                self.logger.debug("Model mapping updated in background")
+
+                # Increase the poling interval after we get our first response
+                poll_interval = 10
+            except Exception:  # pylint: disable=broad-exception-caught
+                # Poll again later
+                pass
+
+    def unload_llms(self, _, __):
+        """
+        Unload the currently loaded LLM.
+        """
+        requests.post(f"http://localhost:{self.port}/api/v0/unload")
+
+    def load_llm(self, _, __, model_name):
+        """Load an LLM model."""
+
+        # Create config for loading
+        config = {"model_name": model_name}
+
+        # Use the executor to make the request asynchronously
+        self.executor.submit(
+            lambda: requests.post(
+                f"http://localhost:{self.port}/api/v0/load", json=config
+            )
+        )
+
+    def show_logs(self, _, __):
+        """
+        Show the log file in a new window.
+        """
+        try:
+            system = platform.system().lower()
+            if system == "darwin":
+                # Use Terminal.app to show live logs on macOS
+                try:
+                    subprocess.Popen(
+                        [
+                            "osascript",
+                            "-e",
+                            f'tell application "Terminal" to do script "tail -f {self.log_file}"',
+                        ]
+                    )
+                except (subprocess.CalledProcessError, FileNotFoundError) as e:
+                    self.logger.error(f"Failed to open Terminal for logs: {e}")
+                    self.show_balloon_notification(
+                        "Error",
+                        f"Failed to open logs in Terminal. Log file: {self.log_file}",
+                    )
+            elif system == "windows":
+                # Use PowerShell on Windows
+                subprocess.Popen(
+                    [
+                        "powershell",
+                        "Start-Process",
+                        "powershell",
+                        "-ArgumentList",
+                        f'"-NoExit", "Get-Content -Wait {self.log_file}"',
+                    ]
+                )
+            else:
+                # Unsupported platform
+                self.logger.error(f"Log viewing not supported on platform: {system}")
+
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            self.logger.error(f"Error opening logs: {str(e)}")
+
+    def open_documentation(self, _, __):
+        """
+        Open the documentation in the default web browser.
+        """
+        webbrowser.open("https://lemonade-server.ai/docs/")
+
+    def open_llm_chat(self, _, __):
+        """
+        Open the LLM chat in the default web browser.
+        """
+        webbrowser.open(f"http://localhost:{self.port}/#llm-chat")
+
+    def open_model_manager(self, _, __):
+        """
+        Open the model manager in the default web browser.
+        """
+        webbrowser.open(f"http://localhost:{self.port}/#model-management")
+
+    def check_server_state(self):
+        """
+        Check the server state using the health endpoint
+        """
+        try:
+            response = requests.get(
+                f"http://localhost:{self.port}/api/v0/health",
+                timeout=0.1,
+            )
+            response.raise_for_status()
+            response_data = response.json()
+            checkpoint = response_data.get("model_loaded")
+
+            # Convert checkpoint to model name if possible
+            if checkpoint:
+                # Create a mapping from checkpoint to model name
+                checkpoint_to_model = {
+                    model_info["checkpoint"]: model_name
+                    for model_name, model_info in self.downloaded_models.items()
+                }
+                # Use the model name if available, otherwise use the checkpoint
+                self.loaded_llm = checkpoint_to_model.get(checkpoint, checkpoint)
+            else:
+                self.loaded_llm = None
+
+            return True  # Successfully checked server state
+
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            # Server might not be ready yet, so just log at debug level
+            self.logger.debug(f"Error checking server state: {str(e)}")
+            return False  # Failed to check server state
+
+    def change_port(self, _, __, new_port):
+        """
+        Change the server port and restart the server.
+        """
+
+        try:
+
+            # Stop the current server
+            if self.server_thread and self.server_thread.is_alive():
+                # Set should_exit flag on the uvicorn server instance
+                if (
+                    hasattr(self.server, "uvicorn_server")
+                    and self.server.uvicorn_server
+                ):
+                    self.server.uvicorn_server.should_exit = True
+                self.server_thread.join(timeout=2)
+
+            # Update the port in both the tray and the server instance
+            self.port = new_port
+
+            # Clear the old server instance to ensure a fresh start
+            # This prevents middleware conflicts when restarting
+            self.server = None
+
+            self.server_thread = threading.Thread(target=self.start_server, daemon=True)
+            self.server_thread.start()
+
+            self.show_balloon_notification(
+                "Port Changed",
+                f"Lemonade Server is now running on port {self.port}",
+            )
+
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            self.logger.error(f"Error changing port: {str(e)}")
+            self.show_balloon_notification("Error", f"Failed to change port: {str(e)}")
+
+    def change_context_size(self, _, __, new_ctx_size):
+        """
+        Change the server context size and restart the server.
+        """
+        try:
+            # Stop the current server
+            if self.server_thread and self.server_thread.is_alive():
+                # Set should_exit flag on the uvicorn server instance
+                if (
+                    hasattr(self.server, "uvicorn_server")
+                    and self.server.uvicorn_server
+                ):
+                    self.server.uvicorn_server.should_exit = True
+                self.server_thread.join(timeout=2)
+            # Update the context size in both the tray and the server instance
+            self.ctx_size = new_ctx_size
+            if self.server:
+                self.server.ctx_size = new_ctx_size
+            # Restart the server
+            self.server_thread = threading.Thread(target=self.start_server, daemon=True)
+            self.server_thread.start()
+            # Show notification
+            ctx_size_label = (
+                f"{new_ctx_size//1024}K" if new_ctx_size >= 1024 else str(new_ctx_size)
+            )
+            self.show_balloon_notification(
+                "Context Size Changed",
+                f"Lemonade Server context size is now {ctx_size_label}",
+            )
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            self.logger.error(f"Error changing context size: {str(e)}")
+            self.show_balloon_notification(
+                "Error", f"Failed to change context size: {str(e)}"
+            )
+
+    def _using_installer(self):
+        """
+        Check if the user is using the NSIS installer by checking for embeddable python
+        """
+        py_home = Path(sys.executable).parent
+        pth_file = (
+            py_home / f"python{sys.version_info.major}{sys.version_info.minor}._pth"
+        )
+        return pth_file.exists()
+
+    def upgrade_to_latest(self, _, __):
+        """
+        Download and launch the Lemonade Server installer if the user is using the NSIS installer
+        Otherwise, simply open the browser to the release page
+        """
+
+        # If the user installed from source, simple open their browser to the release page
+        if not self._using_installer():
+            webbrowser.open("https://github.com/lemonade-sdk/lemonade/releases/latest")
+            return
+
+        # Show notification that download is starting
+        self.show_balloon_notification(
+            "Upgrading Lemonade",
+            "Downloading Lemonade Server Installer. Please wait...",
+        )
+
+        # Create temporary file for the installer
+        installer_path = os.path.join(
+            tempfile.gettempdir(), "Lemonade_Server_Installer.exe"
+        )
+        if os.path.exists(installer_path):
+            os.remove(installer_path)
+
+        # Download the installer
+        response = requests.get(self.latest_version_url, stream=True)
+        response.raise_for_status()
+
+        # Save the installer to disk and force write to disk
+        with open(installer_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+            f.flush()
+            os.fsync(f.fileno())
+
+        # Launch the installer as a completely detached process
+        # subprocess.DETACHED_PROCESS - Creates a process that's not attached to the console
+        # subprocess.CREATE_NEW_PROCESS_GROUP - Creates a new process group
+        # close_fds=True - Closes file descriptors to prevent inheritance
+        subprocess.Popen(
+            [installer_path],
+            creationflags=subprocess.DETACHED_PROCESS
+            | subprocess.CREATE_NEW_PROCESS_GROUP,
+            close_fds=True,
+            shell=True,
+            cwd=tempfile.gettempdir(),
+        )
+
+        # No need to quit the application, the installer will handle it
+
+    def toggle_debug_logs(self, _, __):
+        """
+        Toggle debug logs on and off.
+        """
+        try:
+            new_level = "debug" if not self.debug_logs_enabled else "info"
+            response = requests.post(
+                f"http://localhost:{self.port}/api/v1/log-level",
+                json={"level": new_level},
+            )
+            response.raise_for_status()
+            self.debug_logs_enabled = not self.debug_logs_enabled
+            self.show_balloon_notification(
+                "Debug Logs",
+                f"Debug logs {'enabled' if self.debug_logs_enabled else 'disabled'}",
+            )
+        except (FileNotFoundError, ValueError) as e:
+            self.logger.error(f"Error toggling debug logs: {str(e)}")
+            self.show_balloon_notification("Error", "Failed to toggle debug logs.")
+
+    def create_menu(self):
+        """
+        Create the Lemonade-specific context menu.
+        """
+        # Check server health when menu is opened
+        status_successfully_checked = self.check_server_state()
+
+        items = []
+
+        if not status_successfully_checked:
+            items.append(
+                MenuItem("Server Busy - See Logs for details", None, enabled=False)
+            )
+        elif self.loaded_llm:
+            items.extend(
+                [
+                    MenuItem(f"Loaded: {self.loaded_llm}", None, enabled=False),
+                    MenuItem("Unload LLM", self.unload_llms),
+                ]
+            )
+        else:
+            items.extend(
+                [
+                    MenuItem("No models loaded", None, enabled=False),
+                ]
+            )
+
+        # Create menu items for all downloaded models
+        model_menu_items = []
+        if not self.downloaded_models:
+            model_menu_items.append(
+                MenuItem(
+                    "No models available: Use the Model Manager to pull models",
+                    None,
+                    enabled=False,
+                )
+            )
+        else:
+            for model_name, _ in self.downloaded_models.items():
+                # Create a function that returns the lambda to properly capture the variables
+                def create_handler(mod):
+                    return lambda icon, item: self.load_llm(icon, item, mod)
+
+                model_item = MenuItem(model_name, create_handler(model_name))
+
+                # Set checked property instead of modifying the text
+                model_item.checked = model_name == self.loaded_llm
+                model_menu_items.append(model_item)
+
+        load_submenu = Menu(*model_menu_items)
+
+        # Create port selection submenu with 3 options
+        port_menu_items = []
+        port_options = [
+            8000,
+            8020,
+            8040,
+            8060,
+            8080,
+            9000,
+        ]
+
+        for port_option in port_options:
+            # Create a function that returns the lambda to properly capture the port variable
+            def create_port_handler(port):
+                return lambda icon, item: self.change_port(icon, item, port)
+
+            # Set checked property instead of modifying the text
+            port_item = MenuItem(
+                f"Port {port_option}", create_port_handler(port_option)
+            )
+            port_item.checked = port_option == self.port
+            port_menu_items.append(port_item)
+
+        port_submenu = Menu(*port_menu_items)
+
+        # Create context size selection submenu with 6 options
+        ctx_size_menu_items = []
+        ctx_size_options = [
+            ("4K", 4096),
+            ("8K", 8192),
+            ("16K", 16384),
+            ("32K", 32768),
+            ("64K", 65536),
+            ("128K", 131072),
+        ]
+
+        for ctx_label, ctx_value in ctx_size_options:
+            # Create a function that returns the lambda to properly capture the ctx_size variable
+            def create_ctx_handler(ctx_size):
+                return lambda icon, item: self.change_context_size(icon, item, ctx_size)
+
+            ctx_item = MenuItem(
+                f"Context size {ctx_label}", create_ctx_handler(ctx_value)
+            )
+            ctx_item.checked = ctx_value == self.ctx_size
+            ctx_size_menu_items.append(ctx_item)
+
+        ctx_size_submenu = Menu(*ctx_size_menu_items)
+
+        # Create the Logs submenu
+        debug_log_text = "Enable Debug Logs"
+        debug_log_item = MenuItem(debug_log_text, self.toggle_debug_logs)
+        debug_log_item.checked = self.debug_logs_enabled
+
+        logs_submenu = Menu(
+            MenuItem("Show Logs", self.show_logs),
+            Menu.SEPARATOR,
+            debug_log_item,
+        )
+
+        if status_successfully_checked:
+            items.append(MenuItem("Load Model", None, submenu=load_submenu))
+        items.append(MenuItem("Port", None, submenu=port_submenu))
+        items.append(MenuItem("Context Size", None, submenu=ctx_size_submenu))
+        items.append(Menu.SEPARATOR)
+
+        # Only show upgrade option if newer version is available
+        if parse_version(self.latest_version) > parse_version(self.current_version):
+            items.append(
+                MenuItem(
+                    f"Upgrade to version {self.latest_version}", self.upgrade_to_latest
+                )
+            )
+
+        items.append(MenuItem("Documentation", self.open_documentation))
+        items.append(MenuItem("LLM Chat", self.open_llm_chat))
+        items.append(MenuItem("Model Manager", self.open_model_manager))
+        items.append(MenuItem("Logs", None, submenu=logs_submenu))
+        items.append(Menu.SEPARATOR)
+        items.append(MenuItem("Quit Lemonade", self.exit_app))
+        return Menu(*items)
+
+    def start_server(self):
+        """
+        Start the uvicorn server.
+        """
+        self.server = self.server_factory()
+
+        # Ensure the server uses the current port from the tray
+        # This is important when changing ports
+        self.server.port = self.port
+
+        self.server.uvicorn_server = self.server.run_in_thread(self.server.host)
+        self.server.uvicorn_server.run()
+
+    def run(self):
+        """
+        Run the Lemonade tray application.
+        """
+
+        # Start the background model mapping update thread
+        self.model_update_thread = threading.Thread(
+            target=self.update_downloaded_models_background, daemon=True
+        )
+        self.model_update_thread.start()
+
+        # Start the version check thread
+        self.version_check_thread = threading.Thread(
+            target=self.update_version_background, daemon=True
+        )
+        self.version_check_thread.start()
+
+        # Start the server in a separate thread
+        self.server_thread = threading.Thread(target=self.start_server, daemon=True)
+        self.server_thread.start()
+
+        # Provide an on_ready hook that Windows base tray will call after
+        # the HWND/icon are created. macOS will call it immediately after run.
+        def _on_ready():
+            system = platform.system().lower()
+            if system == "darwin":
+                message = (
+                    "Lemonade Server is running! "
+                    "Click the tray icon above to access options."
+                )
+            else:  # Windows/Linux
+                message = (
+                    "Lemonade Server is running! "
+                    "Right-click the tray icon below to access options."
+                )
+            self.show_balloon_notification("Woohoo!", message)
+
+        # Attach hook for both implementations to invoke after init
+        self.on_ready = _on_ready
+
+        # Call the parent run method which handles platform-specific initialization
+        super().run()
+
+    def exit_app(self, icon, item):
+        """
+        Exit the application.
+        """
+        # Stop the background threads
+        self.stop_model_update.set()
+        self.stop_version_check.set()
+
+        if self.model_update_thread and self.model_update_thread.is_alive():
+            self.model_update_thread.join(timeout=1)
+        if self.version_check_thread and self.version_check_thread.is_alive():
+            self.version_check_thread.join(timeout=1)
+
+        # Platform-specific exit handling
+        system = platform.system().lower()
+        if system == "darwin":  # macOS
+            # For macOS, quit the rumps application
+            import rumps
+
+            rumps.quit_application()
+        else:
+            # Call parent exit method for Windows
+            super().exit_app(icon, item)
+
+        # Stop the server using the CLI stop command to ensure a rigorous cleanup
+        # This must be a subprocess to ensure the cleanup doesnt kill itself
+        subprocess.Popen(
+            [sys.executable, "-m", "lemonade_server.cli", "stop"], shell=True
+        )