osdk_core/model/provider/
huggingface.rs

1use async_trait::async_trait;
2use serde::Deserialize;
3
4use crate::backend::Ctx;
5use crate::error::{Error, Result};
6use crate::model::provider::{get_cached_json, ModelProvider, RemoteModelFile, RemoteSnapshot};
7use crate::model::{ModelRef, ProviderId};
8
9pub struct HuggingFace {
10    token: Option<String>,
11    allow_auth: bool,
12}
13
14impl HuggingFace {
15    pub fn new(allow_auth: bool) -> Self {
16        Self {
17            token: None,
18            allow_auth,
19        }
20    }
21
22    #[cfg(test)]
23    pub fn with_token(token: impl Into<String>) -> Self {
24        Self {
25            token: Some(token.into()),
26            allow_auth: true,
27        }
28    }
29}
30
31impl Default for HuggingFace {
32    fn default() -> Self {
33        Self::new(true)
34    }
35}
36
37#[derive(Debug, Deserialize)]
38struct ModelInfo {
39    sha: String,
40    #[serde(default)]
41    siblings: Vec<RepoSibling>,
42}
43
44#[derive(Debug, Deserialize)]
45struct RepoSibling {
46    rfilename: String,
47    #[serde(default)]
48    size: Option<u64>,
49    #[serde(default, rename = "blobId")]
50    blob_id: Option<String>,
51    #[serde(default)]
52    lfs: Option<LfsFile>,
53}
54
55#[derive(Debug, Deserialize)]
56struct LfsFile {
57    sha256: String,
58    size: u64,
59}
60
61#[async_trait]
62impl ModelProvider for HuggingFace {
63    async fn resolve(
64        &self,
65        ctx: &Ctx,
66        reference: &ModelRef,
67        endpoint: &str,
68    ) -> Result<RemoteSnapshot> {
69        if reference.provider != ProviderId::HuggingFace {
70            return Err(Error::config(format!(
71                "Hugging Face provider cannot resolve {}",
72                reference.provider
73            )));
74        }
75        let endpoint = endpoint.trim_end_matches('/');
76        let metadata_url = metadata_url(endpoint, &reference.repository, &reference.revision)?;
77        let headers = if self.allow_auth {
78            auth_headers(self.token.as_deref())
79        } else {
80            Vec::new()
81        };
82        let cache_identity = format!(
83            "{}:{}:{}@{}",
84            reference.provider, endpoint, reference.repository, reference.revision
85        );
86        let info: ModelInfo = get_cached_json(
87            ctx,
88            reference.provider.as_str(),
89            &cache_identity,
90            &metadata_url,
91            &headers,
92        )
93        .await?;
94        if info.sha.trim().is_empty() {
95            return Err(Error::other("Hugging Face response has no commit revision"));
96        }
97        let mut files = Vec::with_capacity(info.siblings.len());
98        for sibling in info.siblings {
99            crate::model::safe_relative_path(&sibling.rfilename)?;
100            let sha256 = sibling.lfs.as_ref().map(|lfs| lfs.sha256.clone());
101            let size = sibling.lfs.as_ref().map(|lfs| lfs.size).or(sibling.size);
102            files.push(RemoteModelFile {
103                url: file_url(
104                    endpoint,
105                    &reference.repository,
106                    &info.sha,
107                    &sibling.rfilename,
108                )?,
109                path: sibling.rfilename,
110                size,
111                etag: sha256.clone().or(sibling.blob_id),
112                sha256,
113                headers: headers.clone(),
114            });
115        }
116        if files.is_empty() {
117            return Err(Error::other(format!(
118                "Hugging Face repository {}@{} contains no files",
119                reference.repository, reference.revision
120            )));
121        }
122        Ok(RemoteSnapshot {
123            revision: info.sha,
124            endpoint: endpoint.to_string(),
125            files,
126        })
127    }
128}
129
130fn metadata_url(endpoint: &str, repository: &str, revision: &str) -> Result<String> {
131    let mut url = reqwest::Url::parse(endpoint)
132        .map_err(|error| Error::config(format!("invalid Hugging Face endpoint: {error}")))?;
133    {
134        let mut segments = url
135            .path_segments_mut()
136            .map_err(|_| Error::config("Hugging Face endpoint cannot be a base URL"))?;
137        segments.pop_if_empty();
138        segments.extend(["api", "models"]);
139        for part in repository.split('/') {
140            segments.push(part);
141        }
142        segments.extend(["revision", revision]);
143    }
144    url.query_pairs_mut().append_pair("blobs", "true");
145    Ok(url.into())
146}
147
148fn file_url(endpoint: &str, repository: &str, revision: &str, path: &str) -> Result<String> {
149    let mut url = reqwest::Url::parse(endpoint)
150        .map_err(|error| Error::config(format!("invalid Hugging Face endpoint: {error}")))?;
151    {
152        let mut segments = url
153            .path_segments_mut()
154            .map_err(|_| Error::config("Hugging Face endpoint cannot be a base URL"))?;
155        segments.pop_if_empty();
156        for part in repository.split('/') {
157            segments.push(part);
158        }
159        segments.extend(["resolve", revision]);
160        for part in path.split('/') {
161            segments.push(part);
162        }
163    }
164    Ok(url.into())
165}
166
167fn auth_headers(explicit: Option<&str>) -> Vec<(String, String)> {
168    if let Some(value) = explicit {
169        return vec![("Authorization".into(), format!("Bearer {value}"))];
170    }
171    for key in ["OSDK_HF_TOKEN", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"] {
172        if let Ok(value) = std::env::var(key) {
173            let value = value.trim();
174            if !value.is_empty() {
175                return vec![("Authorization".into(), format!("Bearer {value}"))];
176            }
177        }
178    }
179    Vec::new()
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn builds_encoded_metadata_and_file_urls() {
188        let metadata =
189            metadata_url("https://hub.example.test", "owner/repo", "feature/branch").unwrap();
190        assert!(metadata.contains("/api/models/owner/repo/revision/feature%2Fbranch"));
191        assert!(metadata.contains("blobs=true"));
192        let file = file_url(
193            "https://hub.example.test",
194            "owner/repo",
195            "abc123",
196            "weights/model.bin",
197        )
198        .unwrap();
199        assert_eq!(
200            file,
201            "https://hub.example.test/owner/repo/resolve/abc123/weights/model.bin"
202        );
203    }
204}
osdk_core/model/provider/huggingface.rs

osdk_core/model/provider/
huggingface.rs
