osdk_core/model/provider/
mod.rs

1use async_trait::async_trait;
2use serde::de::DeserializeOwned;
3
4use crate::backend::Ctx;
5use crate::error::{Error, Result};
6use crate::model::ModelRef;
7
8pub mod huggingface;
9pub mod modelscope;
10
11#[derive(Debug, Clone)]
12pub struct RemoteModelFile {
13    pub path: String,
14    pub size: Option<u64>,
15    pub sha256: Option<String>,
16    pub etag: Option<String>,
17    pub url: String,
18    pub headers: Vec<(String, String)>,
19}
20
21#[derive(Debug, Clone)]
22pub struct RemoteSnapshot {
23    pub revision: String,
24    pub endpoint: String,
25    pub files: Vec<RemoteModelFile>,
26}
27
28#[async_trait]
29pub trait ModelProvider: Send + Sync {
30    async fn resolve(
31        &self,
32        ctx: &Ctx,
33        reference: &ModelRef,
34        endpoint: &str,
35    ) -> Result<RemoteSnapshot>;
36}
37
38pub async fn get_cached_json<T: DeserializeOwned>(
39    ctx: &Ctx,
40    provider: &str,
41    cache_identity: &str,
42    url: &str,
43    headers: &[(String, String)],
44) -> Result<T> {
45    let hash = blake3::hash(cache_identity.as_bytes()).to_hex().to_string();
46    let cache = ctx
47        .dirs
48        .remote_cache()
49        .join("models")
50        .join(provider)
51        .join(hash);
52    if ctx.config.settings.offline {
53        let bytes = std::fs::read(&cache).map_err(|_| {
54            Error::other(format!(
55                "offline model metadata cache miss for {cache_identity}"
56            ))
57        })?;
58        return Ok(serde_json::from_slice(&bytes)?);
59    }
60
61    let mut request = ctx.client.get(url);
62    for (key, value) in headers {
63        request = request.header(key, value);
64    }
65    match request.send().await {
66        Ok(response) => match response
67            .error_for_status()
68            .map_err(|error| Error::network(url, error))
69        {
70            Ok(response) => {
71                let bytes = response
72                    .bytes()
73                    .await
74                    .map_err(|error| Error::network(url, error))?;
75                let parsed = serde_json::from_slice(&bytes)?;
76                write_atomic(&cache, &bytes)?;
77                Ok(parsed)
78            }
79            Err(error) => read_stale(&cache).or(Err(error)),
80        },
81        Err(error) => {
82            let error = Error::network(url, error);
83            read_stale(&cache).or(Err(error))
84        }
85    }
86}
87
88fn read_stale<T: DeserializeOwned>(path: &std::path::Path) -> Result<T> {
89    let bytes = std::fs::read(path).map_err(|error| Error::io(path, error))?;
90    tracing::warn!(path = %path.display(), "using stale cached model metadata");
91    Ok(serde_json::from_slice(&bytes)?)
92}
93
94fn write_atomic(path: &std::path::Path, bytes: &[u8]) -> Result<()> {
95    if let Some(parent) = path.parent() {
96        std::fs::create_dir_all(parent).map_err(|error| Error::io(parent, error))?;
97    }
98    let temporary = path.with_extension(format!("tmp-{}", std::process::id()));
99    std::fs::write(&temporary, bytes).map_err(|error| Error::io(&temporary, error))?;
100    std::fs::rename(&temporary, path).map_err(|error| Error::io(path, error))
101}
osdk_core/model/provider/mod.rs

osdk_core/model/provider/
mod.rs
