diff --git a/CHANGELOG.md b/CHANGELOG.md index 637bd80..814dc3a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ## [Unreleased] +### Features + +- **Embeddings via any OpenAI-compatible `/v1/embeddings` endpoint.** New `embedding_model: "openai-compat"` option computes embeddings on a server (LM Studio, llama.cpp, vLLM, Ollama's OpenAI shim, or a self-hosted endpoint) instead of a local ONNX model — useful for larger/multilingual embedders such as Qwen3-Embedding, or GPU offload. New `OpenAICompatEmbeddingFunction` in [`mempalace/embedding.py`](mempalace/embedding.py) speaks the standard `/v1/embeddings` protocol over stdlib `urllib` (no new dependency), batches requests, re-sorts the response by `index`, and L2-normalizes for the cosine collection. Endpoint settings are resolved by `MempalaceConfig` as a single source of truth — `embedding_api_url` / `embedding_api_model` / `embedding_api_key` in `config.json`, each overridable via the matching `MEMPALACE_EMBEDDING_API_*` env var. The embedding function's `name()` encodes the model id so changing it forces `mempalace repair rebuild-index` (different vector space). Mirrors the existing `openai-compat` LLM provider naming; stays local when the endpoint is on your machine/LAN. (#1559) + --- ## [3.7.0] — 2026-08-02 diff --git a/README.md b/README.md index 4253acc..f2b231b 100644 --- a/README.md +++ b/README.md @@ -293,6 +293,7 @@ verbatim drawer per user/assistant message, idempotent and resume-safe. - Python 3.9+ - A vector-store backend (ChromaDB by default) - ~300 MB disk for the embedding model. Onboarding (`python -m mempalace.onboarding`) offers `embeddinggemma-300m` (multilingual, 100+ languages, recommended) or `all-MiniLM-L6-v2` (English-only, ~30 MB). See the docstring at [`mempalace/embedding.py`](mempalace/embedding.py) for details and migration notes. +- Optional — compute embeddings on a server instead of locally. Set `embedding_model: "openai-compat"` in `~/.mempalace/config.json` together with `embedding_api_url` / `embedding_api_model` (and `embedding_api_key` if the server needs auth) to use any OpenAI-compatible `/v1/embeddings` endpoint — LM Studio, llama.cpp, vLLM, Ollama's OpenAI shim, or a self-hosted server (e.g. a larger multilingual or GPU-served embedder). Each key is overridable via the matching `MEMPALACE_EMBEDDING_API_*` env var. When the endpoint is on your machine or LAN, no content leaves your network. Switching to it requires `mempalace repair rebuild-index` (different vector space). No API key is required for the core benchmark path. diff --git a/mempalace/config.py b/mempalace/config.py index 21f779f..25d0af3 100644 --- a/mempalace/config.py +++ b/mempalace/config.py @@ -748,7 +748,10 @@ class MempalaceConfig: Values: ``"minilm"`` (ChromaDB's all-MiniLM-L6-v2 — English-only), ``"embeddinggemma"`` (multilingual, 100+ languages, default for - new installs since onboarding writes the choice). Read from env + new installs since onboarding writes the choice), or + ``"openai-compat"`` (embeddings served by an OpenAI-compatible + ``/v1/embeddings`` endpoint — see ``embedding_api_url`` / + ``embedding_api_model`` / ``embedding_api_key``). Read from env ``MEMPALACE_EMBEDDING_MODEL`` first, then ``embedding_model`` in ``config.json``, then ``"minilm"`` as a back-compat fallback for palaces created before onboarding asked the question. @@ -832,6 +835,52 @@ class MempalaceConfig: except (OSError, NotImplementedError): pass + def _resolve_str_setting(self, env_var: str, config_key: str): + """Resolve a string setting: env var > ``config.json`` > ``None``. + + Whitespace-only values are treated as unset, so a blank env var or a + hand-edited empty config key doesn't mask the value below it. Unlike + ``embedding_model`` the result is not lower-cased — URLs, model ids, + and API keys are case-sensitive. + """ + env_val = os.environ.get(env_var) + if env_val and env_val.strip(): + return env_val.strip() + cfg_val = self._file_config.get(config_key) + if isinstance(cfg_val, str) and cfg_val.strip(): + return cfg_val.strip() + return None + + @property + def embedding_api_url(self): + """Base URL of the OpenAI-compatible ``/v1/embeddings`` endpoint. + + Used only when ``embedding_model == "openai-compat"``. Resolved from + env ``MEMPALACE_EMBEDDING_API_URL`` first, then ``embedding_api_url`` + in ``config.json``; ``None`` when unset. Accepts a bare host, a + ``…/v1`` base, or a full endpoint URL. + """ + return self._resolve_str_setting("MEMPALACE_EMBEDDING_API_URL", "embedding_api_url") + + @property + def embedding_api_model(self): + """Server-side model id for the ``openai-compat`` embeddings endpoint. + + Resolved from env ``MEMPALACE_EMBEDDING_API_MODEL`` first, then + ``embedding_api_model`` in ``config.json``; ``None`` when unset. + """ + return self._resolve_str_setting("MEMPALACE_EMBEDDING_API_MODEL", "embedding_api_model") + + @property + def embedding_api_key(self): + """Optional bearer token / API key for the embeddings endpoint. + + Resolved from env ``MEMPALACE_EMBEDDING_API_KEY`` first, then + ``embedding_api_key`` in ``config.json``; ``None`` when unset (for + local endpoints that need no auth). + """ + return self._resolve_str_setting("MEMPALACE_EMBEDDING_API_KEY", "embedding_api_key") + @property def topic_tunnel_min_count(self): """Minimum number of overlapping confirmed topics required to create diff --git a/mempalace/embedding.py b/mempalace/embedding.py index 2ad9c44..7d4eb5c 100644 --- a/mempalace/embedding.py +++ b/mempalace/embedding.py @@ -1,10 +1,12 @@ """Embedding function factory with hardware acceleration. -Returns a ChromaDB-compatible embedding function bound to a user-selected -ONNX Runtime execution provider. +Returns a ChromaDB-compatible embedding function — either a local ONNX model +bound to a user-selected ONNX Runtime execution provider, or an +OpenAI-compatible HTTP ``/v1/embeddings`` endpoint. -Two embedding models are available, selected via ``MEMPALACE_EMBEDDING_MODEL`` -or ``embedding_model`` in ``~/.mempalace/config.json``: +Three embedding-model options are available, selected via +``MEMPALACE_EMBEDDING_MODEL`` or ``embedding_model`` in +``~/.mempalace/config.json``: * ``minilm`` (default) — ``all-MiniLM-L6-v2``, 384-dim, English-only training. ChromaDB's default; what every existing palace was built with. @@ -15,6 +17,16 @@ or ``embedding_model`` in ``~/.mempalace/config.json``: model is lazy-downloaded from HuggingFace on first use. Switching models on an existing palace requires ``mempalace repair rebuild-index`` (different vector space). +* ``openai-compat`` — embeddings served by any OpenAI-compatible + ``/v1/embeddings`` endpoint (LM Studio, llama.cpp, vLLM, Ollama's OpenAI + shim, or a self-hosted server) instead of a local ONNX model. Useful for + larger / multilingual embedders (e.g. Qwen3-Embedding) or GPU offload. + Endpoint settings are read from ``config.json`` as ``embedding_api_url`` / + ``embedding_api_model`` / ``embedding_api_key`` (each overridable via the + matching ``MEMPALACE_EMBEDDING_API_*`` env var). Vectors are L2-normalized + for the cosine collection; the dimension is whatever the server returns, so + switching to/from this backend also requires ``mempalace repair + rebuild-index``. Stays local when the endpoint is on your machine/LAN. Supported devices (env ``MEMPALACE_EMBEDDING_DEVICE`` or ``embedding_device`` in ``~/.mempalace/config.json``): @@ -31,11 +43,14 @@ rather than hard-failing — mining must still work on a laptop without CUDA. from __future__ import annotations +import hashlib import logging import os import threading from typing import Optional +from .version import __version__ + logger = logging.getLogger(__name__) _PROVIDER_MAP = { @@ -447,6 +462,165 @@ class EmbeddinggemmaONNX: return self(input) +# ── OpenAI-compatible embedding API ────────────────────────────────────── +# Fetch embeddings from an OpenAI-compatible ``/v1/embeddings`` server +# (LM Studio, llama.cpp, vLLM, Ollama's OpenAI shim, or any compatible +# endpoint) instead of running a model locally. Selected by +# ``embedding_model == "openai-compat"``. Connection settings (URL, model, +# optional key) are resolved by :class:`~mempalace.config.MempalaceConfig` +# as the single source of truth — see ``embedding_api_url`` / +# ``embedding_api_model`` / ``embedding_api_key`` (each env-overridable). +_EF_API_BATCH = 64 +_EF_API_TIMEOUT = 120 + + +class EmbeddingAPIError(RuntimeError): + """Raised when the embedding API is unreachable or returns an invalid body. + + Module-specific subclass mirroring ``llm_client.LLMError`` so callers can + distinguish embedding-endpoint failures; subclasses ``RuntimeError`` so + existing ``except RuntimeError`` paths still catch it. + """ + + +class OpenAICompatEmbeddingFunction: + """ChromaDB-compatible EF backed by an OpenAI-compatible ``/v1/embeddings`` + endpoint (LM Studio, llama.cpp, vLLM, Ollama's OpenAI shim, etc.). + + Selected via ``embedding_model == "openai-compat"``. Vectors are produced + server-side and fetched over HTTP, which changes the vector space — so + ``name()`` encodes the model id: ChromaDB persists the EF name on the + collection and rejects mismatched reads, the signal to run ``mempalace + repair rebuild-index`` after changing model/endpoint. stdlib ``urllib`` + only, no new dependency. + """ + + def __init__(self, base_url: str, model: str, api_key: Optional[str] = None): + self._url = self._resolve_url(base_url) + self._model = model + self._api_key = api_key + + @staticmethod + def _resolve_url(base_url: str) -> str: + """Accept a base host, a ``/v1`` base, or a full endpoint URL. + + Mirrors ``llm_client.OpenAICompatProvider._resolve_url`` so both sides + treat an ``http://host:port`` endpoint the same way. + """ + url = base_url.rstrip("/") + if url.endswith("/embeddings"): + return url + if url.endswith("/v1"): + return f"{url}/embeddings" + return f"{url}/v1/embeddings" + + def name(self) -> str: + # Encode the model so switching it changes the persisted EF identity + # and forces a rebuild_index (vectors from a different model/space are + # not interchangeable). ChromaDB compares this on every read. + return f"openai_compat_emb_{self._model}".replace("/", "_") + + def embed_query(self, input): # noqa: A002 — ChromaDB EF protocol uses `input` + # ChromaDB 1.5 dispatches query embedding through embed_query (add uses + # __call__). Mirror the EmbeddingFunction protocol default: same path. + return self(input) + + def __call__(self, input): # noqa: A002 — ChromaDB EF protocol uses `input` + import http.client + import json + from urllib.error import HTTPError, URLError + from urllib.request import Request, urlopen + + headers = { + "Content-Type": "application/json", + # Some hosted (Cloudflare-fronted) endpoints 403 the default + # ``Python-urllib`` User-Agent — send our own (see issue #1570). + "User-Agent": f"mempalace/{__version__}", + } + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + + out: list = [] + texts = list(input) + for start in range(0, len(texts), _EF_API_BATCH): + batch = texts[start : start + _EF_API_BATCH] + # encoding_format=float is explicit so a server that defaults to + # base64 doesn't hand back strings we'd mis-parse as vectors. + payload = {"model": self._model, "input": batch, "encoding_format": "float"} + req = Request(self._url, data=json.dumps(payload).encode("utf-8"), headers=headers) + try: + with urlopen(req, timeout=_EF_API_TIMEOUT) as resp: + data = json.loads(resp.read()) + # ValueError covers an invalid/missing URL scheme and json.JSONDecodeError; + # http.client.HTTPException covers low-level protocol faults (BadStatusLine, + # IncompleteRead) common with local/overloaded servers. + except (HTTPError, URLError, OSError, http.client.HTTPException, ValueError) as e: + raise EmbeddingAPIError( + f"Embedding API request to {self._url} failed: {e}. Check that the " + f"server is reachable and MEMPALACE_EMBEDDING_API_URL / embedding_api_url " + f"is correct." + ) from e + out.extend(self._vectors_from_response(data, len(batch))) + return out + + def _vectors_from_response(self, data, n: int) -> list: + """Validate one ``/v1/embeddings`` response and return L2-normed vectors. + + Guards every way a non-conformant server could corrupt the store + silently: a missing/short ``data`` array, response ``index`` values + that aren't the contiguous ``0..n-1`` batch positions (sorting then + zipping positionally would otherwise misalign vectors with texts), and + malformed / ragged / base64 embedding payloads. All failures raise + :class:`EmbeddingAPIError` naming the endpoint rather than a cryptic + numpy error — a silent wrong result would break the 100%-recall promise. + """ + import numpy as np + + if not isinstance(data, dict): + raise EmbeddingAPIError( + f"Embedding API at {self._url} returned a non-object response: {data}" + ) + rows = data.get("data") + if not isinstance(rows, list): + raise EmbeddingAPIError( + f"Embedding API at {self._url} returned no 'data' array: {data.get('error', data)}" + ) + if len(rows) != n: + raise EmbeddingAPIError( + f"Embedding API at {self._url} returned {len(rows)} embeddings for {n} inputs" + ) + # The endpoint may return rows out of order — sort by index, then + # require the indices to be exactly 0..n-1 so positional alignment is + # provably correct (a server using absolute or duplicate indices would + # otherwise pass the count check yet map vectors to the wrong texts). + try: + rows = sorted(rows, key=lambda d: d.get("index", -1)) + indices = [r.get("index") for r in rows] + except AttributeError as e: + raise EmbeddingAPIError( + f"Embedding API at {self._url} returned non-object rows: {e}" + ) from e + if indices != list(range(n)): + raise EmbeddingAPIError( + f"Embedding API at {self._url} returned non-contiguous or duplicate " + f"'index' values; cannot align embeddings with inputs" + ) + try: + arr = np.asarray([r["embedding"] for r in rows], dtype=np.float32) + except (KeyError, TypeError, ValueError) as e: + raise EmbeddingAPIError( + f"Embedding API at {self._url} returned malformed embeddings: {e}" + ) from e + if arr.ndim != 2: + raise EmbeddingAPIError( + f"Embedding API at {self._url} returned non-vector embeddings (shape {arr.shape})" + ) + # L2-normalize so cosine == dot product (collection uses + # hnsw:space=cosine), matching EmbeddinggemmaONNX above. + norms = np.linalg.norm(arr, axis=1, keepdims=True) + 1e-12 + return (arr / norms).tolist() + + def get_embedding_function(device: Optional[str] = None, model: Optional[str] = None): """Return a cached embedding function for the requested device + model. @@ -464,6 +638,41 @@ def get_embedding_function(device: Optional[str] = None, model: Optional[str] = if model is None: model = cfg.embedding_model + # OpenAI-compatible embedding API: bypasses local ONNX entirely. Checked + # before device→provider resolution since it needs no hardware accelerator. + if model == "openai-compat": + from .config import MempalaceConfig + + cfg = MempalaceConfig() + url = cfg.embedding_api_url + if not url: + raise ValueError( + "embedding_model='openai-compat' requires an endpoint — set " + "embedding_api_url in ~/.mempalace/config.json or the " + "MEMPALACE_EMBEDDING_API_URL env var (e.g. http://host:port)" + ) + api_model = cfg.embedding_api_model + if not api_model: + raise ValueError( + "embedding_model='openai-compat' requires a model — set " + "embedding_api_model in ~/.mempalace/config.json or the " + "MEMPALACE_EMBEDDING_API_MODEL env var" + ) + api_key = cfg.embedding_api_key + # Include a fingerprint of the key (never the raw secret) so a token + # rotation busts the cache in long-lived processes (e.g. MCP server). + key_fp = hashlib.sha256((api_key or "").encode("utf-8")).hexdigest()[:16] + cache_key = ("openai-compat", url, api_model, key_fp) + cached = _EF_CACHE.get(cache_key) + if cached is not None: + return cached + ef = OpenAICompatEmbeddingFunction(base_url=url, model=api_model, api_key=api_key) + _EF_CACHE[cache_key] = ef + logger.info( + "Embedding function initialized (openai-compat url=%s model=%s)", url, api_model + ) + return ef + providers, effective = _resolve_providers(device) cache_key = (model, tuple(providers)) cached = _EF_CACHE.get(cache_key) # lock-free fast path; dict.get is GIL-atomic @@ -493,15 +702,21 @@ def get_embedding_function(device: Optional[str] = None, model: Optional[str] = def describe_device(device: Optional[str] = None) -> str: - """Return a short human-readable label for the resolved device. + """Return a short human-readable label for the resolved embedding backend. - Used by the miner CLI header so users can see at a glance whether GPU - acceleration actually engaged. + Used by the miner CLI header / MCP status so users can see at a glance + whether GPU acceleration engaged — or, for the ``openai-compat`` backend, + that embeddings are served by a remote endpoint rather than local hardware + (in which case the ``embedding_device`` accelerator label is irrelevant). """ if device is None: from .config import MempalaceConfig - device = MempalaceConfig().embedding_device + cfg = MempalaceConfig() + if cfg.embedding_model == "openai-compat": + url = cfg.embedding_api_url + return f"openai-compat ({url})" if url else "openai-compat" + device = cfg.embedding_device _, effective = _resolve_providers(device) return effective diff --git a/tests/conftest.py b/tests/conftest.py index b9ad6d9..2f0ba73 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -40,6 +40,7 @@ _TEST_EMBED_DIM = 384 _TEST_TOKEN_RE = re.compile(r"\w+", re.UNICODE) _REAL_EMBEDDING_TEST_MODULES = { "test_embedding", + "test_embedding_api", "test_embeddinggemma", } diff --git a/tests/test_embedding_api.py b/tests/test_embedding_api.py new file mode 100644 index 0000000..0f07b30 --- /dev/null +++ b/tests/test_embedding_api.py @@ -0,0 +1,352 @@ +"""Tests for the OpenAI-compatible embedding API backend (issue #1559). + +Covers ``OpenAICompatEmbeddingFunction``, the ``embedding_model == +"openai-compat"`` selection branch in ``get_embedding_function``, and the +``MempalaceConfig`` properties that are the single source of truth for the +endpoint settings. No server required — ``urllib.request.urlopen`` is mocked. +""" + +import json + +import pytest + +import mempalace.embedding as embedding +from mempalace.config import MempalaceConfig + + +@pytest.fixture(autouse=True) +def isolate_embedding_cache(monkeypatch): + monkeypatch.setattr(embedding, "_EF_CACHE", {}) + + +# ── Fake HTTP layer ─────────────────────────────────────────────────────── + + +class _FakeResp: + def __init__(self, body: bytes): + self._body = body + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def read(self): + return self._body + + +def _fake_urlopen(*, dim=4, one_hot=False, shuffle=False, captured=None): + """Return a urlopen stand-in that echoes one embedding per input text. + + ``one_hot`` makes each vector a unit vector at its own index (so order is + observable); ``shuffle`` reverses the returned rows to prove the EF + re-sorts by ``index``; ``captured`` collects the Request objects. + """ + + def fake(req, timeout=None): + if captured is not None: + captured.append(req) + body = json.loads(req.data.decode()) + n = len(body["input"]) + rows = [] + for i in range(n): + if one_hot: + vec = [0.0] * max(dim, n) + vec[i] = 1.0 + else: + vec = [float(i + 1)] * dim + rows.append({"index": i, "embedding": vec}) + if shuffle: + rows = list(reversed(rows)) + return _FakeResp(json.dumps({"data": rows, "model": body["model"]}).encode()) + + return fake + + +# ── OpenAICompatEmbeddingFunction ───────────────────────────────────────── + + +def test_resolve_url_variants(): + ef = embedding.OpenAICompatEmbeddingFunction + assert ef("http://h:8420", "m")._url == "http://h:8420/v1/embeddings" + assert ef("http://h:8420/", "m")._url == "http://h:8420/v1/embeddings" + assert ef("http://h:8420/v1", "m")._url == "http://h:8420/v1/embeddings" + assert ef("http://h:8420/v1/embeddings", "m")._url == "http://h:8420/v1/embeddings" + + +def test_name_encodes_model(): + ef = embedding.OpenAICompatEmbeddingFunction + assert ef("http://h", "small").name() == "openai_compat_emb_small" + # HF-style ids with slashes are flattened to a safe identifier + assert ef("http://h", "Qwen/Qwen3-Embedding-0.6B").name() == ( + "openai_compat_emb_Qwen_Qwen3-Embedding-0.6B" + ) + + +def test_embeds_and_l2_normalizes(monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(dim=4)) + ef = embedding.OpenAICompatEmbeddingFunction("http://h:8420", "small") + out = ef(["a", "b"]) + assert len(out) == 2 + assert len(out[0]) == 4 + for vec in out: + assert abs(sum(x * x for x in vec) ** 0.5 - 1.0) < 1e-6 + + +def test_sorts_response_by_index(monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(one_hot=True, shuffle=True)) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + out = ef(["x", "y", "z"]) + # Server returned rows reversed; the EF must realign by index so out[i] + # is the one-hot vector for position i. + for i, vec in enumerate(out): + assert max(range(len(vec)), key=lambda j: vec[j]) == i + + +def test_sends_bearer_header_when_key_set(monkeypatch): + captured = [] + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(captured=captured)) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m", api_key="sk-secret") + ef(["a"]) + assert captured[0].get_header("Authorization") == "Bearer sk-secret" + + +def test_no_auth_header_without_key(monkeypatch): + captured = [] + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(captured=captured)) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + ef(["a"]) + assert captured[0].get_header("Authorization") is None + + +def test_batches_large_input(monkeypatch): + captured = [] + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(captured=captured)) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + out = ef([f"t{i}" for i in range(130)]) # > _EF_API_BATCH (64) + assert len(out) == 130 + assert len(captured) == 3 # 64 + 64 + 2 + + +def test_embed_query_delegates_to_call(monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(dim=4)) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + assert ef.embed_query(["q"]) == ef(["q"]) + + +def test_raises_on_count_mismatch(monkeypatch): + def short(req, timeout=None): + return _FakeResp(json.dumps({"data": [{"index": 0, "embedding": [1.0]}]}).encode()) + + monkeypatch.setattr("urllib.request.urlopen", short) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(RuntimeError, match="embeddings for"): + ef(["a", "b"]) + + +def test_raises_on_transport_error(monkeypatch): + from urllib.error import URLError + + def boom(req, timeout=None): + raise URLError("connection refused") + + monkeypatch.setattr("urllib.request.urlopen", boom) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(RuntimeError, match="failed"): + ef(["a"]) + + +# ── get_embedding_function selection branch ─────────────────────────────── + + +class _FakeCfg: + def __init__(self, url=None, model=None, key=None, embedding_model="openai-compat"): + self.embedding_api_url = url + self.embedding_api_model = model + self.embedding_api_key = key + self.embedding_model = embedding_model + + +def test_get_embedding_function_selects_openai_compat(monkeypatch): + monkeypatch.setattr( + "mempalace.config.MempalaceConfig", lambda *a, **k: _FakeCfg("http://h:8420", "small") + ) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(dim=4)) + ef = embedding.get_embedding_function(device="cpu", model="openai-compat") + assert isinstance(ef, embedding.OpenAICompatEmbeddingFunction) + assert ef.name() == "openai_compat_emb_small" + assert len(ef(["hi"])[0]) == 4 + + +def test_openai_compat_requires_url(monkeypatch): + monkeypatch.setattr("mempalace.config.MempalaceConfig", lambda *a, **k: _FakeCfg(None, "small")) + with pytest.raises(ValueError, match="requires an endpoint"): + embedding.get_embedding_function(device="cpu", model="openai-compat") + + +def test_openai_compat_requires_model(monkeypatch): + monkeypatch.setattr( + "mempalace.config.MempalaceConfig", lambda *a, **k: _FakeCfg("http://h", None) + ) + with pytest.raises(ValueError, match="requires a model"): + embedding.get_embedding_function(device="cpu", model="openai-compat") + + +# ── MempalaceConfig endpoint settings (single source of truth) ──────────── + + +def test_config_api_url_from_file(tmp_path, monkeypatch): + monkeypatch.delenv("MEMPALACE_EMBEDDING_API_URL", raising=False) + (tmp_path / "config.json").write_text(json.dumps({"embedding_api_url": "http://host:8420"})) + assert MempalaceConfig(config_dir=str(tmp_path)).embedding_api_url == "http://host:8420" + + +def test_config_api_env_overrides_file(tmp_path, monkeypatch): + (tmp_path / "config.json").write_text(json.dumps({"embedding_api_url": "http://from-config"})) + monkeypatch.setenv("MEMPALACE_EMBEDDING_API_URL", " http://from-env ") + assert MempalaceConfig(config_dir=str(tmp_path)).embedding_api_url == "http://from-env" + + +def test_config_api_unset_is_none(tmp_path, monkeypatch): + for var in ("MEMPALACE_EMBEDDING_API_URL", "MEMPALACE_EMBEDDING_API_MODEL"): + monkeypatch.delenv(var, raising=False) + cfg = MempalaceConfig(config_dir=str(tmp_path)) + assert cfg.embedding_api_url is None + assert cfg.embedding_api_model is None + + +def test_config_api_blank_value_is_none(tmp_path, monkeypatch): + monkeypatch.delenv("MEMPALACE_EMBEDDING_API_MODEL", raising=False) + (tmp_path / "config.json").write_text(json.dumps({"embedding_api_model": " "})) + assert MempalaceConfig(config_dir=str(tmp_path)).embedding_api_model is None + + +def test_config_api_model_and_key_preserve_case(tmp_path, monkeypatch): + for var in ("MEMPALACE_EMBEDDING_API_MODEL", "MEMPALACE_EMBEDDING_API_KEY"): + monkeypatch.delenv(var, raising=False) + (tmp_path / "config.json").write_text( + json.dumps({"embedding_api_model": "Qwen3-Embedding", "embedding_api_key": "AbC-XyZ"}) + ) + cfg = MempalaceConfig(config_dir=str(tmp_path)) + assert cfg.embedding_api_model == "Qwen3-Embedding" + assert cfg.embedding_api_key == "AbC-XyZ" + + +def test_config_api_blank_env_falls_through_to_file(tmp_path, monkeypatch): + (tmp_path / "config.json").write_text(json.dumps({"embedding_api_url": "http://from-config"})) + monkeypatch.setenv("MEMPALACE_EMBEDDING_API_URL", " ") # blank must not mask the file value + assert MempalaceConfig(config_dir=str(tmp_path)).embedding_api_url == "http://from-config" + + +# ── request shape + malformed-response hardening (review findings) ──────── + + +def test_request_targets_v1_embeddings_with_expected_body(monkeypatch): + captured = [] + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen(captured=captured)) + embedding.OpenAICompatEmbeddingFunction("http://h:8420", "small")(["a", "b"]) + req = captured[0] + assert req.full_url == "http://h:8420/v1/embeddings" + assert req.get_header("Content-type") == "application/json" + # Custom User-Agent so Cloudflare-fronted endpoints don't 403 us (#1570). + assert req.get_header("User-agent", "").startswith("mempalace/") + assert json.loads(req.data) == { + "model": "small", + "input": ["a", "b"], + "encoding_format": "float", + } + + +def test_embedding_api_error_is_runtimeerror(): + assert issubclass(embedding.EmbeddingAPIError, RuntimeError) + + +def test_raises_on_missing_embedding_key(monkeypatch): + def bad(req, timeout=None): + return _FakeResp(json.dumps({"data": [{"index": 0}]}).encode()) # no "embedding" + + monkeypatch.setattr("urllib.request.urlopen", bad) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(embedding.EmbeddingAPIError, match="malformed embeddings"): + ef(["a"]) + + +def test_raises_on_non_contiguous_indices(monkeypatch): + # Count matches (2 rows for 2 inputs) but the indices are absolute, not + # 0..n-1 — sort+positional-zip would silently misalign vectors with texts. + def bad(req, timeout=None): + rows = [{"index": 64, "embedding": [1.0]}, {"index": 65, "embedding": [2.0]}] + return _FakeResp(json.dumps({"data": rows}).encode()) + + monkeypatch.setattr("urllib.request.urlopen", bad) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(embedding.EmbeddingAPIError, match="non-contiguous"): + ef(["a", "b"]) + + +def test_raises_on_http_protocol_exception(monkeypatch): + # BadStatusLine / IncompleteRead — common with local/overloaded servers. + from http.client import HTTPException + + def boom(req, timeout=None): + raise HTTPException("incomplete read") + + monkeypatch.setattr("urllib.request.urlopen", boom) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(embedding.EmbeddingAPIError, match="failed"): + ef(["a"]) + + +def test_raises_on_value_error_from_urlopen(monkeypatch): + # urlopen raises ValueError on an invalid/missing URL scheme. + def boom(req, timeout=None): + raise ValueError("unknown url type") + + monkeypatch.setattr("urllib.request.urlopen", boom) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(embedding.EmbeddingAPIError, match="failed"): + ef(["a"]) + + +def test_raises_on_non_object_response(monkeypatch): + def bad(req, timeout=None): + return _FakeResp(json.dumps([1, 2, 3]).encode()) # JSON list, not an object + + monkeypatch.setattr("urllib.request.urlopen", bad) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(embedding.EmbeddingAPIError, match="non-object response"): + ef(["a"]) + + +def test_raises_and_surfaces_server_error_body(monkeypatch): + # HTTP 200 with an OpenAI-style error envelope (no "data") — surface it. + def err(req, timeout=None): + return _FakeResp(json.dumps({"error": {"message": "model not found"}}).encode()) + + monkeypatch.setattr("urllib.request.urlopen", err) + ef = embedding.OpenAICompatEmbeddingFunction("http://h", "m") + with pytest.raises(embedding.EmbeddingAPIError, match="model not found"): + ef(["a"]) + + +def test_get_embedding_function_caches_instance(monkeypatch): + monkeypatch.setattr( + "mempalace.config.MempalaceConfig", lambda *a, **k: _FakeCfg("http://h:8420", "small", "k") + ) + a = embedding.get_embedding_function(device="cpu", model="openai-compat") + b = embedding.get_embedding_function(device="cpu", model="openai-compat") + assert a is b + + +def test_describe_device_reports_openai_compat_endpoint(monkeypatch): + monkeypatch.setattr( + "mempalace.config.MempalaceConfig", + lambda *a, **k: _FakeCfg("http://10.0.0.1:8420", "small"), + ) + assert embedding.describe_device() == "openai-compat (http://10.0.0.1:8420)" + + +def test_describe_device_openai_compat_without_url(monkeypatch): + monkeypatch.setattr("mempalace.config.MempalaceConfig", lambda *a, **k: _FakeCfg(None, "small")) + assert embedding.describe_device() == "openai-compat"