Skip to content

fix(embedder): safe 3072-token default + token-aware truncation - #2190

Open
RerankerGuo wants to merge 2 commits into
MemTensor:mainfrom
RerankerGuo:fix/issue-2121-embedding-3072-token-truncate
Open

fix(embedder): safe 3072-token default + token-aware truncation#2190
RerankerGuo wants to merge 2 commits into
MemTensor:mainfrom
RerankerGuo:fix/issue-2121-embedding-3072-token-truncate

Conversation

@RerankerGuo

Copy link
Copy Markdown
Contributor

Background

Issue #2121 reports that memos-local-plugin ingestion passes source text directly to the embedding API, and text-embedding-3 enforces a hard 3072-token input cap. Previously, the embedder's default max_tokens was 8192 (which exceeds the API limit) and the truncation was a plain character-slice that did not account for token counts at all. This caused out-of-budget embedding errors for longer documents and knowledge chunks.

Changes

  • memos/embedders/base.py:
    • Added _SAFE_EMBEDDING_MAX_TOKENS = 3072 matching text-embedding-3's input limit.
    • Added BaseEmbedder._effective_max_tokens(): returns an explicitly configured positive max_tokens, otherwise the 3072 safe default (also used when the user configures max_tokens=None or max_tokens=0).
    • Rewrote _truncate_texts() to delegate to the existing token-aware _truncate_text_to_tokens() binary-search helper. A fast-path length check keeps the cheap case cheap (text length <= max tokens means no token count required).
  • memos/configs/embedder.py:
    • Changed BaseEmbedderConfig.max_tokens default from 8192 to None so that the new 3072 safe-limit takes effect for any user not overriding the field explicitly; the docstring was updated.
  • tests/embedders/test_base.py:
    • Added coverage for _effective_max_tokens() behaviour across None, 0, and explicit positive values.
    • Added coverage for _truncate_text_to_tokens(): short text unchanged, empty/None/0 boundaries, long CJK text stays within budget.
    • Added coverage for _truncate_texts(): short texts untouched, very long CJK text truncated within the 3072-token default, explicit 10-token override honoured.

Verification

  • Ruff: ruff format + ruff check --fix applied (2 pre-existing BLE001 broad-excepts in _count_tokens_for_embedding left intact).
  • python3 -m py_compile succeeds for modified files.

Impact

  • Backwards compatible for callers that set an explicit max_tokens; behaviour changes only when the field was left at the previous 8192 default.
  • Eliminates the silent "send 8192 chars and hope for the best" behaviour for text-embedding-3 endpoints; the default now stays inside the API budget.

…3072-token default

- embedders/base.py: add _SAFE_EMBEDDING_MAX_TOKENS=3072, new
  BaseEmbedder._effective_max_tokens() fallback, and make
  _truncate_texts() token-aware by calling _truncate_text_to_tokens
  (binary search over prefix length) instead of a naive character-slice.
- configs/embedder.py: BaseEmbedderConfig.max_tokens default changed
  from 8192 to None so the new 3072-token safe-limit kicks in for
  providers such as text-embedding-3 that enforce a hard input cap.
- tests/embedders/test_base.py: new tests covering effective_max_tokens
  for None/0/explicit overrides, _truncate_text_to_tokens bounds on
  long CJK text, and _truncate_texts with the new 3072 default.
Closes MemTensor#2121.
@Memtensor-AI Memtensor-AI added area:core MOS 编排层 / 框架底座 / 跨模块问题 area:model llm + embedder + reranker status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Jul 31, 2026
@Memtensor-AI

Memtensor-AI commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

🤖 Open Code Review

Target: PR #2190
Task: 61ebe3a4627d2a33
Base: main
Head: fix/issue-2121-embedding-3072-token-truncate

🔍 OpenCodeReview found 5 issue(s) in this PR.

⚠️ 1 warning(s) occurred during review.


1. src/memos/embedders/base.py (L171-L173)

The fast-path assumes that len(t) <= max_tokens guarantees the token count is also within budget, citing CJK as the worst case (~1 char/token). This assumption is incorrect for emoji and rare Unicode.

With tiktoken's cl100k_base, emoji such as U+1F600 (😀) encode to 2 tokens each, and rare code points that fall back to byte-level tokens can produce even more. A text composed of 3072 emoji characters has len(t) == 3072, passes the fast-path unchanged, yet has ~6144 actual tokens — far above the max_tokens limit.

Suggested fix: use a tighter bound, or skip the fast-path entirely and always call _count_tokens_for_embedding when you need a correct answer:

# Conservative: only skip counting when the text is very short
# (well under half of max_tokens to account for high-token-density chars)
if len(t) * 2 <= max_tokens:   # even emoji won't exceed budget
    truncated.append(t)
    continue

