Skip to content

refactor(rust): split token counter backends - #42165

Merged
yujonglee-berri merged 6 commits into
mainfrom
refactor-tokenizer-rs
Sep 20, 2026
Merged

yujonglee-berri merged 6 commits into
mainfrom
refactor-tokenizer-rs

Conversation

@devin-ai-integration

@devin-ai-integration devin-ai-integration Bot commented Sep 20, 2026 •

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • Tokenizer backends are coupled inside one Rust crate
  • Reserved merge ranks can silently produce incorrect token counts

How it solves it:

  • Splits fast, Hugging Face, and tiktoken into backend crates
  • Keeps shared request accounting behind the Tokenizer interface
  • Gives every tokenizer crate its own error module
  • Rejects reserved ranks before counting
  • Checks backend parity and all eight feature combinations

User Flow

Before: a Rust caller loads a malformed rank file and receives an incorrect count

  1. They run the rank-file client below with the reserved-rank input
  2. It accepts the file and prints input_tokens=3
  3. They repeat with the valid input and see input_tokens=2

After: the same caller gets an explicit error for the malformed file

  1. They run the rank-file client below with the reserved-rank input
  2. It prints error=failed to load tokenizer: tiktoken rank file: rank 4294967295 is reserved
  3. They repeat with the valid input and see input_tokens=2

Relevant issues

Affected release

Linear ticket

Pre-Submission checklist

  • I have added meaningful tests
  • The focused tokenizer tests and feature matrix pass locally
  • My PR passes all required CI/CD checks at the current tip
  • My PR's scope is limited to tokenizer backend extraction and validation
  • Greptile confidence is at least 4/5 at the current tip

Validation

The four tokenizer crates pass 168 tests. The counter tests and Python bridge checks pass with each of the eight feature combinations. Workspace Clippy passes with warnings denied, including benchmark targets. Full workspace test compilation, Rust formatting, and the CI shell syntax checks pass

The regression test fails without the reserved-rank validation and passes with it. The shared cl100k/o200k fixtures now run against both fast and tiktoken, including requests with tools and large prompts. Every named tiktoken encoding is checked against its reference encoder on ordinary text, Unicode, whitespace, and special-token spellings. Encoding failures preserve the backend error source

Dependency inspection confirms that the default Python bridge includes only fast, and the counter with default features disabled has no built-in tokenizer dependency

Screenshots / Proof of Fix

These are real calls through the public Rust library, without mocks. Token counting is local and does not call an LLM provider

Shared setup and complete client

Run from the PR checkout. The before snapshot uses the merge base, and the after client uses the current checkout

mkdir -p /tmp/pr42165-audit/before
git archive 0e61ed51c0a00b66d5c1e878ac2b42a4c505cfd0 litellm-rust .cargo rust-toolchain.toml | tar -x -C /tmp/pr42165-audit/before
python3 - <<'PYSETUP'
from pathlib import Path
import base64
root = Path('/tmp/pr42165-audit')
ranks = ''.join(f'{base64.b64encode(bytes([b])).decode()} {b}\n' for b in range(256))
(root / 'reserved.tiktoken').write_text(ranks + 'YWI= 4294967295\n')
(root / 'valid.tiktoken').write_text(ranks + 'YWI= 256\n')
for label, checkout in [('before', root / 'before'), ('after', Path.cwd())]:
    project = root / f'{label}-client'
    (project / 'src').mkdir(parents=True, exist_ok=True)
    (project / 'Cargo.toml').write_text(f'''[package]
name = "token-counter-audit-{label}"
version = "0.1.0"
edition = "2024"
[dependencies]
litellm-token-counter = {{ path = "{checkout}/litellm-rust/crates/token-counter" }}
''')
PYSETUP

Save this as /tmp/pr42165-audit/before-client/src/main.rs and copy it to /tmp/pr42165-audit/after-client/src/main.rs

use litellm_token_counter::{CountableRequest, TokenCounter};

fn main() {
    let path = std::env::args().nth(1).expect("rank file argument");
    let ranks = std::fs::read_to_string(path).unwrap();
    match TokenCounter::from_cl100k_ranks(&ranks) {
        Ok(counter) => {
            let request = CountableRequest::parse(br#"{"prompt":"aab"}"#).unwrap();
            let count = counter.count_request(&request).unwrap();
            println!("input_tokens={}", count.input_tokens);
        }
        Err(error) => println!("error={error}"),
    }
}

Before (0e61ed5)

Reserved rank

  1. Run cargo run --quiet --manifest-path /tmp/pr42165-audit/before-client/Cargo.toml -- /tmp/pr42165-audit/reserved.tiktoken
  2. Output: input_tokens=3

Valid rank

  1. Run cargo run --quiet --manifest-path /tmp/pr42165-audit/before-client/Cargo.toml -- /tmp/pr42165-audit/valid.tiktoken
  2. Output: input_tokens=2

After (661da87)

Reserved rank

  1. Run cargo run --quiet --manifest-path /tmp/pr42165-audit/after-client/Cargo.toml -- /tmp/pr42165-audit/reserved.tiktoken
  2. Output: error=failed to load tokenizer: tiktoken rank file: rank 4294967295 is reserved

Valid rank

  1. Run cargo run --quiet --manifest-path /tmp/pr42165-audit/after-client/Cargo.toml -- /tmp/pr42165-audit/valid.tiktoken
  2. Output: input_tokens=2

Type

Refactoring and bug fix

Caveats (if any)

Medium

  • Python wheels still enable only the fast backend
  • Current-tip CI and automated reviews are not yet confirmed

Low

  • Fresh wheel runtime and size checks remain with CI
  • Earlier Greptile and Codecov reports predate the current changes
  • No Veria or Bugbot review was available during this audit

Final Attestation

  • Tests cover backend parity, invalid ranks, errors, and feature combinations

yujonglee-berri and others added 2 commits September 20, 2026 14:10
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".

  • Disable automatic comment, CI, and merge conflict monitoring

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@codspeed

codspeed Bot commented Sep 20, 2026 •

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing refactor-tokenizer-rs (661da87) with main (e7c8928)

Open in CodSpeed

@greptile-apps

greptile-apps Bot commented Sep 20, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

The PR appears safe to merge, with no new actionable defects identified and no previous findings outstanding

Summary

Splits Rust tokenization into fast, Hugging Face, and tiktoken backend crates while retaining shared request accounting through the Tokenizer interface

  • Rejects reserved merge ranks before counting
  • Keeps Python wheels on the fast backend and preserves existing Python tokenizer routing
  • Restores benchmark dependencies and expands backend parity coverage and CI checks across all eight feature combinations

Reviews (2) · Last reviewed commit: "fix(rust): validate tokenizer ranks and ..."

Comment thread litellm-rust/crates/token-counter/Cargo.toml
Comment thread litellm-rust/crates/token-counter/tests/token_counter.rs Outdated
Comment thread tests/test_litellm/rust_bridge/test_token_counter.py Outdated
@greptile-apps

This comment has been minimized.

@codecov

codecov Bot commented Sep 20, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

yujonglee-berri and others added 3 commits September 20, 2026 21:49
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@yujonglee-berri

Copy link
Copy Markdown
Contributor

@greptile-apps wdyt

@yujonglee-berri
yujonglee-berri merged commit 65c6495 into main Sep 20, 2026
97 checks passed
@yujonglee-berri
yujonglee-berri deleted the refactor-tokenizer-rs branch September 20, 2026 22:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant