From 054e1d5414a0537f86922061db8baa2a0af669ad Mon Sep 17 00:00:00 2001 From: akushonkamen Date: Thu, 8 Oct 2026 21:51:51 +0800 Subject: [PATCH 1/2] fix(embedding): truncate plain list inputs to max_input_tokens (#5760) The OpenAI-compatible embedder's _prepare_embedding_input passed plain list[str] inputs straight through to the provider, skipping the max_input_tokens guard entirely, while the base prepare_embedding_input collapsed the same input to an empty string for text-only embedders and left items unbounded for multimodal embedders. The embedding queue contract explicitly allows list inputs, so any producer could submit oversized items that fail with provider errors (e.g. bigmodel 1210). Both guard layers now keep the list shape and bound every item individually via truncate_embedding_input; max_input_tokens=None keeps the historical passthrough. Single-string and dict-part multimodal behavior is unchanged. --- openviking/models/embedder/base.py | 36 +++++- .../models/embedder/openai_embedders.py | 7 +- tests/unit/test_openai_embedder.py | 108 ++++++++++++++++++ 3 files changed, 145 insertions(+), 6 deletions(-) diff --git a/openviking/models/embedder/base.py b/openviking/models/embedder/base.py index efb9ec1369..3036a1f516 100644 --- a/openviking/models/embedder/base.py +++ b/openviking/models/embedder/base.py @@ -10,7 +10,18 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from threading import Lock -from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union +from typing import ( + Any, + AsyncIterator, + Awaitable, + Callable, + Dict, + List, + Optional, + TypeVar, + Union, + cast, +) from openviking.telemetry import get_current_telemetry from openviking.utils.embedding_input import ( @@ -309,11 +320,26 @@ def __init__(self, model_name: str, config: Optional[Dict[str, Any]] = None): def prepare_embedding_input(self, content: "EmbeddingInput") -> "EmbeddingInput": """Apply this embedder's input guard before provider calls. - Plain text is truncated to ``max_input_tokens``. For embedders that - support images, multimodal inputs (a list of content parts) are kept as a - list. Embedders that do not support multimodal input get the - text parts extracted, so image parts are safely dropped. + Plain text is truncated to ``max_input_tokens``. A plain text batch (a + list of strings) keeps its list shape and every item is truncated + individually, matching the provider's per-item input limits. For + embedders that support images, multimodal inputs (a list of content + parts) are kept as a list. Embedders that do not support multimodal + input get the text parts extracted, so image parts are safely dropped. """ + if isinstance(content, list) and all(isinstance(item, str) for item in content): + if self.max_input_tokens is None: + return content + # mypy cannot narrow ``EmbeddingInput`` items via ``all(isinstance(...))``, + # so rebind them as plain strings before the per-item truncation. + text_items: List[str] = [] + for item in content: + if isinstance(item, str): + text_items.append(item) + return cast( + "EmbeddingInput", + [truncate_embedding_input(item, self.max_input_tokens) for item in text_items], + ) if isinstance(content, list) and self.supports_multimodal: if self.max_input_tokens is None: return content diff --git a/openviking/models/embedder/openai_embedders.py b/openviking/models/embedder/openai_embedders.py index c4c122d201..c084c684a5 100644 --- a/openviking/models/embedder/openai_embedders.py +++ b/openviking/models/embedder/openai_embedders.py @@ -22,6 +22,7 @@ from openviking.models.vlm.registry import DEFAULT_AZURE_API_VERSION from openviking.telemetry import get_current_telemetry from openviking.utils.async_client_cache import LoopScopedAsyncClientCache +from openviking.utils.embedding_input import truncate_embedding_input logger = get_logger(__name__) @@ -321,7 +322,11 @@ def supports_multimodal(self) -> bool: def _prepare_embedding_input(self, text: EmbeddingInput | List[str]) -> EmbeddingInput | List[str]: if isinstance(text, list) and all(isinstance(item, str) for item in text): - return text + if self.max_input_tokens is None: + return text + # mypy cannot narrow via ``all(isinstance(...))``; rebind as strings. + text_items: List[str] = [item for item in text if isinstance(item, str)] + return [truncate_embedding_input(item, self.max_input_tokens) for item in text_items] return self.prepare_embedding_input(text) def _build_kwargs( diff --git a/tests/unit/test_openai_embedder.py b/tests/unit/test_openai_embedder.py index 5ca3090ef0..5f97031b30 100644 --- a/tests/unit/test_openai_embedder.py +++ b/tests/unit/test_openai_embedder.py @@ -5,6 +5,10 @@ from unittest.mock import MagicMock, patch from openviking.models.embedder import OpenAIDenseEmbedder +from openviking.utils.embedding_input import ( + EMBEDDING_TRUNCATION_SUFFIX, + estimate_embedding_input_tokens, +) class TestOpenAIDenseEmbedder: @@ -277,3 +281,107 @@ def test_telemetry_called_when_module_available(self, mock_openai_class): assert result.dense_vector is not None mock_telemetry.add_token_usage_by_source.assert_called_once_with("embedding", 8, 0) + + +class TestEmbeddingInputTruncationGuard: + """Plain list[str] inputs must respect max_input_tokens at both guard layers. + + Regression tests for the truncation guard being skipped for plain text + batches: the OpenAI-compatible lower guard passed them through untouched, + while the base guard collapsed them to an empty string for text-only + embedders and left them unbounded for multimodal embedders. + """ + + OVERSIZED_TEXT = "x" * 5000 + TOKEN_BUDGET = 64 + + def _make_embedder(self, input_type=None, max_input_tokens=TOKEN_BUDGET): + config = {} + if max_input_tokens is not None: + config["max_input_tokens"] = max_input_tokens + kwargs = { + "model_name": "embedding-3", + "api_key": "test-api-key", + "dimension": 1536, + "config": config, + } + if input_type is not None: + kwargs["input_type"] = input_type + return OpenAIDenseEmbedder(**kwargs) + + def _assert_within_budget(self, item): + """Each truncated item keeps its content within the budget (suffix overhead allowed).""" + assert item.endswith(EMBEDDING_TRUNCATION_SUFFIX.lstrip()) + suffix_budget = estimate_embedding_input_tokens(EMBEDDING_TRUNCATION_SUFFIX) + assert estimate_embedding_input_tokens(item) <= self.TOKEN_BUDGET + suffix_budget + + def test_lower_guard_truncates_plain_list_items(self): + """_prepare_embedding_input should bound every item of a plain text batch.""" + embedder = self._make_embedder() + result = embedder._prepare_embedding_input([self.OVERSIZED_TEXT, "short text"]) + + assert isinstance(result, list) + assert len(result) == 2 + self._assert_within_budget(result[0]) + assert result[1] == "short text" + + def test_lower_guard_keeps_list_when_no_budget(self): + """Without max_input_tokens the plain text batch is returned as-is.""" + embedder = self._make_embedder(max_input_tokens=None) + batch = ["first text", "second text"] + + assert embedder._prepare_embedding_input(batch) == batch + + def test_upper_guard_does_not_collapse_plain_list_for_text_only(self): + """Text-only embedders must not fold a plain text batch into an empty string.""" + embedder = self._make_embedder() + assert embedder.supports_multimodal is False + + result = embedder.prepare_embedding_input([self.OVERSIZED_TEXT]) + + assert isinstance(result, list) + assert len(result) == 1 + self._assert_within_budget(result[0]) + + def test_upper_guard_truncates_plain_list_for_multimodal(self): + """Multimodal embedders must not pass plain text batch items through unbounded.""" + embedder = self._make_embedder(input_type="multimodal") + assert embedder.supports_multimodal is True + + result = embedder.prepare_embedding_input([self.OVERSIZED_TEXT]) + + assert isinstance(result, list) + assert len(result) == 1 + self._assert_within_budget(result[0]) + + def test_embed_request_carries_truncated_list(self): + """embed() must send the truncated batch, not the raw oversized one.""" + mock_client = MagicMock() + with patch( + "openviking.models.embedder.openai_embedders.openai.OpenAI" + ) as mock_openai_class: + mock_openai_class.return_value = mock_client + mock_embedding = MagicMock() + mock_embedding.embedding = [0.1] * 1536 + mock_response = MagicMock() + mock_response.data = [mock_embedding] + mock_client.embeddings.create.return_value = mock_response + + embedder = self._make_embedder() + embedder.embed([self.OVERSIZED_TEXT, "short text"]) + + call_kwargs = mock_client.embeddings.create.call_args[1] + sent_input = call_kwargs["input"] + assert isinstance(sent_input, list) + assert len(sent_input) == 2 + self._assert_within_budget(sent_input[0]) + assert sent_input[1] == "short text" + + def test_single_string_truncation_unchanged(self): + """The single-string path keeps its existing truncation behavior.""" + embedder = self._make_embedder() + + result = embedder.prepare_embedding_input(self.OVERSIZED_TEXT) + + assert isinstance(result, str) + self._assert_within_budget(result) From 4b238836803b468baf27211b0b712db4883fd799 Mon Sep 17 00:00:00 2001 From: qin-ctx Date: Thu, 15 Oct 2026 21:13:30 +0800 Subject: [PATCH 2/2] fix(embedding): narrow batch truncation to OpenAI input guard --- openviking/models/embedder/base.py | 36 +----- .../models/embedder/openai_embedders.py | 7 +- tests/unit/test_openai_embedder.py | 116 +----------------- 3 files changed, 12 insertions(+), 147 deletions(-) diff --git a/openviking/models/embedder/base.py b/openviking/models/embedder/base.py index 3036a1f516..efb9ec1369 100644 --- a/openviking/models/embedder/base.py +++ b/openviking/models/embedder/base.py @@ -10,18 +10,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from threading import Lock -from typing import ( - Any, - AsyncIterator, - Awaitable, - Callable, - Dict, - List, - Optional, - TypeVar, - Union, - cast, -) +from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union from openviking.telemetry import get_current_telemetry from openviking.utils.embedding_input import ( @@ -320,26 +309,11 @@ def __init__(self, model_name: str, config: Optional[Dict[str, Any]] = None): def prepare_embedding_input(self, content: "EmbeddingInput") -> "EmbeddingInput": """Apply this embedder's input guard before provider calls. - Plain text is truncated to ``max_input_tokens``. A plain text batch (a - list of strings) keeps its list shape and every item is truncated - individually, matching the provider's per-item input limits. For - embedders that support images, multimodal inputs (a list of content - parts) are kept as a list. Embedders that do not support multimodal - input get the text parts extracted, so image parts are safely dropped. + Plain text is truncated to ``max_input_tokens``. For embedders that + support images, multimodal inputs (a list of content parts) are kept as a + list. Embedders that do not support multimodal input get the + text parts extracted, so image parts are safely dropped. """ - if isinstance(content, list) and all(isinstance(item, str) for item in content): - if self.max_input_tokens is None: - return content - # mypy cannot narrow ``EmbeddingInput`` items via ``all(isinstance(...))``, - # so rebind them as plain strings before the per-item truncation. - text_items: List[str] = [] - for item in content: - if isinstance(item, str): - text_items.append(item) - return cast( - "EmbeddingInput", - [truncate_embedding_input(item, self.max_input_tokens) for item in text_items], - ) if isinstance(content, list) and self.supports_multimodal: if self.max_input_tokens is None: return content diff --git a/openviking/models/embedder/openai_embedders.py b/openviking/models/embedder/openai_embedders.py index c084c684a5..e7169468b5 100644 --- a/openviking/models/embedder/openai_embedders.py +++ b/openviking/models/embedder/openai_embedders.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: AGPL-3.0 """OpenAI Embedder Implementation""" -from typing import Any, Dict, List, Literal, Optional +from typing import Any, Dict, List, Literal, Optional, cast import httpx import openai @@ -324,9 +324,8 @@ def _prepare_embedding_input(self, text: EmbeddingInput | List[str]) -> Embeddin if isinstance(text, list) and all(isinstance(item, str) for item in text): if self.max_input_tokens is None: return text - # mypy cannot narrow via ``all(isinstance(...))``; rebind as strings. - text_items: List[str] = [item for item in text if isinstance(item, str)] - return [truncate_embedding_input(item, self.max_input_tokens) for item in text_items] + text = cast(List[str], text) + return [truncate_embedding_input(item, self.max_input_tokens) for item in text] return self.prepare_embedding_input(text) def _build_kwargs( diff --git a/tests/unit/test_openai_embedder.py b/tests/unit/test_openai_embedder.py index 5f97031b30..e812ce0f5f 100644 --- a/tests/unit/test_openai_embedder.py +++ b/tests/unit/test_openai_embedder.py @@ -5,10 +5,6 @@ from unittest.mock import MagicMock, patch from openviking.models.embedder import OpenAIDenseEmbedder -from openviking.utils.embedding_input import ( - EMBEDDING_TRUNCATION_SUFFIX, - estimate_embedding_input_tokens, -) class TestOpenAIDenseEmbedder: @@ -91,7 +87,7 @@ def test_default_input_type_embed_downgrades_multimodal_parts(self, mock_openai_ @patch("openviking.models.embedder.openai_embedders.openai.OpenAI") def test_embed_preserves_openai_batch_text_input(self, mock_openai_class): - """OpenAI batch text input should keep the historical list[str] passthrough.""" + """OpenAI batches keep item boundaries and truncate each oversized item.""" mock_client = MagicMock() mock_openai_class.return_value = mock_client @@ -105,13 +101,13 @@ def test_embed_preserves_openai_batch_text_input(self, mock_openai_class): embedder = OpenAIDenseEmbedder( model_name="text-embedding-3-small", api_key="test-api-key", - dimension=1536, + config={"max_input_tokens": 1}, ) - embedder.embed(["first text", "second text"]) + embedder.embed(["first text", "x"]) call_kwargs = mock_client.embeddings.create.call_args[1] - assert call_kwargs["input"] == ["first text", "second text"] + assert call_kwargs["input"] == ["firs\n...(truncated for embedding)", "x"] @patch("openviking.models.embedder.openai_embedders.openai.OpenAI") def test_multimodal_input_type_passes_content_parts(self, mock_openai_class): @@ -281,107 +277,3 @@ def test_telemetry_called_when_module_available(self, mock_openai_class): assert result.dense_vector is not None mock_telemetry.add_token_usage_by_source.assert_called_once_with("embedding", 8, 0) - - -class TestEmbeddingInputTruncationGuard: - """Plain list[str] inputs must respect max_input_tokens at both guard layers. - - Regression tests for the truncation guard being skipped for plain text - batches: the OpenAI-compatible lower guard passed them through untouched, - while the base guard collapsed them to an empty string for text-only - embedders and left them unbounded for multimodal embedders. - """ - - OVERSIZED_TEXT = "x" * 5000 - TOKEN_BUDGET = 64 - - def _make_embedder(self, input_type=None, max_input_tokens=TOKEN_BUDGET): - config = {} - if max_input_tokens is not None: - config["max_input_tokens"] = max_input_tokens - kwargs = { - "model_name": "embedding-3", - "api_key": "test-api-key", - "dimension": 1536, - "config": config, - } - if input_type is not None: - kwargs["input_type"] = input_type - return OpenAIDenseEmbedder(**kwargs) - - def _assert_within_budget(self, item): - """Each truncated item keeps its content within the budget (suffix overhead allowed).""" - assert item.endswith(EMBEDDING_TRUNCATION_SUFFIX.lstrip()) - suffix_budget = estimate_embedding_input_tokens(EMBEDDING_TRUNCATION_SUFFIX) - assert estimate_embedding_input_tokens(item) <= self.TOKEN_BUDGET + suffix_budget - - def test_lower_guard_truncates_plain_list_items(self): - """_prepare_embedding_input should bound every item of a plain text batch.""" - embedder = self._make_embedder() - result = embedder._prepare_embedding_input([self.OVERSIZED_TEXT, "short text"]) - - assert isinstance(result, list) - assert len(result) == 2 - self._assert_within_budget(result[0]) - assert result[1] == "short text" - - def test_lower_guard_keeps_list_when_no_budget(self): - """Without max_input_tokens the plain text batch is returned as-is.""" - embedder = self._make_embedder(max_input_tokens=None) - batch = ["first text", "second text"] - - assert embedder._prepare_embedding_input(batch) == batch - - def test_upper_guard_does_not_collapse_plain_list_for_text_only(self): - """Text-only embedders must not fold a plain text batch into an empty string.""" - embedder = self._make_embedder() - assert embedder.supports_multimodal is False - - result = embedder.prepare_embedding_input([self.OVERSIZED_TEXT]) - - assert isinstance(result, list) - assert len(result) == 1 - self._assert_within_budget(result[0]) - - def test_upper_guard_truncates_plain_list_for_multimodal(self): - """Multimodal embedders must not pass plain text batch items through unbounded.""" - embedder = self._make_embedder(input_type="multimodal") - assert embedder.supports_multimodal is True - - result = embedder.prepare_embedding_input([self.OVERSIZED_TEXT]) - - assert isinstance(result, list) - assert len(result) == 1 - self._assert_within_budget(result[0]) - - def test_embed_request_carries_truncated_list(self): - """embed() must send the truncated batch, not the raw oversized one.""" - mock_client = MagicMock() - with patch( - "openviking.models.embedder.openai_embedders.openai.OpenAI" - ) as mock_openai_class: - mock_openai_class.return_value = mock_client - mock_embedding = MagicMock() - mock_embedding.embedding = [0.1] * 1536 - mock_response = MagicMock() - mock_response.data = [mock_embedding] - mock_client.embeddings.create.return_value = mock_response - - embedder = self._make_embedder() - embedder.embed([self.OVERSIZED_TEXT, "short text"]) - - call_kwargs = mock_client.embeddings.create.call_args[1] - sent_input = call_kwargs["input"] - assert isinstance(sent_input, list) - assert len(sent_input) == 2 - self._assert_within_budget(sent_input[0]) - assert sent_input[1] == "short text" - - def test_single_string_truncation_unchanged(self): - """The single-string path keeps its existing truncation behavior.""" - embedder = self._make_embedder() - - result = embedder.prepare_embedding_input(self.OVERSIZED_TEXT) - - assert isinstance(result, str) - self._assert_within_budget(result)