From c3572b0dd219c39cac5a5c1fdf628159a1cbdcac Mon Sep 17 00:00:00 2001 From: autodev-bot Date: Thu, 8 Oct 2026 16:04:44 +0800 Subject: [PATCH 1/2] fix: split memory items that exceed the embedder per-item limit Items whose token count fell in (embedder limit, chat_window_max_tokens] were neither split nor embeddable, so /product/add returned 200 while the item was stored with metadata.embedding=None and vector_sync != "success", permanently invisible to semantic search (issue #2461). - expose chat_window_max_tokens as an explicit, env-configurable field on BaseMemReaderConfig (MEM_READER_CHAT_WINDOW_MAX_TOKENS), default 1024 - derive the split budget as min(window, embedder per-item limit) - re-check chunker output and hard-split any over-budget chunk, preferring a punctuation boundary and falling back to a binary-search character window; a chunker that fails or yields nothing now hard-splits instead of passing the over-budget item through - truncate to budget before the per-item embedding retry so one over-limit item cannot leave a whole batch without vectors --- src/memos/api/config.py | 29 ++ src/memos/configs/mem_reader.py | 17 + src/memos/mem_reader/multi_modal_struct.py | 167 +++++--- src/memos/mem_reader/simple_struct.py | 153 +++++++- tests/mem_reader/test_embed_budget_split.py | 355 ++++++++++++++++++ .../mem_reader/test_project_id_propagation.py | 4 +- 6 files changed, 673 insertions(+), 52 deletions(-) create mode 100644 tests/mem_reader/test_embed_budget_split.py diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 0247dfa42..e25024786 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -685,6 +685,32 @@ def get_reader_config() -> dict[str, Any]: }, } + @staticmethod + def get_chat_window_max_tokens() -> int: + """Get the mem_reader per-item token budget. + + Items above this size are split before embedding. Deployments whose embedder has a + lower per-item limit than the default (1024) should lower this too, otherwise + content in between is stored without a vector (issue #2461). + """ + raw = os.getenv("MEM_READER_CHAT_WINDOW_MAX_TOKENS", "") + if not raw.strip(): + return 1024 + try: + value = int(raw) + except ValueError: + logger.warning( + "Invalid MEM_READER_CHAT_WINDOW_MAX_TOKENS=%r; falling back to 1024", raw + ) + return 1024 + if value <= 0: + logger.warning( + "MEM_READER_CHAT_WINDOW_MAX_TOKENS must be positive (got %r); falling back to 1024", + raw, + ) + return 1024 + return value + @staticmethod def get_oss_config() -> dict[str, Any] | None: """Get OSS configuration and validate connection.""" @@ -775,6 +801,7 @@ def get_internet_config() -> dict[str, Any]: }, }, "chat_chunker": reader_config, + "chat_window_max_tokens": APIConfig.get_chat_window_max_tokens(), }, }, }, @@ -1054,6 +1081,7 @@ def get_product_default_config() -> dict[str, Any]: }, }, "chat_chunker": reader_config, + "chat_window_max_tokens": APIConfig.get_chat_window_max_tokens(), "direct_markdown_hostnames": [ h.strip() for h in os.getenv( @@ -1189,6 +1217,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple["MOSConfig", "Gene }, }, "chat_chunker": reader_config, + "chat_window_max_tokens": APIConfig.get_chat_window_max_tokens(), }, }, "enable_textual_memory": True, diff --git a/src/memos/configs/mem_reader.py b/src/memos/configs/mem_reader.py index 5bbcfa0c0..12dc51f73 100644 --- a/src/memos/configs/mem_reader.py +++ b/src/memos/configs/mem_reader.py @@ -59,6 +59,23 @@ def parse_datetime(cls, value): chat_chunker: dict[str, Any] = Field( default=None, description="Configuration for the MemReader chat chunk strategy" ) + chat_window_max_tokens: int = Field( + default=1024, + description=( + "Maximum number of tokens per memory item. Items above this size are split " + "before embedding. The effective split budget is the minimum of this value " + "and the embedder's per-item limit, so texts that the embedder would reject " + "are never handed to it whole (issue #2461)." + ), + ) + + @field_validator("chat_window_max_tokens") + @classmethod + def validate_chat_window_max_tokens(cls, value: int) -> int: + """Reject non-positive windows; they would make every item unsplittable.""" + if value <= 0: + raise ValueError("chat_window_max_tokens must be a positive integer") + return value class SimpleStructMemReaderConfig(BaseMemReaderConfig): diff --git a/src/memos/mem_reader/multi_modal_struct.py b/src/memos/mem_reader/multi_modal_struct.py index c46331f1d..ac3b53609 100644 --- a/src/memos/mem_reader/multi_modal_struct.py +++ b/src/memos/mem_reader/multi_modal_struct.py @@ -92,7 +92,9 @@ def _embed_memory_items(self, items: list[TextualMemoryItem]) -> None: """Compute embeddings for a list of memory items in-place. Attempts a single batch call first; falls back to per-item calls if the - batch fails. Errors are logged but never raised so callers always + batch fails. The per-item retry truncates each text to the embedder budget before + calling, so a single over-limit item cannot leave the whole batch without + embeddings (issue #2461). Errors are logged but never raised so callers always continue normally. """ valid = [w for w in items if w and w.memory] @@ -106,9 +108,18 @@ def _embed_memory_items(self, items: list[TextualMemoryItem]) -> None: except Exception as e: logger.error(f"[MultiModalStruct] Error batch computing embeddings: {e}") logger.warning("[EMBED_FALLBACK] batch_size=%d", len(texts)) + budget = self._get_embed_budget() for w in valid: + text = w.memory or "" + payload = self._truncate_to_budget(text, budget) + if payload != text: + logger.warning( + "[EMBED_TRUNCATE] fallback truncated %d-token text to budget %d", + self._count_tokens_safe(text), + budget, + ) try: - w.metadata.embedding = self.embedder.embed([w.memory])[0] + w.metadata.embedding = self.embedder.embed([payload])[0] except Exception as e2: logger.error(f"[MultiModalStruct] Error computing embedding for item: {e2}") @@ -118,6 +129,12 @@ def _split_large_memory_item( """ Split a single memory item that exceeds max_tokens into multiple chunks. + The chunker is the primary splitter, but its output is not trusted blindly: some + backends (e.g. ``sentence``/chonkie) do not recognise CJK sentence terminators and + can hand back a chunk that is still larger than the budget. Every produced chunk is + therefore re-checked and hard-split when it still exceeds the budget, guaranteeing + that nothing over the embedder's per-item limit is ever sent to it (issue #2461). + Args: item: TextualMemoryItem to split max_tokens: Maximum tokens per chunk @@ -129,68 +146,98 @@ def _split_large_memory_item( if not item_text: return [item] - item_tokens = self._count_tokens(item_text) + item_tokens = self._count_tokens_safe(item_text) if item_tokens <= max_tokens: return [item] - # Use chunker to split the text - try: - chunks = self.chunker.chunk(item_text) - split_items = [] - source_info = dict(item.metadata.info or {}) - source_internal_info = dict(item.metadata.internal_info or {}) - ingest_batch_id = str(source_internal_info.get("ingest_batch_id") or uuid.uuid4()) - chunk_total = len(chunks) - - def _create_chunk_item(chunk_idx: int, chunk): - # Different chunkers are not fully consistent: - # some return Chunk-like objects with `.text`, while others return raw strings. - chunk_text = chunk.text if hasattr(chunk, "text") else chunk - if not chunk_text or not chunk_text.strip(): - return None - chunk_info = { - "user_id": item.metadata.user_id, - "session_id": item.metadata.session_id, - **source_info, - } - chunk_internal_info = { - **source_internal_info, - "ingest_batch_id": ingest_batch_id, - "chunk_index": chunk_idx, - "chunk_total": chunk_total, - } - # Create a new memory item for each chunk, preserving original metadata - split_item = self._make_memory_item( - value=chunk_text, - info=chunk_info, - memory_type=item.metadata.memory_type, - tags=item.metadata.tags or [], - key=item.metadata.key, - sources=item.metadata.sources or [], - background=item.metadata.background or "", - need_embed=False, - ) - split_item.metadata.internal_info = chunk_internal_info - return split_item + # The chunker is best-effort: any failure here (unsupported backend API, empty + # output) must NOT fall back to the original over-budget item, which is exactly the + # silent no-vector write of issue #2461. Hard-split instead. + raw_texts = self._chunk_within_budget(item_text, max_tokens) + if not raw_texts: + return [item] + split_items = [] + source_info = dict(item.metadata.info or {}) + source_internal_info = dict(item.metadata.internal_info or {}) + ingest_batch_id = str(source_internal_info.get("ingest_batch_id") or uuid.uuid4()) + chunk_total = len(raw_texts) + + def _create_chunk_item(chunk_idx: int, chunk_text: str): + chunk_info = { + "user_id": item.metadata.user_id, + "session_id": item.metadata.session_id, + **source_info, + } + chunk_internal_info = { + **source_internal_info, + "ingest_batch_id": ingest_batch_id, + "chunk_index": chunk_idx, + "chunk_total": chunk_total, + } + # Create a new memory item for each chunk, preserving original metadata + split_item = self._make_memory_item( + value=chunk_text, + info=chunk_info, + memory_type=item.metadata.memory_type, + tags=item.metadata.tags or [], + key=item.metadata.key, + sources=item.metadata.sources or [], + background=item.metadata.background or "", + need_embed=False, + ) + split_item.metadata.internal_info = chunk_internal_info + return split_item + + try: # Use thread pool to parallel process chunks, but keep the original order with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: futures = [ - executor.submit(_create_chunk_item, chunk_idx, chunk) - for chunk_idx, chunk in enumerate(chunks) + executor.submit(_create_chunk_item, chunk_idx, chunk_text) + for chunk_idx, chunk_text in enumerate(raw_texts) ] for future in futures: split_item = future.result() if split_item is not None: split_items.append(split_item) - - return split_items if split_items else [item] except Exception as e: logger.warning( - f"[MultiModalStruct] Failed to split large memory item: {e}. Returning original item." + f"[MultiModalStruct] Failed building split item: {e}. Returning original item." ) return [item] + return split_items if split_items else [item] + + def _chunk_within_budget(self, text: str, max_tokens: int) -> list[str]: + """Split ``text`` with the chunker, guaranteeing every piece fits ``max_tokens``. + + Chunker output is re-checked per chunk and hard-split when it still exceeds the + budget. If the chunker raises or produces nothing usable, the hard split runs over + the raw text so the caller never receives an over-budget item (issue #2461). + """ + raw_texts: list[str] = [] + try: + for chunk in self.chunker.chunk(text): + # Different chunkers are not fully consistent: some return Chunk-like + # objects with `.text`, while others return raw strings. + chunk_text = chunk.text if hasattr(chunk, "text") else chunk + if not chunk_text or not chunk_text.strip(): + continue + if self._count_tokens_safe(chunk_text) > max_tokens: + raw_texts.extend(self._hard_split_text(chunk_text, max_tokens)) + else: + raw_texts.append(chunk_text) + except Exception as e: + logger.warning( + f"[MultiModalStruct] Chunker failed ({e}); falling back to a hard split." + ) + return self._hard_split_text(text, max_tokens) + + if not raw_texts: + # The chunker cannot segment this text at all; hard-split it directly. + return self._hard_split_text(text, max_tokens) + return raw_texts + def _concat_multi_modal_memories( self, all_memory_items: list[TextualMemoryItem], max_tokens=None, overlap=200 ) -> list[TextualMemoryItem]: @@ -206,7 +253,7 @@ def _concat_multi_modal_memories( if not all_memory_items: return [] - max_tokens = max_tokens or self.chat_window_max_tokens + max_tokens = max_tokens or self._get_embed_budget() # Split large memory items before processing processed_items = [] @@ -219,7 +266,7 @@ def _concat_multi_modal_memories( # Create a list to hold futures with their original index futures = [] for idx, item in enumerate(all_memory_items): - if (item.memory or "") and self._count_tokens(item.memory) > max_tokens: + if (item.memory or "") and self._count_tokens_safe(item.memory) > max_tokens: future = executor.submit(self._split_large_memory_item, item, max_tokens) futures.append( (idx, future, True) @@ -245,7 +292,7 @@ def _concat_multi_modal_memories( # serial chunk large memory items for item in all_memory_items: item_text = item.memory or "" - item_tokens = self._count_tokens(item_text) + item_tokens = self._count_tokens_safe(item_text) if item_tokens > max_tokens: # Split the large item into multiple chunks split_items = self._split_large_memory_item(item, max_tokens) @@ -253,6 +300,26 @@ def _concat_multi_modal_memories( else: processed_items.append(item) + # Guard: nothing over the embed budget may proceed. If an item is still oversized + # after splitting (chunker produced nothing usable and the hard split could not + # segment it either), emit the items directly instead of aggregating them into an + # even larger window — that aggregation is what produced the silent no-vector write + # in issue #2461. The per-item embedding step truncates as a final safety net. + oversize = [ + it for it in processed_items if self._count_tokens_safe(it.memory or "") > max_tokens + ] + if oversize: + logger.warning( + "[MultiModalStruct] %d item(s) still exceed the embed budget (%d tokens) " + "after splitting; emitting them directly so the embedder can truncate " + "rather than silently dropping the vector", + len(oversize), + max_tokens, + ) + with timed_stage("add", "embedding", window_count=len(processed_items)): + self._embed_memory_items(processed_items) + return processed_items + # If only one item after processing, compute embedding and return if len(processed_items) == 1: single_item = processed_items[0] diff --git a/src/memos/mem_reader/simple_struct.py b/src/memos/mem_reader/simple_struct.py index 2804bac23..a67e9e616 100644 --- a/src/memos/mem_reader/simple_struct.py +++ b/src/memos/mem_reader/simple_struct.py @@ -86,6 +86,13 @@ def from_config(_config): logger = log.get_logger(__name__) + +# Sentence terminators used when picking a hard-split boundary. CJK terminators are +# included because the default ``sentence`` chunker (chonkie) does not recognise them, +# so long Chinese paragraphs can reach this fallback intact. +_SPLIT_PUNCT_CHARS = frozenset("。！？；!?;\n…") +_SPLIT_TRAILING_CHARS = frozenset("\"'”’）)】」』]") + PROMPT_DICT = { "chat": { "en": SIMPLE_STRUCT_MEM_READER_PROMPT, @@ -217,6 +224,119 @@ def set_graph_db(self, graph_db: "BaseGraphDB | None") -> None: def set_searcher(self, searcher: "Searcher | None") -> None: self.searcher = searcher + def _count_tokens_safe(self, text: str) -> int: + """Token-count helper that tolerates partially-mocked readers. + + Tests and embedders sometimes construct a reader without a live tokenizer; falling + back to a cheap character estimate keeps the budget logic usable instead of + raising ``AttributeError``. + """ + counter = getattr(self, "_count_tokens", None) + if counter is None: + return len(text or "") + return counter(text) + + def _get_embed_budget(self) -> int: + """Return the largest token count that is safe to hand to the embedder. + + The split budget is the minimum of the configured chat window and the embedder's + own per-item limit. Without this, an item in the + ``(embedder limit, chat_window_max_tokens]`` band is neither split nor embeddable + and ends up stored without a vector (issue #2461). + + The embedder limit is only honoured when it is a real positive integer, which keeps + mocked embedders (``MagicMock``) behaving as if no limit were configured. + """ + window = getattr(self, "chat_window_max_tokens", None) + budget = window if isinstance(window, int) and window > 0 else 1024 + + embedder_config = getattr(self.embedder, "config", None) + embedder_limit = getattr(embedder_config, "max_tokens", None) + if ( + isinstance(embedder_limit, int) + and not isinstance(embedder_limit, bool) + and embedder_limit > 0 + ): + budget = min(budget, embedder_limit) + return budget + + def _find_hard_split_index(self, text: str, budget: int) -> int: + """Return the length of the first piece of ``text`` that fits within ``budget``. + + Binary-searches the largest prefix whose token count is ``<= budget``, then prefers + to cut at the last sentence terminator inside that prefix so the split lands on a + natural boundary instead of mid-word. Always returns a value in ``[1, len(text))`` + so callers make progress and terminate. + """ + if budget <= 0: + return 1 + + low, high = 1, len(text) + best = 0 + while low <= high: + mid = (low + high) // 2 + if self._count_tokens_safe(text[:mid]) <= budget: + best = mid + low = mid + 1 + else: + high = mid - 1 + + if best <= 0: + # Even one character exceeds the budget; still make progress. + return 1 + if best >= len(text): + return len(text) + + # Prefer the last punctuation boundary within the largest half of the window. The + # search walks one position past ``best`` so a terminator that lands just outside + # the budget-sized prefix can still be kept on the left side. + scan_end = min(best + 1, len(text)) + lower_bound = max(best // 2, 1) + for idx in range(scan_end, lower_bound - 1, -1): + if idx <= 0 or idx > len(text): + continue + char = text[idx - 1] + if char in _SPLIT_PUNCT_CHARS: + return idx + # Trailing quote/bracket right after a terminator: keep it on the left side. + if char in _SPLIT_TRAILING_CHARS and idx >= 2 and text[idx - 2] in _SPLIT_PUNCT_CHARS: + return idx + return best + + def _hard_split_text(self, text: str, budget: int) -> list[str]: + """Split ``text`` into pieces each within ``budget`` tokens. + + This is the last-resort guarantee behind :meth:`_split_large_memory_item`: it makes + no assumption about the chunker having honoured the budget. Pieces are cut at + punctuation where possible and by binary-searched character window otherwise. + """ + if not text: + return [] + budget = budget if budget and budget > 0 else 1 + + pieces: list[str] = [] + remaining = text + # Each iteration consumes at least one character, so this terminates. + while remaining: + if self._count_tokens_safe(remaining) <= budget: + pieces.append(remaining) + break + cut = self._find_hard_split_index(remaining, budget) + cut = max(1, min(cut, len(remaining) - 1)) + pieces.append(remaining[:cut]) + remaining = remaining[cut:] + + return [piece for piece in pieces if piece.strip()] + + def _truncate_to_budget(self, text: str, budget: int) -> str: + """Truncate ``text`` to at most ``budget`` tokens (best-effort prefix).""" + if not text or budget <= 0: + return text + if self._count_tokens_safe(text) <= budget: + return text + cut = self._find_hard_split_index(text, budget) + return text[: max(1, cut)] + def _make_memory_item( self, value: str, @@ -244,7 +364,7 @@ def _make_memory_item( status="activated", tags=tags or [], key=key if key is not None else derive_key(value), - embedding=self.embedder.embed([value])[0] if need_embed else None, + embedding=self._embed_text(value) if need_embed else None, usage=[], sources=sources or [], background=background, @@ -255,7 +375,38 @@ def _make_memory_item( ), ) + def _embed_text(self, value: str) -> list[float] | None: + """Embed a single text, truncating it to the embedder budget first. + + Returns ``None`` when the embedder fails even on the truncated text, so callers + store the item without a vector instead of raising. The truncation is a safety net + for texts that slipped past the splitter; it is logged so the degradation is + observable rather than silent (issue #2461). + """ + text = value or "" + budget = self._get_embed_budget() + payload = self._truncate_to_budget(text, budget) + if payload != text: + logger.warning( + "[EMBED_TRUNCATE] text exceeded embedder budget (%d tokens > budget %d); " + "embedding truncated prefix of %d chars", + self._count_tokens_safe(text), + budget, + len(payload), + ) + try: + return self.embedder.embed([payload])[0] + except Exception as e: + logger.error( + "[EMBED_FAIL] Failed to embed text of %d tokens (budget %d): %s", + self._count_tokens_safe(text), + budget, + e, + ) + return None + def _safe_generate(self, messages: list[dict]) -> str | None: + """Generate a response, swallowing provider errors as ``None``.""" try: return self.llm.generate(messages) except Exception: diff --git a/tests/mem_reader/test_embed_budget_split.py b/tests/mem_reader/test_embed_budget_split.py new file mode 100644 index 000000000..c2f2813c7 --- /dev/null +++ b/tests/mem_reader/test_embed_budget_split.py @@ -0,0 +1,355 @@ +"""Regression tests for issue #2461. + +A memory item whose token count sits in ``(embedder per-item limit, +chat_window_max_tokens]`` used to be neither split (the splitter only fired above +``chat_window_max_tokens``) nor embeddable (it exceeded the embedder's limit). The item +was still persisted, but with ``metadata.embedding is None`` and ``vector_sync != +"success"``, so semantic search could never recall it — a silent no-vector write. + +These tests pin the fix: + +1. ``chat_window_max_tokens`` is an explicit, env-configurable config field. +2. The effective split budget is ``min(window, embedder limit)``. +3. ``_split_large_memory_item`` re-checks chunker output and hard-splits anything still + over budget, so every text handed to the embedder fits. +4. ``_embed_memory_items`` truncates before per-item retry, so one over-limit item cannot + leave a whole batch without embeddings. +""" + +import unittest + +from unittest.mock import MagicMock, patch + +from pydantic import ValidationError + +from memos.api.config import APIConfig +from memos.chunkers import ChunkerFactory +from memos.configs.mem_reader import ( + SimpleStructMemReaderConfig, +) +from memos.embedders.factory import EmbedderFactory +from memos.llms.factory import LLMFactory +from memos.mem_reader.multi_modal_struct import MultiModalStructMemReader +from memos.mem_reader.simple_struct import SimpleStructMemReader +from memos.memories.textual.item import ( + TextualMemoryItem, + TreeNodeTextualMemoryMetadata, +) + + +class _FakeEmbedder: + """Embedder stub whose configured limit mirrors the real one. + + ``embed`` raises for any text whose character length exceeds ``limit``, which mimics a + provider rejecting over-limit input, and returns a deterministic vector otherwise. + """ + + def __init__(self, limit: int = 10_000): + self.config = MagicMock() + self.config.max_tokens = limit + self.limit = limit + self.calls: list[list[str]] = [] + + def embed(self, texts: list[str]) -> list[list[float]]: + self.calls.append(list(texts)) + for text in texts: + if len(text) > self.limit: + raise ValueError("input exceeds the embedder per-item limit") + return [[float(len(t))] * 4 for t in texts] + + +def _make_item(memory: str, **internal_info) -> TextualMemoryItem: + return TextualMemoryItem( + memory=memory, + metadata=TreeNodeTextualMemoryMetadata( + user_id="u1", + session_id="s1", + memory_type="LongTermMemory", + tags=[], + key=None, + sources=[], + background="", + internal_info=internal_info or None, + ), + ) + + +def _reader_config_dict(window: int) -> dict: + """Minimal valid mem_reader config dict (no optional providers).""" + return { + "llm": { + "backend": "openai", + "config": {"model_name_or_path": "gpt-4o", "api_key": "test-key"}, + }, + "embedder": { + "backend": "ollama", + "config": {"model_name_or_path": "nomic-embed-text:latest"}, + }, + "chunker": { + "backend": "sentence", + "config": { + "tokenizer_or_token_counter": "gpt2", + "chunk_size": 512, + "chunk_overlap": 128, + "min_sentences_per_chunk": 1, + "save_rawfile": False, + }, + }, + "chat_window_max_tokens": window, + } + + +def _build_reader( + embedder: _FakeEmbedder, + window: int = 1024, + chunker=None, + cls=MultiModalStructMemReader, +): + """Construct a reader with a mocked embedder/chunker and a lightweight config stub. + + ``MultiModalStructMemReader`` rebuilds a ``SimpleStructMemReaderConfig`` from + ``config.model_dump()``, so the stub must return a dict the base config accepts. + """ + cfg = MagicMock() + cfg.model_dump.return_value = _reader_config_dict(window) + cfg.chat_window_max_tokens = window + cfg.remove_prompt_example = False + cfg.general_llm = None + cfg.preference_extractor_llm = None + cfg.image_parser_llm = None + cfg.document_parser_llm = None + cfg.direct_markdown_hostnames = None + cfg.oss_config = None + cfg.skills_dir_config = None + cfg.memory_version_switch = "off" + + with ( + patch.object(LLMFactory, "from_config", return_value=MagicMock()), + patch.object(EmbedderFactory, "from_config", return_value=embedder), + patch.object(ChunkerFactory, "from_config", return_value=chunker or MagicMock()), + ): + reader = cls(cfg) + + # MagicMock auto-attributes would make every budget look unlimited; be explicit. + reader.embedder = embedder + reader.chunker = chunker or MagicMock() + reader.chat_window_max_tokens = window + return reader + + +class TestEffectiveBudget(unittest.TestCase): + """R2: the split budget honors the embedder's per-item limit.""" + + def test_budget_capped_by_embedder_limit(self): + reader = _build_reader(_FakeEmbedder(limit=512), window=1024) + self.assertEqual(reader._get_embed_budget(), 512) + + def test_budget_uses_window_when_embedder_is_more_permissive(self): + reader = _build_reader(_FakeEmbedder(limit=8192), window=1024) + self.assertEqual(reader._get_embed_budget(), 1024) + + def test_budget_defaults_to_1024_without_config(self): + reader = _build_reader(_FakeEmbedder(limit=8192)) + reader.chat_window_max_tokens = 1024 + self.assertEqual(reader._get_embed_budget(), 1024) + + def test_mock_embedder_has_no_limit(self): + """A MagicMock config must not be mistaken for a real positive limit.""" + embedder = MagicMock() + reader = _build_reader(embedder, window=1024) + reader.embedder = embedder + self.assertEqual(reader._get_embed_budget(), 1024) + + +class TestHardSplit(unittest.TestCase): + """R3: post-chunker hard split guarantees every chunk fits.""" + + def test_punctuation_free_cjk_paragraph_is_split_within_budget(self): + reader = _build_reader(_FakeEmbedder(limit=10_000), window=512) + # Punctuation-free CJK paragraph: the sentence chunker does not split it. + text = "中" * 1170 + pieces = reader._hard_split_text(text, 512) + + self.assertGreaterEqual(len(pieces), 2) + self.assertEqual("".join(pieces), text) + for piece in pieces: + self.assertLessEqual(reader._count_tokens(piece), 512) + self.assertGreater(len(piece), 0) + + def test_split_prefers_punctuation_boundary(self): + reader = _build_reader(_FakeEmbedder(limit=10_000), window=10) + # "。" should be the cut point rather than an arbitrary character position. + text = "一二三四五六七八九十。" + "x" * 40 + pieces = reader._hard_split_text(text, 10) + self.assertTrue(pieces[0].endswith("。"), pieces[0]) + + def test_empty_chunker_output_still_gets_split(self): + """A chunker that yields nothing must not put the over-budget item through as-is. + + Returning the original item here would reproduce the silent no-vector bug, so the + hard split takes over instead. + """ + chunker = MagicMock() + chunker.chunk.return_value = [] + reader = _build_reader(_FakeEmbedder(limit=10_000), window=5, chunker=chunker) + reader._count_tokens = len + item = _make_item("abcdefghij") + + result = reader._split_large_memory_item(item, max_tokens=5) + + self.assertGreater(len(result), 1) + self.assertEqual("".join(it.memory for it in result), "abcdefghij") + for it in result: + self.assertLessEqual(len(it.memory), 5) + + def test_item_within_budget_is_untouched(self): + chunker = MagicMock() + reader = _build_reader(_FakeEmbedder(limit=10_000), window=1024, chunker=chunker) + reader._count_tokens = len + item = _make_item("short") + + result = reader._split_large_memory_item(item, max_tokens=1024) + + self.assertEqual(result, [item]) + chunker.chunk.assert_not_called() + + def test_chunker_output_expanded_and_indexed_contiguously(self): + """A chunk the chunker failed to size must be hard-split and re-indexed.""" + chunker = MagicMock() + # First chunk fits, second one is oversized and must be hard-split. + chunker.chunk.return_value = ["a" * 4, "b" * 12] + reader = _build_reader(_FakeEmbedder(limit=10_000), window=5, chunker=chunker) + reader._count_tokens = len + reader._make_memory_item = MagicMock( + side_effect=lambda *, value, **kwargs: _make_item(value) + ) + + result = reader._split_large_memory_item(_make_item("x" * 20), max_tokens=5) + + self.assertGreater(len(result), 2) + self.assertEqual( + [it.metadata.internal_info["chunk_index"] for it in result], + list(range(len(result))), + ) + self.assertEqual({it.metadata.internal_info["chunk_total"] for it in result}, {len(result)}) + batch_ids = {it.metadata.internal_info["ingest_batch_id"] for it in result} + self.assertEqual(len(batch_ids), 1) + for it in result: + self.assertLessEqual(len(it.memory), 5) + + +class TestConcatUsesEffectiveBudget(unittest.TestCase): + """R4: the split trigger uses the effective budget, not the raw window.""" + + def test_item_in_the_gap_is_split(self): + chunker = MagicMock() + chunker.chunk.return_value = [] # force the deterministic hard-split path + reader = _build_reader(_FakeEmbedder(limit=512), window=1024, chunker=chunker) + reader._count_tokens = len + reader._make_memory_item = MagicMock( + side_effect=lambda *, value, **kwargs: _make_item(value) + ) + # 700 "tokens" of content: inside (512, 1024], previously never split. + item = _make_item("z" * 700) + + result = reader._concat_multi_modal_memories([item]) + + self.assertGreaterEqual(len(result), 2) + for it in result: + self.assertLessEqual(len(it.memory), 512) + self.assertIsNotNone(it.metadata.embedding) + + +class TestEmbedFallback(unittest.TestCase): + """R5: per-item retry truncates, so one bad item cannot zero out a batch.""" + + def test_over_limit_item_still_gets_an_embedding(self): + embedder = _FakeEmbedder(limit=10) + reader = _build_reader(embedder, window=10) + reader._count_tokens = len + + good = _make_item("short") + bad = _make_item("q" * 100) + + calls: list[list[str]] = [] + + def embed_impl(texts: list[str]) -> list[list[float]]: + calls.append(list(texts)) + if len(texts) > 1: + raise ValueError("batch call rejected by the provider") + return [[1.0] * 4 for _ in texts] + + embedder.embed = MagicMock(side_effect=embed_impl) + reader._embed_memory_items([good, bad]) + + self.assertIsNotNone(good.metadata.embedding) + self.assertIsNotNone(bad.metadata.embedding) + # The retried text must have been truncated to the budget. + retried = [c[0] for c in calls if len(c) == 1] + self.assertTrue(retried) + self.assertTrue(all(len(text) <= 10 for text in retried)) + self.assertTrue(any(len(text) == 10 for text in retried)) + + +class TestConfigPlumbing(unittest.TestCase): + """R1: config field default + env override reachable from the API layer.""" + + def test_env_var_overrides_window(self): + with patch.dict("os.environ", {"MEM_READER_CHAT_WINDOW_MAX_TOKENS": "512"}): + self.assertEqual(APIConfig.get_chat_window_max_tokens(), 512) + + def test_env_var_defaults_to_1024(self): + with patch.dict("os.environ", {}, clear=False): + import os + + os.environ.pop("MEM_READER_CHAT_WINDOW_MAX_TOKENS", None) + self.assertEqual(APIConfig.get_chat_window_max_tokens(), 1024) + + def test_invalid_env_var_falls_back_to_1024(self): + for raw in ("abc", "0", "-5", ""): + with patch.dict("os.environ", {"MEM_READER_CHAT_WINDOW_MAX_TOKENS": raw}): + self.assertEqual(APIConfig.get_chat_window_max_tokens(), 1024) + + def test_reader_config_field_has_default(self): + field = SimpleStructMemReaderConfig.model_fields["chat_window_max_tokens"] + self.assertEqual(field.default, 1024) + + def test_reader_config_rejects_non_positive_window(self): + base = _reader_config_dict(1024) + with self.assertRaises(ValidationError): + SimpleStructMemReaderConfig(**{**base, "chat_window_max_tokens": 0}) + + +class TestSimpleStructReaderUsesBudget(unittest.TestCase): + """R5/P1: the base reader's single-text embedding path respects the budget.""" + + def test_make_memory_item_truncates_over_budget_text(self): + embedder = _FakeEmbedder(limit=10) + reader = _build_reader(embedder, window=10, cls=SimpleStructMemReader) + reader._count_tokens = len + + item = reader._make_memory_item( + value="y" * 500, + info={"user_id": "u1", "session_id": "s1"}, + memory_type="LongTermMemory", + ) + + self.assertIsNotNone(item.metadata.embedding) + self.assertEqual(embedder.calls[-1][0], "y" * 10) + + def test_make_memory_item_within_budget_is_verbatim(self): + embedder = _FakeEmbedder(limit=10_000) + reader = _build_reader(embedder, window=1024, cls=SimpleStructMemReader) + reader._count_tokens = len + + reader._make_memory_item( + value="verbatim text", + info={"user_id": "u1", "session_id": "s1"}, + memory_type="LongTermMemory", + ) + + self.assertEqual(embedder.calls[-1][0], "verbatim text") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/mem_reader/test_project_id_propagation.py b/tests/mem_reader/test_project_id_propagation.py index bf55aca46..7ead96346 100644 --- a/tests/mem_reader/test_project_id_propagation.py +++ b/tests/mem_reader/test_project_id_propagation.py @@ -260,7 +260,9 @@ def test_build_window_picks_first_nonempty(self): _assert_fields(self, result) def test_split_large_memory_item_assigns_shared_ingest_batch_id(self): - self.reader._count_tokens = MagicMock(return_value=999) + # Token count proportional to length so chunker output survives the new + # post-chunker budget re-check (see issue #2461). + self.reader._count_tokens = lambda text: len(text) self.reader.chunker.chunk.return_value = ["chunk one", "chunk two"] def fake_make_memory_item( From 243ef50068787b1d00a6a45599d3d3b60b3f9481 Mon Sep 17 00:00:00 2001 From: autodev-bot Date: Thu, 8 Oct 2026 16:49:54 +0800 Subject: [PATCH 2/2] fix: keep split budget decisions on the safe token counter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the review findings on the #2461 fix: - multi_modal_struct: the window loop and the overlap trim still called `_count_tokens` directly, so a reader built without a live tokenizer made its budget decision with `_count_tokens_safe` and then raised `AttributeError` a few lines later, dropping the whole scene. Both call sites now use the safe counter. - multi_modal_struct: `_chunk_within_budget` wrapped the per-chunk token check and hard split in the same `try` as the chunker, so a tokenizer error was reported as "Chunker failed" and the already-validated chunks were thrown away. The guard now covers only `chunker.chunk(text)`. - simple_struct: `_find_hard_split_index` could return `best + 1`, letting a terminator just past the budget be kept on the left side. That made the reported prefix `budget + 1` tokens; `_truncate_to_budget` used it directly, so the last guard before `embedder.embed([payload])` could still hand the provider an over-limit payload — the exact rejection this change set out to remove. The scan now stays inside the fitting prefix. - simple_struct: drop the unreachable range guard in the scan loop. --- src/memos/mem_reader/multi_modal_struct.py | 28 ++++++++++++--------- src/memos/mem_reader/simple_struct.py | 9 +++---- tests/mem_reader/test_embed_budget_split.py | 21 ++++++++++++++-- 3 files changed, 39 insertions(+), 19 deletions(-) diff --git a/src/memos/mem_reader/multi_modal_struct.py b/src/memos/mem_reader/multi_modal_struct.py index ac3b53609..6e9f7fcc3 100644 --- a/src/memos/mem_reader/multi_modal_struct.py +++ b/src/memos/mem_reader/multi_modal_struct.py @@ -217,22 +217,25 @@ def _chunk_within_budget(self, text: str, max_tokens: int) -> list[str]: """ raw_texts: list[str] = [] try: - for chunk in self.chunker.chunk(text): - # Different chunkers are not fully consistent: some return Chunk-like - # objects with `.text`, while others return raw strings. - chunk_text = chunk.text if hasattr(chunk, "text") else chunk - if not chunk_text or not chunk_text.strip(): - continue - if self._count_tokens_safe(chunk_text) > max_tokens: - raw_texts.extend(self._hard_split_text(chunk_text, max_tokens)) - else: - raw_texts.append(chunk_text) + # Materialise inside the guard: some chunkers raise lazily while iterating. + chunks = list(self.chunker.chunk(text)) except Exception as e: logger.warning( f"[MultiModalStruct] Chunker failed ({e}); falling back to a hard split." ) return self._hard_split_text(text, max_tokens) + for chunk in chunks: + # Different chunkers are not fully consistent: some return Chunk-like + # objects with `.text`, while others return raw strings. + chunk_text = chunk.text if hasattr(chunk, "text") else chunk + if not chunk_text or not chunk_text.strip(): + continue + if self._count_tokens_safe(chunk_text) > max_tokens: + raw_texts.extend(self._hard_split_text(chunk_text, max_tokens)) + else: + raw_texts.append(chunk_text) + if not raw_texts: # The chunker cannot segment this text at all; hard-split it directly. return self._hard_split_text(text, max_tokens) @@ -347,7 +350,7 @@ def _concat_multi_modal_memories( # Check if adding this item would exceed max_tokens (same logic as _iter_chat_windows) # Note: After splitting large items, each item should be <= max_tokens, # but we still check to handle edge cases - if self._count_tokens(cur_text + line) > max_tokens and cur_text: + if self._count_tokens_safe(cur_text + line) > max_tokens and cur_text: # Yield current window window = self._build_window_from_items(buf_items, info) if window: @@ -357,7 +360,8 @@ def _concat_multi_modal_memories( # (same logic as _iter_chat_windows) while ( buf_items - and self._count_tokens("".join([it.memory or "" for it in buf_items])) > overlap + and self._count_tokens_safe("".join([it.memory or "" for it in buf_items])) + > overlap ): buf_items.pop(0) # Recalculate cur_text from remaining items diff --git a/src/memos/mem_reader/simple_struct.py b/src/memos/mem_reader/simple_struct.py index a67e9e616..631f89062 100644 --- a/src/memos/mem_reader/simple_struct.py +++ b/src/memos/mem_reader/simple_struct.py @@ -288,13 +288,12 @@ def _find_hard_split_index(self, text: str, budget: int) -> int: return len(text) # Prefer the last punctuation boundary within the largest half of the window. The - # search walks one position past ``best`` so a terminator that lands just outside - # the budget-sized prefix can still be kept on the left side. - scan_end = min(best + 1, len(text)) + # search stays inside the fitting prefix so the returned index never slips a + # terminator just past the budget — a +1 overshoot is enough to make the + # provider reject the payload (issue #2461). + scan_end = min(best, len(text)) lower_bound = max(best // 2, 1) for idx in range(scan_end, lower_bound - 1, -1): - if idx <= 0 or idx > len(text): - continue char = text[idx - 1] if char in _SPLIT_PUNCT_CHARS: return idx diff --git a/tests/mem_reader/test_embed_budget_split.py b/tests/mem_reader/test_embed_budget_split.py index c2f2813c7..8f8d309f9 100644 --- a/tests/mem_reader/test_embed_budget_split.py +++ b/tests/mem_reader/test_embed_budget_split.py @@ -178,10 +178,27 @@ def test_punctuation_free_cjk_paragraph_is_split_within_budget(self): def test_split_prefers_punctuation_boundary(self): reader = _build_reader(_FakeEmbedder(limit=10_000), window=10) - # "。" should be the cut point rather than an arbitrary character position. - text = "一二三四五六七八九十。" + "x" * 40 + # The "。" fits inside the budget-sized prefix, so it must be preferred as the + # cut point over an arbitrary character position. + text = "一二三四五。六七八九十" + "x" * 40 pieces = reader._hard_split_text(text, 10) self.assertTrue(pieces[0].endswith("。"), pieces[0]) + self.assertLessEqual(reader._count_tokens(pieces[0]), 10) + + def test_split_index_never_exceeds_budget(self): + """A terminator just past the budget must not be kept on the left side. + + Returning ``best + 1`` here produced a prefix of ``budget + 1`` tokens, which the + provider rejects outright when the embedder limit is the binding constraint — the + item then gets persisted without a vector (issue #2461). + """ + reader = _build_reader(_FakeEmbedder(limit=10_000), window=10) + # "。" lands at index 10, one past the fitting prefix (whose length is 10). + text = "一二三四五六七八九十" + "。" + "x" * 40 + + cut = reader._find_hard_split_index(text, 10) + + self.assertLessEqual(reader._count_tokens(text[:cut]), 10, repr(text[:cut])) def test_empty_chunker_output_still_gets_split(self): """A chunker that yields nothing must not put the over-budget item through as-is.