use serde::{Deserialize, Serialize};
use typed_builder::TypedBuilder;
use super::StorageConfig;
use crate::Result;
pub const HF_TOKEN: &str = "hf.token";
pub const HF_ENDPOINT: &str = "hf.endpoint";
pub const HF_REVISION: &str = "hf.revision";
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, TypedBuilder)]
pub struct HfConfig {
#[builder(default, setter(strip_option, into))]
pub token: Option<String>,
#[builder(default, setter(strip_option, into))]
pub endpoint: Option<String>,
#[builder(default, setter(strip_option, into))]
pub revision: Option<String>,
}
impl TryFrom<&StorageConfig> for HfConfig {
type Error = crate::Error;
fn try_from(config: &StorageConfig) -> Result<Self> {
let props = config.props();
let mut cfg = HfConfig::default();
if let Some(token) = props.get(HF_TOKEN) {
cfg.token = Some(token.clone());
}
if let Some(endpoint) = props.get(HF_ENDPOINT) {
cfg.endpoint = Some(endpoint.clone());
}
if let Some(revision) = props.get(HF_REVISION) {
cfg.revision = Some(revision.clone());
}
Ok(cfg)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hf_config_builder() {
let cfg = HfConfig::builder()
.token("hf_mytoken")
.endpoint("https://huggingface.co")
.revision("dev")
.build();
assert_eq!(cfg.token.as_deref(), Some("hf_mytoken"));
assert_eq!(cfg.endpoint.as_deref(), Some("https://huggingface.co"));
assert_eq!(cfg.revision.as_deref(), Some("dev"));
}
#[test]
fn test_hf_config_from_storage_config() {
let storage_config = StorageConfig::new()
.with_prop(HF_TOKEN, "hf_abc123")
.with_prop(HF_ENDPOINT, "https://huggingface.co");
let cfg = HfConfig::try_from(&storage_config).unwrap();
assert_eq!(cfg.token.as_deref(), Some("hf_abc123"));
assert_eq!(cfg.endpoint.as_deref(), Some("https://huggingface.co"));
}
#[test]
fn test_hf_config_empty() {
let cfg = HfConfig::try_from(&StorageConfig::new()).unwrap();
assert_eq!(cfg.token, None);
assert_eq!(cfg.endpoint, None);
}
}