Skip to content

Commit 02edbe0

Browse files
JustinBacherclaude
andcommitted
feat(embedding): optionally truncate oversized embeddings to configured dims
Matryoshka (MRL) embedding models served behind OpenAI-compatible endpoints (llama.cpp, some vLLM builds) ignore the `dimensions=` request parameter and always return their native dimension. Honcho then hard-fails on the dimension mismatch, making these models unusable even though truncating an MRL vector to a smaller size is valid. Add an opt-in `EMBEDDING_TRUNCATE_TO_DIMENSIONS` flag. When enabled and the provider returns a vector larger than `EMBEDDING_VECTOR_DIMENSIONS`, keep the leading components and L2-renormalize so cosine/inner-product similarity stays well-defined. Off by default, so a genuinely misconfigured model still fails loudly. Also keeps embeddings at or below pgvector's 2000-dimension HNSW index cap when the native model dimension exceeds it. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 055d73b commit 02edbe0

4 files changed

Lines changed: 128 additions & 3 deletions

File tree

.env.template

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,12 @@ PERFORMANCE_LOG_FORMAT=compact # compact|rich
1919
# EMBEDDING_VECTOR_DIMENSIONS=1536
2020
# EMBEDDING_MAX_INPUT_TOKENS=8192
2121
# EMBEDDING_MAX_TOKENS_PER_REQUEST=300000
22+
# Truncate + L2-renormalize oversized embeddings down to EMBEDDING_VECTOR_DIMENSIONS
23+
# instead of erroring on a dimension mismatch. Use with Matryoshka (MRL) models
24+
# served behind OpenAI-compatible endpoints (llama.cpp, some vLLM builds) that
25+
# ignore the `dimensions=` request parameter and always emit their native size.
26+
# Off by default so a genuinely misconfigured model still fails loudly.
27+
# EMBEDDING_TRUNCATE_TO_DIMENSIONS=false
2228
# EMBEDDING_MODEL_CONFIG__TRANSPORT=openai
2329
# EMBEDDING_MODEL_CONFIG__MODEL=text-embedding-3-small
2430
# EMBEDDING_MODEL_CONFIG__OVERRIDES__BASE_URL=

src/config.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -758,6 +758,14 @@ def _MODEL_CONFIG_DEFAULT() -> ConfiguredEmbeddingModelSettings:
758758
# saturated, message creation skips the fast path entirely and the
759759
# reconciler embeds on its next cycle. 0 disables the fast path.
760760
MAX_PENDING_EMBED_TASKS: Annotated[int, Field(default=50, ge=0)] = 50
761+
# When the embedding provider returns a vector larger than
762+
# VECTOR_DIMENSIONS, truncate it to VECTOR_DIMENSIONS and L2-renormalize
763+
# instead of raising a dimension-mismatch error. Intended for MRL-trained
764+
# models (e.g. Qwen3-Embedding) served behind OpenAI-compatible endpoints
765+
# (llama.cpp, some vLLM builds) that ignore the `dimensions=` request
766+
# parameter and always emit their native dimension. Off by default so a
767+
# genuinely misconfigured model still fails loudly.
768+
TRUNCATE_TO_DIMENSIONS: bool = False
761769

762770
@model_validator(mode="before")
763771
@classmethod

src/embedding_client.py

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import asyncio
22
import logging
3+
import math
34
import threading
45
import time
56
from collections import defaultdict
@@ -173,11 +174,13 @@ def __init__(
173174
max_input_tokens: int,
174175
max_tokens_per_request: int,
175176
send_dimensions: bool,
177+
truncate_to_dimensions: bool = False,
176178
):
177179
self.transport: str = config.transport
178180
self.model: str = config.model
179181
self.vector_dimensions: int = vector_dimensions
180182
self.send_dimensions: bool = send_dimensions
183+
self.truncate_to_dimensions: bool = truncate_to_dimensions
181184

182185
if self.transport == "gemini":
183186
if not config.api_key:
@@ -223,6 +226,20 @@ def _validate_embedding_dimensions(self, embedding: list[float]) -> list[float]:
223226
)
224227
return embedding
225228

229+
def _fit_embedding_dimensions(self, embedding: list[float]) -> list[float]:
230+
"""Coerce a provider embedding to ``vector_dimensions``, then validate.
231+
232+
When ``truncate_to_dimensions`` is enabled and the provider returned a
233+
larger vector than configured (common with MRL-trained models served
234+
behind OpenAI-compatible endpoints that ignore ``dimensions=``), keep
235+
the leading ``vector_dimensions`` components and L2-renormalize so
236+
cosine similarity stays well-defined. Otherwise, fall through to the
237+
strict validator, which raises on any mismatch.
238+
"""
239+
if self.truncate_to_dimensions and len(embedding) > self.vector_dimensions:
240+
embedding = _l2_truncate(embedding, self.vector_dimensions)
241+
return self._validate_embedding_dimensions(embedding)
242+
226243
async def embed(self, query: str) -> list[float]:
227244
token_count = len(self.encoding.encode(query))
228245

@@ -264,7 +281,7 @@ async def _call_openai() -> list[float]:
264281
if self.send_dimensions:
265282
openai_kwargs["dimensions"] = self.vector_dimensions
266283
response = await openai_client.embeddings.create(**openai_kwargs)
267-
return self._validate_embedding_dimensions(response.data[0].embedding)
284+
return self._fit_embedding_dimensions(response.data[0].embedding)
268285

