Aryan SinghandClaude Opus 4.8
031a051687
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS
78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
"""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
|