File size: 11,378 Bytes
d61821a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 | """Strict client for the local LM Studio server.
The client uses LM Studio's OpenAI-compatible chat-completions endpoint because
that endpoint supports custom tools. Model discovery is performed before
inference, and model identity mismatches are fatal.
"""
from __future__ import annotations
from dataclasses import asdict, dataclass
import json
import os
import re
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
from .specs import ModelSpec
class LMStudioError(RuntimeError):
"""Raised when discovery, identity validation, or inference fails."""
class LMStudioTransportError(LMStudioError):
"""Raised when no valid server response was observed and one retry is safe."""
def normalize_identity(value: str) -> str:
return re.sub(r"[^a-z0-9]+", "", value.lower())
def _record_identity(record: dict[str, Any]) -> str:
values = [
str(record.get("id", "")),
str(record.get("key", "")),
str(record.get("display_name", "")),
str(record.get("name", "")),
]
return " ".join(item for item in values if item)
@dataclass(frozen=True, slots=True)
class DiscoveryResult:
openai_models: tuple[dict[str, Any], ...]
native_models: tuple[dict[str, Any], ...]
endpoint_errors: tuple[str, ...]
def to_dict(self) -> dict[str, Any]:
return asdict(self)
@dataclass(frozen=True, slots=True)
class ResolvedModel:
inference_key: str
openai_record: dict[str, Any]
native_record: dict[str, Any] | None
def to_dict(self) -> dict[str, Any]:
return asdict(self)
def select_expected_model(
spec: ModelSpec,
openai_models: tuple[dict[str, Any], ...] | list[dict[str, Any]],
native_models: tuple[dict[str, Any], ...] | list[dict[str, Any]] = (),
) -> ResolvedModel:
"""Resolve one and only one inference-visible model matching the fixed identity."""
expected = normalize_identity(spec.expected_identity)
matches = [
record
for record in openai_models
if expected in normalize_identity(_record_identity(record))
]
if not matches:
visible = sorted(filter(None, (_record_identity(item) for item in openai_models)))
raise LMStudioError(
f"Expected {spec.canonical_name}, but no matching model is visible through "
f"{spec.discovery_endpoint}. Visible models: {visible or ['<none>']}"
)
if len(matches) > 1:
exact = [
record
for record in matches
if normalize_identity(str(record.get("id", ""))) == expected
or normalize_identity(str(record.get("key", ""))) == expected
]
if len(exact) == 1:
matches = exact
else:
raise LMStudioError(
"Model identity is ambiguous; refusing to select a quantization or variant silently: "
+ ", ".join(_record_identity(item) for item in matches)
)
record = matches[0]
inference_key = str(record.get("id") or record.get("key") or "")
if not inference_key:
raise LMStudioError("Matching LM Studio model record has no inference identifier")
if inference_key != spec.expected_inference_key:
raise LMStudioError(
f"Expected inference key {spec.expected_inference_key}, but LM Studio exposed {inference_key}"
)
native_match: dict[str, Any] | None = None
native_candidates = [
item for item in native_models if expected in normalize_identity(_record_identity(item))
]
if len(native_candidates) == 1:
native_match = native_candidates[0]
elif native_candidates:
selected = [item for item in native_candidates if item.get("selected_variant")]
if len(selected) == 1:
native_match = selected[0]
if native_match is None:
raise LMStudioError(
"The matching model has no unique native /api/v1/models record; "
"variant and runtime metadata cannot be verified"
)
quantization = native_match.get("quantization", {})
quantization_name = quantization.get("name") if isinstance(quantization, dict) else quantization
expected_runtime = {
"selected_variant": spec.expected_variant,
"format": spec.expected_format,
"quantization": spec.expected_quantization,
}
actual_runtime = {
"selected_variant": native_match.get("selected_variant"),
"format": native_match.get("format"),
"quantization": quantization_name,
}
mismatches = [
f"{field}: expected {expected_runtime[field]!r}, observed {actual_runtime[field]!r}"
for field in expected_runtime
if expected_runtime[field] != actual_runtime[field]
]
loaded_instances = native_match.get("loaded_instances", [])
loaded_contexts = {
item.get("config", {}).get("context_length")
for item in loaded_instances
if isinstance(item, dict) and isinstance(item.get("config"), dict)
}
if loaded_contexts != {spec.context_length}:
observed_contexts = sorted(loaded_contexts, key=lambda value: str(value))
mismatches.append(
f"loaded context length: expected only {spec.context_length}, observed {observed_contexts}"
)
capabilities = native_match.get("capabilities", {})
reasoning = capabilities.get("reasoning", {}) if isinstance(capabilities, dict) else {}
observed_reasoning = reasoning.get("default") if isinstance(reasoning, dict) else None
if observed_reasoning is None:
observed_reasoning = "none"
if observed_reasoning != spec.reasoning_mode:
mismatches.append(
f"reasoning mode: expected {spec.reasoning_mode!r}, observed {observed_reasoning!r}"
)
if mismatches:
raise LMStudioError(
f"LM Studio runtime does not match {spec.model_id}: " + "; ".join(mismatches)
)
return ResolvedModel(
inference_key=inference_key,
openai_record=dict(record),
native_record=dict(native_match),
)
class LMStudioClient:
def __init__(self, spec: ModelSpec, timeout_seconds: float = 10.0):
self.spec = spec
self.timeout_seconds = timeout_seconds
def _headers(self) -> dict[str, str]:
headers = {"Content-Type": "application/json"}
token = os.environ.get(self.spec.api_token_env, "").strip()
if token:
headers["Authorization"] = f"Bearer {token}"
return headers
def _request(
self,
method: str,
endpoint: str,
payload: dict[str, Any] | None = None,
) -> dict[str, Any]:
data = None if payload is None else json.dumps(payload).encode("utf-8")
request = Request(
self.spec.base_url + endpoint,
data=data,
method=method,
headers=self._headers(),
)
try:
with urlopen(request, timeout=self.timeout_seconds) as response:
body = response.read().decode("utf-8")
except HTTPError as exc:
detail = exc.read().decode("utf-8", errors="replace")
error_type = (
LMStudioTransportError
if exc.code in {408, 429, 500, 502, 503, 504}
else LMStudioError
)
raise error_type(
f"LM Studio returned HTTP {exc.code} for {endpoint}: {detail}"
) from exc
except URLError as exc:
raise LMStudioTransportError(
f"Cannot connect to LM Studio at {self.spec.base_url}. "
"Start the server on port 1234 and load the frozen model. "
f"Underlying error: {exc.reason}"
) from exc
try:
decoded = json.loads(body)
except json.JSONDecodeError as exc:
raise LMStudioError(f"LM Studio returned non-JSON data for {endpoint}") from exc
if not isinstance(decoded, dict):
raise LMStudioError(f"LM Studio returned an unexpected response for {endpoint}")
return decoded
def discover(self) -> DiscoveryResult:
errors: list[str] = []
openai_models: tuple[dict[str, Any], ...] = ()
native_models: tuple[dict[str, Any], ...] = ()
try:
response = self._request("GET", self.spec.discovery_endpoint)
data = response.get("data", [])
if isinstance(data, list):
openai_models = tuple(item for item in data if isinstance(item, dict))
except LMStudioError as exc:
errors.append(str(exc))
try:
response = self._request("GET", self.spec.native_discovery_endpoint)
data = response.get("models", [])
if isinstance(data, list):
native_models = tuple(item for item in data if isinstance(item, dict))
except LMStudioError as exc:
errors.append(str(exc))
if not openai_models and not native_models:
raise LMStudioError("; ".join(errors) or "LM Studio returned no model records")
return DiscoveryResult(openai_models, native_models, tuple(errors))
def resolve(self, discovery: DiscoveryResult | None = None) -> tuple[DiscoveryResult, ResolvedModel]:
result = discovery or self.discover()
resolved = select_expected_model(self.spec, result.openai_models, result.native_models)
return result, resolved
def chat_completions(
self,
model_key: str,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
max_tokens: int | None = None,
seed: int | None = None,
) -> dict[str, Any]:
payload: dict[str, Any] = {
"model": model_key,
"messages": messages,
"temperature": self.spec.temperature,
"top_p": self.spec.top_p,
"max_tokens": max_tokens or self.spec.max_tokens,
"seed": self.spec.seed if seed is None else seed,
"stream": False,
}
if tools is not None:
payload["tools"] = tools
return self._request("POST", self.spec.inference_endpoint, payload)
def inference_probe(self, model_key: str) -> dict[str, Any]:
response = self.chat_completions(
model_key,
messages=[
{
"role": "user",
"content": "Reply with exactly MODEL_OK and no other text.",
}
],
max_tokens=256,
)
try:
message = response["choices"][0]["message"]
content = message["content"]
except (KeyError, IndexError, TypeError) as exc:
raise LMStudioError("Inference probe returned an invalid chat-completions response") from exc
if not isinstance(content, str) or content.strip() != "MODEL_OK":
finish_reason = response.get("choices", [{}])[0].get("finish_reason")
reasoning_content = message.get("reasoning_content", "")
raise LMStudioError(
"Inference probe did not produce the required MODEL_OK marker "
f"(finish_reason={finish_reason!r}, visible={content!r}, "
f"reasoning_chars={len(reasoning_content)})"
)
return response
|