| """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 |
|
|
