fix(embedding): bound truncation estimate above real tokenizer costs

openviking/utils/embedding_input.py

@@ -8,44 +8,73 @@
 
 EMBEDDING_TRUNCATION_SUFFIX = "\n...(truncated for embedding)"
 _CJK_RUN = re.compile("[\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]+")
+_CYRILLIC_RUN = re.compile("[\u0400-\u052f]+")
+_SCRIPT_RUN = re.compile(f"({_CJK_RUN.pattern})|({_CYRILLIC_RUN.pattern})")
+
+# Costs are tracked in tenths of a token so every per-character rate stays an
+# exact integer. Tokenizers of the embedding models commonly served behind
+# OpenAI-compatible endpoints spend far more than four characters per token on
+# non-English text: measured against llama-server's /tokenize, Cyrillic prose
+# needs roughly 2.7 characters per token and even English prose and shell code
+# land near 3. A flat "4 characters per token" estimate therefore lets text
+# truncated to ``max_input_tokens`` arrive with 30-50% more real tokens and get
+# rejected by the provider. The rates below sit above those observed costs, so
+# a document truncated under this estimate stays within the real limit.
+_TENTHS_PER_TOKEN = 10
+_CJK_TENTHS = 10  # CJK: one token per character.
+_CYRILLIC_TENTHS = 5  # Cyrillic: two characters per token.
+_ORDINARY_TENTHS = 4  # Everything else: 2.5 characters per token.
+_MAX_TENTHS_PER_CHAR = _CJK_TENTHS
+
+
+def _run_tenths(match: re.Match[str]) -> int:
+    """Return the per-character cost of a script run matched by ``_SCRIPT_RUN``."""
+    return _CJK_TENTHS if match.group(1) is not None else _CYRILLIC_TENTHS
 
 
 def truncate_embedding_input(
     text: str,
     max_tokens: int,
     suffix: str = EMBEDDING_TRUNCATION_SUFFIX,
 ) -> str:
-    """Keep the original estimated-token prefix budget, then append the marker."""
+    """Keep the longest prefix whose estimated tokens plus ``suffix`` fit the budget."""
     if not text:
         return text
     if max_tokens <= 0:
         return suffix.lstrip()
-    if len(text) <= max_tokens:
-        return text
 
-    # Quarter-token units preserve the original estimate without rounding each
-    # run. No prefix longer than 4 * max_tokens can fit, even with only ASCII.
-    max_units = max_tokens * 4
+    # Reserve room for the marker so the final payload still fits the budget.
+    suffix_units = len(suffix) * _ORDINARY_TENTHS
+    max_units = max(max_tokens * _TENTHS_PER_TOKEN - suffix_units, 0)
+    if len(text) * _MAX_TENTHS_PER_CHAR <= max_units:
+        return text
     if text.isascii():
