1use wasm_bindgen::prelude::{wasm_bindgen, JsError};
11
12#[wasm_bindgen]
17pub struct Encoding {
18 name: &'static str,
20 bpe: &'static tiktoken::CoreBpe,
22}
23
24#[wasm_bindgen]
25impl Encoding {
26 pub fn encode(&self, text: &str) -> Vec<u32> {
31 self.bpe.encode(text)
32 }
33
34 #[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 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 pub fn count(&self, text: &str) -> usize {
55 self.bpe.count(text)
56 }
57
58 #[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 #[wasm_bindgen(js_name = vocabSize, getter)]
69 pub fn vocab_size(&self) -> usize {
70 self.bpe.vocab_size()
71 }
72
73 #[wasm_bindgen(js_name = numSpecialTokens, getter)]
75 pub fn num_special_tokens(&self) -> usize {
76 self.bpe.num_special_tokens()
77 }
78
79 #[wasm_bindgen(getter)]
81 pub fn name(&self) -> String {
82 self.name.to_string()
83 }
84}
85
86#[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#[wasm_bindgen(js_name = getEncoding)]
114pub fn get_encoding(name: &str) -> Result<Encoding, JsError> {
115 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#[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#[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#[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#[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#[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#[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#[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}