"""Strict LM Studio client for the pinned dense-retrieval embedding model.""" from __future__ import annotations from dataclasses import asdict, dataclass import json import math import os from typing import Any from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen from .specs import EmbeddingSpec class EmbeddingStudioError(RuntimeError): """Raised when embedding discovery, identity checks, or inference fails.""" def validate_embedding_record(spec: EmbeddingSpec, record: dict[str, Any]) -> None: quantization = record.get("quantization", {}) quantization_name = quantization.get("name") if isinstance(quantization, dict) else quantization expected = { "type": "embedding", "key": spec.model_key, "display_name": spec.expected_display_name, "format": spec.expected_format, "quantization": spec.expected_quantization, "size_bytes": spec.expected_size_bytes, "max_context_length": spec.max_context_length, } actual = { "type": record.get("type"), "key": record.get("key"), "display_name": record.get("display_name"), "format": record.get("format"), "quantization": quantization_name, "size_bytes": record.get("size_bytes"), "max_context_length": record.get("max_context_length"), } mismatches = [ f"{field}: expected {expected[field]!r}, observed {actual[field]!r}" for field in expected if expected[field] != actual[field] ] loaded_instances = record.get("loaded_instances", []) if isinstance(loaded_instances, list) and 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.loaded_context_length}: mismatches.append( f"loaded context length: expected only {spec.loaded_context_length}, " f"observed {sorted(loaded_contexts, key=lambda value: str(value))}" ) if mismatches: raise EmbeddingStudioError( "LM Studio embedding runtime does not match EMB001: " + "; ".join(mismatches) ) @dataclass(frozen=True, slots=True) class EmbeddingProbeResult: model_key: str vector_count: int vector_dimension: int l2_norms: tuple[float, ...] pairwise_cosine: float usage: dict[str, Any] def to_dict(self) -> dict[str, Any]: return asdict(self) class LMStudioEmbeddingClient: def __init__(self, spec: EmbeddingSpec, timeout_seconds: float = 30.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") raise EmbeddingStudioError( f"LM Studio returned HTTP {exc.code} for {endpoint}: {detail}" ) from exc except URLError as exc: raise EmbeddingStudioError( f"Cannot connect to LM Studio embeddings at {self.spec.base_url}: {exc.reason}" ) from exc try: decoded = json.loads(body) except json.JSONDecodeError as exc: raise EmbeddingStudioError(f"LM Studio returned non-JSON data for {endpoint}") from exc if not isinstance(decoded, dict): raise EmbeddingStudioError(f"LM Studio returned an unexpected response for {endpoint}") return decoded def resolve(self) -> dict[str, Any]: response = self._request("GET", self.spec.discovery_endpoint) models = response.get("models", []) if not isinstance(models, list): raise EmbeddingStudioError("LM Studio model discovery response has no models list") matches = [ item for item in models if isinstance(item, dict) and item.get("key") == self.spec.model_key ] if len(matches) != 1: visible = [ item.get("key") for item in models if isinstance(item, dict) and item.get("type") == "embedding" ] raise EmbeddingStudioError( f"Expected exactly one {self.spec.model_key!r} record; visible embeddings: {visible}" ) record = dict(matches[0]) validate_embedding_record(self.spec, record) return record def loaded_model_keys(self) -> tuple[str, ...]: response = self._request("GET", self.spec.discovery_endpoint) models = response.get("models", []) if not isinstance(models, list): raise EmbeddingStudioError("LM Studio model discovery response has no models list") return tuple( str(item.get("key")) for item in models if isinstance(item, dict) and isinstance(item.get("loaded_instances"), list) and item.get("loaded_instances") ) def embed(self, inputs: list[str]) -> tuple[tuple[float, ...], ...]: if not inputs or any(not isinstance(item, str) or not item for item in inputs): raise ValueError("embedding inputs must be non-empty strings") response = self._request( "POST", self.spec.inference_endpoint, {"model": self.spec.model_key, "input": inputs}, ) return self._vectors_from_response(response, len(inputs)) def _vectors_from_response( self, response: dict[str, Any], expected_count: int, ) -> tuple[tuple[float, ...], ...]: if response.get("model") != self.spec.model_key: raise EmbeddingStudioError( f"Embedding response model mismatch: {response.get('model')!r}" ) data = response.get("data", []) if not isinstance(data, list) or len(data) != expected_count: raise EmbeddingStudioError("Embedding response vector count does not match input count") ordered = sorted(data, key=lambda item: item.get("index", -1) if isinstance(item, dict) else -1) vectors: list[tuple[float, ...]] = [] for expected_index, item in enumerate(ordered): if not isinstance(item, dict) or item.get("index") != expected_index: raise EmbeddingStudioError("Embedding response indices are malformed") raw_vector = item.get("embedding") if not isinstance(raw_vector, list) or len(raw_vector) != self.spec.vector_dimension: raise EmbeddingStudioError( f"Expected {self.spec.vector_dimension}-dimensional embedding at index " f"{expected_index}" ) try: vector = tuple(float(value) for value in raw_vector) except (TypeError, ValueError) as exc: raise EmbeddingStudioError("Embedding vector contains non-numeric data") from exc if not all(math.isfinite(value) for value in vector): raise EmbeddingStudioError("Embedding vector contains non-finite data") norm = math.sqrt(sum(value * value for value in vector)) if self.spec.normalized and not math.isclose(norm, 1.0, abs_tol=1e-4): raise EmbeddingStudioError( f"Embedding at index {expected_index} is not L2-normalized: {norm}" ) vectors.append(vector) return tuple(vectors) def probe(self) -> EmbeddingProbeResult: inputs = [ "def binary_search(items, target): return -1", "class HttpClient: pass", ] response = self._request( "POST", self.spec.inference_endpoint, {"model": self.spec.model_key, "input": inputs}, ) vectors = self._vectors_from_response(response, len(inputs)) if vectors[0] == vectors[1]: raise EmbeddingStudioError("Distinct probe inputs produced identical vectors") norms = tuple(math.sqrt(sum(value * value for value in vector)) for vector in vectors) pairwise_cosine = sum(left * right for left, right in zip(vectors[0], vectors[1])) usage = response.get("usage", {}) return EmbeddingProbeResult( model_key=self.spec.model_key, vector_count=len(vectors), vector_dimension=self.spec.vector_dimension, l2_norms=norms, pairwise_cosine=pairwise_cosine, usage=dict(usage) if isinstance(usage, dict) else {}, )