or simply remove the fast-path and rely on the early-exit inside _truncate_text_to_tokens (line 110–111) which already skips truncation when current_tokens <= max_tokens.


2. src/memos/embedders/base.py (L191)

This fallback unconditionally applies a 3072-token truncation limit to every embedder that has not explicitly set max_tokens in its config — including models that support 8192 or more tokens (e.g. text-embedding-ada-002, many open-source models). Previously the code skipped truncation entirely when max_tokens was None, so this is a silent breaking change that can degrade embedding quality for longer documents with no warning to the user.

A safer approach is to make this opt-in rather than opt-out, or at least document clearly that subclasses for higher-capacity models should override _effective_max_tokens or always set max_tokens in their config:

def _effective_max_tokens(self) -> int | None:
    config = getattr(self, "config", None)
    if config is not None:
        configured = getattr(config, "max_tokens", None)
        if configured is not None and configured > 0:
            return configured
    # Return None to opt out of truncation for unconfigured embedders;
    # subclasses for token-limited APIs (e.g. text-embedding-3) should
    # override this or set max_tokens in their config.
    return None

3. tests/embedders/test_base.py (L22-L23)

_ConcreteEmbedder.init does nothing but delegate to super().init(config) with the same signature. Python inherits the parent init automatically when no override is defined, making this method dead code. Remove it entirely; the class only needs the embed override.

💡 Suggested Change

Before:

    def __init__(self, config):
        super().__init__(config)

After:

class _ConcreteEmbedder(BaseEmbedder):
    def embed(self, texts: list[str]) -> list[list[float]]:
        return [[0.0] for _ in texts]

4. tests/embedders/test_base.py (L70-L71)

The lower-bound check measures characters, not tokens, implicitly relying on a 1-char-per-token CJK assumption. If the tokeniser encodes any CJK subsequence as multiple tokens, truncation may correctly stop below limit characters while still satisfying the token budget, causing a spurious test failure. Both assertions should use the same counting function.

💡 Suggested Change

Before:

        assert _count_tokens_for_embedding(truncated) <= limit
        assert len(truncated) >= limit

After:

        assert _count_tokens_for_embedding(truncated) <= limit
        assert _count_tokens_for_embedding(truncated) >= int(limit * 0.9)

5. tests/embedders/test_base.py (L36-L37)

This test conflates two independent concerns: (a) that None max_tokens falls back to the safe constant, and (b) that the constant equals exactly 3072. If the constant is intentionally changed, the failure will point at a test named 'falls_back_to_safe_default', making the root cause harder to diagnose. Move the constant-value assertion to its own dedicated test.

💡 Suggested Change

Before:

        assert emb._effective_max_tokens() == _SAFE_EMBEDDING_MAX_TOKENS
        assert _SAFE_EMBEDDING_MAX_TOKENS == 3072

After:

        assert emb._effective_max_tokens() == _SAFE_EMBEDDING_MAX_TOKENS


def test_safe_embedding_max_tokens_constant_value():
    assert _SAFE_EMBEDDING_MAX_TOKENS == 3072

🧹 Filtered 1 low-confidence OCR finding(s) before posting/fix-loop (existing_code_mismatch: 1).

Generated by cloud-assistant via Open Code Review.

@Memtensor-AI

Copy link
Copy Markdown
Collaborator

⚠️ Automated Test Results: INCONCLUSIVE

Automated tests inconclusive (auto-generated test defect); treated as non-blocking. Manual review recommended. Details: The AI-generated test file defines a _ConcreteEmbedder subclass of BaseEmbedder but fails to implement the abstract __init__ method, causing all instantiation-based tests to fail with TypeError. Additionally, one CJK truncation test has an off-by-one assertion that expects the truncated length to be >= the token limit, but token-aware truncation correctly produces a length <= limit. [advisory, non-gating] AI-generated tests on branch test/auto-gen-bc3c30a4c517a246-20260731105109: 24/65 passed, 41 failed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/issue-2121-embedding-3072-token-truncate

@Memtensor-AI

Copy link
Copy Markdown
Collaborator

⚠️ Automated Test Results: INCONCLUSIVE

Automated tests inconclusive (auto-generated test defect); treated as non-blocking. Manual review recommended. Details: The test asserts that a token-truncated CJK string has a character length >= the token limit (500), but token-aware truncation legitimately returns fewer characters than the token budget when the tokenizer produces a boundary just below the limit. [advisory, non-gating] AI-generated tests on branch test/auto-gen-61ebe3a4627d2a33-20260803120352: 58/58 passed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/issue-2121-embedding-3072-token-truncate

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core MOS 编排层 / 框架底座 / 跨模块问题 area:model llm + embedder + reranker status:in-progress Someone or AI is working on it | 人工或 AI 正在处理

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants