feat(llm): oMLX-backed chat client + server control + model-id resolution
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS
This commit is contained in:
1 parent
52450c0835
commit
031a051687
3 files changed
+164
No files matched your search
Whitespace-only changes.
@@ -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
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user