diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index aba8e1c47..ac2ced1f5 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -1,7 +1,9 @@ import os from typing import Any from typing import Callable +from typing import List from typing import Optional +from typing import Tuple from typing import Union import httpx @@ -30,6 +32,89 @@ from botocore.config import Config +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. + """ + + 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.""" + + def _chunks(self, text: str) -> List[str]: + 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)] + + @staticmethod + def _average(vectors: List[List[float]], weights: List[int]) -> List[float]: + total = float(sum(weights)) or 1.0 + dim = len(vectors[0]) + avg = [0.0] * dim + for vec, w in zip(vectors, weights): + for k in range(dim): + avg[k] += vec[k] * w + avg = [x / total for x in avg] + norm = sum(x * x for x in avg) ** 0.5 + if norm > 0: + avg = [x / norm for x in avg] + return avg + + def _plan(self, texts: List[str]) -> Tuple[List[str], List[int]]: + flat: List[str] = [] + owner: List[int] = [] + for i, text in enumerate(texts): + for chunk in self._chunks(text): + flat.append(chunk) + owner.append(i) + return flat, owner + + def _reduce( + self, + num_texts: int, + owner: List[int], + flat: List[str], + embeddings: List[List[float]], + ) -> List[List[float]]: + out: List[List[float]] = [] + for i in range(num_texts): + idxs = [j for j, o in enumerate(owner) if o == i] + if len(idxs) == 1: + out.append(embeddings[idxs[0]]) + else: + out.append( + self._average( + [embeddings[j] for j in idxs], + [max(1, len(flat[j])) for j in idxs], + ), + ) + return out + + 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) + + 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) + + def SingleStoreEmbeddingsFactory( model_name: str, api_key: Optional[str] = None, @@ -152,7 +237,23 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: ) if http_client is not None: openai_kwargs['http_client'] = http_client - return OpenAIEmbeddings( + + 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). + kwargs.setdefault('check_embedding_ctx_length', True) + return OpenAIEmbeddings( + **openai_kwargs, + **kwargs, + ) + + # 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). + kwargs.setdefault('check_embedding_ctx_length', False) + return _ChunkedOpenAIEmbeddings( **openai_kwargs, **kwargs, ) diff --git a/singlestoredb/tests/test_embeddings.py b/singlestoredb/tests/test_embeddings.py new file mode 100644 index 000000000..4ff60ca5a --- /dev/null +++ b/singlestoredb/tests/test_embeddings.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python +# type: ignore +"""SingleStoreDB embeddings testing.""" +import asyncio +import importlib.util +import math +import os +import sys +import types +import unittest + + +class MockOpenAIEmbeddings: + + def __init__(self, **kwargs): + self.kwargs = kwargs + for key, value in kwargs.items(): + setattr(self, key, value) + self.seen_documents = [] + self.async_seen_documents = [] + + @staticmethod + def _embedding_for(text): + if text.startswith('a'): + return [1.0, 0.0] + if text.startswith('b'): + return [0.0, 1.0] + return [0.0, -1.0] + + def embed_documents(self, texts, chunk_size=None, **kwargs): + self.seen_documents.extend(texts) + self.seen_chunk_size = chunk_size + self.seen_kwargs = kwargs + return [self._embedding_for(text) for text in texts] + + async def aembed_documents(self, texts, chunk_size=None, **kwargs): + self.async_seen_documents.extend(texts) + self.async_seen_chunk_size = chunk_size + self.async_seen_kwargs = kwargs + return [self._embedding_for(text) for text in texts] + + +class MockBedrockEmbeddings: + + def __init__(self, **kwargs): + self.kwargs = kwargs + + +class MockConfig: + + def __init__(self, **kwargs): + self.kwargs = kwargs + + +class MockClient: + pass + + +class MockTimeout: + + def __init__(self, connect=None, read=None): + self.connect = connect + self.read = read + + +class MockBoto3(types.ModuleType): + + def client(self, *args, **kwargs): + return types.SimpleNamespace( + _endpoint=types.SimpleNamespace( + _event_emitter=types.SimpleNamespace( + register_first=lambda *args, **kwargs: None, + ), + ), + ) + + +class TestEmbeddings(unittest.TestCase): + + @classmethod + def setUpClass(cls): + sys.modules.pop('singlestoredb.ai.embeddings', None) + sys.modules.pop('_test_embeddings_module', None) + + httpx = types.ModuleType('httpx') + httpx.Client = MockClient + httpx.Timeout = MockTimeout + sys.modules['httpx'] = httpx + + langchain_openai = types.ModuleType('langchain_openai') + langchain_openai.OpenAIEmbeddings = MockOpenAIEmbeddings + sys.modules['langchain_openai'] = langchain_openai + + langchain_aws = types.ModuleType('langchain_aws') + langchain_aws.BedrockEmbeddings = MockBedrockEmbeddings + sys.modules['langchain_aws'] = langchain_aws + + botocore = types.ModuleType('botocore') + botocore.UNSIGNED = 'unsigned' + sys.modules['botocore'] = botocore + + botocore_config = types.ModuleType('botocore.config') + botocore_config.Config = MockConfig + sys.modules['botocore.config'] = botocore_config + + sys.modules['boto3'] = MockBoto3('boto3') + + 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) + sys.modules['_test_embeddings_module'] = module + assert spec.loader is not None, spec + 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', + api_key='token', + base_url='http://localhost:8000', + hosting_platform='NovaMultiTenant', + ) + + 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 + + embedding.embed_documents(['a' * 24001]) + assert [len(x) for x in embedding.seen_documents] == [24000, 1] + + def test_azure_factory_keeps_langchain_tokenization(self): + embedding = self.embeddings.SingleStoreEmbeddingsFactory( + model_name='text-embedding-3-small', + api_key='token', + base_url='http://localhost:8000', + hosting_platform='Azure', + ) + + assert isinstance(embedding, MockOpenAIEmbeddings) + assert not isinstance(embedding, self.embeddings._ChunkedOpenAIEmbeddings) + assert embedding.kwargs['check_embedding_ctx_length'] is True + + def test_chunked_embedding_reduces_weighted_average_to_one_vector(self): + embedding = self.embeddings._ChunkedOpenAIEmbeddings( + model='shared-qwen3-embed-0-6b', + api_key='token', + base_url='http://localhost:8000', + check_embedding_ctx_length=False, + ) + embedding.max_chunk_chars = 3 + + out = embedding.embed_documents(['aaabbb', 'zz']) + + assert embedding.seen_documents == ['aaa', 'bbb', 'zz'] + assert len(out) == 2, out + assert math.isclose(out[0][0], 1.0 / math.sqrt(2.0)), out[0] + assert math.isclose(out[0][1], 1.0 / math.sqrt(2.0)), out[0] + assert out[1] == [0.0, -1.0], out[1] + + def test_async_chunked_embedding_uses_same_reduction(self): + async def run(): + embedding = self.embeddings._ChunkedOpenAIEmbeddings( + model='shared-qwen3-embed-0-6b', + api_key='token', + base_url='http://localhost:8000', + check_embedding_ctx_length=False, + ) + embedding.max_chunk_chars = 3 + + out = await embedding.aembed_documents(['aaabbb']) + + assert embedding.async_seen_documents == ['aaa', 'bbb'] + assert len(out) == 1, out + assert math.isclose(out[0][0], 1.0 / math.sqrt(2.0)), out[0] + assert math.isclose(out[0][1], 1.0 / math.sqrt(2.0)), out[0] + + asyncio.run(run()) + + +if __name__ == '__main__': + unittest.main()