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