diff --git a/reme/components/embedding_store/base_embedding_store.py b/reme/components/embedding_store/base_embedding_store.py index 500ece5b..9075b1de 100644 --- a/reme/components/embedding_store/base_embedding_store.py +++ b/reme/components/embedding_store/base_embedding_store.py @@ -1,6 +1,7 @@ """Base embedding store with abstract interface for caching and retrieval.""" from abc import abstractmethod +import unicodedata import numpy as np @@ -31,6 +32,28 @@ class BaseEmbeddingStore(BaseComponent): self.max_retries = max_retries self.is_healthy: bool = True + def _truncate(self, text: str) -> str: + """Truncate text using a CJK-aware character budget. + + ASCII text keeps its historical character limit. For non-ASCII text, + narrow characters cost one unit while CJK and other full-width + characters cost 1.5 units because they commonly consume more + embedding tokens. The estimate reserves a 5% safety margin, and + integer half-units avoid floating-point boundary errors. + """ + limit = max(0, self.max_input_length) + if text.isascii(): + return text[:limit] + + # Reserve a 5% safety margin for token estimation. + budget = limit * 2 * 95 // 100 + used = 0 + for index, char in enumerate(text): + used += 3 if unicodedata.east_asian_width(char) in {"W", "F"} else 2 + if used > budget: + return text[:index] + return text + @abstractmethod async def health_check(self, timeout: float = 2.0) -> bool: """Probe the provider; sets and returns is_healthy.""" diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index dc023a13..adf8676a 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -84,9 +84,6 @@ class LocalEmbeddingStore(BaseEmbeddingStore): # -- Batching -- - def _truncate(self, text: str) -> str: - return text if len(text) <= self.max_input_length else text[: self.max_input_length] - def _partition_by_cache(self, texts: list[str]) -> tuple[list[np.ndarray | None], list[Miss]]: results: list[np.ndarray | None] = [None] * len(texts) misses: list[Miss] = [] diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 48181541..e3934f39 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -46,6 +46,19 @@ def run(coro): return asyncio.run(coro) +def test_truncate_uses_cjk_aware_integer_budget(): + """Truncation should preserve ASCII behavior and budget non-ASCII text.""" + store = BadNodeEmbeddingStore(name="t_base_embedding_truncate", max_input_length=10) + + assert store._truncate("abcdefghijk") == "abcdefghij" + assert store._truncate("中文中文中文中文") == "中文中文中文" + assert store._truncate("éabcdefghij") == "éabcdefgh" + + store.max_input_length = -1 + assert store._truncate("text") == "" + assert store._truncate("中文") == "" + + def test_compute_batch_rejects_embeddings_with_wrong_dimension(): """Provider results with wrong dimensions are not padded, truncated, or cached."""