mistralrs_core/utils/
tokens.rs

1use std::{env, fs};
2use thiserror::Error;
3
4use anyhow::Result;
5use tracing::info;
6
7use crate::pipeline::TokenSource;
8
9#[derive(Error, Debug)]
10enum TokenRetrievalError {
11    #[error("No home directory.")]
12    HomeDirectoryMissing,
13}
14
15/// This reads a token from a specified source. If the token cannot be read, a warning is logged with `tracing`
16/// and *no token is used*.
17pub(crate) fn get_token(source: &TokenSource) -> Result<Option<String>> {
18    fn skip_token(input: &str) -> Option<String> {
19        info!("Could not load token at {input:?}, using no HF token.");
20        None
21    }
22
23    let token = match source {
24        TokenSource::Literal(data) => Some(data.clone()),
25        TokenSource::EnvVar(envvar) => env::var(envvar).ok().or_else(|| skip_token(envvar)),
26        TokenSource::Path(path) => fs::read_to_string(path).ok().or_else(|| skip_token(path)),
27        TokenSource::CacheToken => {
28            let home = format!(
29                "{}/.cache/huggingface/token",
30                dirs::home_dir()
31                    .ok_or(TokenRetrievalError::HomeDirectoryMissing)?
32                    .display()
33            );
34
35            fs::read_to_string(home.clone())
36                .ok()
37                .or_else(|| skip_token(&home))
38        }
39        TokenSource::None => None,
40    };
41
42    Ok(token.map(|s| s.trim().to_string()))
43}