fix(embedding): bound vectorizer text input

Author: r266-tech

openviking/storage/vectordb/vectorize/vectorizer.py

@@ -1,7 +1,10 @@
 # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
 # SPDX-License-Identifier: AGPL-3.0
+import logging
 from typing import Any, Dict, List, Optional, Tuple, TypedDict
 
+logger = logging.getLogger(__name__)
+
 
 class DenseMeta(TypedDict, total=False):
     ModelName: str
@@ -22,6 +25,45 @@ class VectorizeMeta(TypedDict, total=False):
     Sparse: SparseMeta
 
 
+def _safe_vectorizer_config(vectorizer: Any) -> Optional[Dict[str, Any]]:
+    try:
+        config = getattr(vectorizer, "config", None)
+    except AttributeError:
+        return None
+    except Exception as exc:
+        raise RuntimeError("failed to read vectorizer config") from exc
+    if config is None:
+        return None
+    return config if isinstance(config, dict) else None
+
+
+def _resolve_runtime_max_input_tokens(config: Optional[Dict[str, Any]]) -> Optional[int]:
+    if config is None:
+        return None
+    raw_value = config.get("max_input_tokens")
+    if raw_value is None or isinstance(raw_value, bool):
+        if isinstance(raw_value, bool):
+            logger.warning("Invalid vectorizer max_input_tokens=%r; limit disabled", raw_value)
+        return None
+    if isinstance(raw_value, int):
+        max_input_tokens = raw_value
+    elif isinstance(raw_value, str) and raw_value.strip().isdigit():
+        max_input_tokens = int(raw_value)
+    else:
+        logger.warning("Invalid vectorizer max_input_tokens=%r; limit disabled", raw_value)
+        return None
+    return max_input_tokens if max_input_tokens > 0 else None
+
+
+def _truncate_provider_text(text: str, max_input_tokens: int) -> str:
+    try:
+        from openviking.utils.embedding_input import truncate_embedding_input
+    except ImportError as exc:
+        raise RuntimeError("embedding input truncation is unavailable") from exc
+
+    return truncate_embedding_input(text, max_input_tokens)
+
+
 class VectorizerAdapter:
     """Adapter for vectorizer to handle data vectorization.
 
@@ -42,6 +84,9 @@ def __init__(self, vectorizer: Any, vectorize_meta: VectorizeMeta):
         self.image_field = dense_meta.get("ImageField", "")
         self.video_field = dense_meta.get("VideoField", "")
         self.vectorizer = vectorizer
+        self.max_input_tokens = _resolve_runtime_max_input_tokens(
+            _safe_vectorizer_config(vectorizer)
+        )
         sparse_meta = vectorize_meta.get("Sparse", {})
         self.dense_model = {
             "name": dense_meta.get("ModelName", ""),
@@ -59,6 +104,14 @@ def __init__(self, vectorizer: Any, vectorize_meta: VectorizeMeta):
         )
         self.dim = self.vectorizer.get_dense_vector_dim(self.dense_model, self.sparse_model)
 
+    def _prepare_text(self, text: Any) -> Any:
+        if self.max_input_tokens is None or not isinstance(text, str):
+            return text
+        # The vectorizer provider request has one text slot shared by dense and
+        # sparse generation. Keep stored raw fields untouched, but bound the
+        # provider-facing text so one oversized file cannot fail the whole write.
+        return _truncate_provider_text(text, self.max_input_tokens)
+
     def get_dim(self) -> int:
         """Get the dimension of the dense vector.
 
