Skip to main content

opendal_service_hf/
config.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use super::backend::HfBuilder;
19use super::core::HfDownloadMode;
20use super::core::HfRepoType;
21use super::core::HfUri;
22use serde::Deserialize;
23use serde::Serialize;
24use std::fmt::Debug;
25
26use super::HUGGINGFACE_SCHEME;
27
28/// Configuration for Hugging Face service support.
29#[derive(Default, Serialize, Deserialize, Clone, PartialEq, Eq)]
30#[serde(default)]
31#[non_exhaustive]
32pub struct HfConfig {
33    /// Repo type of this backend. Required.
34    pub repo_type: Option<HfRepoType>,
35    /// Repo id of this backend.
36    ///
37    /// This is required.
38    pub repo_id: Option<String>,
39    /// Revision of this backend.
40    ///
41    /// Default is main.
42    pub revision: Option<String>,
43    /// Root of this backend. Can be "/path/to/dir".
44    ///
45    /// Default is "/".
46    pub root: Option<String>,
47    /// Token of this backend.
48    ///
49    /// This is optional.
50    pub token: Option<String>,
51    /// Endpoint of the Hugging Face Hub.
52    ///
53    /// The default is `https://huggingface.co`.
54    pub endpoint: Option<String>,
55    /// Download mode. Either `xet` (default) or `http`.
56    ///
57    /// When unset, the mode is resolved from the `HF_HUB_DISABLE_XET`
58    /// environment variable: a non-empty value forces `http`, otherwise it
59    /// defaults to `xet`. An explicit value here takes precedence.
60    ///
61    /// See <https://huggingface.co/docs/huggingface_hub/package_reference/environment_variables#hfhubdisablexet>.
62    pub download_mode: Option<HfDownloadMode>,
63}
64
65impl Debug for HfConfig {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("HfConfig")
68            .field(
69                "repo_type",
70                &self.repo_type.as_ref().map(HfRepoType::as_str),
71            )
72            .field("repo_id", &self.repo_id)
73            .field("revision", &self.revision)
74            .field("root", &self.root)
75            .field("download_mode", &self.download_mode)
76            .finish_non_exhaustive()
77    }
78}
79
80impl opendal_core::Configurator for HfConfig {
81    type Builder = HfBuilder;
82
83    fn from_uri(uri: &opendal_core::OperatorUri) -> opendal_core::Result<Self> {
84        let opts = uri.options();
85
86        // Reconstruct the full path from authority (name) and root.
87        // OperatorUri splits "hf://datasets/user/repo@rev/path" into
88        // name="datasets" and root="user/repo@rev/path".
89        let mut path = String::new();
90        if let Some(name) = uri.name()
91            && !name.is_empty()
92        {
93            path.push_str(name);
94        }
95        if let Some(root) = uri.root()
96            && !root.is_empty()
97        {
98            if !path.is_empty() {
99                path.push('/');
100            }
101            path.push_str(root);
102        }
103
104        let download_mode = opts
105            .get("download_mode")
106            .map(|s| HfDownloadMode::parse(s))
107            .transpose()?;
108
109        if !path.is_empty() {
110            // Full URI like "hf://datasets/user/repo@rev/path"
111            let parsed = HfUri::parse(&path)?;
112            Ok(Self {
113                repo_type: Some(parsed.repo.repo_type),
114                repo_id: Some(parsed.repo.repo_id),
115                revision: parsed.repo.revision,
116                root: opts.get("root").cloned(),
117                token: opts.get("token").cloned(),
118                endpoint: opts.get("endpoint").cloned(),
119                download_mode,
120            })
121        } else {
122            // Bare scheme from via_iter, all config is in options.
123            let repo_type = opts
124                .get("repo_type")
125                .ok_or_else(|| {
126                    opendal_core::Error::new(
127                        opendal_core::ErrorKind::ConfigInvalid,
128                        "repo_type is required",
129                    )
130                    .with_context("service", HUGGINGFACE_SCHEME)
131                })
132                .and_then(|s| HfRepoType::parse(s))?;
133            Ok(Self {
134                repo_type: Some(repo_type),
135                repo_id: opts.get("repo_id").cloned(),
136                revision: opts.get("revision").cloned(),
137                root: opts.get("root").cloned(),
138                token: opts.get("token").cloned(),
139                endpoint: opts.get("endpoint").cloned(),
140                download_mode,
141            })
142        }
143    }
144
145    fn into_builder(self) -> Self::Builder {
146        HfBuilder { config: self }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use opendal_core::Configurator;
154    use opendal_core::OperatorUri;
155
156    #[test]
157    fn from_uri_with_all_components() {
158        let uri = OperatorUri::new(
159            "hf://datasets/username/my_dataset@dev/train/data.csv",
160            Vec::<(String, String)>::new(),
161        )
162        .unwrap();
163
164        let cfg = HfConfig::from_uri(&uri).unwrap();
165        assert_eq!(cfg.repo_type, Some(HfRepoType::Dataset));
166        assert_eq!(cfg.repo_id.as_deref(), Some("username/my_dataset"));
167        assert_eq!(cfg.revision.as_deref(), Some("dev"));
168        assert!(cfg.root.is_none());
169    }
170
171    #[test]
172    fn from_uri_via_iter_options() {
173        // Simulates the via_iter path: bare scheme with options map.
174        let uri = OperatorUri::new(
175            "huggingface",
176            vec![
177                ("repo_type".to_string(), "dataset".to_string()),
178                (
179                    "repo_id".to_string(),
180                    "opendal/huggingface-testdata".to_string(),
181                ),
182                ("revision".to_string(), "main".to_string()),
183                ("root".to_string(), "/testdata/".to_string()),
184            ],
185        )
186        .unwrap();
187
188        let cfg = HfConfig::from_uri(&uri).unwrap();
189        assert_eq!(cfg.repo_type, Some(HfRepoType::Dataset));
190        assert_eq!(cfg.repo_id.as_deref(), Some("opendal/huggingface-testdata"));
191        assert_eq!(cfg.revision.as_deref(), Some("main"));
192        assert_eq!(cfg.root.as_deref(), Some("/testdata/"));
193    }
194
195    #[test]
196    fn from_uri_download_mode_http() {
197        let uri = OperatorUri::new(
198            "huggingface",
199            vec![
200                ("repo_type".to_string(), "dataset".to_string()),
201                ("repo_id".to_string(), "user/repo".to_string()),
202                ("download_mode".to_string(), "http".to_string()),
203            ],
204        )
205        .unwrap();
206
207        let cfg = HfConfig::from_uri(&uri).unwrap();
208        assert_eq!(cfg.download_mode, Some(HfDownloadMode::Http));
209    }
210
211    #[test]
212    fn from_uri_download_mode_defaults_to_xet() {
213        let uri = OperatorUri::new(
214            "huggingface",
215            vec![
216                ("repo_type".to_string(), "model".to_string()),
217                ("repo_id".to_string(), "user/repo".to_string()),
218            ],
219        )
220        .unwrap();
221
222        let cfg = HfConfig::from_uri(&uri).unwrap();
223        assert_eq!(cfg.download_mode.unwrap_or_default(), HfDownloadMode::Xet);
224    }
225}