Skip to content

feat(embedding): declare a model's width without sending dimensions #1857

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
NullSense wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(embedding): declare a model's width without sending dimensions #1857

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 106 additions & 36 deletions packages/core/src/repowise/core/providers/embedding/openai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,13 @@
text-embedding-ada-002 → 1536 dims

REPOWISE_EMBEDDING_DIMS (or the ``dimensions`` arg) overrides the width and
is passed to the API so the returned vectors match. See ``OpenAIEmbedder``.
is passed to the API so the returned vectors match — for a model that
supports a variable (Matryoshka) width.

REPOWISE_EMBEDDING_DECLARED_DIMS (or the ``declared_dimensions`` arg)
overrides the width *without* passing it to the API, for a model whose
width is fixed and that rejects the ``dimensions`` parameter outright.
See ``OpenAIEmbedder``.
"""

from __future__ import annotations
Expand All @@ -44,9 +50,18 @@ class OpenAIEmbedder:
base_url: Optional custom base URL for OpenAI-compatible endpoints.
dimensions: Output width for a model not in ``_DIMS`` (e.g. a local
OpenAI-compatible embedder). Falls back to REPOWISE_EMBEDDING_DIMS,
then to the known-model table, then 1536. An overridden width is
also sent to the API so the returned vectors match the declaration;
an endpoint that does not implement the parameter ignores it.
then to ``declared_dimensions``, then to the known-model table,
then 1536. Also sent to the API as the ``dimensions`` parameter,
for a model that supports reshaping its output (Matryoshka); an
endpoint that does not implement the parameter ignores it.
declared_dimensions: Output width for a model whose width is fixed
and that rejects the ``dimensions`` parameter — e.g. NVIDIA's
Nemotron-3-Embed-1B, which 400s with "does not support Matryoshka
embeddings; dimensions must be unset" if it's sent at all. Falls
back to REPOWISE_EMBEDDING_DECLARED_DIMS. Unlike ``dimensions``,
never sent to the API — it only tells this adapter, and therefore
the vector store, what width to expect. Ignored when ``dimensions``
(or REPOWISE_EMBEDDING_DIMS) is also set.
"""

