diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index ac2ced1f..b701c069 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -1,6 +1,18 @@ +"""LangChain embeddings models for SingleStore-hosted inference APIs. + +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 +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,278 @@ from botocore import UNSIGNED from botocore.config import Config +_DEFAULT_TOKEN_ID_PLATFORMS = frozenset({'Nova', 'NovaMultiTenant'}) + +# Throwaway text used to compare a bare encode against a wrapped one. +_AFFIX_PROBE = 'x' + +# 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 + + +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: + """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 + """Platforms whose route accepts token IDs for this model. + + 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 when the factory looks the model up through the management API. +# +# 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 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, + send_token_ids=True, + ), +} + + +@lru_cache(maxsize=None) +def _load_tokenizer(tokenizer_name: str) -> Any: + """Load a HuggingFace tokenizer, once per name per process. + + 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) + + +def _derive_special_affixes(tokenizer: Any) -> Tuple[List[int], List[int]]: + """Return the token IDs a tokenizer adds before and after content. + + 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(_AFFIX_PROBE, add_special_tokens=False)) + wrapped = list(tokenizer.encode(_AFFIX_PROBE, add_special_tokens=True)) + if not bare: + 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):] + 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 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 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 + # 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) + ] + + +def _resolve_max_input_tokens( + policy: _ModelPolicy, + info: Any, + override: Optional[int], +) -> int: + """Resolve the context window: caller override first, then server, then policy. + + ``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) + 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 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. + """ + 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: + # 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 + + +_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 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`` 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 = 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.""" + """Characters per chunk when ``token_chunker`` is None. - def _chunks(self, text: str) -> List[str]: + 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 + """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: + """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) + 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 +329,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]]: + """Chunk every input, 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 +359,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 +367,23 @@ 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 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] + ) + 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 +393,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). @@ -239,9 +511,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, @@ -249,10 +521,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 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, + 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..00269b86 100644 --- a/singlestoredb/tests/test_embeddings.py +++ b/singlestoredb/tests/test_embeddings.py @@ -9,9 +9,27 @@ import types import unittest +# 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 + +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 +37,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 +97,57 @@ 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 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.""" + + 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 +174,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,19 +186,61 @@ def setUpClass(cls): spec.loader.exec_module(module) cls.embeddings = module - def test_non_azure_factory_sends_raw_strings_and_uses_chunk_cap(self): - embedding = self.embeddings.SingleStoreEmbeddingsFactory( - model_name='shared-qwen3-embed-0-6b', + @classmethod + def tearDownClass(cls): + # 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: + 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): + # 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 assert embedding.max_chunk_chars == 24000, embedding.max_chunk_chars 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] == [24000, 1] def test_azure_factory_keeps_langchain_tokenization(self): @@ -176,6 +291,155 @@ 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] + # 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 == 2 + + embedding.embed_documents(['abcdefg']) + + sent = embedding.seen_documents + assert len(sent) == 4, sent + for chunk in sent: + assert isinstance(chunk, list), chunk + assert all(isinstance(token, int) for token in chunk), chunk + # 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'), 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=6) + + assert embedding.token_chunker.budget == 3 + + out = embedding.embed_documents(['aaab']) + + # 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] + + 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]) + # 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]) + + # 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 + + def test_tokenizer_load_failure_falls_back_and_warns(self): + self.fail_tokenizer_load(RuntimeError('huggingface.co unreachable')) + + with self.assertWarns( + self.embeddings.TokenizationFallbackWarning, + ) 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 == 24000 + + 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] == [24000, 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 + 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..97e4db3e --- /dev/null +++ b/singlestoredb/tests/test_embeddings_live.py @@ -0,0 +1,157 @@ +#!/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) + # 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__': + unittest.main()