xai_rust/api/

tokenizer.rs

//! Tokenizer API for counting and getting tokens.

use crate::client::XaiClient;
use crate::models::tokenizer::{TokenizeRequest, TokenizeResponse};
use crate::{Error, Result};

/// Tokenizer API client.
#[derive(Debug, Clone)]
pub struct TokenizerApi {
    client: XaiClient,
}

impl TokenizerApi {
    pub(crate) fn new(client: XaiClient) -> Self {
        Self { client }
    }

    /// Tokenize text using a specific model.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use xai_rust::{XaiClient, TokenizeRequest};
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = XaiClient::from_env()?;
    ///
    /// let request = TokenizeRequest::new("grok-4", "Hello, world!");
    /// let response = client.tokenizer().tokenize(request).await?;
    ///
    /// println!("Token count: {}", response.count());
    /// println!("Tokens: {:?}", response.tokens);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn tokenize(&self, request: TokenizeRequest) -> Result<TokenizeResponse> {
        let url = format!("{}/tokenize-text", self.client.base_url());

        let response = self
            .client
            .send(self.client.http().post(&url).json(&request))
            .await?;

        if !response.status().is_success() {
            return Err(Error::from_response(response).await);
        }

        Ok(response.json().await?)
    }

    /// Tokenize text using a model name and text string.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use xai_rust::XaiClient;
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = XaiClient::from_env()?;
    ///
    /// let response = client.tokenizer()
    ///     .tokenize_text("grok-4", "Hello, world!")
    ///     .await?;
    ///
    /// println!("Tokens: {:?}", response.tokens);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn tokenize_text(
        &self,
        model: impl Into<String>,
        text: impl Into<String>,
    ) -> Result<TokenizeResponse> {
        self.tokenize(TokenizeRequest::new(model, text)).await
    }

    /// Count tokens in text using a specific model.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use xai_rust::XaiClient;
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = XaiClient::from_env()?;
    ///
    /// let count = client.tokenizer()
    ///     .count_tokens("grok-4", "Hello, world!")
    ///     .await?;
    ///
    /// println!("Token count: {}", count);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn count_tokens(
        &self,
        model: impl Into<String>,
        text: impl Into<String>,
    ) -> Result<usize> {
        let response = self.tokenize_text(model, text).await?;
        Ok(response.count())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use wiremock::matchers::{method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    #[tokio::test]
    async fn tokenize_forwards_model_and_text_payload() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/tokenize-text"))
            .respond_with(move |req: &wiremock::Request| {
                let body = serde_json::from_slice::<serde_json::Value>(&req.body).unwrap();
                assert_eq!(body["model"], "grok-4");
                assert_eq!(body["text"], "Hello tokenizer");
                ResponseTemplate::new(200).set_body_json(json!({
                    "tokens": [10, 20, 30],
                    "token_count": 3
                }))
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let response = client
            .tokenizer()
            .tokenize(TokenizeRequest::new("grok-4", "Hello tokenizer"))
            .await
            .unwrap();
        assert_eq!(response.tokens, vec![10, 20, 30]);
        assert_eq!(response.count(), 3);
    }

    #[tokio::test]
    async fn count_tokens_prefers_explicit_token_count() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/tokenize-text"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "tokens": [1, 2],
                "token_count": 9
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let count = client
            .tokenizer()
            .count_tokens("grok-4", "count this")
            .await
            .unwrap();
        assert_eq!(count, 9);
    }

    #[tokio::test]
    async fn count_tokens_falls_back_to_token_vector_length() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/tokenize-text"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "tokens": [7, 8, 9, 10]
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let count = client
            .tokenizer()
            .count_tokens("grok-4", "fallback count")
            .await
            .unwrap();
        assert_eq!(count, 4);
    }

    #[tokio::test]
    async fn tokenize_returns_api_error_for_non_success_response() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/tokenize-text"))
            .respond_with(ResponseTemplate::new(400).set_body_json(json!({
                "error": {
                    "message": "bad tokenize request",
                    "type": "invalid_request_error"
                }
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let err = client
            .tokenizer()
            .tokenize_text("grok-4", "bad input")
            .await
            .unwrap_err();
        match err {
            Error::Api {
                status,
                message,
                error_type,
            } => {
                assert_eq!(status, 400);
                assert_eq!(message, "bad tokenize request");
                assert_eq!(error_type.as_deref(), Some("invalid_request_error"));
            }
            other => panic!("expected api error, got {other:?}"),
        }
    }
}