269286
return await _emit_embedding_call(
270287
provider=self.transport,
@@ -469,7 +486,7 @@ async def _call_provider() -> dict[str, dict[int, list[float]]]:
469486
response = await self.client.embeddings.create(**openai_kwargs)
470487
for item, embedding_data in zip(batch, response.data, strict=True):
471488
result[item.text_id][item.chunk_index] = (
472-
self._validate_embedding_dimensions(embedding_data.embedding)
489+
self._fit_embedding_dimensions(embedding_data.embedding)
473490
)
474491
return result
475492

@@ -531,6 +548,21 @@ def _accumulate_embeddings(
531548
}
532549

533550

551+
def _l2_truncate(embedding: list[float], target_dimensions: int) -> list[float]:
552+
"""Keep the leading ``target_dimensions`` components and L2-renormalize.
553+
554+
Valid for Matryoshka (MRL) embeddings, where the leading subvector is a
555+
trained lower-dimensional representation. Renormalizing restores unit norm
556+
so cosine/inner-product similarity is preserved. A zero-norm slice is
557+
returned unchanged (degenerate input; nothing to normalize).
558+
"""
559+
sliced = embedding[:target_dimensions]
560+
norm = math.sqrt(sum(value * value for value in sliced))
561+
if norm == 0.0:
562+
return sliced
563+
return [value / norm for value in sliced]
564+
565+
534566
def _chunk_text_with_tokens(
535567
text: str,
536568
encoded_tokens: list[int],
@@ -603,6 +635,7 @@ def _get_client(self) -> _EmbeddingClient:
603635
max_input_tokens=settings.EMBEDDING.MAX_INPUT_TOKENS,
604636
max_tokens_per_request=settings.EMBEDDING.MAX_TOKENS_PER_REQUEST,
605637
send_dimensions=settings.EMBEDDING.resolve_send_dimensions(),
638+
truncate_to_dimensions=settings.EMBEDDING.TRUNCATE_TO_DIMENSIONS,
606639
)
607640
self._instance_signature = signature
608641
logger.debug(
@@ -627,6 +660,7 @@ def _get_settings_signature(self) -> tuple[object, ...]:
627660
settings.EMBEDDING.MAX_INPUT_TOKENS,
628661
settings.EMBEDDING.MAX_TOKENS_PER_REQUEST,
629662
settings.EMBEDDING.resolve_send_dimensions(),
663+
settings.EMBEDDING.TRUNCATE_TO_DIMENSIONS,
630664
)
631665

632666
async def embed(self, query: str) -> list[float]:

tests/llm/test_embedding_client.py

Lines changed: 78 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
1+
import math
12
from types import SimpleNamespace
23
from typing import Any
34

45
import pytest
56

67
from src.config import EmbeddingModelConfig
7-
from src.embedding_client import _EmbeddingClient # pyright: ignore[reportPrivateUsage]
8+
from src.embedding_client import ( # pyright: ignore[reportPrivateUsage]
9+
_EmbeddingClient,
10+
_l2_truncate,
11+
)
812

913

1014
class FakeOpenAIEmbeddingsAPI:
@@ -92,6 +96,79 @@ def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
9296
await client.embed("hello world")
9397

9498

99+
def test_l2_truncate_slices_and_renormalizes() -> None:
100+
truncated = _l2_truncate([0.1] * 16, 8)
101+
102+
assert len(truncated) == 8
103+
assert math.isclose(math.sqrt(sum(v * v for v in truncated)), 1.0, abs_tol=1e-9)
104+
105+
106+
def test_l2_truncate_leaves_zero_norm_slice_untouched() -> None:
107+
assert _l2_truncate([0.0] * 4, 2) == [0.0, 0.0]
108+
109+
110+
@pytest.mark.asyncio
111+
async def test_openai_embedding_client_truncates_oversized_when_enabled(
112+
monkeypatch: pytest.MonkeyPatch,
113+
) -> None:
114+
# Provider returns a native 16-dim vector (e.g. an MRL model behind a server
115+
# that ignores `dimensions=`); we configured 8 and enabled truncation.
116+
fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 16)
117+
118+
class FakeOpenAIClient:
119+
def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
120+
self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
121+
122+
monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient)
123+
124+
client = _EmbeddingClient(
125+
EmbeddingModelConfig(
126+
transport="openai",
127+
model="qwen3-embedding",
128+
api_key="test-key",
129+
),
130+
vector_dimensions=8,
131+
max_input_tokens=8192,
132+
max_tokens_per_request=300_000,
133+
send_dimensions=True,
134+
truncate_to_dimensions=True,
135+
)
136+
137+
embedding = await client.embed("hello world")
138+
139+
assert len(embedding) == 8
140+
assert math.isclose(math.sqrt(sum(v * v for v in embedding)), 1.0, abs_tol=1e-9)
141+
142+
143+
@pytest.mark.asyncio
144+
async def test_openai_embedding_client_rejects_oversized_without_truncation(
145+
monkeypatch: pytest.MonkeyPatch,
146+
) -> None:
147+
fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 16)
148+
149+
class FakeOpenAIClient:
150+
def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
151+
self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
152+
153+
monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient)
154+
155+
client = _EmbeddingClient(
156+
EmbeddingModelConfig(
157+
transport="openai",
158+
model="qwen3-embedding",
159+
api_key="test-key",
160+
),
161+
vector_dimensions=8,
162+
max_input_tokens=8192,
163+
max_tokens_per_request=300_000,
164+
send_dimensions=True,
165+
truncate_to_dimensions=False,
166+
)
167+
168+
with pytest.raises(ValueError, match="Embedding dimension mismatch"):
169+
await client.embed("hello world")
170+
171+
95172
@pytest.mark.asyncio
96173
async def test_gemini_embedding_client_uses_output_dimensionality(
97174
monkeypatch: pytest.MonkeyPatch,

0 commit comments

Comments
 (0)