trading-agent/llm/server.py

"""Ensure the oMLX server is up for a run.

oMLX runs as a managed background service (auto-starts on login) and loads models
on-demand with LRU eviction, so there's usually nothing to do but health-check. For the
nightly launchd job we still `omlx start` defensively and warm the active model so the
first agent turn isn't paying cold-load latency.
"""

from __future__ import annotations

import os
import shutil
import subprocess
import time

import httpx
from loguru import logger

from config.settings import get_settings

_OMLX_CANDIDATES = ["omlx", os.path.expanduser("~/.omlx/bin/omlx")]


def _omlx_bin() -> str | None:
    for cand in _OMLX_CANDIDATES:
        path = shutil.which(cand) if os.path.sep not in cand else (cand if os.path.exists(cand) else None)
        if path:
            return path
    return None


def server_up() -> bool:
    s = get_settings()
    try:
        r = httpx.get(
            f"{s.llm_base_url}/models",
            headers={"Authorization": f"Bearer {s.llm_api_key}"},
            timeout=5,
        )
        return r.status_code == 200
    except Exception:
        return False


def start_server() -> None:
    if server_up():
        logger.info("oMLX already up at {}", get_settings().llm_base_url)
        return
    bin_ = _omlx_bin()
    if not bin_:
        logger.warning("omlx CLI not found; cannot auto-start. Start the oMLX app manually.")
        return
    proc = subprocess.run([bin_, "start"], capture_output=True, text=True, timeout=120)
    logger.info("omlx start -> {} {}", proc.returncode, (proc.stdout + proc.stderr).strip()[:300])


def wait_healthy(timeout_s: int = 120) -> bool:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        if server_up():
            return True
        time.sleep(3)
    return False


def warm_model() -> None:
    """Fire a tiny completion so the active model is resident before the first agent turn."""
    from llm.client import build_llm

    try:
        build_llm().invoke("ok")
        logger.info("warmed model")
    except Exception as e:
        logger.warning("warm_model failed (non-fatal): {}", e)


def ensure_ready(warm: bool = True) -> bool:
    start_server()
    ok = wait_healthy(120)
    if ok and warm:
        warm_model()
    logger.info("oMLX ready: {}", ok)
    return ok


if __name__ == "__main__":
    ensure_ready()