_DIMS: ClassVar[dict[str, int]] = {
Expand All @@ -66,6 +81,7 @@ def __init__(
timeout: float | None = None,
base_url: str | None = None,
dimensions: int | None = None,
declared_dimensions: int | None = None,
) -> None:
self._api_key = api_key or os.environ.get("OPENAI_API_KEY")
if not self._api_key:
Expand All @@ -77,47 +93,70 @@ def __init__(
self._timeout = resolve_embedding_timeout(
timeout, self._DEFAULT_TIMEOUT, provider_env="OPENAI_EMBEDDING_TIMEOUT"
)
# When the user overrides the width, request that width from the API too,
# so the returned vectors match the declaration instead of the model's
# default — otherwise the store is sized to a width the vectors never
# have. Servers that don't implement the parameter ignore it (returning
# their native width, which the override then correctly declares); one
# that can't honour it rejects the request loudly rather than silently
# returning a mismatched width. No model is special-cased here: the
# endpoint, not a hardcoded name table, decides what it accepts.
self._dimensions, self._request_dimensions = self._resolve_dimensions(dimensions, model)
# When the user overrides the width via `dimensions`, request that width
# from the API too, so the returned vectors match the declaration instead
# of the model's default — otherwise the store is sized to a width the
# vectors never have. Servers that don't implement the parameter ignore
# it (returning their native width, which the override then correctly
# declares); one that can't honour it rejects the request loudly rather
# than silently returning a mismatched width.
#
# `declared_dimensions` exists for the model that can't honour it and
# doesn't reject it loudly either — it rejects the request outright,
# because the parameter asks for reshaping (Matryoshka) support the
# model doesn't have. There, the only way to tell repowise the real
# width is to declare it without ever sending it. Neither knob hardcodes
# a model name: the operator declares the width for whatever model
# they're pointing at, and the endpoint decides what it accepts.
self._dimensions, self._request_dimensions, self._dims_declared_only = (
self._resolve_dimensions(dimensions, declared_dimensions, model)
)
self._client: object | None = None # cached; created once on first embed()

@classmethod
def _resolve_dimensions(cls, dimensions: int | None, model: str) -> tuple[int, int | None]:
"""Resolve ``(declared_width, override)``.
def _resolve_dimensions(
cls, dimensions: int | None, declared_dimensions: int | None, model: str
) -> tuple[int, int | None, bool]:
"""Resolve ``(declared_width, request_override, declared_only)``.

``override`` is the user-chosen width — explicit arg or
``REPOWISE_EMBEDDING_DIMS`` — once validated, or ``None`` when the width
falls back to the known-model table. It doubles as the value to request
from the API: ``None`` means send no ``dimensions`` (keep the stock
request byte-identical). Precedence for the declared width:
explicit arg > REPOWISE_EMBEDDING_DIMS > known-model table > 1536.
``request_override`` is the value to send the API as the ``dimensions``
parameter, or ``None`` to keep the request byte-identical to the stock
call. ``declared_only`` is True when the width came from
``declared_dimensions`` / ``REPOWISE_EMBEDDING_DECLARED_DIMS`` rather
than from ``dimensions`` / ``REPOWISE_EMBEDDING_DIMS`` or the ``_DIMS``
table — kept only so :meth:`embed`'s width-mismatch error can name the
setting that actually chose the number.

Precedence for the declared width: explicit ``dimensions=`` >
``REPOWISE_EMBEDDING_DIMS`` > explicit ``declared_dimensions=`` >
``REPOWISE_EMBEDDING_DECLARED_DIMS`` > the ``_DIMS`` table > 1536.
The first two are also requested from the API; the last three never
are, ``_DIMS`` because it is not a user override at all, and
``declared_dimensions`` because it exists specifically for a model
that must never receive the parameter.

A local OpenAI-compatible embedder (e.g. a self-hosted model) is not in
``_DIMS``; without an override its width would silently default to 1536
and mismatch the store. Mirrors the Ollama/Gemini embedders, which
``_DIMS``; without a declared width its width would silently default to
1536 and mismatch the store. Mirrors the Ollama/Gemini embedders, which
already honour REPOWISE_EMBEDDING_DIMS.
"""
if dimensions is None:
env = os.environ.get("REPOWISE_EMBEDDING_DIMS")
if env:
try:
dimensions = int(env)
except ValueError:
# Match the message every other bad value raises below,
# instead of a raw "invalid literal for int()".
raise ValueError("dimensions must be a positive integer") from None
if dimensions is None:
return cls._DIMS.get(model, 1536), None
if isinstance(dimensions, bool) or not isinstance(dimensions, int) or dimensions <= 0:
raise ValueError("dimensions must be a positive integer")
return dimensions, dimensions
dimensions = _parse_dimensions_env(env)
if dimensions is not None:
width = _validate_dimensions(dimensions)
return width, width, False

if declared_dimensions is None:
env = os.environ.get("REPOWISE_EMBEDDING_DECLARED_DIMS")
if env:
declared_dimensions = _parse_dimensions_env(env)
if declared_dimensions is not None:
width = _validate_dimensions(declared_dimensions)
return width, None, True

return cls._DIMS.get(model, 1536), None, False

@property
def dimensions(self) -> int:
Expand All @@ -142,6 +181,7 @@ async def embed(self, texts: list[str]) -> list[list[float]]:
timeout = self._timeout
request_dimensions = self._request_dimensions
expected_dimensions = self._dimensions
dims_declared_only = self._dims_declared_only

def _embed_sync() -> list[list[float]]:
import openai # type: ignore[import-untyped]
Expand All @@ -166,22 +206,52 @@ def _embed_sync() -> list[list[float]]:
f"Set REPOWISE_EMBEDDING_DIMS={actual} to match the server's"
f" native output, or remove the override to use the model's default."
)
elif dims_declared_only:
hint = (
f"The width {expected_dimensions} came from REPOWISE_EMBEDDING_DECLARED_DIMS"
f" (or declared_dimensions=), which is deliberately never sent as the API's"
f" 'dimensions' parameter. Set it to {actual} to match what the server"
f" actually returns, or — if this endpoint does accept the parameter after"
f" all — switch to REPOWISE_EMBEDDING_DIMS={actual} instead."
)
else:
hint = (
f"The width {expected_dimensions} came from the built-in _DIMS table for"
f" {model!r}. Add or update OpenAIEmbedder._DIMS[{model!r}] = {actual},"
f" or set REPOWISE_EMBEDDING_DIMS={actual}."
)
reason = (
"The endpoint likely ignored the 'dimensions' parameter."
if request_dimensions is not None
else "Nothing told this endpoint to produce that width — it never was."
)
raise ValueError(
f"OpenAIEmbedder declared {expected_dimensions}-dimensional vectors but the"
f" API returned {actual} (model={model!r}). The endpoint likely ignored"
f" the 'dimensions' parameter. {hint}"
f" API returned {actual} (model={model!r}). {reason} {hint}"
)
return [_l2_normalize(v) for v in raw_vectors]

return await asyncio.to_thread(_embed_sync)


def _validate_dimensions(value: int) -> int:
"""Raise the one message every bad width — explicit or from env — shares."""
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
raise ValueError("dimensions must be a positive integer")
return value


def _parse_dimensions_env(raw: str) -> int:
"""Parse an env var's dimensions value, or raise the same message a bad
explicit ``dimensions=`` / ``declared_dimensions=`` does.
"""
try:
parsed = int(raw)
except ValueError:
raise ValueError("dimensions must be a positive integer") from None
return _validate_dimensions(parsed)


def _l2_normalize(vec: list[float]) -> list[float]:
"""L2-normalize a vector to unit length (cosine similarity = dot product)."""
norm = math.sqrt(sum(x * x for x in vec))
Expand Down
106 changes: 106 additions & 0 deletions tests/unit/test_persistence/test_openai_embedder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,61 @@ def test_malformed_env_raises_the_same_message(monkeypatch):
OpenAIEmbedder(api_key="k", model="local-embedder")


# ---------------------------------------------------------------------------
# declared_dimensions / REPOWISE_EMBEDDING_DECLARED_DIMS — declare a width
# without ever sending it, for a model that rejects the 'dimensions' param.
# ---------------------------------------------------------------------------


def test_declared_dimensions_sets_width_without_a_dimensions_override():
emb = OpenAIEmbedder(api_key="k", model="nemotron-embed-1b", declared_dimensions=2048)
assert emb.dimensions == 2048
assert emb._request_dimensions is None


def test_declared_dimensions_from_env(monkeypatch):
monkeypatch.setenv("REPOWISE_EMBEDDING_DECLARED_DIMS", "2048")
emb = OpenAIEmbedder(api_key="k", model="nemotron-embed-1b")
assert emb.dimensions == 2048
assert emb._request_dimensions is None


def test_dimensions_arg_beats_declared_dimensions_arg():
# A model that DOES accept the parameter still takes the send-it path when
# both are given — declared_dimensions is the fallback for one that can't.
emb = OpenAIEmbedder(
api_key="k", model="local-embedder", dimensions=1024, declared_dimensions=2048
)
assert emb.dimensions == 1024
assert emb._request_dimensions == 1024


def test_embedding_dims_env_beats_declared_dimensions_env(monkeypatch):
monkeypatch.setenv("REPOWISE_EMBEDDING_DIMS", "1024")
monkeypatch.setenv("REPOWISE_EMBEDDING_DECLARED_DIMS", "2048")
emb = OpenAIEmbedder(api_key="k", model="local-embedder")
assert emb.dimensions == 1024
assert emb._request_dimensions == 1024


def test_declared_dimensions_arg_beats_declared_dimensions_env(monkeypatch):
monkeypatch.setenv("REPOWISE_EMBEDDING_DECLARED_DIMS", "4096")
emb = OpenAIEmbedder(api_key="k", model="nemotron-embed-1b", declared_dimensions=2048)
assert emb.dimensions == 2048


@pytest.mark.parametrize("bad", [0, -5, True])
def test_invalid_declared_dimensions_raises(bad):
with pytest.raises(ValueError, match="dimensions must be a positive integer"):
OpenAIEmbedder(api_key="k", model="nemotron-embed-1b", declared_dimensions=bad)


def test_malformed_declared_dimensions_env_raises_the_same_message(monkeypatch):
monkeypatch.setenv("REPOWISE_EMBEDDING_DECLARED_DIMS", "abc")
with pytest.raises(ValueError, match="dimensions must be a positive integer"):
OpenAIEmbedder(api_key="k", model="nemotron-embed-1b")


def test_timeout_from_shared_env(monkeypatch):
monkeypatch.setenv("REPOWISE_EMBEDDING_TIMEOUT", "180")
assert OpenAIEmbedder(api_key="k", model="local-embedder")._timeout == 180.0
Expand Down Expand Up @@ -230,6 +285,37 @@ async def test_default_request_omits_dimensions():
assert "dimensions" not in kwargs


async def test_declared_only_width_is_never_sent():
emb = OpenAIEmbedder(api_key="k", model="nemotron-embed-1b", declared_dimensions=2048)
kwargs = await _capture_create_kwargs(emb)
assert "dimensions" not in kwargs


async def test_declared_only_width_is_never_sent_even_to_a_strict_endpoint():
# nvidia/Nemotron-3-Embed-1B's own behaviour, reproduced: a fake endpoint
# that 400s the instant 'dimensions' shows up in the request at all, since
# the model isn't Matryoshka-capable and has no default width to fall back
# to. This is the case _DIMS/REPOWISE_EMBEDDING_DIMS cannot serve — sending
# the override at all breaks the request, so the only usable path is one
# that declares the width without ever asking the API to produce it.
emb = OpenAIEmbedder(api_key="k", model="nemotron-embed-1b", declared_dimensions=2048)

def fake_create(**kwargs):
if "dimensions" in kwargs:
raise ValueError(
"Model 'nemotron-embed-1b' does not support Matryoshka embeddings; "
"dimensions must be unset"
)
return _make_mock_response([[1.0] + [0.0] * 2047])

with patch("openai.OpenAI") as mock_client:
mock_client.return_value.embeddings.create.side_effect = fake_create
result = await emb.embed(["hello"])

assert len(result) == 1
assert len(result[0]) == 2048


# ---------------------------------------------------------------------------
# Width verification — the check added in fix/embedder-width-verification
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -273,6 +359,26 @@ async def test_embed_raises_when_api_returns_wrong_width_with_user_override(monk
assert "REPOWISE_EMBEDDING_DIMS" in msg


async def test_embed_raises_when_api_returns_wrong_width_with_declared_only_override(monkeypatch):
# Declared via declared_dimensions= (never sent). The error message should
# reference REPOWISE_EMBEDDING_DECLARED_DIMS, not _DIMS or the sent-width
# hint — neither would tell this user what to change.
emb = OpenAIEmbedder(api_key="k", model="nemotron-embed-1b", declared_dimensions=2048)
assert emb.dimensions == 2048

with patch("openai.OpenAI") as mock_client:
mock_client.return_value.embeddings.create.return_value = _make_mock_response(
[[1.0, 0.0, 0.0]]
)
with pytest.raises(ValueError, match="2048") as exc_info:
await emb.embed(["hello"])

msg = str(exc_info.value)
assert "3" in msg
assert "REPOWISE_EMBEDDING_DECLARED_DIMS" in msg
assert "_DIMS table" not in msg


async def test_embed_width_check_is_skipped_for_empty_response():
# embed([]) short-circuits before the API call and returns []; the guard
# must not fire on the empty list itself.
Expand Down