diff --git a/openviking/models/embedder/openai_embedders.py b/openviking/models/embedder/openai_embedders.py index a0a76e9e8f..909061bec3 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 @@ -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 from openviking_cli.utils import get_logger logger = get_logger(__name__) @@ -322,7 +323,10 @@ 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 + 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 5ca3090ef0..e812ce0f5f 100644 --- a/tests/unit/test_openai_embedder.py +++ b/tests/unit/test_openai_embedder.py @@ -87,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 @@ -101,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):