diff --git a/application/core/settings.py b/application/core/settings.py index 65b9314c..a5d5b8ec 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -48,6 +48,7 @@ class Settings(BaseSettings): EMBEDDINGS_NAME: str = "huggingface_sentence-transformers/all-mpnet-base-v2" EMBEDDINGS_BASE_URL: Optional[str] = None # Remote embeddings API URL (OpenAI-compatible) EMBEDDINGS_KEY: Optional[str] = None # api key for embeddings (if using openai, just copy API_KEY) + EMBEDDINGS_MAX_INPUT_TOKENS: Optional[int] = None # truncate each remote embed input to N tokens (overflow lost) # Optional directory of operator-supplied model YAMLs, loaded after the # built-in catalog under application/core/models/. Later wins on # duplicate model id. See application/core/models/README.md. diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py index cd7a0e6b..142f75f8 100644 --- a/application/vectorstore/base.py +++ b/application/vectorstore/base.py @@ -6,6 +6,7 @@ import requests from langchain_openai import OpenAIEmbeddings from application.core.settings import settings +from application.utils import get_encoding class RemoteEmbeddings: @@ -23,8 +24,51 @@ class RemoteEmbeddings: self.headers["Authorization"] = f"Bearer {api_key}" self.dimension = 768 + def _truncate_inputs(self, inputs): + """Clip each input to ``EMBEDDINGS_MAX_INPUT_TOKENS`` tokens. + + The remote server (e.g. llama.cpp) hard-rejects any single input + larger than its physical batch size with a 500. When the setting is + configured, each input is truncated to that many tokens before the + request and the overflow is dropped (lossy by design). Token counts + use the shared tiktoken encoding, which differs from the server's + tokenizer, so set the limit with headroom under the server's true + limit to absorb tokenizer skew. + + Args: + inputs: A single string or a list of strings to embed. + + Returns: + The inputs with each string clipped to the token limit, or the + inputs unchanged when the limit is unset or non-positive. + """ + limit = settings.EMBEDDINGS_MAX_INPUT_TOKENS + if not limit or limit <= 0: + return inputs + + encoding = get_encoding() + + def clip(text): + if not isinstance(text, str): + return text + tokens = encoding.encode(text) + if len(tokens) <= limit: + return text + logging.warning( + "Truncating remote embeddings input from %d to %d tokens (%d dropped)", + len(tokens), + limit, + len(tokens) - limit, + ) + return encoding.decode(tokens[:limit]) + + if isinstance(inputs, list): + return [clip(text) for text in inputs] + return clip(inputs) + def _embed(self, inputs): """Send embedding request to remote API in OpenAI-compatible format.""" + inputs = self._truncate_inputs(inputs) payload = {"input": inputs} if self.model_name: payload["model"] = self.model_name diff --git a/application/worker.py b/application/worker.py index 6727f693..fea14c75 100755 --- a/application/worker.py +++ b/application/worker.py @@ -1217,7 +1217,7 @@ def remote_worker( min_tokens=cfg.chunking.min_tokens, duplicate_headers=cfg.chunking.duplicate_headers, ) - docs = chunker.chunk(documents=raw_docs) + raw_docs = chunker.chunk(documents=raw_docs) docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs] tokens = count_tokens_docs(docs) logging.info("Total tokens calculated: %d", tokens) diff --git a/tests/vectorstore/test_remote_embeddings_truncation.py b/tests/vectorstore/test_remote_embeddings_truncation.py new file mode 100644 index 00000000..5d1bf700 --- /dev/null +++ b/tests/vectorstore/test_remote_embeddings_truncation.py @@ -0,0 +1,85 @@ +"""Tests for the ``EMBEDDINGS_MAX_INPUT_TOKENS`` truncation net. + +The remote embeddings server (e.g. llama.cpp) hard-rejects any single input +larger than its physical batch size with a 500. When the setting is +configured, ``RemoteEmbeddings`` clips each input to that many tokens before +the request; the overflow is dropped (lossy by design). +""" + +from unittest.mock import MagicMock + +from application.core.settings import settings +from application.utils import get_encoding +from application.vectorstore import base +from application.vectorstore.base import RemoteEmbeddings + + +def _capture_post(monkeypatch): + """Patch ``requests.post`` and return a dict recording the sent payload.""" + captured = {} + + def fake_post(url, headers=None, json=None, timeout=None): + captured["payload"] = json + n_inputs = len(json["input"]) if isinstance(json["input"], list) else 1 + resp = MagicMock() + resp.raise_for_status.return_value = None + resp.json.return_value = { + "data": [{"index": i, "embedding": [0.0]} for i in range(n_inputs)] + } + return resp + + monkeypatch.setattr(base.requests, "post", fake_post) + return captured + + +def test_truncates_oversized_input_to_limit(monkeypatch): + monkeypatch.setattr(settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 10) + captured = _capture_post(monkeypatch) + enc = get_encoding() + + long_text = " ".join(["word"] * 100) # ~100 tokens, far over the limit of 10 + emb = RemoteEmbeddings(api_url="https://example.test", model_name="m") + emb.embed_documents([long_text]) + + sent = captured["payload"]["input"][0] + assert sent == enc.decode(enc.encode(long_text)[:10]) + assert len(enc.encode(sent)) <= 10 + + +def test_short_input_is_unchanged(monkeypatch): + monkeypatch.setattr(settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 10) + captured = _capture_post(monkeypatch) + + short_text = "hello world" + emb = RemoteEmbeddings(api_url="https://example.test", model_name="m") + emb.embed_documents([short_text]) + + assert captured["payload"]["input"][0] == short_text + + +def test_no_truncation_when_setting_unset(monkeypatch): + monkeypatch.setattr(settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None) + captured = _capture_post(monkeypatch) + enc = get_encoding() + + long_text = " ".join(["word"] * 100) + emb = RemoteEmbeddings(api_url="https://example.test", model_name="m") + emb.embed_documents([long_text]) + + sent = captured["payload"]["input"][0] + assert sent == long_text + assert len(enc.encode(sent)) > 10 + + +def test_query_path_is_truncated(monkeypatch): + """``embed_query`` passes a bare string through the same net.""" + monkeypatch.setattr(settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 10) + captured = _capture_post(monkeypatch) + enc = get_encoding() + + long_text = " ".join(["word"] * 100) + emb = RemoteEmbeddings(api_url="https://example.test", model_name="m") + emb.embed_query(long_text) + + sent = captured["payload"]["input"] + assert sent == enc.decode(enc.encode(long_text)[:10])