1use anyhow::Result;
 2use language_model_core::{LanguageModelRequest, Role};
 3
 4use crate::Model;
 5
 6/// Count tokens for an xAI model using tiktoken. This is synchronous;
 7/// callers should spawn it on a background thread if needed.
 8pub fn count_xai_tokens(request: LanguageModelRequest, model: Model) -> Result<u64> {
 9    let messages = request
10        .messages
11        .into_iter()
12        .map(|message| tiktoken_rs::ChatCompletionRequestMessage {
13            role: match message.role {
14                Role::User => "user".into(),
15                Role::Assistant => "assistant".into(),
16                Role::System => "system".into(),
17            },
18            content: Some(message.string_contents()),
19            name: None,
20            function_call: None,
21        })
22        .collect::<Vec<_>>();
23
24    let model_name = if model.max_token_count() >= 100_000 {
25        "gpt-4o"
26    } else {
27        "gpt-4"
28    };
29    tiktoken_rs::num_tokens_from_messages(model_name, &messages).map(|tokens| tokens as u64)
30}