Skip to main content

tiktoken_wasm/
lib.rs

1//! WebAssembly bindings for the tiktoken BPE tokenizer.
2//!
3//! Provides browser-compatible wrappers around the core `tiktoken` crate,
4//! enabling high-performance token encoding, decoding, counting, and
5//! cost estimation directly in JavaScript/TypeScript applications.
6//!
7//! All encoding instances are cached globally via `OnceLock`, so repeated
8//! calls to `getEncoding()` with the same name return the same underlying data.
9
10use wasm_bindgen::prelude::{wasm_bindgen, JsError};
11
12/// WASM wrapper around a tiktoken encoding instance.
13///
14/// Created via [`get_encoding`] or [`encoding_for_model`].
15/// Call `.free()` when done to release WASM memory.
16#[wasm_bindgen]
17pub struct Encoding {
18    /// encoding name (e.g. "cl100k_base") — always a static string
19    name: &'static str,
20    /// reference to the globally cached CoreBpe instance
21    bpe: &'static tiktoken::CoreBpe,
22}
23
24#[wasm_bindgen]
25impl Encoding {
26    /// Encode text into token ids (returns `Uint32Array` in JS).
27    ///
28    /// Special tokens like `<|endoftext|>` are treated as ordinary text.
29    /// Use `encodeWithSpecialTokens()` to recognize them.
30    pub fn encode(&self, text: &str) -> Vec<u32> {
31        self.bpe.encode(text)
32    }
33
34    /// Encode text into token ids, recognizing special tokens.
35    ///
36    /// Special tokens (e.g. `<|endoftext|>`) are encoded as their designated ids
37    /// instead of being split into sub-word pieces.
38    #[wasm_bindgen(js_name = encodeWithSpecialTokens)]
39    pub fn encode_with_special_tokens(&self, text: &str) -> Vec<u32> {
40        self.bpe.encode_with_special_tokens(text)
41    }
42
43    /// Decode token ids back to a UTF-8 string.
44    ///
45    /// Uses lossy UTF-8 conversion — invalid byte sequences are replaced with U+FFFD.
46    pub fn decode(&self, tokens: &[u32]) -> String {
47        let bytes = self.bpe.decode(tokens);
48        String::from_utf8_lossy(&bytes).into_owned()
49    }
50
51    /// Count tokens without building the full token id array.
52    ///
53    /// Faster than `encode(text).length` for cases where you only need the count.
54    pub fn count(&self, text: &str) -> usize {
55        self.bpe.count(text)
56    }
57
58    /// Count tokens, recognizing special tokens.
59    ///
60    /// Like `count()` but special tokens (e.g. `<|endoftext|>`) are counted
61    /// as single tokens instead of being split into sub-word pieces.
62    #[wasm_bindgen(js_name = countWithSpecialTokens)]
63    pub fn count_with_special_tokens(&self, text: &str) -> usize {
64        self.bpe.count_with_special_tokens(text)
65    }
66
67    /// Get the number of regular (non-special) tokens in the vocabulary.
68    #[wasm_bindgen(js_name = vocabSize, getter)]
69    pub fn vocab_size(&self) -> usize {
70        self.bpe.vocab_size()
71    }
72
73    /// Get the number of special tokens in the vocabulary.
74    #[wasm_bindgen(js_name = numSpecialTokens, getter)]
75    pub fn num_special_tokens(&self) -> usize {
76        self.bpe.num_special_tokens()
77    }
78
79    /// Get the encoding name (e.g. `"cl100k_base"`).
80    #[wasm_bindgen(getter)]
81    pub fn name(&self) -> String {
82        self.name.to_string()
83    }
84}
85
86/// List all available encoding names.
87///
88/// Returns an array of strings: `["cl100k_base", "o200k_base", ...]`
89#[wasm_bindgen(js_name = listEncodings)]
90pub fn list_encodings() -> Vec<String> {
91    tiktoken::list_encodings()
92        .iter()
93        .map(|s| s.to_string())
94        .collect()
95}
96
97/// Get an encoding by name.
98///
99/// Supported encodings:
100/// - `"cl100k_base"` — GPT-4, GPT-3.5-turbo
101/// - `"o200k_base"` — GPT-4o, GPT-4.1, o1, o3
102/// - `"o200k_harmony"` — gpt-oss (harmony chat format)
103/// - `"p50k_base"` — text-davinci-002/003
104/// - `"p50k_edit"` — text-davinci-edit
105/// - `"r50k_base"` — GPT-3 (davinci, curie, etc.)
106/// - `"gpt2"` — GPT-2 (alias for r50k_base)
107/// - `"llama3"` — Meta Llama 3/4
108/// - `"deepseek_v3"` — DeepSeek V3/R1
109/// - `"qwen2"` — Qwen 2/2.5/3
110/// - `"mistral_v3"` — Mistral/Codestral/Pixtral
111///
112/// Throws `Error` for unknown encoding names.
113#[wasm_bindgen(js_name = getEncoding)]
114pub fn get_encoding(name: &str) -> Result<Encoding, JsError> {
115    // look up the static name from tiktoken's canonical list (single source of truth)
116    let static_name = tiktoken::list_encodings()
117        .iter()
118        .find(|&&n| n == name)
119        .ok_or_else(|| JsError::new(&format!("unknown encoding: {name}")))?;
120    let bpe = tiktoken::get_encoding(name)
121        .ok_or_else(|| JsError::new(&format!("unknown encoding: {name}")))?;
122    Ok(Encoding {
123        name: static_name,
124        bpe,
125    })
126}
127
128/// Get an encoding for a model name (e.g. `"gpt-4o"`, `"o3-mini"`, `"llama-4"`, `"deepseek-r1"`).
129///
130/// Supports models from OpenAI, Meta, DeepSeek, Qwen, and Mistral.
131/// Automatically resolves the model name to the correct encoding.
132/// Throws `Error` for unknown model names.
133#[wasm_bindgen(js_name = encodingForModel)]
134pub fn encoding_for_model(model: &str) -> Result<Encoding, JsError> {
135    let name = tiktoken::model_to_encoding(model)
136        .ok_or_else(|| JsError::new(&format!("unknown model: {model}")))?;
137    let bpe = tiktoken::get_encoding(name)
138        .ok_or_else(|| JsError::new(&format!("unknown encoding: {name}")))?;
139    Ok(Encoding { name, bpe })
140}
141
142/// Map a model name to its encoding name without loading the encoding.
143///
144/// Returns the encoding name string (e.g. `"o200k_base"`) or `null` for unknown models.
145#[wasm_bindgen(js_name = modelToEncoding)]
146pub fn model_to_encoding(model: &str) -> Option<String> {
147    tiktoken::model_to_encoding(model).map(|s| s.to_string())
148}
149
150/// Estimate cost in USD for a given model, input token count, and output token count.
151///
152/// Supports OpenAI, Anthropic Claude, Google Gemini, Meta Llama, DeepSeek, Qwen, and Mistral models.
153/// Throws `Error` for unknown model ids.
154#[wasm_bindgen(js_name = estimateCost)]
155pub fn estimate_cost(
156    model_id: &str,
157    input_tokens: u32,
158    output_tokens: u32,
159) -> Result<f64, JsError> {
160    tiktoken::pricing::estimate_cost(model_id, input_tokens as u64, output_tokens as u64)
161        .ok_or_else(|| JsError::new(&format!("unknown model: {model_id}")))
162}
163
164/// Get model pricing and metadata.
165///
166/// Returns a typed object with: `id`, `provider`, `inputPer1m`, `outputPer1m`,
167/// `cachedInputPer1m`, `contextWindow`, `maxOutput`.
168///
169/// Throws `Error` for unknown model ids.
170#[wasm_bindgen(js_name = getModelInfo)]
171pub fn get_model_info(model_id: &str) -> Result<ModelInfo, JsError> {
172    let model = tiktoken::pricing::get_model(model_id)
173        .ok_or_else(|| JsError::new(&format!("unknown model: {model_id}")))?;
174    Ok(convert_model(model))
175}
176
177/// List all supported models with pricing info.
178///
179/// Returns an array of `ModelInfo` objects.
180#[wasm_bindgen(js_name = allModels)]
181pub fn all_models() -> Vec<ModelInfo> {
182    tiktoken::pricing::all_models()
183        .iter()
184        .map(convert_model)
185        .collect()
186}
187
188/// List models filtered by provider name.
189///
190/// Provider names: `"OpenAI"`, `"Anthropic"`, `"Google"`, `"Meta"`, `"DeepSeek"`, `"Alibaba"`, `"Mistral"`.
191/// Returns an empty array for unknown providers.
192#[wasm_bindgen(js_name = modelsByProvider)]
193pub fn models_by_provider(provider: &str) -> Vec<ModelInfo> {
194    let Some(provider) = parse_provider(provider) else {
195        return Vec::new();
196    };
197
198    tiktoken::pricing::models_by_provider(provider)
199        .iter()
200        .map(|m| convert_model(m))
201        .collect()
202}
203
204fn convert_model(m: &tiktoken::pricing::Model) -> ModelInfo {
205    ModelInfo {
206        id: m.id,
207        provider: m.provider.to_string(),
208        input_per_1m: m.pricing.input_per_1m,
209        output_per_1m: m.pricing.output_per_1m,
210        cached_input_per_1m: m.pricing.cached_input_per_1m,
211        context_window: m.context_window,
212        max_output: m.max_output,
213    }
214}
215
216fn parse_provider(s: &str) -> Option<tiktoken::pricing::Provider> {
217    match s {
218        "OpenAI" => Some(tiktoken::pricing::Provider::OpenAI),
219        "Anthropic" => Some(tiktoken::pricing::Provider::Anthropic),
220        "Google" => Some(tiktoken::pricing::Provider::Google),
221        "Meta" => Some(tiktoken::pricing::Provider::Meta),
222        "DeepSeek" => Some(tiktoken::pricing::Provider::DeepSeek),
223        "Alibaba" => Some(tiktoken::pricing::Provider::Alibaba),
224        "Mistral" => Some(tiktoken::pricing::Provider::Mistral),
225        _ => None,
226    }
227}
228
229/// Model pricing and metadata.
230#[wasm_bindgen]
231#[derive(Clone)]
232pub struct ModelInfo {
233    id: &'static str,
234    provider: String,
235    input_per_1m: f64,
236    output_per_1m: f64,
237    cached_input_per_1m: Option<f64>,
238    context_window: u32,
239    max_output: u32,
240}
241
242#[wasm_bindgen]
243impl ModelInfo {
244    #[wasm_bindgen(getter)]
245    pub fn id(&self) -> String {
246        self.id.to_string()
247    }
248    #[wasm_bindgen(getter)]
249    pub fn provider(&self) -> String {
250        self.provider.clone()
251    }
252    #[wasm_bindgen(getter, js_name = inputPer1m)]
253    pub fn input_per_1m(&self) -> f64 {
254        self.input_per_1m
255    }
256    #[wasm_bindgen(getter, js_name = outputPer1m)]
257    pub fn output_per_1m(&self) -> f64 {
258        self.output_per_1m
259    }
260    #[wasm_bindgen(getter, js_name = cachedInputPer1m)]
261    pub fn cached_input_per_1m(&self) -> Option<f64> {
262        self.cached_input_per_1m
263    }
264    #[wasm_bindgen(getter, js_name = contextWindow)]
265    pub fn context_window(&self) -> u32 {
266        self.context_window
267    }
268    #[wasm_bindgen(getter, js_name = maxOutput)]
269    pub fn max_output(&self) -> u32 {
270        self.max_output
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277
278    #[test]
279    fn all_encodings_roundtrip() {
280        for &name in tiktoken::list_encodings() {
281            let enc = get_encoding(name).unwrap();
282            let text = "hello world 你好 🚀";
283            let tokens = enc.encode(text);
284            let decoded = enc.decode(&tokens);
285            assert_eq!(decoded, text, "roundtrip failed for {name}");
286        }
287    }
288
289    #[test]
290    fn encoding_for_known_models() {
291        let models = [
292            "gpt-4o", "gpt-4", "gpt-3.5-turbo", "llama-4", "deepseek-r1", "qwen3", "mistral-large",
293        ];
294        for model in models {
295            let enc = encoding_for_model(model);
296            assert!(enc.is_ok(), "encoding_for_model failed for {model}");
297        }
298    }
299
300    #[test]
301    fn list_encodings_count() {
302        let names = list_encodings();
303        assert_eq!(names.len(), 9);
304    }
305
306    #[test]
307    fn all_models_count() {
308        let models = all_models();
309        assert_eq!(models.len(), tiktoken::pricing::all_models().len());
310    }
311
312    #[test]
313    fn models_by_valid_provider() {
314        let openai = models_by_provider("OpenAI");
315        assert!(!openai.is_empty());
316        for m in &openai {
317            assert_eq!(m.provider, "OpenAI");
318        }
319    }
320
321    #[test]
322    fn models_by_invalid_provider() {
323        let unknown = models_by_provider("NonExistent");
324        assert!(unknown.is_empty());
325    }
326
327    #[test]
328    fn estimate_cost_known_model() {
329        let cost = estimate_cost("gpt-4o", 1000, 1000).unwrap();
330        assert!(cost > 0.0);
331    }
332
333    #[test]
334    fn estimate_cost_unknown_model() {
335        assert!(estimate_cost("fake-model", 1000, 1000).is_err());
336    }
337
338    #[test]
339    fn get_model_info_known() {
340        let info = get_model_info("gpt-4o").unwrap();
341        assert_eq!(info.id(), "gpt-4o");
342        assert_eq!(info.provider(), "OpenAI");
343        assert!(info.context_window() > 0);
344    }
345
346    #[test]
347    fn get_model_info_unknown() {
348        assert!(get_model_info("fake-model").is_err());
349    }
350
351    #[test]
352    fn unknown_encoding_error() {
353        assert!(get_encoding("nonexistent").is_err());
354    }
355
356    #[test]
357    fn unknown_model_encoding_error() {
358        assert!(encoding_for_model("nonexistent-model-xyz").is_err());
359    }
360
361    #[test]
362    fn model_to_encoding_known() {
363        let name = model_to_encoding("gpt-4o");
364        assert_eq!(name.as_deref(), Some("o200k_base"));
365    }
366
367    #[test]
368    fn model_to_encoding_unknown() {
369        assert!(model_to_encoding("fake-model").is_none());
370    }
371
372    #[test]
373    fn parse_provider_all_variants() {
374        assert!(parse_provider("OpenAI").is_some());
375        assert!(parse_provider("Anthropic").is_some());
376        assert!(parse_provider("Google").is_some());
377        assert!(parse_provider("Meta").is_some());
378        assert!(parse_provider("DeepSeek").is_some());
379        assert!(parse_provider("Alibaba").is_some());
380        assert!(parse_provider("Mistral").is_some());
381        assert!(parse_provider("Unknown").is_none());
382    }
383}