From 031a05168735b10e94affffe224ddc387d444deb Mon Sep 17 00:00:00 2001 From: Aryan Singh Date: Wed, 1 Jul 2026 16:29:29 -0500 Subject: [PATCH] feat(llm): oMLX-backed chat client + server control + model-id resolution Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS --- llm/__init__.py | 0 llm/client.py | 77 +++++++++++++++++++++++++++++++++++++++++++ llm/server.py | 87 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 164 insertions(+) create mode 100644 llm/__init__.py create mode 100644 llm/client.py create mode 100644 llm/server.py diff --git a/llm/__init__.py b/llm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/llm/client.py b/llm/client.py new file mode 100644 index 0000000..c163fc8 --- /dev/null +++ b/llm/client.py @@ -0,0 +1,77 @@ +"""LangChain chat model pointed at the local oMLX server (Apple MLX, OpenAI-compatible). + +oMLX parses the model's native tool-call format into standard OpenAI `tool_calls`, so we +rely on that rather than doing our own parsing. `resolve_model_id` matches the active +profile against whatever id oMLX serves, so the MODEL_PROFILE toggle survives id changes. +oMLX enforces an API key (see settings.llm_api_key, auto-discovered from ~/.omlx). +""" + +from __future__ import annotations + +import httpx +from langchain_openai import ChatOpenAI + +from config.settings import get_settings + + +def _auth_header() -> dict[str, str]: + return {"Authorization": f"Bearer {get_settings().llm_api_key}"} + + +def list_loaded_models() -> list[str]: + """Model ids currently served by oMLX (GET /v1/models). Empty if unreachable.""" + s = get_settings() + try: + r = httpx.get(f"{s.llm_base_url}/models", headers=_auth_header(), timeout=10) + r.raise_for_status() + return [m["id"] for m in r.json().get("data", [])] + except Exception: + return [] + + +def resolve_model_id() -> str: + """Pick the served model id for the active profile. + + Priority: exact match of the configured id → keyword match from the profile → + first served model → the configured id as-is (lets errors surface at call time). + """ + s = get_settings() + served = list_loaded_models() + target = s.effective_model_id + if target in served: + return target + for mid in served: + low = mid.lower() + if any(kw in low for kw in s.active_profile.match_keywords): + return mid + return served[0] if served else target + + +def build_llm( + tools: list | None = None, + temperature: float | None = None, + json_mode: bool = False, + max_tokens: int | None = None, +) -> ChatOpenAI: + """Build a ChatOpenAI bound to oMLX. + + json_mode forces a JSON object response (used for the finalize step so structured picks + don't drift into prose). max_tokens overrides the configured default when a step needs + more room. Don't combine json_mode with tools. + """ + s = get_settings() + model_kwargs: dict = {} + if json_mode: + model_kwargs["response_format"] = {"type": "json_object"} + llm = ChatOpenAI( + base_url=s.llm_base_url, + api_key=s.llm_api_key, + model=resolve_model_id(), + temperature=s.llm_temperature if temperature is None else temperature, + max_tokens=max_tokens or s.llm_max_tokens, + timeout=180, + model_kwargs=model_kwargs, + ) + if tools: + llm = llm.bind_tools(tools) + return llm diff --git a/llm/server.py b/llm/server.py new file mode 100644 index 0000000..1fe6a34 --- /dev/null +++ b/llm/server.py @@ -0,0 +1,87 @@ +"""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()