From 8c14429a11cdb29784739a970e2a94def0a85569 Mon Sep 17 00:00:00 2001 From: Karish Date: Wed, 19 Aug 2026 15:54:41 +0530 Subject: [PATCH 1/5] [backend:feature-improvement] Add model based registry for client-side tokenization MCDB-98680 --- singlestoredb/ai/embeddings.py | 272 ++++++++++++++++++--- singlestoredb/tests/test_embeddings.py | 313 ++++++++++++++++++++++++- 2 files changed, 550 insertions(+), 35 deletions(-) diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index ac2ced1f..c9adbe87 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -1,6 +1,18 @@ +"""LangChain embeddings models for SingleStore-hosted inference APIs. + +Models served on the 'Nova' platforms are not OpenAI models, so tiktoken is the wrong +tokenizer for them. Where this module knows a model's real tokenizer it chunks and +encodes client-side and puts model-native token IDs on the wire; otherwise it falls +back to sending raw text in character-bounded chunks. +""" import os +import warnings +from dataclasses import dataclass +from functools import lru_cache from typing import Any from typing import Callable +from typing import Dict +from typing import FrozenSet from typing import List from typing import Optional from typing import Tuple @@ -31,31 +43,216 @@ from botocore import UNSIGNED from botocore.config import Config +_DEFAULT_TOKEN_ID_PLATFORMS = frozenset({'Nova', 'NovaMultiTenant'}) + + +@dataclass(frozen=True) +class _ModelPolicy: + """Tokenization policy for one model.""" + + max_input_tokens: int + send_token_ids: bool = False + tokenizer_name: Optional[str] = None + """HuggingFace repo to load the tokenizer from. None means use the model name.""" + + token_id_platforms: FrozenSet[str] = _DEFAULT_TOKEN_ID_PLATFORMS + """Hosting platforms whose route accepts token IDs for this model. + + This is per-model-and-route rather than global: the 'Amazon' route decodes integer + ``input`` arrays with tiktoken, so model-native IDs sent that way are decoded into + unrelated text and embedded without error. + """ + + +# Keyed on a lowercased HuggingFace repo id, which is what InferenceAPIInfo.model_name +# resolves to, so entries survive deployment alias renames. +# +# Entries are explicit opt-in: there is no family-prefix or wildcard matching, even +# for models known to share a tokenizer. A model that reaches the token path without a +# parity run is the failure mode this registry exists to prevent -- mismatched special +# tokens return a well-formed, unit-norm, plausible-looking vector and raise nothing. +# To add a model: +# +# 1. Confirm the serving stack accepts token IDs on that route. vLLM types the +# embeddings input as list[int] | list[list[int]] | str | list[str]; Bedrock +# does not. +# 2. Confirm the context window actually served, including any --max-model-len +# override at launch, rather than the window the model card advertises. +# 3. Run test_live_token_id_parity from singlestoredb/tests/test_embeddings.py +# against a real deployment and require cosine > 0.9999. +# 4. Add the entry plus a unit test asserting its resolved budget and affixes. +_MODEL_POLICIES: Dict[str, _ModelPolicy] = { + 'qwen/qwen3-embedding-0.6b': _ModelPolicy( + max_input_tokens=32768, + send_token_ids=True, + ), +} + + +@lru_cache(maxsize=None) +def _load_tokenizer(tokenizer_name: str) -> Any: + """Load and memoize a HuggingFace tokenizer. + + Memoized because parsing Qwen3's ~11 MB tokenizer.json is slow, and because there + is no baked tokenizer cache in the serving image, so the first load in a container + fetches from huggingface.co over the network. + """ + from transformers import AutoTokenizer # type: ignore[import-not-found] + return AutoTokenizer.from_pretrained(tokenizer_name) + + +def _derive_special_affixes(tokenizer: Any) -> Tuple[List[int], List[int]]: + """Return the token IDs a tokenizer adds before and after content. + + Derived by diffing a probe encode rather than hardcoded, so this holds for + BOS-style models too and self-corrects if a model revision or a ``transformers`` + upgrade changes special-token handling. + """ + bare = list(tokenizer.encode('x', add_special_tokens=False)) + wrapped = list(tokenizer.encode('x', add_special_tokens=True)) + if not bare: + return [], [] + for start in range(len(wrapped) - len(bare) + 1): + if wrapped[start:start + len(bare)] == bare: + return wrapped[:start], wrapped[start + len(bare):] + return [], [] + + +@dataclass(frozen=True) +class _TokenChunker: + """Splits text into model-native token ID chunks that fit the context window.""" + + tokenizer: Any + max_input_tokens: int + prefix: List[int] + suffix: List[int] + + @property + def budget(self) -> int: + """Content tokens allowed per chunk, after reserving room for the affixes.""" + return max(1, self.max_input_tokens - len(self.prefix) - len(self.suffix)) + + def chunks(self, text: str) -> List[List[int]]: + """Encode ``text`` into wrapped, in-budget token ID chunks.""" + content = list(self.tokenizer.encode(text, add_special_tokens=False)) + budget = self.budget + # Every chunk is wrapped individually. Slicing an already-wrapped encoding + # would put the special suffix on the last chunk only, leaving every earlier + # chunk pooled at the wrong position. + return [ + self.prefix + content[i:i + budget] + self.suffix + for i in range(0, max(len(content), 1), budget) + ] + + +def _resolve_max_input_tokens( + policy: _ModelPolicy, + info: Any, + override: Optional[int], +) -> int: + """Resolve the context window, preferring caller and server values over the policy. + + ``max_input_tokens`` is read off ``info`` defensively so that the registry constant + is superseded automatically if the inference API ever starts reporting the window, + without needing another SDK release. + """ + if override is not None: + return int(override) + from_info = getattr(info, 'max_input_tokens', None) + if from_info: + return int(from_info) + return policy.max_input_tokens + + +def _token_chunker_for( + model_name: str, + hosting_platform: Optional[str], + info: Any = None, + max_input_tokens: Optional[int] = None, +) -> Optional[_TokenChunker]: + """Build a token chunker for a model, or None to keep character chunking. + + Returns None when the model is not in the registry, when its route does not accept + token IDs, or when the tokenizer cannot be loaded. + """ + policy = _MODEL_POLICIES.get(model_name.strip().lower()) + if policy is None or not policy.send_token_ids: + return None + if hosting_platform not in policy.token_id_platforms: + return None + + tokenizer_name = policy.tokenizer_name or model_name + try: + tokenizer = _load_tokenizer(tokenizer_name) + except Exception as exc: + # Any failure -- transformers missing, blocked egress, hub outage, renamed + # repo -- degrades to character chunking with text on the wire. That is + # correct, just coarser. Warn so the degradation is not silent: an egress + # change switching this feature off invisibly is the failure class this + # tokenization work exists to eliminate. + warnings.warn( + f'Could not load tokenizer {tokenizer_name!r} for model ' + f'{model_name!r} ({type(exc).__name__}: {exc}). Falling back to ' + 'character-based chunking with raw text; long inputs may be chunked ' + 'less precisely.', + ) + return None + + prefix, suffix = _derive_special_affixes(tokenizer) + return _TokenChunker( + tokenizer=tokenizer, + max_input_tokens=_resolve_max_input_tokens(policy, info, max_input_tokens), + prefix=prefix, + suffix=suffix, + ) + + +_Chunk = Union[str, List[int]] + class _ChunkedOpenAIEmbeddings(OpenAIEmbeddings): """OpenAIEmbeddings for non-OpenAI models behind an OpenAI-compatible endpoint. - These models (e.g. Qwen served on the 'Nova' platform) tokenize server-side with - their own tokenizer, so inputs are sent as raw text (``check_embedding_ctx_length`` - should be False). Because the server rejects (or silently truncates) inputs longer - than its context window, this class splits long inputs into character-bounded chunks - itself, embeds each chunk, and length-weighted-averages them back into a single - vector per input -- irrespective of the flag -- so long texts never hit the server's - hard limit. + tiktoken is the wrong tokenizer for these models (e.g. Qwen served on the 'Nova' + platforms), so ``check_embedding_ctx_length`` should be False to keep langchain + from encoding with it. Because the server rejects (or silently truncates) inputs + longer than its context window, this class splits long inputs into chunks itself, + embeds each chunk, and weighted-averages them back into a single vector per input + -- irrespective of the flag -- so long texts never hit the server's hard limit. + + With a ``token_chunker`` set, chunking uses the model's own tokenizer and sends + token IDs. Without one, it falls back to a coarse character split and sends text + for the server to tokenize. """ - max_chunk_chars: int = 24000 - """Maximum characters per chunk. This is a coarse character-based guard for - models whose exact tokenizer/context metadata is not yet available to the client. - Override per model if the deployment's context window is known to be smaller or - larger.""" + max_chunk_chars: int = 6000 + """Maximum characters per chunk, used only when ``token_chunker`` is None. + + Coarse character-based guard used because the client does not have the model's + tokenizer. Sized to stay under common ~8k-token Nova embedding windows even when + characters map roughly 1:1 to tokens (code / CJK). Models with larger windows + (e.g. Qwen3 Embedding ~32k) can raise this; models with smaller windows (e.g. + ~4k) should lower it. + """ - def _chunks(self, text: str) -> List[str]: + token_chunker: Optional[Any] = None + """A :class:`_TokenChunker`, or None to chunk by characters and send raw text.""" + + def _chunks(self, text: str) -> List[_Chunk]: + if self.token_chunker is not None: + return list(self.token_chunker.chunks(text)) n = max(1, self.max_chunk_chars) if len(text) <= n: return [text] return [text[i:i + n] for i in range(0, len(text), n)] + def _weight(self, chunk: _Chunk) -> int: + """Weight of a chunk in the reduction, in units of content.""" + if self.token_chunker is None: + return max(1, len(chunk)) + affix_len = len(self.token_chunker.prefix) + len(self.token_chunker.suffix) + return max(1, len(chunk) - affix_len) + @staticmethod def _average(vectors: List[List[float]], weights: List[int]) -> List[float]: total = float(sum(weights)) or 1.0 @@ -70,20 +267,25 @@ def _average(vectors: List[List[float]], weights: List[int]) -> List[float]: avg = [x / norm for x in avg] return avg - def _plan(self, texts: List[str]) -> Tuple[List[str], List[int]]: - flat: List[str] = [] + def _plan( + self, texts: List[str], + ) -> Tuple[List[_Chunk], List[int], List[int]]: + """Split every input into chunks, tracking which input each chunk came from.""" + flat: List[_Chunk] = [] owner: List[int] = [] + weights: List[int] = [] for i, text in enumerate(texts): for chunk in self._chunks(text): flat.append(chunk) owner.append(i) - return flat, owner + weights.append(self._weight(chunk)) + return flat, owner, weights def _reduce( self, num_texts: int, owner: List[int], - flat: List[str], + weights: List[int], embeddings: List[List[float]], ) -> List[List[float]]: out: List[List[float]] = [] @@ -95,7 +297,7 @@ def _reduce( out.append( self._average( [embeddings[j] for j in idxs], - [max(1, len(flat[j])) for j in idxs], + [weights[j] for j in idxs], ), ) return out @@ -103,16 +305,22 @@ def _reduce( def embed_documents( self, texts: List[str], chunk_size: Optional[int] = None, **kwargs: Any, ) -> List[List[float]]: - flat, owner = self._plan(texts) - embeddings = super().embed_documents(flat, chunk_size=chunk_size, **kwargs) - return self._reduce(len(texts), owner, flat, embeddings) + flat, owner, weights = self._plan(texts) + # langchain forwards batch elements untouched when check_embedding_ctx_length + # is False, so token ID lists reach the wire as-is despite the str signature. + embeddings = super().embed_documents( + flat, chunk_size=chunk_size, **kwargs, # type: ignore[arg-type] + ) + return self._reduce(len(texts), owner, weights, embeddings) async def aembed_documents( self, texts: List[str], chunk_size: Optional[int] = None, **kwargs: Any, ) -> List[List[float]]: - flat, owner = self._plan(texts) - embeddings = await super().aembed_documents(flat, chunk_size=chunk_size, **kwargs) - return self._reduce(len(texts), owner, flat, embeddings) + flat, owner, weights = self._plan(texts) + embeddings = await super().aembed_documents( + flat, chunk_size=chunk_size, **kwargs, # type: ignore[arg-type] + ) + return self._reduce(len(texts), owner, weights, embeddings) def SingleStoreEmbeddingsFactory( @@ -122,6 +330,7 @@ def SingleStoreEmbeddingsFactory( obo_token_getter: Optional[Callable[[], Optional[str]]] = None, base_url: Optional[str] = None, hosting_platform: Optional[str] = None, + max_input_tokens: Optional[int] = None, **kwargs: Any, ) -> Union[OpenAIEmbeddings, BedrockEmbeddings]: """Return an embeddings model instance (OpenAIEmbeddings or BedrockEmbeddings). @@ -249,10 +458,19 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: ) # Non-OpenAI models (e.g. Qwen on 'Nova'): tiktoken would send OpenAI token IDs the - # model can't interpret -> nonsensical embeddings. Send raw text so the server - # tokenizes with the model's own tokenizer, and chunk long inputs ourselves (the - # server otherwise rejects or silently truncates over-context input). + # model can't interpret -> nonsensical embeddings. Either encode with the model's + # own tokenizer client-side, or send raw text and let the server tokenize. Either + # way chunk long inputs ourselves, since the server otherwise rejects or silently + # truncates over-context input. kwargs.setdefault('check_embedding_ctx_length', False) + token_chunker = _token_chunker_for( + info.model_name, + info.hosting_platform, + info=info, + max_input_tokens=max_input_tokens, + ) + if token_chunker is not None: + kwargs['token_chunker'] = token_chunker return _ChunkedOpenAIEmbeddings( **openai_kwargs, **kwargs, diff --git a/singlestoredb/tests/test_embeddings.py b/singlestoredb/tests/test_embeddings.py index 4ff60ca5..6531937e 100644 --- a/singlestoredb/tests/test_embeddings.py +++ b/singlestoredb/tests/test_embeddings.py @@ -9,9 +9,29 @@ import types import unittest +# Arbitrary IDs outside the fake tokenizers' character-derived range, standing in for +# a model's BOS/EOS. The real Qwen3 EOS (151643) is deliberately not used here: the +# affixes are derived from the tokenizer, so no test should know a real special ID. +FAKE_BOS = 900001 +FAKE_EOS = 900002 + +LIVE_MODEL_ENV = 'SINGLESTOREDB_EMBEDDINGS_LIVE_MODEL' + +INJECTED_MODULES = ( + 'httpx', + 'langchain_openai', + 'langchain_aws', + 'botocore', + 'botocore.config', + 'boto3', + 'transformers', +) + class MockOpenAIEmbeddings: + specials = (FAKE_BOS, FAKE_EOS) + def __init__(self, **kwargs): self.kwargs = kwargs for key, value in kwargs.items(): @@ -19,11 +39,15 @@ def __init__(self, **kwargs): self.seen_documents = [] self.async_seen_documents = [] - @staticmethod - def _embedding_for(text): - if text.startswith('a'): + @classmethod + def _embedding_for(cls, item): + if not isinstance(item, str): + item = ''.join( + chr(token) for token in item if token not in cls.specials + ) + if item.startswith('a'): return [1.0, 0.0] - if text.startswith('b'): + if item.startswith('b'): return [0.0, 1.0] return [0.0, -1.0] @@ -75,10 +99,41 @@ def client(self, *args, **kwargs): ) +class FakeTokenizer: + """One token per character, plus configurable special-token affixes.""" + + def __init__(self, prefix=(), suffix=()): + self.prefix = list(prefix) + self.suffix = list(suffix) + + def encode(self, text, add_special_tokens=True): + tokens = [ord(char) for char in text] + if add_special_tokens: + return self.prefix + tokens + self.suffix + return tokens + + +class FakeAutoTokenizer: + """Stands in for ``transformers.AutoTokenizer`` so unit tests stay offline.""" + + tokenizer = None + error = None + + @classmethod + def from_pretrained(cls, name, **kwargs): + if cls.error is not None: + raise cls.error + return cls.tokenizer + + class TestEmbeddings(unittest.TestCase): @classmethod def setUpClass(cls): + cls.saved_modules = { + name: sys.modules.get(name) for name in INJECTED_MODULES + } + sys.modules.pop('singlestoredb.ai.embeddings', None) sys.modules.pop('_test_embeddings_module', None) @@ -105,6 +160,10 @@ def setUpClass(cls): sys.modules['boto3'] = MockBoto3('boto3') + transformers = types.ModuleType('transformers') + transformers.AutoTokenizer = FakeAutoTokenizer + sys.modules['transformers'] = transformers + path = os.path.join(os.path.dirname(__file__), '..', 'ai', 'embeddings.py') spec = importlib.util.spec_from_file_location('_test_embeddings_module', path) module = importlib.util.module_from_spec(spec) @@ -113,7 +172,41 @@ def setUpClass(cls): spec.loader.exec_module(module) cls.embeddings = module - def test_non_azure_factory_sends_raw_strings_and_uses_chunk_cap(self): + @classmethod + def tearDownClass(cls): + # The live check below, and anything else importing these for real, must not + # inherit the fakes. + sys.modules.pop('_test_embeddings_module', None) + for name, module in cls.saved_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + def tearDown(self): + # _load_tokenizer memoizes, so each test must start from an empty cache. + self.embeddings._load_tokenizer.cache_clear() + FakeAutoTokenizer.tokenizer = None + FakeAutoTokenizer.error = None + + def use_tokenizer(self, prefix=(), suffix=()): + FakeAutoTokenizer.tokenizer = FakeTokenizer(prefix=prefix, suffix=suffix) + + def fail_tokenizer_load(self, error): + FakeAutoTokenizer.error = error + + def qwen_embedding(self, **kwargs): + return self.embeddings.SingleStoreEmbeddingsFactory( + model_name='Qwen/Qwen3-Embedding-0.6B', + api_key='token', + base_url='http://localhost:8000', + hosting_platform='NovaMultiTenant', + **kwargs, + ) + + def test_unregistered_model_sends_raw_strings_and_uses_chunk_cap(self): + # A model with no registry entry keeps the pre-tokenization behavior: text on + # the wire, split on characters. embedding = self.embeddings.SingleStoreEmbeddingsFactory( model_name='shared-qwen3-embed-0-6b', api_key='token', @@ -123,10 +216,12 @@ def test_non_azure_factory_sends_raw_strings_and_uses_chunk_cap(self): assert isinstance(embedding, self.embeddings._ChunkedOpenAIEmbeddings) assert embedding.kwargs['check_embedding_ctx_length'] is False - assert embedding.max_chunk_chars == 24000, embedding.max_chunk_chars + assert embedding.token_chunker is None + assert embedding.max_chunk_chars == 6000, embedding.max_chunk_chars - embedding.embed_documents(['a' * 24001]) - assert [len(x) for x in embedding.seen_documents] == [24000, 1] + embedding.embed_documents(['a' * 6001]) + assert all(isinstance(x, str) for x in embedding.seen_documents) + assert [len(x) for x in embedding.seen_documents] == [6000, 1] def test_azure_factory_keeps_langchain_tokenization(self): embedding = self.embeddings.SingleStoreEmbeddingsFactory( @@ -176,6 +271,208 @@ async def run(): asyncio.run(run()) + def test_registry_entry_for_qwen3_embedding(self): + self.use_tokenizer(suffix=[FAKE_EOS]) + + chunker = self.embeddings._token_chunker_for( + 'Qwen/Qwen3-Embedding-0.6B', 'Nova', + ) + + assert chunker is not None + assert chunker.max_input_tokens == 32768, chunker.max_input_tokens + assert chunker.prefix == [] + assert chunker.suffix == [FAKE_EOS] + assert chunker.budget == 32767, chunker.budget + + def test_token_path_wraps_every_chunk_and_stays_within_budget(self): + self.use_tokenizer(suffix=[FAKE_EOS]) + embedding = self.qwen_embedding(max_input_tokens=4) + + assert embedding.token_chunker is not None + assert embedding.token_chunker.budget == 3 + + embedding.embed_documents(['abcdefg']) + + sent = embedding.seen_documents + assert len(sent) == 3, sent + for chunk in sent: + assert isinstance(chunk, list), chunk + assert all(isinstance(token, int) for token in chunk), chunk + assert len(chunk) <= 4, chunk + # Every chunk carries the affix, not just the last one. Slicing a wrapped + # encoding instead would mispool every chunk but the final one. + assert chunk[-1] == FAKE_EOS, chunk + assert sent == [ + [ord('a'), ord('b'), ord('c'), FAKE_EOS], + [ord('d'), ord('e'), ord('f'), FAKE_EOS], + [ord('g'), FAKE_EOS], + ], sent + + def test_token_path_weights_reduction_by_content_tokens_only(self): + self.use_tokenizer(prefix=[FAKE_BOS], suffix=[FAKE_EOS]) + embedding = self.qwen_embedding(max_input_tokens=5) + + assert embedding.token_chunker.budget == 3 + + out = embedding.embed_documents(['aaab']) + + # Chunks weigh 3 and 1 content tokens; counting the two affix tokens as well + # would weigh them 5 and 3 and pull the result toward the shorter chunk. + assert len(out) == 1, out + assert math.isclose(out[0][0], 3.0 / math.sqrt(10.0)), out[0] + assert math.isclose(out[0][1], 1.0 / math.sqrt(10.0)), out[0] + + def test_affix_derivation_covers_prefix_suffix_and_neither(self): + derive = self.embeddings._derive_special_affixes + + assert derive(FakeTokenizer(suffix=[FAKE_EOS])) == ([], [FAKE_EOS]) + assert derive(FakeTokenizer(prefix=[FAKE_BOS])) == ([FAKE_BOS], []) + assert derive( + FakeTokenizer(prefix=[FAKE_BOS], suffix=[FAKE_EOS]), + ) == ([FAKE_BOS], [FAKE_EOS]) + assert derive(FakeTokenizer()) == ([], []) + + def test_token_ids_refused_on_platforms_outside_the_allowlist(self): + self.use_tokenizer(suffix=[FAKE_EOS]) + + # The Bedrock route decodes integer inputs with tiktoken, so model-native IDs + # would be silently decoded into unrelated text and embedded. + assert self.embeddings._token_chunker_for( + 'Qwen/Qwen3-Embedding-0.6B', 'Amazon', + ) is None + assert self.embeddings._token_chunker_for( + 'Qwen/Qwen3-Embedding-0.6B', 'Nova', + ) is not None + + def test_tokenizer_load_failure_falls_back_and_warns(self): + self.fail_tokenizer_load(RuntimeError('huggingface.co unreachable')) + + with self.assertWarns(UserWarning) as caught: + embedding = self.qwen_embedding() + + assert 'huggingface.co unreachable' in str(caught.warning) + assert embedding.token_chunker is None + assert embedding.max_chunk_chars == 6000 + + embedding.embed_documents(['a' * 6001]) + assert all(isinstance(x, str) for x in embedding.seen_documents) + assert [len(x) for x in embedding.seen_documents] == [6000, 1] + + def test_max_input_tokens_override_beats_registry_and_info(self): + self.use_tokenizer(suffix=[FAKE_EOS]) + chunker_for = self.embeddings._token_chunker_for + info = types.SimpleNamespace(max_input_tokens=1024) + + assert chunker_for( + 'Qwen/Qwen3-Embedding-0.6B', 'Nova', max_input_tokens=512, + ).max_input_tokens == 512 + assert chunker_for( + 'Qwen/Qwen3-Embedding-0.6B', 'Nova', info=info, + ).max_input_tokens == 1024 + assert chunker_for( + 'Qwen/Qwen3-Embedding-0.6B', 'Nova', info=info, max_input_tokens=256, + ).max_input_tokens == 256 + + +@unittest.skipUnless( + os.environ.get(LIVE_MODEL_ENV), + f'set {LIVE_MODEL_ENV} to a deployed embedding model name to run the live ' + 'server-contract check', +) +class TestLiveTokenIdParity(unittest.TestCase): + """Server-contract tripwire against a real deployment; never runs in CI. + + Now that the client owns tokenization, this is what detects a vLLM upgrade, an + ``--hf-overrides`` change, or a model revision that alters tokenization or + pooling. Needs ``SINGLESTOREDB_USER_TOKEN`` plus either an org context or + ``SINGLESTOREDB_INFERENCE_API_BASE_URL`` and + ``SINGLESTOREDB_INFERENCE_API_HOSTING_PLATFORM``. + """ + + text = ( + 'SingleStore is a distributed SQL database that supports both ' + 'transactional and analytical workloads over the same data, with ' + 'vector search built in.' + ) + + @staticmethod + def cosine(left, right): + dot = sum(a * b for a, b in zip(left, right)) + left_norm = sum(a * a for a in left) ** 0.5 + right_norm = sum(b * b for b in right) ** 0.5 + return dot / (left_norm * right_norm) + + def native_embedding(self): + """An embeddings model on the token path, as the factory built it.""" + from singlestoredb.ai.embeddings import SingleStoreEmbeddingsFactory + + model_name = os.environ[LIVE_MODEL_ENV] + embedding = SingleStoreEmbeddingsFactory(model_name=model_name) + assert embedding.token_chunker is not None, ( + f'{model_name} did not take the token path; check its registry entry ' + 'and that the tokenizer loaded' + ) + return embedding + + def text_embedding(self): + """An embeddings model that puts raw text on the wire, as the baseline.""" + embedding = self.native_embedding() + embedding.token_chunker = None + return embedding + + def retokenized_embedding(self, tokenizer, prefix, suffix): + """An embeddings model that puts ``tokenizer``'s IDs on the wire.""" + import dataclasses + + embedding = self.native_embedding() + embedding.token_chunker = dataclasses.replace( + embedding.token_chunker, + tokenizer=tokenizer, + prefix=prefix, + suffix=suffix, + ) + return embedding + + def cosine_against_text(self, embedding): + return self.cosine( + self.text_embedding().embed_documents([self.text])[0], + embedding.embed_documents([self.text])[0], + ) + + def test_native_token_ids_match_raw_text(self): + cos = self.cosine_against_text(self.native_embedding()) + assert cos > 0.9999, cos + + def test_dropping_special_affixes_breaks_parity(self): + # Guards the LAST-pooling assumption: without the trailing special token the + # sentence vector becomes the hidden state of the last content token instead. + unwrapped = self.retokenized_embedding( + self.native_embedding().token_chunker.tokenizer, [], [], + ) + + cos = self.cosine_against_text(unwrapped) + assert cos < 0.99, ( + f'dropping the special affixes still matched raw text (cos={cos}); the ' + 'server-side tokenization or pooling contract has changed' + ) + + def test_tiktoken_ids_are_not_equivalent(self): + import tiktoken + + encoding = tiktoken.get_encoding('cl100k_base') + + class TiktokenShim: + def encode(self, text, add_special_tokens=True): + return encoding.encode(text) + + cos = self.cosine_against_text( + self.retokenized_embedding(TiktokenShim(), [], []), + ) + assert cos < 0.9, ( + f'tiktoken IDs matched raw text (cos={cos}); the server is no longer ' + 'interpreting the input as model-native token IDs' + ) + if __name__ == '__main__': unittest.main() From 86f541797f9d61dfec1ffb00eab68f0aead7d8ef Mon Sep 17 00:00:00 2001 From: Karish Date: Wed, 19 Aug 2026 18:25:49 +0530 Subject: [PATCH 2/5] [backend:feature-improvement] Minor refactor, limit handling and improved tests MCDB-98680 --- singlestoredb/ai/embeddings.py | 246 +++++++++++++------- singlestoredb/tests/test_embeddings.py | 231 ++++++++---------- singlestoredb/tests/test_embeddings_live.py | 153 ++++++++++++ 3 files changed, 408 insertions(+), 222 deletions(-) create mode 100644 singlestoredb/tests/test_embeddings_live.py diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index c9adbe87..9aac8a32 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -1,9 +1,9 @@ """LangChain embeddings models for SingleStore-hosted inference APIs. -Models served on the 'Nova' platforms are not OpenAI models, so tiktoken is the wrong -tokenizer for them. Where this module knows a model's real tokenizer it chunks and -encodes client-side and puts model-native token IDs on the wire; otherwise it falls -back to sending raw text in character-bounded chunks. +Nova-hosted models are not OpenAI models, so tiktoken is the wrong tokenizer for them. +For models in the registry below, this module encodes with the model's own tokenizer +and sends token IDs. For everything else it sends raw text in character-sized chunks +and lets the server tokenize. """ import os import warnings @@ -45,6 +45,36 @@ _DEFAULT_TOKEN_ID_PLATFORMS = frozenset({'Nova', 'NovaMultiTenant'}) +_AFFIX_PROBE = 'x' +"""Throwaway text used to compare a bare encode against a wrapped one.""" + +_WINDOW_SAFETY_MARGIN = 1 +"""Tokens held back so a full chunk stays strictly under the context window. + +Without it a full chunk is exactly ``max_input_tokens`` long, so an off-by-one in the +server's length check would reject only the longest inputs. +""" + + +class TokenizationFallbackWarning(UserWarning): + """A model could not use client-side tokenization. + + Embeddings stay correct, since the server tokenizes the raw text itself, but long + inputs are split on a coarse character budget instead of the real context window. + Every fallback warns, because a silent one looks exactly like success. Silence with + ``warnings.filterwarnings('ignore', category=TokenizationFallbackWarning)``. + """ + + +def _warn_fallback(model_name: str, reason: str) -> None: + warnings.warn( + f'Using character-based chunking with raw text for model {model_name!r} ' + f'because {reason}. Embeddings remain correct, but long inputs are split on a ' + f'coarse character budget rather than on the model context window.', + TokenizationFallbackWarning, + stacklevel=3, + ) + @dataclass(frozen=True) class _ModelPolicy: @@ -56,31 +86,32 @@ class _ModelPolicy: """HuggingFace repo to load the tokenizer from. None means use the model name.""" token_id_platforms: FrozenSet[str] = _DEFAULT_TOKEN_ID_PLATFORMS - """Hosting platforms whose route accepts token IDs for this model. + """Platforms whose route accepts token IDs for this model. - This is per-model-and-route rather than global: the 'Amazon' route decodes integer - ``input`` arrays with tiktoken, so model-native IDs sent that way are decoded into - unrelated text and embedded without error. + Default-deny, so a platform added later does not inherit the token path untested. + 'Amazon' and 'Azure' return earlier in the factory and never reach this check. + + 'NovaMultiTenant' is verified by a live parity run (Qwen3-Embedding-0.6B). 'Nova' + is inferred from serving the same image: tenancy changes routing and auth, not the + tokenizer inside the container. """ # Keyed on a lowercased HuggingFace repo id, which is what InferenceAPIInfo.model_name -# resolves to, so entries survive deployment alias renames. +# resolves to when the factory looks the model up through the management API. # -# Entries are explicit opt-in: there is no family-prefix or wildcard matching, even -# for models known to share a tokenizer. A model that reaches the token path without a -# parity run is the failure mode this registry exists to prevent -- mismatched special -# tokens return a well-formed, unit-norm, plausible-looking vector and raise nothing. -# To add a model: +# Opt-in only: no prefix or wildcard matching, even between models that share a +# tokenizer. Wrong token IDs do not raise -- they return a well-formed, unit-norm +# vector -- so no model reaches the token path unverified. To add one: # -# 1. Confirm the serving stack accepts token IDs on that route. vLLM types the -# embeddings input as list[int] | list[list[int]] | str | list[str]; Bedrock -# does not. -# 2. Confirm the context window actually served, including any --max-model-len -# override at launch, rather than the window the model card advertises. -# 3. Run test_live_token_id_parity from singlestoredb/tests/test_embeddings.py -# against a real deployment and require cosine > 0.9999. -# 4. Add the entry plus a unit test asserting its resolved budget and affixes. +# 1. Confirm the route accepts token IDs. vLLM does; Bedrock decodes them with +# tiktoken instead, turning them into unrelated text. +# 2. Confirm the window actually served, including any --max-model-len override at +# launch, not the one on the model card. Record the full window; the affixes and +# _WINDOW_SAFETY_MARGIN are subtracted from it. +# 3. Run tests/test_embeddings_live.py against a real deployment (set +# SINGLESTOREDB_EMBEDDINGS_LIVE_MODEL) and require cosine > 0.9999. +# 4. Add a unit test for the resolved budget and affixes. _MODEL_POLICIES: Dict[str, _ModelPolicy] = { 'qwen/qwen3-embedding-0.6b': _ModelPolicy( max_input_tokens=32768, @@ -91,11 +122,10 @@ class _ModelPolicy: @lru_cache(maxsize=None) def _load_tokenizer(tokenizer_name: str) -> Any: - """Load and memoize a HuggingFace tokenizer. + """Load a HuggingFace tokenizer, once per name per process. - Memoized because parsing Qwen3's ~11 MB tokenizer.json is slow, and because there - is no baked tokenizer cache in the serving image, so the first load in a container - fetches from huggingface.co over the network. + Memoized because parsing Qwen3's ~11 MB tokenizer.json is slow, and the serving + image bakes in no tokenizer cache, so the first load fetches from huggingface.co. """ from transformers import AutoTokenizer # type: ignore[import-not-found] return AutoTokenizer.from_pretrained(tokenizer_name) @@ -104,41 +134,71 @@ def _load_tokenizer(tokenizer_name: str) -> Any: def _derive_special_affixes(tokenizer: Any) -> Tuple[List[int], List[int]]: """Return the token IDs a tokenizer adds before and after content. - Derived by diffing a probe encode rather than hardcoded, so this holds for - BOS-style models too and self-corrects if a model revision or a ``transformers`` - upgrade changes special-token handling. + Measured by diffing a bare encode against a wrapped one instead of hardcoded, so + it covers BOS-style models too and follows any change in the model revision or in + ``transformers``. + + Raises: + ValueError: if the tokenizer does not wrap content in a fixed prefix and + suffix. Assuming no affixes here would cause the exact mispooling this + function prevents, so the caller must fall back instead. """ - bare = list(tokenizer.encode('x', add_special_tokens=False)) - wrapped = list(tokenizer.encode('x', add_special_tokens=True)) + bare = list(tokenizer.encode(_AFFIX_PROBE, add_special_tokens=False)) + wrapped = list(tokenizer.encode(_AFFIX_PROBE, add_special_tokens=True)) if not bare: - return [], [] + raise ValueError( + f'tokenizer encoded the probe {_AFFIX_PROBE!r} to no tokens, so its ' + 'special-token affixes cannot be derived', + ) for start in range(len(wrapped) - len(bare) + 1): if wrapped[start:start + len(bare)] == bare: return wrapped[:start], wrapped[start + len(bare):] - return [], [] + raise ValueError( + f'could not locate the bare probe encoding {bare} inside its wrapped form ' + f'{wrapped}, so this tokenizer does not simply surround content with a fixed ' + 'prefix and suffix; per-chunk special tokens cannot be reproduced safely', + ) @dataclass(frozen=True) class _TokenChunker: - """Splits text into model-native token ID chunks that fit the context window.""" + """Splits text into token ID chunks that fit the model's context window.""" tokenizer: Any max_input_tokens: int prefix: List[int] suffix: List[int] + def __post_init__(self) -> None: + if self.budget < 1: + raise ValueError( + f'max_input_tokens={self.max_input_tokens} leaves no room for content ' + f'after {len(self.prefix) + len(self.suffix)} special token(s) and a ' + f'{_WINDOW_SAFETY_MARGIN}-token safety margin', + ) + @property def budget(self) -> int: - """Content tokens allowed per chunk, after reserving room for the affixes.""" - return max(1, self.max_input_tokens - len(self.prefix) - len(self.suffix)) + """Content tokens per chunk, after the affixes and the safety margin. + + Not clamped on purpose: a window too small for the affixes means the caller or + the registry is wrong, and clamping would emit chunks larger than the window + it was asked to respect. + """ + return ( + self.max_input_tokens + - len(self.prefix) + - len(self.suffix) + - _WINDOW_SAFETY_MARGIN + ) def chunks(self, text: str) -> List[List[int]]: """Encode ``text`` into wrapped, in-budget token ID chunks.""" content = list(self.tokenizer.encode(text, add_special_tokens=False)) budget = self.budget - # Every chunk is wrapped individually. Slicing an already-wrapped encoding - # would put the special suffix on the last chunk only, leaving every earlier - # chunk pooled at the wrong position. + # Wrap each chunk on its own. Slicing an already-wrapped encoding would leave + # the suffix on the last chunk only, so every earlier chunk pools at the wrong + # position. return [ self.prefix + content[i:i + budget] + self.suffix for i in range(0, max(len(content), 1), budget) @@ -150,11 +210,10 @@ def _resolve_max_input_tokens( info: Any, override: Optional[int], ) -> int: - """Resolve the context window, preferring caller and server values over the policy. + """Resolve the context window: caller override first, then server, then policy. - ``max_input_tokens`` is read off ``info`` defensively so that the registry constant - is superseded automatically if the inference API ever starts reporting the window, - without needing another SDK release. + ``info`` is read defensively so the registry constant gives way automatically if + the inference API ever starts reporting the window, with no new SDK release. """ if override is not None: return int(override) @@ -172,40 +231,49 @@ def _token_chunker_for( ) -> Optional[_TokenChunker]: """Build a token chunker for a model, or None to keep character chunking. - Returns None when the model is not in the registry, when its route does not accept - token IDs, or when the tokenizer cannot be loaded. + Returns None if the model is not registered, if its route does not accept token + IDs, or if the tokenizer cannot be loaded or inspected. All three warn, since none + of them is visible in the embeddings themselves. """ - policy = _MODEL_POLICIES.get(model_name.strip().lower()) + key = model_name.strip().lower() + policy = _MODEL_POLICIES.get(key) if policy is None or not policy.send_token_ids: + _warn_fallback( + model_name, + f'no tokenization policy is registered under {key!r}. If this is a ' + 'deployment alias rather than a HuggingFace repo id, the registry cannot ' + 'match it', + ) return None if hosting_platform not in policy.token_id_platforms: + _warn_fallback( + model_name, + f'hosting platform {hosting_platform!r} is not known to accept token IDs ' + f'for this model (allowed: {sorted(policy.token_id_platforms)})', + ) return None tokenizer_name = policy.tokenizer_name or model_name try: tokenizer = _load_tokenizer(tokenizer_name) + prefix, suffix = _derive_special_affixes(tokenizer) + return _TokenChunker( + tokenizer=tokenizer, + max_input_tokens=_resolve_max_input_tokens(policy, info, max_input_tokens), + prefix=prefix, + suffix=suffix, + ) except Exception as exc: - # Any failure -- transformers missing, blocked egress, hub outage, renamed - # repo -- degrades to character chunking with text on the wire. That is - # correct, just coarser. Warn so the degradation is not silent: an egress - # change switching this feature off invisibly is the failure class this - # tokenization work exists to eliminate. - warnings.warn( - f'Could not load tokenizer {tokenizer_name!r} for model ' - f'{model_name!r} ({type(exc).__name__}: {exc}). Falling back to ' - 'character-based chunking with raw text; long inputs may be chunked ' - 'less precisely.', + # transformers missing, blocked egress, hub outage, renamed repo, unreadable + # special tokens, unusable window: all degrade to character chunking with text + # on the wire, which is correct but coarser. + _warn_fallback( + model_name, + f'tokenizer {tokenizer_name!r} could not be prepared ' + f'({type(exc).__name__}: {exc})', ) return None - prefix, suffix = _derive_special_affixes(tokenizer) - return _TokenChunker( - tokenizer=tokenizer, - max_input_tokens=_resolve_max_input_tokens(policy, info, max_input_tokens), - prefix=prefix, - suffix=suffix, - ) - _Chunk = Union[str, List[int]] @@ -213,26 +281,23 @@ def _token_chunker_for( class _ChunkedOpenAIEmbeddings(OpenAIEmbeddings): """OpenAIEmbeddings for non-OpenAI models behind an OpenAI-compatible endpoint. - tiktoken is the wrong tokenizer for these models (e.g. Qwen served on the 'Nova' - platforms), so ``check_embedding_ctx_length`` should be False to keep langchain - from encoding with it. Because the server rejects (or silently truncates) inputs - longer than its context window, this class splits long inputs into chunks itself, - embeds each chunk, and weighted-averages them back into a single vector per input - -- irrespective of the flag -- so long texts never hit the server's hard limit. + tiktoken is the wrong tokenizer for these models (e.g. Qwen on the 'Nova' + platforms), so ``check_embedding_ctx_length`` should be False to stop langchain + encoding with it. That also turns off langchain's own long-input handling, so this + class always chunks inputs itself, embeds each chunk, and weighted-averages them + back into one vector per input. Otherwise the server rejects, or silently + truncates, anything over its context window. - With a ``token_chunker`` set, chunking uses the model's own tokenizer and sends - token IDs. Without one, it falls back to a coarse character split and sends text - for the server to tokenize. + With a ``token_chunker`` it chunks by real tokens and sends token IDs. Without one + it splits on characters and sends text for the server to tokenize. """ max_chunk_chars: int = 6000 - """Maximum characters per chunk, used only when ``token_chunker`` is None. + """Characters per chunk when ``token_chunker`` is None. - Coarse character-based guard used because the client does not have the model's - tokenizer. Sized to stay under common ~8k-token Nova embedding windows even when - characters map roughly 1:1 to tokens (code / CJK). Models with larger windows - (e.g. Qwen3 Embedding ~32k) can raise this; models with smaller windows (e.g. - ~4k) should lower it. + A coarse stand-in for a token count, since the client has no tokenizer here. Sized + to stay under a common ~8k-token Nova window even when characters map nearly 1:1 + to tokens, as in code or CJK. Raise it for larger windows, lower it for ~4k ones. """ token_chunker: Optional[Any] = None @@ -247,7 +312,7 @@ def _chunks(self, text: str) -> List[_Chunk]: return [text[i:i + n] for i in range(0, len(text), n)] def _weight(self, chunk: _Chunk) -> int: - """Weight of a chunk in the reduction, in units of content.""" + """How much this chunk counts for in the average, in units of content.""" if self.token_chunker is None: return max(1, len(chunk)) affix_len = len(self.token_chunker.prefix) + len(self.token_chunker.suffix) @@ -270,7 +335,7 @@ def _average(vectors: List[List[float]], weights: List[int]) -> List[float]: def _plan( self, texts: List[str], ) -> Tuple[List[_Chunk], List[int], List[int]]: - """Split every input into chunks, tracking which input each chunk came from.""" + """Chunk every input, tracking which input each chunk came from.""" flat: List[_Chunk] = [] owner: List[int] = [] weights: List[int] = [] @@ -306,8 +371,9 @@ def embed_documents( self, texts: List[str], chunk_size: Optional[int] = None, **kwargs: Any, ) -> List[List[float]]: flat, owner, weights = self._plan(texts) - # langchain forwards batch elements untouched when check_embedding_ctx_length - # is False, so token ID lists reach the wire as-is despite the str signature. + # langchain passes batch elements through untouched when + # check_embedding_ctx_length is False, so token ID lists reach the wire as-is + # despite the str signature. embeddings = super().embed_documents( flat, chunk_size=chunk_size, **kwargs, # type: ignore[arg-type] ) @@ -448,9 +514,9 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: openai_kwargs['http_client'] = http_client if info.hosting_platform == 'Azure': - # Genuine OpenAI (Azure) models: tiktoken is the correct tokenizer, and the - # model name is passed above so it selects the right encoding. Keep langchain's - # client-side tokenization + long-input chunking (all correct for these models). + # Real OpenAI models: tiktoken is the right tokenizer, and the model name passed + # above picks the right encoding. Keep langchain's own tokenization and + # long-input chunking, which are both correct here. kwargs.setdefault('check_embedding_ctx_length', True) return OpenAIEmbeddings( **openai_kwargs, @@ -458,10 +524,10 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: ) # Non-OpenAI models (e.g. Qwen on 'Nova'): tiktoken would send OpenAI token IDs the - # model can't interpret -> nonsensical embeddings. Either encode with the model's - # own tokenizer client-side, or send raw text and let the server tokenize. Either - # way chunk long inputs ourselves, since the server otherwise rejects or silently - # truncates over-context input. + # model cannot read, giving meaningless embeddings. So either encode with the + # model's own tokenizer, or send raw text for the server to tokenize. Either way we + # chunk long inputs here, since the server rejects or silently truncates anything + # over its window. kwargs.setdefault('check_embedding_ctx_length', False) token_chunker = _token_chunker_for( info.model_name, diff --git a/singlestoredb/tests/test_embeddings.py b/singlestoredb/tests/test_embeddings.py index 6531937e..cb535887 100644 --- a/singlestoredb/tests/test_embeddings.py +++ b/singlestoredb/tests/test_embeddings.py @@ -9,14 +9,12 @@ import types import unittest -# Arbitrary IDs outside the fake tokenizers' character-derived range, standing in for -# a model's BOS/EOS. The real Qwen3 EOS (151643) is deliberately not used here: the -# affixes are derived from the tokenizer, so no test should know a real special ID. +# Stand-ins for a model's BOS/EOS, outside the fake tokenizers' ord()-derived range. +# Not the real Qwen3 EOS (151643): affixes come from the tokenizer, so no test should +# know a real special ID. FAKE_BOS = 900001 FAKE_EOS = 900002 -LIVE_MODEL_ENV = 'SINGLESTOREDB_EMBEDDINGS_LIVE_MODEL' - INJECTED_MODULES = ( 'httpx', 'langchain_openai', @@ -113,6 +111,22 @@ def encode(self, text, add_special_tokens=True): return tokens +class RewritingTokenizer: + """Rewrites content when adding specials instead of wrapping it.""" + + def encode(self, text, add_special_tokens=True): + if add_special_tokens: + return [FAKE_BOS, FAKE_EOS] + return [ord(char) for char in text] + + +class EmptyTokenizer: + """Encodes the probe to nothing, so no affixes can be located.""" + + def encode(self, text, add_special_tokens=True): + return [] + + class FakeAutoTokenizer: """Stands in for ``transformers.AutoTokenizer`` so unit tests stay offline.""" @@ -174,8 +188,8 @@ def setUpClass(cls): @classmethod def tearDownClass(cls): - # The live check below, and anything else importing these for real, must not - # inherit the fakes. + # Any later real import of these must not get the fakes. test_embeddings_live + # would pass against a constant vector if it did. sys.modules.pop('_test_embeddings_module', None) for name, module in cls.saved_modules.items(): if module is None: @@ -205,15 +219,21 @@ def qwen_embedding(self, **kwargs): ) def test_unregistered_model_sends_raw_strings_and_uses_chunk_cap(self): - # A model with no registry entry keeps the pre-tokenization behavior: text on - # the wire, split on characters. - embedding = self.embeddings.SingleStoreEmbeddingsFactory( - model_name='shared-qwen3-embed-0-6b', - api_key='token', - base_url='http://localhost:8000', - hosting_platform='NovaMultiTenant', - ) + # No registry entry keeps the old behavior: text on the wire, split on + # characters. A deployment alias lands here, so it must warn rather than look + # like success. + with self.assertWarns( + self.embeddings.TokenizationFallbackWarning, + ) as caught: + embedding = self.embeddings.SingleStoreEmbeddingsFactory( + model_name='shared-qwen3-embed-0-6b', + api_key='token', + base_url='http://localhost:8000', + hosting_platform='NovaMultiTenant', + ) + assert 'no tokenization policy is registered' in str(caught.warning) + assert 'deployment alias' in str(caught.warning) assert isinstance(embedding, self.embeddings._ChunkedOpenAIEmbeddings) assert embedding.kwargs['check_embedding_ctx_length'] is False assert embedding.token_chunker is None @@ -282,42 +302,46 @@ def test_registry_entry_for_qwen3_embedding(self): assert chunker.max_input_tokens == 32768, chunker.max_input_tokens assert chunker.prefix == [] assert chunker.suffix == [FAKE_EOS] - assert chunker.budget == 32767, chunker.budget + # 32768 window - 1 suffix token - 1 safety margin. + assert chunker.budget == 32766, chunker.budget def test_token_path_wraps_every_chunk_and_stays_within_budget(self): self.use_tokenizer(suffix=[FAKE_EOS]) embedding = self.qwen_embedding(max_input_tokens=4) assert embedding.token_chunker is not None - assert embedding.token_chunker.budget == 3 + assert embedding.token_chunker.budget == 2 embedding.embed_documents(['abcdefg']) sent = embedding.seen_documents - assert len(sent) == 3, sent + assert len(sent) == 4, sent for chunk in sent: assert isinstance(chunk, list), chunk assert all(isinstance(token, int) for token in chunk), chunk - assert len(chunk) <= 4, chunk - # Every chunk carries the affix, not just the last one. Slicing a wrapped - # encoding instead would mispool every chunk but the final one. + # Strictly under the window, so an off-by-one in the server's length check + # cannot reject the longest chunks. + assert len(chunk) < 4, chunk + # Every chunk carries the affix. Slicing a wrapped encoding would mispool + # all but the last one. assert chunk[-1] == FAKE_EOS, chunk assert sent == [ - [ord('a'), ord('b'), ord('c'), FAKE_EOS], - [ord('d'), ord('e'), ord('f'), FAKE_EOS], + [ord('a'), ord('b'), FAKE_EOS], + [ord('c'), ord('d'), FAKE_EOS], + [ord('e'), ord('f'), FAKE_EOS], [ord('g'), FAKE_EOS], ], sent def test_token_path_weights_reduction_by_content_tokens_only(self): self.use_tokenizer(prefix=[FAKE_BOS], suffix=[FAKE_EOS]) - embedding = self.qwen_embedding(max_input_tokens=5) + embedding = self.qwen_embedding(max_input_tokens=6) assert embedding.token_chunker.budget == 3 out = embedding.embed_documents(['aaab']) - # Chunks weigh 3 and 1 content tokens; counting the two affix tokens as well - # would weigh them 5 and 3 and pull the result toward the shorter chunk. + # Content weights are 3 and 1. Counting the two affix tokens too would make + # them 5 and 3, pulling the result toward the shorter chunk. assert len(out) == 1, out assert math.isclose(out[0][0], 3.0 / math.sqrt(10.0)), out[0] assert math.isclose(out[0][1], 1.0 / math.sqrt(10.0)), out[0] @@ -330,16 +354,57 @@ def test_affix_derivation_covers_prefix_suffix_and_neither(self): assert derive( FakeTokenizer(prefix=[FAKE_BOS], suffix=[FAKE_EOS]), ) == ([FAKE_BOS], [FAKE_EOS]) + # A tokenizer that adds nothing is a real match, not a failure. assert derive(FakeTokenizer()) == ([], []) + def test_affix_derivation_refuses_unrecognizable_tokenizers(self): + # Empty affixes here would send token IDs with no special tokens, the exact + # mispooling this derivation prevents. + derive = self.embeddings._derive_special_affixes + + with self.assertRaises(ValueError): + derive(RewritingTokenizer()) + with self.assertRaises(ValueError): + derive(EmptyTokenizer()) + + def test_unrecognizable_tokenizer_falls_back_and_warns(self): + self.embeddings._load_tokenizer.cache_clear() + FakeAutoTokenizer.tokenizer = RewritingTokenizer() + + with self.assertWarns( + self.embeddings.TokenizationFallbackWarning, + ) as caught: + embedding = self.qwen_embedding() + + assert 'could not be prepared' in str(caught.warning) + assert embedding.token_chunker is None + + def test_window_too_small_for_affixes_falls_back_and_warns(self): + # Clamping the budget instead would quietly emit chunks larger than the window. + self.use_tokenizer(prefix=[FAKE_BOS], suffix=[FAKE_EOS]) + + with self.assertWarns( + self.embeddings.TokenizationFallbackWarning, + ) as caught: + embedding = self.qwen_embedding(max_input_tokens=3) + + assert 'leaves no room for content' in str(caught.warning) + assert embedding.token_chunker is None + def test_token_ids_refused_on_platforms_outside_the_allowlist(self): self.use_tokenizer(suffix=[FAKE_EOS]) - # The Bedrock route decodes integer inputs with tiktoken, so model-native IDs - # would be silently decoded into unrelated text and embedded. - assert self.embeddings._token_chunker_for( - 'Qwen/Qwen3-Embedding-0.6B', 'Amazon', - ) is None + # Bedrock decodes integer inputs with tiktoken, so model-native IDs would + # quietly become unrelated text and get embedded. + with self.assertWarns( + self.embeddings.TokenizationFallbackWarning, + ) as caught: + refused = self.embeddings._token_chunker_for( + 'Qwen/Qwen3-Embedding-0.6B', 'Amazon', + ) + + assert refused is None + assert 'not known to accept token IDs' in str(caught.warning) assert self.embeddings._token_chunker_for( 'Qwen/Qwen3-Embedding-0.6B', 'Nova', ) is not None @@ -347,7 +412,9 @@ def test_token_ids_refused_on_platforms_outside_the_allowlist(self): def test_tokenizer_load_failure_falls_back_and_warns(self): self.fail_tokenizer_load(RuntimeError('huggingface.co unreachable')) - with self.assertWarns(UserWarning) as caught: + with self.assertWarns( + self.embeddings.TokenizationFallbackWarning, + ) as caught: embedding = self.qwen_embedding() assert 'huggingface.co unreachable' in str(caught.warning) @@ -374,105 +441,5 @@ def test_max_input_tokens_override_beats_registry_and_info(self): ).max_input_tokens == 256 -@unittest.skipUnless( - os.environ.get(LIVE_MODEL_ENV), - f'set {LIVE_MODEL_ENV} to a deployed embedding model name to run the live ' - 'server-contract check', -) -class TestLiveTokenIdParity(unittest.TestCase): - """Server-contract tripwire against a real deployment; never runs in CI. - - Now that the client owns tokenization, this is what detects a vLLM upgrade, an - ``--hf-overrides`` change, or a model revision that alters tokenization or - pooling. Needs ``SINGLESTOREDB_USER_TOKEN`` plus either an org context or - ``SINGLESTOREDB_INFERENCE_API_BASE_URL`` and - ``SINGLESTOREDB_INFERENCE_API_HOSTING_PLATFORM``. - """ - - text = ( - 'SingleStore is a distributed SQL database that supports both ' - 'transactional and analytical workloads over the same data, with ' - 'vector search built in.' - ) - - @staticmethod - def cosine(left, right): - dot = sum(a * b for a, b in zip(left, right)) - left_norm = sum(a * a for a in left) ** 0.5 - right_norm = sum(b * b for b in right) ** 0.5 - return dot / (left_norm * right_norm) - - def native_embedding(self): - """An embeddings model on the token path, as the factory built it.""" - from singlestoredb.ai.embeddings import SingleStoreEmbeddingsFactory - - model_name = os.environ[LIVE_MODEL_ENV] - embedding = SingleStoreEmbeddingsFactory(model_name=model_name) - assert embedding.token_chunker is not None, ( - f'{model_name} did not take the token path; check its registry entry ' - 'and that the tokenizer loaded' - ) - return embedding - - def text_embedding(self): - """An embeddings model that puts raw text on the wire, as the baseline.""" - embedding = self.native_embedding() - embedding.token_chunker = None - return embedding - - def retokenized_embedding(self, tokenizer, prefix, suffix): - """An embeddings model that puts ``tokenizer``'s IDs on the wire.""" - import dataclasses - - embedding = self.native_embedding() - embedding.token_chunker = dataclasses.replace( - embedding.token_chunker, - tokenizer=tokenizer, - prefix=prefix, - suffix=suffix, - ) - return embedding - - def cosine_against_text(self, embedding): - return self.cosine( - self.text_embedding().embed_documents([self.text])[0], - embedding.embed_documents([self.text])[0], - ) - - def test_native_token_ids_match_raw_text(self): - cos = self.cosine_against_text(self.native_embedding()) - assert cos > 0.9999, cos - - def test_dropping_special_affixes_breaks_parity(self): - # Guards the LAST-pooling assumption: without the trailing special token the - # sentence vector becomes the hidden state of the last content token instead. - unwrapped = self.retokenized_embedding( - self.native_embedding().token_chunker.tokenizer, [], [], - ) - - cos = self.cosine_against_text(unwrapped) - assert cos < 0.99, ( - f'dropping the special affixes still matched raw text (cos={cos}); the ' - 'server-side tokenization or pooling contract has changed' - ) - - def test_tiktoken_ids_are_not_equivalent(self): - import tiktoken - - encoding = tiktoken.get_encoding('cl100k_base') - - class TiktokenShim: - def encode(self, text, add_special_tokens=True): - return encoding.encode(text) - - cos = self.cosine_against_text( - self.retokenized_embedding(TiktokenShim(), [], []), - ) - assert cos < 0.9, ( - f'tiktoken IDs matched raw text (cos={cos}); the server is no longer ' - 'interpreting the input as model-native token IDs' - ) - - if __name__ == '__main__': unittest.main() diff --git a/singlestoredb/tests/test_embeddings_live.py b/singlestoredb/tests/test_embeddings_live.py new file mode 100644 index 00000000..f38b5984 --- /dev/null +++ b/singlestoredb/tests/test_embeddings_live.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python +# type: ignore +"""Live server-contract checks for singlestoredb.ai.embeddings. + +Not unit tests. They need a real deployment, real credentials, and network access, +and a green run is a fact about one model on one platform at one moment, not an +invariant of this code. They exist because no offline test can verify what the +token path assumes: that the server reads our integers as model-native token IDs. +A fake tokenizer agrees with whatever we assert. + +Kept out of test_embeddings.py on purpose. That module swaps fakes into sys.modules +for langchain_openai and transformers, so if its teardown were ever skipped these +would import the fakes and pass against a constant vector -- the one false green +this file cannot afford. + +Marked ``management`` like the other tests that need real cloud credentials, so +CI's ``-m 'not management'`` excludes them by policy rather than by whether an env +var happens to be unset. The skipUnless on top of that is just so a local run +without a deployment says what to set. + +Run against a live deployment, named by its deployment name not its HF id:: + + SINGLESTOREDB_EMBEDDINGS_LIVE_MODEL=shared-qwen3-embed-0-6b + SINGLESTOREDB_MANAGEMENT_TOKEN=... # resolves that name to the HF id + SINGLESTOREDB_PROJECT=... + SINGLESTOREDB_USER_TOKEN=... # authenticates the embeddings call + SINGLESTOREDB_URL=... # any value; skips the test container + +Keep SINGLESTOREDB_INFERENCE_API_BASE_URL unset. With it the factory builds the +model info itself, so the registry key and the wire ``model`` collapse into one +string and no value satisfies both. +""" +import os +import unittest + +import pytest + +LIVE_MODEL_ENV = 'SINGLESTOREDB_EMBEDDINGS_LIVE_MODEL' + + +@pytest.mark.management +@unittest.skipUnless( + os.environ.get(LIVE_MODEL_ENV), + f'set {LIVE_MODEL_ENV} to a deployed embedding model name to run the live ' + 'server-contract check', +) +class TestLiveTokenIdParity(unittest.TestCase): + """Catches a vLLM upgrade, an ``--hf-overrides`` change, or a model revision + that shifts tokenization or pooling.""" + + text = ( + 'SingleStore is a distributed SQL database that supports both ' + 'transactional and analytical workloads over the same data, with ' + 'vector search built in.' + ) + + @staticmethod + def cosine(left, right): + dot = sum(a * b for a, b in zip(left, right)) + left_norm = sum(a * a for a in left) ** 0.5 + right_norm = sum(b * b for b in right) ** 0.5 + return dot / (left_norm * right_norm) + + def native_embedding(self): + """An embeddings model on the token path, as the factory built it.""" + from singlestoredb.ai.embeddings import SingleStoreEmbeddingsFactory + + model_name = os.environ[LIVE_MODEL_ENV] + embedding = SingleStoreEmbeddingsFactory(model_name=model_name) + assert embedding.token_chunker is not None, ( + f'{model_name} did not take the token path; check its registry entry ' + 'and that the tokenizer loaded' + ) + return embedding + + def text_embedding(self): + """An embeddings model that puts raw text on the wire, as the baseline.""" + embedding = self.native_embedding() + embedding.token_chunker = None + return embedding + + def retokenized_embedding(self, tokenizer, prefix, suffix): + """An embeddings model that puts ``tokenizer``'s IDs on the wire.""" + import dataclasses + + embedding = self.native_embedding() + embedding.token_chunker = dataclasses.replace( + embedding.token_chunker, + tokenizer=tokenizer, + prefix=prefix, + suffix=suffix, + ) + return embedding + + def cosine_against_text(self, embedding): + return self.cosine( + self.text_embedding().embed_documents([self.text])[0], + embedding.embed_documents([self.text])[0], + ) + + def test_native_token_ids_match_raw_text(self): + cos = self.cosine_against_text(self.native_embedding()) + assert cos > 0.9999, cos + + def test_dropping_special_affixes_breaks_parity(self): + # Guards the LAST-pooling assumption: with no trailing special token, the + # sentence vector becomes the last content token's hidden state instead. + unwrapped = self.retokenized_embedding( + self.native_embedding().token_chunker.tokenizer, [], [], + ) + + cos = self.cosine_against_text(unwrapped) + assert cos < 0.99, ( + f'dropping the special affixes still matched raw text (cos={cos}); the ' + 'server-side tokenization or pooling contract has changed' + ) + + def test_tiktoken_ids_are_not_equivalent(self): + import tiktoken + + encoding = tiktoken.get_encoding('cl100k_base') + + class TiktokenShim: + def encode(self, text, add_special_tokens=True): + return encoding.encode(text) + + cos = self.cosine_against_text( + self.retokenized_embedding(TiktokenShim(), [], []), + ) + assert cos < 0.9, ( + f'tiktoken IDs matched raw text (cos={cos}); the server is no longer ' + 'interpreting the input as model-native token IDs' + ) + + def test_batched_inputs_match_single_inputs(self): + # Several inputs in one call put a list of token-ID arrays on the wire. The + # other tests send one input, so they only ever cover a one-element list. + embedding = self.native_embedding() + texts = [ + 'Vector search over transactional data.', + 'Distributed SQL with both columnstore and rowstore tables.', + 'Client-side tokenization for Nova-hosted embedding models.', + ] + + batched = embedding.embed_documents(texts) + assert len(batched) == len(texts), len(batched) + for text, got in zip(texts, batched): + cos = self.cosine(embedding.embed_documents([text])[0], got) + assert cos > 0.9999, (text, cos) + + +if __name__ == '__main__': + unittest.main() From ef1ae8310909e72484fddb4e79ad21a061e46403 Mon Sep 17 00:00:00 2001 From: Karish Date: Thu, 20 Aug 2026 15:59:52 +0530 Subject: [PATCH 3/5] [backend:feature-improvement] ci issues MCDB-98680 --- singlestoredb/ai/embeddings.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index 9aac8a32..34f07975 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -45,15 +45,13 @@ _DEFAULT_TOKEN_ID_PLATFORMS = frozenset({'Nova', 'NovaMultiTenant'}) +# Throwaway text used to compare a bare encode against a wrapped one. _AFFIX_PROBE = 'x' -"""Throwaway text used to compare a bare encode against a wrapped one.""" +# Tokens held back so a full chunk stays strictly under the context window. Without it +# a full chunk is exactly max_input_tokens long, so an off-by-one in the server's +# length check would reject only the longest inputs. _WINDOW_SAFETY_MARGIN = 1 -"""Tokens held back so a full chunk stays strictly under the context window. - -Without it a full chunk is exactly ``max_input_tokens`` long, so an off-by-one in the -server's length check would reject only the longest inputs. -""" class TokenizationFallbackWarning(UserWarning): From 11f6560ab5ca1bd43e03bafa42612d74459d70f4 Mon Sep 17 00:00:00 2001 From: Karish Date: Thu, 20 Aug 2026 16:31:55 +0530 Subject: [PATCH 4/5] [backend:feature-improvement] revert max_chunk_chars MCDB-98680 --- singlestoredb/ai/embeddings.py | 7 +++---- singlestoredb/tests/test_embeddings.py | 12 ++++++------ 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index 34f07975..b701c069 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -290,12 +290,11 @@ class always chunks inputs itself, embeds each chunk, and weighted-averages them it splits on characters and sends text for the server to tokenize. """ - max_chunk_chars: int = 6000 + max_chunk_chars: int = 24000 """Characters per chunk when ``token_chunker`` is None. - A coarse stand-in for a token count, since the client has no tokenizer here. Sized - to stay under a common ~8k-token Nova window even when characters map nearly 1:1 - to tokens, as in code or CJK. Raise it for larger windows, lower it for ~4k ones. + A coarse stand-in for a token count, since the client has no tokenizer here. + Override per model if the deployment's context window is known. """ token_chunker: Optional[Any] = None diff --git a/singlestoredb/tests/test_embeddings.py b/singlestoredb/tests/test_embeddings.py index cb535887..00269b86 100644 --- a/singlestoredb/tests/test_embeddings.py +++ b/singlestoredb/tests/test_embeddings.py @@ -237,11 +237,11 @@ def test_unregistered_model_sends_raw_strings_and_uses_chunk_cap(self): assert isinstance(embedding, self.embeddings._ChunkedOpenAIEmbeddings) assert embedding.kwargs['check_embedding_ctx_length'] is False assert embedding.token_chunker is None - assert embedding.max_chunk_chars == 6000, embedding.max_chunk_chars + assert embedding.max_chunk_chars == 24000, embedding.max_chunk_chars - embedding.embed_documents(['a' * 6001]) + embedding.embed_documents(['a' * 24001]) assert all(isinstance(x, str) for x in embedding.seen_documents) - assert [len(x) for x in embedding.seen_documents] == [6000, 1] + assert [len(x) for x in embedding.seen_documents] == [24000, 1] def test_azure_factory_keeps_langchain_tokenization(self): embedding = self.embeddings.SingleStoreEmbeddingsFactory( @@ -419,11 +419,11 @@ def test_tokenizer_load_failure_falls_back_and_warns(self): assert 'huggingface.co unreachable' in str(caught.warning) assert embedding.token_chunker is None - assert embedding.max_chunk_chars == 6000 + assert embedding.max_chunk_chars == 24000 - embedding.embed_documents(['a' * 6001]) + embedding.embed_documents(['a' * 24001]) assert all(isinstance(x, str) for x in embedding.seen_documents) - assert [len(x) for x in embedding.seen_documents] == [6000, 1] + assert [len(x) for x in embedding.seen_documents] == [24000, 1] def test_max_input_tokens_override_beats_registry_and_info(self): self.use_tokenizer(suffix=[FAKE_EOS]) From 364e0ac781d895813d2b4526ea11dd16034fba6e Mon Sep 17 00:00:00 2001 From: Karish Date: Thu, 20 Aug 2026 17:17:25 +0530 Subject: [PATCH 5/5] [backend:feature-improvement] Relax live batch cosine threshold MCDB-98680 Prod shared-qwen3-embed-0-6b returned 0.99987 for batched vs single inputs; vLLM padded GEMM noise, not a tokenizer contract miss. Co-authored-by: Cursor --- singlestoredb/tests/test_embeddings_live.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/singlestoredb/tests/test_embeddings_live.py b/singlestoredb/tests/test_embeddings_live.py index f38b5984..97e4db3e 100644 --- a/singlestoredb/tests/test_embeddings_live.py +++ b/singlestoredb/tests/test_embeddings_live.py @@ -146,7 +146,11 @@ def test_batched_inputs_match_single_inputs(self): assert len(batched) == len(texts), len(batched) for text, got in zip(texts, batched): cos = self.cosine(embedding.embed_documents([text])[0], got) - assert cos > 0.9999, (text, cos) + # Looser than the native-vs-text check: a batched vLLM forward pads + # mixed-length sequences into one GEMM, so the vectors are not + # bit-identical to three solo calls. Wrong token IDs would land far + # below this, as in test_tiktoken_ids_are_not_equivalent. + assert cos > 0.999, (text, cos) if __name__ == '__main__':