-        return text if len(text) <= max_units else text[:max_units].rstrip() + suffix
-    scan_end = min(len(text), max_units + 1)
+        # ASCII is charged exactly _ORDINARY_TENTHS per character.
+        return (
+            text
+            if len(text) * _ORDINARY_TENTHS <= max_units
+            else text[: max_units // _ORDINARY_TENTHS].rstrip() + suffix
+        )
+
     units = position = 0
-    for match in _CJK_RUN.finditer(text, 0, scan_end):
+    for match in _SCRIPT_RUN.finditer(text):
         start, end = match.span()
-        # Ordinary characters preceding the CJK run cost one unit each.
-        if units + start - position > max_units:
-            prefix_end = position + max_units - units
+        # Ordinary characters preceding the script run cost 0.4 token each.
+        ordinary_span = start - position
+        if units + ordinary_span * _ORDINARY_TENTHS > max_units:
+            prefix_end = position + (max_units - units) // _ORDINARY_TENTHS
             return text[:prefix_end].rstrip() + suffix
-        units += start - position
-        if units + (end - start) * 4 > max_units:
-            prefix_end = start + (max_units - units) // 4
+        units += ordinary_span * _ORDINARY_TENTHS
+        per_char = _run_tenths(match)
+        if units + (end - start) * per_char > max_units:
+            prefix_end = start + (max_units - units) // per_char
             return text[:prefix_end].rstrip() + suffix
-        units += (end - start) * 4
+        units += (end - start) * per_char
         position = end
 
-    if units + scan_end - position <= max_units:
+    if units + (len(text) - position) * _ORDINARY_TENTHS <= max_units:
         return text
-    prefix_end = position + max_units - units
+    prefix_end = position + (max_units - units) // _ORDINARY_TENTHS
     return text[:prefix_end].rstrip() + suffix
 
 

tests/unit/test_embedding_input.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: AGPL-3.0
+"""Tests for embedding input truncation bounds against real tokenizer costs.
+
+The mock tokenizer below is calibrated against the measurements reported in
+issue #6129 (llama-server ``POST /tokenize`` on the configured embedding
+model): Cyrillic prose costs up to ~0.372 tokens per character (2.69
+characters per token), and the worst non-Cyrillic sample — PowerShell source —
+costs ~0.36 tokens per character, well above the historical "4 characters per
+token" assumption. CJK characters are charged one token each.
+"""
+
+import math
+
+from openviking.utils.embedding_input import (
+    EMBEDDING_TRUNCATION_SUFFIX,
+    resolve_embedding_max_input_tokens,
+    truncate_embedding_input,
+)
+
+_MOCK_MILLI_PER_CHAR = {
+    # thousandths of a token per character, from the issue #6129 measurements
+    "cjk": 1000,  # one token per character
+    "cyrillic": 372,  # measured 2.69 characters per token (worst Russian sample)
+    "other": 360,  # measured 2.78 characters per token (worst non-Cyrillic sample)
+}
+_CYRILLIC = "\u0400"
+_MAX_INPUT_TOKENS = 4096
+
+RUSSIAN_MARKDOWN = "Русский текст для проверки памяти. " * 400
+ENGLISH_MARKDOWN = "# Section\n\nSome plain English prose sentence here. " * 500
+POWERSHELL_SOURCE = (
+    'Get-ChildItem -Path $env:SystemRoot -Filter "*.log" | ForEach-Object { $_.Name }' * 250
+)
+
+
+def _mock_tokenizer_tokens(text: str) -> int:
+    """Count tokens like the measured embedding-model tokenizer would."""
+    cjk = sum(1 for char in text if "\u4e00" <= char <= "\u9fff" or "\uac00" <= char <= "\ud7af")
+    cyrillic = sum(1 for char in text if _CYRILLIC <= char <= "\u052f")
+    other = len(text) - cjk - cyrillic
+    total_milli = (
+        cjk * _MOCK_MILLI_PER_CHAR["cjk"]
+        + cyrillic * _MOCK_MILLI_PER_CHAR["cyrillic"]
+        + other * _MOCK_MILLI_PER_CHAR["other"]
+    )
+    return math.ceil(total_milli / 1000)
+
+
+def test_truncated_input_stays_within_real_token_budget():
+    for sample in (RUSSIAN_MARKDOWN, ENGLISH_MARKDOWN, POWERSHELL_SOURCE):
+        truncated = truncate_embedding_input(sample, _MAX_INPUT_TOKENS)
+        assert truncated.endswith(EMBEDDING_TRUNCATION_SUFFIX)
+        assert _mock_tokenizer_tokens(truncated) <= _MAX_INPUT_TOKENS
+
+
+def test_truncation_respects_mixed_script_budget():
+    # 10 Cyrillic + 10 CJK + 10 ASCII characters repeat: 0.5 + 1 + 0.4 = 1.9
+    # tokens per repetition under the billing rates, so 2400 repetitions are
+    # charged above a 4096-token budget and must be truncated.
+    sample = "Ж中a" * 2400
+    assert _mock_tokenizer_tokens(sample) > _MAX_INPUT_TOKENS
+    truncated = truncate_embedding_input(sample, _MAX_INPUT_TOKENS)
+    assert truncated.endswith(EMBEDDING_TRUNCATION_SUFFIX)
+    assert _mock_tokenizer_tokens(truncated) <= _MAX_INPUT_TOKENS
+
+
+def test_truncate_appends_marker_within_budget():
+    text = "word " * 3000
+    truncated = truncate_embedding_input(text, 100)
+    assert truncated.endswith(EMBEDDING_TRUNCATION_SUFFIX)
+    assert _mock_tokenizer_tokens(truncated) <= 100
+    assert len(truncated) < len(text)
+
+
+def test_truncate_passthrough_and_degenerate_budgets():
+    text = "short"
+    assert truncate_embedding_input(text, 100) == text
+    assert truncate_embedding_input("", 100) == ""
+    assert truncate_embedding_input(text, 0) == EMBEDDING_TRUNCATION_SUFFIX.lstrip()
+
+
+def test_resolve_embedding_max_input_tokens():
+    assert resolve_embedding_max_input_tokens(None) is None
+    assert resolve_embedding_max_input_tokens({}, default=512) == 512
+    assert resolve_embedding_max_input_tokens({"max_input_tokens": 2048}) == 2048
+    assert resolve_embedding_max_input_tokens({"max_input_tokens": "2048"}) == 2048
+    assert resolve_embedding_max_input_tokens({"max_input_tokens": 0}) is None
+    assert resolve_embedding_max_input_tokens({"max_input_tokens": "bad"}) is None

tests/unit/test_vectorize_file_strategy.py

@@ -997,8 +997,13 @@ async def test_vectorize_content_only_reads_embedding_input(monkeypatch):
 @pytest.mark.parametrize(
     "content, expected",
     [
-        ("a1👍" * 100, "a1👍" * 26 + "a1" + EMBEDDING_TRUNCATION_SUFFIX),
-        ("中文あ가Ｌ、" * 100, "中文あ가Ｌ、" * 4 + "中文" + EMBEDDING_TRUNCATION_SUFFIX),
+        # 0.4 token per ordinary character, and the truncation marker itself is
+        # reserved from the budget, so the final message stays within the bound.
+        ("a1👍" * 100, "a1👍" * 15 + "a" + EMBEDDING_TRUNCATION_SUFFIX),
+        (
+            "中文あ가Ｌ、" * 100,
+            "中文あ가Ｌ、" * 3 + "中文あ가" + EMBEDDING_TRUNCATION_SUFFIX,
+        ),
         ("a1中👍" * 5, "a1中👍" * 5),
     ],
     ids=["ordinary-characters", "wide-text", "within-budget"],
@@ -1012,7 +1017,7 @@ async def test_vectorize_file_bounds_content_before_enqueue(monkeypatch, content
         embedding_utils,
         "get_openviking_config",
         lambda: types.SimpleNamespace(
-            embedding=types.SimpleNamespace(text_source="content_only", max_input_tokens=20)
+            embedding=types.SimpleNamespace(text_source="content_only", max_input_tokens=30)
         ),
     )
     await embedding_utils.vectorize_file(