@@ -84,7 +137,7 @@ def vectorize_raw_data(
         for raw_data in raw_data_list:
             data = {}
             if self.text_field in raw_data:
-                data["text"] = raw_data[self.text_field]
+                data["text"] = self._prepare_text(raw_data[self.text_field])
             if self.image_field in raw_data:
                 data["image"] = raw_data[self.image_field]
             if self.video_field in raw_data:
@@ -110,7 +163,7 @@ def vectorize_one(
         """
         data = {}
         if text:
-            data["text"] = text
+            data["text"] = self._prepare_text(text)
         if image:
             data["image"] = image
         if video:

tests/unit/test_vectorize_file_strategy.py

@@ -1,9 +1,14 @@
+import builtins
+import logging
 import types
 
 import pytest
 
 from openviking.core.context import Context
-from openviking.utils import embedding_utils
+from openviking.storage.vectordb.vectorize.base import VectorizeResult
+from openviking.storage.vectordb.vectorize.vectorizer import VectorizerAdapter
+from openviking.utils import embedding_input, embedding_utils
+from openviking.utils.embedding_input import EMBEDDING_TRUNCATION_SUFFIX
 
 
 class DummyQueue:
@@ -64,6 +69,226 @@ def __init__(self):
         self.account_id = "default"
 
 
+class DummyVectorizer:
+    config = {"max_input_tokens": 20}
+
+    def __init__(self):
+        self.config = type(self).config
+        self.data = None
+
+    def get_dense_vector_dim(self, _dense_model, _sparse_model):
+        return 1
+
+    def vectorize_document(self, data, _dense_model, _sparse_model):
+        self.data = data
+        return VectorizeResult(dense_vectors=[[1.0] for _ in data])
+
+
+class DummyVectorizerWithoutLimit(DummyVectorizer):
+    config = {}
+
+
+class DummyVectorizerDisabledLimit(DummyVectorizer):
+    config = {"max_input_tokens": 0}
+
+
+class DummyVectorizerMalformedConfig(DummyVectorizer):
+    config = []
+
+
+class DummyVectorizerMalformedLimit(DummyVectorizer):
+    config = {"max_input_tokens": object()}
+
+
+class DummyVectorizerStringLimit(DummyVectorizer):
+    config = {"max_input_tokens": "20"}
+
+
+class DummyVectorizerBoolLimit(DummyVectorizer):
+    config = {"max_input_tokens": True}
+
+
+class DummyVectorizerFloatLimit(DummyVectorizer):
+    config = {"max_input_tokens": 20.5}
+
+
+class DummyVectorizerRaisingConfig(DummyVectorizer):
+    @property
+    def config(self):
+        raise RuntimeError("config unavailable")
+
+    def __init__(self):
+        self.data = None
+
+
+def test_vectorizer_adapter_truncates_provider_text_without_mutating_raw_data():
+    vectorizer = DummyVectorizer()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+    )
+    raw_text = "oversized memory content " * 80
+    raw_data = [{"content": raw_text, "uri": "viking://user/default/resources/big.md"}]
+
+    dense, sparse = adapter.vectorize_raw_data(raw_data)
+
+    assert dense == [[1.0]]
+    assert sparse == []
+    provider_text = vectorizer.data[0]["text"]
+    assert provider_text.endswith(EMBEDDING_TRUNCATION_SUFFIX)
+    assert len(provider_text) < len(raw_text)
+    assert raw_data[0]["content"] == raw_text
+
+
+def test_vectorizer_adapter_accepts_integer_string_limit():
+    vectorizer = DummyVectorizerStringLimit()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+    )
+    raw_text = "oversized memory content " * 80
+
+    adapter.vectorize_raw_data([{"content": raw_text}])
+
+    assert vectorizer.data[0]["text"].endswith(EMBEDDING_TRUNCATION_SUFFIX)
+
+
+def test_vectorizer_adapter_import_failure_fails_closed(monkeypatch):
+    real_import = builtins.__import__
+
+    def fail_embedding_input_import(name, *args, **kwargs):
+        if name == "openviking.utils.embedding_input":
+            raise ImportError("import failed")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", fail_embedding_input_import)
+    vectorizer = DummyVectorizer()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+    )
+    raw_text = "oversized memory content " * 80
+
+    with pytest.raises(RuntimeError, match="truncation is unavailable"):
+        adapter.vectorize_raw_data([{"content": raw_text}])
+
+
+def test_vectorizer_adapter_truncation_failure_fails_closed(monkeypatch):
+    def fail_truncation(_text, _max_input_tokens):
+        raise RuntimeError("truncate failed")
+
+    monkeypatch.setattr(embedding_input, "truncate_embedding_input", fail_truncation)
+    vectorizer = DummyVectorizer()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+    )
+    raw_text = "oversized memory content " * 80
+
+    with pytest.raises(RuntimeError, match="truncate failed"):
+        adapter.vectorize_raw_data([{"content": raw_text}])
+
+
+def test_vectorizer_adapter_raising_config_fails_closed():
+    with pytest.raises(RuntimeError, match="failed to read vectorizer config"):
+        VectorizerAdapter(
+            DummyVectorizerRaisingConfig(),
+            {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+        )
+
+
+def test_vectorizer_adapter_preserves_media_fields_when_truncating_text():
+    vectorizer = DummyVectorizer()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {
+            "Dense": {
+                "TextField": "content",
+                "ImageField": "image",
+                "VideoField": "video",
+                "ModelName": "dummy",
+            }
+        },
+    )
+    raw_text = "oversized memory content " * 80
+    image = {"uri": "data:image/png;base64,aaa"}
+    video = {"uri": "data:video/mp4;base64,bbb"}
+
+    adapter.vectorize_raw_data([{"content": raw_text, "image": image, "video": video}])
+
+    provider_data = vectorizer.data[0]
+    assert provider_data["text"].endswith(EMBEDDING_TRUNCATION_SUFFIX)
+    assert provider_data["image"] is image
+    assert provider_data["video"] is video
+
+
+def test_vectorizer_adapter_vectorize_one_truncates_text_and_preserves_media():
+    vectorizer = DummyVectorizer()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+    )
+    raw_text = "oversized memory content " * 80
+    image = {"uri": "data:image/png;base64,aaa"}
+    video = {"uri": "data:video/mp4;base64,bbb"}
+
+    adapter.vectorize_one(text=raw_text, image=image, video=video)
+
+    provider_data = vectorizer.data[0]
+    assert provider_data["text"].endswith(EMBEDDING_TRUNCATION_SUFFIX)
+    assert provider_data["image"] is image
+    assert provider_data["video"] is video
+
+
+@pytest.mark.parametrize(
+    "vectorizer_cls",
+    [
+        DummyVectorizerWithoutLimit,
+        DummyVectorizerDisabledLimit,
+        DummyVectorizerMalformedConfig,
+        DummyVectorizerMalformedLimit,
+        DummyVectorizerBoolLimit,
+        DummyVectorizerFloatLimit,
+    ],
+)
+def test_vectorizer_adapter_leaves_text_unchanged_without_explicit_limit(vectorizer_cls, caplog):
+    if vectorizer_cls in {
+        DummyVectorizerMalformedLimit,
+        DummyVectorizerBoolLimit,
+        DummyVectorizerFloatLimit,
+    }:
+        caplog.set_level(logging.WARNING)
+    vectorizer = vectorizer_cls()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+    )
+    raw_text = "oversized memory content " * 80
+
+    adapter.vectorize_raw_data([{"content": raw_text}])
+
+    assert vectorizer.data[0]["text"] == raw_text
+    if vectorizer_cls in {
+        DummyVectorizerMalformedLimit,
+        DummyVectorizerBoolLimit,
+        DummyVectorizerFloatLimit,
+    }:
+        assert "max_input_tokens disabled" in caplog.text
+
+
+def test_vectorizer_adapter_vectorize_one_leaves_text_unchanged_without_limit():
+    vectorizer = DummyVectorizerWithoutLimit()
+    adapter = VectorizerAdapter(
+        vectorizer,
+        {"Dense": {"TextField": "content", "ModelName": "dummy"}},
+    )
+    raw_text = "oversized memory content " * 80
+
+    adapter.vectorize_one(text=raw_text)
+
+    assert vectorizer.data[0]["text"] == raw_text
+
+
 @pytest.mark.asyncio
 async def test_vectorize_file_uses_summary_first(monkeypatch):
     queue = DummyQueue()