Skip to content

fix: batch token counting during indexing to prevent unbounded plugin requests - #39571

Open
mesrop2008 wants to merge 1 commit into
langgenius:mainfrom
mesrop2008:fix/batch-token-counting-indexing
Open

fix: batch token counting during indexing to prevent unbounded plugin requests#39571
mesrop2008 wants to merge 1 commit into
langgenius:mainfrom
mesrop2008:fix/batch-token-counting-indexing

Conversation

@mesrop2008

@mesrop2008 mesrop2008 commented Jul 25, 2026

Copy link
Copy Markdown

Summary

Fixes #39560

Token counting during indexing sent the entire document's chunks to the
embedding plugin in a single, unbounded request via
get_text_embedding_num_tokens, while the embedding call itself was already
batched via MAX_CHUNKS in cached_embedding.py. For large documents this
produced a single multi-megabyte request that could fail with 413 Payload Too Large from an intermediate proxy, OOM-kill a memory-constrained
plugin-daemon, or exceed a tokenizer backend's own batch limits (e.g. TEI, as
previously seen in #15035) — always before a single embedding request was
even made.

Changes

  • api/core/rag/embedding/token_counter.py:
    • calculate_segment_token_counts now batches texts using the same
      ModelPropertyKey.MAX_CHUNKS logic already used by cached_embedding.py,
      concatenating per-batch results in order (the function already returned
      one token count per document, so no summing is needed).
    • Added a second, independent cap on serialized UTF-8 byte size per batch
      (_iter_batches), so a single oversized chunk — or text that expands
      after encoding — can't produce an oversized batch even when the chunk
      count alone would allow it.

Testing

  • Added unit tests covering: batching by MAX_CHUNKS, batching by byte-size
    cap independent of chunk count, and the existing economy/empty-document/
    high-quality code paths (updated to account for the new schema lookup).
  • Manually verified against a real (non-mocked) indexing run: built a local
    dify-api image from this branch, swapped it into an existing Docker
    Compose stack via docker-compose.override.yaml, and indexed a ~17 MB
    document split into ~100-character segments (18,500 chunks) with
    high_quality indexing. Worker logs show hundreds of small, successful
    .../text_embedding/num_tokens requests instead of one unbounded request,
    and indexing completed without any 413/OOM/timeout failure at the
    token-counting stage.

Screenshots

N/A — backend-only change; see worker log excerpt in the linked issue
discussion for before/after request pattern.

Checklist

  • This change requires a documentation update, included: Dify Document
  • I understand that this PR may be closed in case there was no previous discussion or issues. (This doesn't apply to typos!)
  • I've added a test for each change that was introduced, and I tried as much as possible to make a single atomic change.
  • I've updated the documentation accordingly.
  • I ran uv run ruff format / uv run ruff check on the changed files and uv run pytest on the affected test file to appease the lint gods

@dosubot dosubot Bot added the size:M This PR changes 30-99 lines, ignoring generated files. label Jul 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:M This PR changes 30-99 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Indexing: token counting sends the whole document in one plugin request (no max_chunks batching)

1 participant