From e7a561e86f80a10c13a4f1ac2745382860890aad Mon Sep 17 00:00:00 2001 From: akushonkamen Date: Thu, 15 Oct 2026 21:43:21 +0800 Subject: [PATCH] fix(embedding): bound truncation estimate above real tokenizer costs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit estimate_embedding_input_tokens charged every non-CJK character 0.25 tokens (4 chars/token). Tokenizers of the embedding models served behind OpenAI-compatible endpoints spend far more characters per token: measured against llama-server POST /tokenize (issue #6129), Cyrillic prose needs ~2.7 chars/token and even English prose and shell code land near 3, so text truncated to max_input_tokens arrived with 1.3-1.5x the configured budget and the provider rejected it — failing the whole session_commit task and losing that session's memory. Charge per character in tenths of a token: CJK 1.0 (unchanged), Cyrillic 0.5, everything else 0.4 — rates at or above the measured worst cases. truncate_embedding_input now also reserves the budget that the truncation marker itself consumes, so the final payload stays within max_input_tokens instead of exceeding it by the marker's cost. The estimate/truncate expectations in test_vectorize_file_strategy are updated for the tighter rates, and tests/unit/test_embedding_input.py adds a mock tokenizer calibrated to the issue's measurements asserting that truncated input stays within the real-token budget. Fixes #6129 Signed-off-by: akushonkamen --- openviking/utils/embedding_input.py | 65 +++++++++++----- tests/unit/test_embedding_input.py | 89 ++++++++++++++++++++++ tests/unit/test_vectorize_file_strategy.py | 11 ++- 3 files changed, 144 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_embedding_input.py diff --git a/openviking/utils/embedding_input.py b/openviking/utils/embedding_input.py index ad4bb3ebe5..00b3fe9462 100644 --- a/openviking/utils/embedding_input.py +++ b/openviking/utils/embedding_input.py @@ -8,6 +8,28 @@ 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( @@ -15,37 +37,44 @@ def truncate_embedding_input( 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 diff --git a/tests/unit/test_embedding_input.py b/tests/unit/test_embedding_input.py new file mode 100644 index 0000000000..9f060d2072 --- /dev/null +++ b/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 diff --git a/tests/unit/test_vectorize_file_strategy.py b/tests/unit/test_vectorize_file_strategy.py index 036e3a98cb..dfac6cae62 100644 --- a/tests/unit/test_vectorize_file_strategy.py +++ b/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(