litellm/litellm-rust/crates/python-bridge/src/token_counter.rs

use std::sync::Arc;

#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
use std::{num::NonZero, thread::available_parallelism};

#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
use litellm_host_python::release_gil;
use litellm_host_python::run_async;
use litellm_token_counter::{
    CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
};
use pyo3::{
    exceptions::{PyRuntimeError, PyValueError},
    prelude::*,
    types::PyAny,
};
use tokio::sync::Semaphore;

use crate::errors::RustBridgeDeclined;

/// Counts the input tokens of a raw request body off the Python event loop with
/// the GIL released. Python owns which requests get here and what to do with
/// the count. At most one encode per core runs at a time; the rest wait in the
/// async task, where a cancelled Python awaiter drops them before any blocking
/// work is scheduled.
#[pyclass(frozen)]
pub(crate) struct TokenCounter {
    inner: Arc<CoreTokenCounter>,
    encode_slots: Arc<Semaphore>,
}

#[pymethods]
impl TokenCounter {
    #[new]
    fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
        #[cfg(feature = "fast")]
        {
            Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
        }
        #[cfg(all(not(feature = "fast"), feature = "huggingface"))]
        {
            Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
        }
        #[cfg(not(any(feature = "fast", feature = "huggingface")))]
        {
            let _ = (py, tokenizer_json);
            Err(RustBridgeDeclined::new_err(
                "tokenizer backend requires the fast or huggingface feature",
            ))
        }
    }

    #[staticmethod]
    fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
        #[cfg(feature = "fast")]
        {
            Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
        }
        #[cfg(not(feature = "fast"))]
        {
            let _ = (py, rank_file);
            Err(RustBridgeDeclined::new_err(
                "tokenizer backend requires the fast feature",
            ))
        }
    }

    #[staticmethod]
    fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
        #[cfg(feature = "fast")]
        {
            Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
        }
        #[cfg(not(feature = "fast"))]
        {
            let _ = (py, rank_file);
            Err(RustBridgeDeclined::new_err(
                "tokenizer backend requires the fast feature",
            ))
        }
    }

    #[staticmethod]
    fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
        #[cfg(feature = "tiktoken")]
        {
            Self::load(py, || CoreTokenCounter::from_tiktoken(encoding))
        }
        #[cfg(not(feature = "tiktoken"))]
        {
            let _ = (py, encoding);
            Err(RustBridgeDeclined::new_err(
                "tokenizer backend requires the tiktoken feature",
            ))
        }
    }

    fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
        let counter = Arc::clone(&self.inner);
        let encode_slots = Arc::clone(&self.encode_slots);
        let body = body.to_vec();
        run_async(
            py,
            async move {
                let _slot = encode_slots
                    .acquire_owned()
                    .await
                    .map_err(|error| Error::Task(error.to_string()))?;
                tokio::task::spawn_blocking(move || count_body(&counter, &body))
                    .await
                    .map_err(|error| Error::Task(error.to_string()))?
            },
            token_count_error_to_pyerr,
        )
    }
}

impl TokenCounter {
    #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
    fn load(
        py: Python<'_>,
        load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
    ) -> PyResult<Self> {
        let inner = release_gil(py, load).map_err(token_count_error_to_pyerr)?;
        Ok(Self {
            inner: Arc::new(inner),
            encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
        })
    }
}

#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
fn encode_parallelism() -> usize {
    available_parallelism().map_or(1, NonZero::get)
}

fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount, Error> {
    let request = CountableRequest::parse(body)?;
    counter.count_request(&request)
}

fn token_count_error_to_pyerr(error: Error) -> PyErr {
    let message = error.to_string();
    match error {
        Error::Load(_)
        | Error::Ranks(_)
        | Error::UnicodeClasses
        | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message),
        Error::RequestParse(_)
        | Error::MissingInput
        | Error::FloatText
        | Error::ContentBlock
        | Error::ArrayItems
        | Error::JsonSerialization(_)
        | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message),
        Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message),
    }
}