Merge pull request #1671 from maximilize/feat/openai-compat-embeddings
feat(embedding): add OpenAI-compatible /v1/embeddings backend
This commit is contained in:
Igor Lins e Silva
GitHub
commit
9bc8cdddc6
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
Loading…
Reference in New Issue