fix(embedding): configurable tokenizer for non-OpenAI embedding models

[Vansh-Sharma27]

Problem

When a non-OpenAI embedding model is configured via EMBEDDING_MODEL_CONFIG__MODEL (made possible by #678), the embedding client still counts tokens with tiktoken and silently falls back to cl100k_base for any model name tiktoken doesn't recognize (src/embedding_client.py:208-211). For models like BAAI/bge-m3 (XLM-RoBERTa tokenizer), cl100k_base undercounts vs the provider's real tokenizer, so prepare_chunks emits "within-limit" chunks the provider then rejects with HTTP 400. The reconciler retries the unchanged payload 20 times (~3h, MAX_SYNC_ATTEMPTS=20 × SYNC_BACKOFF=10min) and marks MessageEmbedding.sync_state='failed', permanently excluding the message from vector search (src/utils/search.py:166 filters embedding IS NOT NULL).

Cause

The chunk-size decision is made locally with tiktoken before the provider call, and the chosen chunk text is persisted. Retry re-embeds that stored text without re-chunking, so a deterministic size error can't self-heal. There was no way to tell Honcho which tokenizer the configured model actually uses.

Runtime-measured undercount (cl100k_base vs BAAI/bge-m3, via the tokenizers package):

• technical text: +44.0% (matches the issue reporter's 30-42% range)
• mixed ru/en: +18.6%
• pure CJK/Russian: over-counts (harmless, just extra chunking)

End-to-end repro without a live provider: 24,360 chars of technical text → prepare_chunks returns 1 chunk (cl100k est 5,801 < 8,192) → bge-m3 real count 8,355 > 8,192 → provider 400 → 20 retries → failed → invisible to vector search.

Solution

Add EMBEDDING_MODEL_CONFIG__TOKENIZER (config + .env.template + config.toml.example + docs/v3/contributing/configuration.mdx + docs/v3/contributing/changing-embeddings.mdx):

• unset (default): tiktoken auto-detection — current behavior, fully backwards compatible
• tiktoken:<encoding>: explicit tiktoken encoding
• hf:<repo-id>: HuggingFace tokenizer via the optional honcho[tokenizers] extra
• file:<path>: local HF tokenizer JSON

The TokenizerLike protocol (encode/decode) decouples the embedding pipeline from tiktoken.Encoding. The _HuggingFaceTokenizer adapter encodes without special tokens and reserves the special-token overhead ([CLS]/[SEP]) from the chunk budget so provider-side counts stay exactly within limit. Unknown models now log a warning pointing at the new setting. Invalid specs raise ValidationException (repo-standard, per src/exceptions.py). The singleton rebuild signature (_get_settings_signature) includes tokenizer so runtime config changes rebuild the client.

Runtime-verified the fix: same 24,360-char input with tokenizer=hf:BAAI/bge-m3 → 2 chunks (8,192 / 1,803 real tokens, both ≤ 8,192), zero over-limit.

Testing

• uv run pytest tests/llm/test_embedding_client.py — 30 passed (16 pre-existing + 14 new covering: spec resolution, tiktoken/hf/file branches, unknown-encoding/invalid-prefix/empty-name ValidationExceptions, missing tokenizers package, unknown-model warning, HF special-token overhead, chunking follows configured tokenizer, singleton signature invalidation).
• uv run pytest tests/llm/ — 252 passed (includes test_model_config.py, schema callers).
• uv run ruff check + uv run ruff format --check — clean.
• uv run basedpyright — 0 errors, 0 warnings.
• DB-backed suites (tests/deriver/, tests/integration/) need Postgres+pgvector, which isn't available in this environment; they fail identically on a pristine checkout (connection refused). Those tests mock the embedding client entirely (tests/conftest.py:484-536), so they don't exercise this diff. CI runs the full matrix with a real pgvector service.

Risk

• Default behavior unchanged: unset TOKENIZER preserves tiktoken auto-detection; existing OpenAI deployments see no change.
• uv.lock: adds tokenizers, huggingface-hub, hf-xet, fsspec; click 8.3.3 → 8.4.2 is forced by huggingface-hub>=1.x's click dependency (satisfies existing typer/uvicorn constraints).
• Special-token budget: the HF adapter subtracts the special-token overhead from max_embedding_tokens. For pathological configs where the limit is ≤ the overhead, the budget could go non-positive — such a config is already inoperable (limit smaller than 2 tokens) and no guard is added.

Out of scope (follow-ups)

The RCA identified several adjacent issues. They're intentionally not bundled here to keep this PR focused on the reported bug; happy to open separate issues/PRs for any the maintainers want:

• Recovery/reindex command for rows already stuck in sync_state='failed' (a future-only fix can't heal them; the reconciler never re-chunks).
• scripts/generate_message_embeddings.py stores full message.content per chunk embedding rather than the actual chunk text.
• ConclusionCreate validates the embedding limit with o200k_base while the client may use a different encoding (src/schemas/api.py:506-519).
• Typed exceptions separating dimension mismatch from token overflow (currently both raise ValueError, misclassified as "token limit" in src/utils/search.py:390-394).
• Live-embedding CI matrix (.github/workflows/live-llm-tests.yml path filters don't include src/embedding_client.py).

Fixes #827

Summary by CodeRabbit

• New Features
  • Added configurable tokenizers for embedding models via EMBEDDING_MODEL_CONFIG__TOKENIZER.
  • Embedding token counting and text chunking now use the selected tokenizer for more accurate limits.
  • Supports tiktoken encodings, Hugging Face tokenizers, and local tokenizer files (with an optional tokenizers package).
  • Embedding tokenizer/client is warmed during app startup to reduce first-request delays.
• Documentation
  • Updated configuration examples and embedding guidance with tokenizer options, required setup for non-OpenAI tokenizers, fallback behavior, and tokenizer-mismatch warnings.

[coderabbitai]

Walkthrough

Embedding configuration now accepts tokenizer specifications, the embedding client resolves and uses them for token-aware chunking, and application startup eagerly warms the client. Documentation, optional dependencies, validation tests, and startup wiring are included.

Changes

Embedding tokenizer configuration

Layer / File(s)	Summary
Tokenizer configuration contract and setup .env.template, config.toml.example, docs/v3/contributing/*.mdx, pyproject.toml, src/config.py	Embedding settings carry an optional tokenizer specification, documented through configuration examples, with a tokenizers optional dependency extra.
Tokenizer resolution and adapters src/embedding_client.py	Tokenizer resolution supports explicit tiktoken encodings, HuggingFace repositories, local tokenizer files, validation, and fallback behavior.
Tokenizer-aware embedding chunking and validation src/embedding_client.py, tests/llm/test_embedding_client.py	Embedding initialization and chunking use the resolved tokenizer, account for special-token overhead, reinitialize when tokenizer settings change, and test supported and invalid specifications.
Startup tokenizer warmup src/main.py, src/deriver/__main__.py	Application and deriver startup paths eagerly warm the embedding client before later embedding work.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Operator
  participant ApplicationStartup
  participant EmbeddingClient
  participant TokenizerResolver
  participant EmbeddingProvider
  Operator->>ApplicationStartup: configure model and tokenizer
  ApplicationStartup->>EmbeddingClient: warmup()
  EmbeddingClient->>TokenizerResolver: resolve tokenizer specification
  TokenizerResolver-->>EmbeddingClient: return tokenizer
  EmbeddingClient->>EmbeddingClient: count tokens and split text
  EmbeddingClient->>EmbeddingProvider: send tokenizer-sized chunks

Loading

Possibly related PRs

• plastic-labs/honcho#647: Both changes modify embedding tokenizer selection and token counting behavior.

Suggested reviewers: rajat-ahuja1997

Poem

I’m a rabbit with tokens to spare,
Splitting long messages with care.
Tiktoken or HF,
Local files join the quest,
So embeddings hop safely through air.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 18.18% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding a configurable tokenizer for non-OpenAI embedding models.
Linked Issues check	✅ Passed	The changes implement configurable embedding tokenizers, preserve fallback behavior, add validation/tests/docs, and address the non-OpenAI token mismatch in #827.
Out of Scope Changes check	✅ Passed	No clear out-of-scope changes are present; docs, tests, dependency updates, and startup warmups all support the tokenizer fix.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

src/embedding_client.py (1)

191-213: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Avoid hf: downloads on the async embedding path. Tokenizer.from_pretrained(...) can fetch from the Hugging Face Hub on first use, and _get_client() constructs _EmbeddingClient synchronously under a class lock. That means the first embed/batch_embed call for a new tokenizer config can stall the event loop and block other callers behind the download. Consider pre-warming at startup, keeping hf: tokenizers in an offline cache, or moving the load off the request path.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/embedding_client.py` around lines 191 - 213, Move Hugging Face tokenizer
loading out of the synchronous request path used by _get_client and
_EmbeddingClient, or pre-warm/cache hf: tokenizers during startup so
Tokenizer.from_pretrained is never invoked on the first embed or batch_embed
call. Preserve the existing validation and file: loading behavior in
_load_huggingface_tokenizer while ensuring concurrent callers are not blocked by
a Hub download under the class lock.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@src/embedding_client.py`:
- Around line 191-213: Move Hugging Face tokenizer loading out of the
synchronous request path used by _get_client and _EmbeddingClient, or
pre-warm/cache hf: tokenizers during startup so Tokenizer.from_pretrained is
never invoked on the first embed or batch_embed call. Preserve the existing
validation and file: loading behavior in _load_huggingface_tokenizer while
ensuring concurrent callers are not blocked by a Hub download under the class
lock.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 5cb89f41-9417-4ffc-a8d1-e175e3207aeb

📥 Commits

Reviewing files that changed from the base of the PR and between 055d73b and 2097325.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (8)

• .env.template
• config.toml.example
• docs/v3/contributing/changing-embeddings.mdx
• docs/v3/contributing/configuration.mdx
• pyproject.toml
• src/config.py
• src/embedding_client.py
• tests/llm/test_embedding_client.py

[coderabbitai]

🧹 Nitpick comments (1)

src/main.py (1)

146-149: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Offload blocking initialization to a worker thread.

The warmup() method is synchronous and can perform blocking network I/O (such as downloading HuggingFace tokenizer files). Calling a synchronous I/O function directly inside an async def function blocks the asyncio event loop. Even during application startup, it is considered best practice to offload blocking operations to prevent delaying or stalling background tasks (such as signal handlers or telemetry).

• src/main.py#L146-L149: Consider wrapping this call with await asyncio.to_thread(embedding_client.warmup). You may need to import asyncio if it is not already imported.
• src/deriver/__main__.py#L74-L78: Consider wrapping this call with await asyncio.to_thread(embedding_client.warmup).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/main.py` around lines 146 - 149, Offload the synchronous embedding
initialization from the asyncio event loop by awaiting asyncio.to_thread around
embedding_client.warmup() in src/main.py lines 146-149 and
src/deriver/__main__.py lines 74-78; add the asyncio import where needed and
preserve the existing startup flow.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@src/main.py`:
- Around line 146-149: Offload the synchronous embedding initialization from the
asyncio event loop by awaiting asyncio.to_thread around
embedding_client.warmup() in src/main.py lines 146-149 and
src/deriver/__main__.py lines 74-78; add the asyncio import where needed and
preserve the existing startup flow.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: ea21182e-d9f3-4d1e-890d-9c1685d7381a

📥 Commits

Reviewing files that changed from the base of the PR and between 2097325 and bfa06a5.

📒 Files selected for processing (3)

• src/deriver/__main__.py
• src/embedding_client.py
• src/main.py

🚧 Files skipped from review as they are similar to previous changes (1)

• src/embedding_client.py

[karaokedjodua]

Finn-loop review of bfa06a5

CI: not configured (stale branch protection)
Mergeability: clean (MERGEABLE)

Review

Adds EMBEDDING_MODEL_CONFIG__TOKENIZER to let non-OpenAI embedding models use their actual tokenizer for chunk-size decisions instead of tiktoken's cl100k_base fallback, preventing silent provider rejections that permanently exclude messages from vector search. Clean implementation with TokenizerLike protocol, _HuggingFaceTokenizer adapter, special-token budget reservation, and 14 new tests covering all branches. Default behavior unchanged (unset → tiktoken auto-detect).

1. Must fix before merge

• [CI] Stale branch protection rule blocks merge: mergeStateStatus: BLOCKED but mergeable: MERGEABLE and gh pr checks --required reports "no required checks configured on the 'fix/827-embedding-tokenizer' branch." No required checks are configured on this branch. An admin must update or remove the stale branch-protection rule to unblock merges.

2. Should fix soon

None.

3. Safe to merge

Code is solid — comprehensive tests, clean backwards-compatible design, well-scoped to the reported issue. Blocked only by the stale branch-protection rule (repo config, not code).

Labels to apply: loop-changes-requested (for [CI] must-fix), needs-human-review (no configured required CI)