aria_engine/
setup.rs

1//! Instance-level Engine setup (in-memory; does not write engine.yml).
2
3use thiserror::Error;
4
5pub const INTL_SITE: &str = "https://ariacompute.com";
6pub const INTL_UPGRADE: &str = "https://github.com/ariacompute";
7pub const CN_SITE: &str = "https://ariacompute.cn";
8pub const CN_UPGRADE: &str = "https://gitee.com/ariacompute";
9
10#[derive(Debug, Error, Clone, PartialEq, Eq)]
11pub enum SetupError {
12    #[error("invalid compute: {0}")]
13    InvalidCompute(String),
14    #[error("{0}")]
15    InvalidKey(String),
16}
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct SetupConfig {
20    pub router: String,
21    pub router_api_key: String,
22    pub site_url: String,
23    pub upgrade_url: String,
24    pub compute: String,
25    pub hf_token: String,
26    pub modelscope_api_token: String,
27}
28
29impl Default for SetupConfig {
30    fn default() -> Self {
31        Self {
32            router: String::new(),
33            router_api_key: String::new(),
34            site_url: String::new(),
35            upgrade_url: String::new(),
36            compute: "auto".into(),
37            hf_token: String::new(),
38            modelscope_api_token: String::new(),
39        }
40    }
41}
42
43/// Partial merge. `None` fields are omitted.
44#[derive(Debug, Clone, Default)]
45pub struct SetupUpdates {
46    pub router: Option<String>,
47    pub router_api_key: Option<String>,
48    pub site_url: Option<String>,
49    pub upgrade_url: Option<String>,
50    pub compute: Option<String>,
51    pub hf_token: Option<String>,
52    pub modelscope_api_token: Option<String>,
53}
54
55fn gateway_region(url: &str) -> Option<&'static str> {
56    let lower = url.to_ascii_lowercase();
57    if lower.contains("ariacompute.cn") || lower.contains("gitee.com/ariacompute") {
58        Some("cn")
59    } else if lower.contains("ariacompute.com") || lower.contains("github.com/ariacompute") {
60        Some("intl")
61    } else {
62        None
63    }
64}
65
66fn pair_urls(region: &str) -> (&'static str, &'static str) {
67    if region == "cn" {
68        (CN_SITE, CN_UPGRADE)
69    } else {
70        (INTL_SITE, INTL_UPGRADE)
71    }
72}
73
74/// Fill missing site/upgrade URLs from a provided TLD.
75pub fn fill_setup_urls(mut cfg: SetupConfig) -> SetupConfig {
76    let region = gateway_region(&cfg.site_url).or_else(|| gateway_region(&cfg.upgrade_url));
77    let Some(region) = region else {
78        return cfg;
79    };
80    let (site, upgrade) = pair_urls(region);
81    if cfg.site_url.is_empty() {
82        cfg.site_url = site.into();
83    }
84    if cfg.upgrade_url.is_empty() {
85        cfg.upgrade_url = upgrade.into();
86    }
87    cfg
88}
89
90fn validate_router_api_key(key: &str) -> Result<(), SetupError> {
91    let t = key.trim();
92    if t.is_empty() {
93        return Ok(());
94    }
95    if t.starts_with("sk-aria_") || t.starts_with("bfvk-") {
96        return Ok(());
97    }
98    Err(SetupError::InvalidKey(
99        "router_api_key must start with sk-aria_ or bfvk-".into(),
100    ))
101}
102
103/// Merge `updates` into `existing`. Validates; does not mutate `existing`.
104pub fn apply_setup(existing: &SetupConfig, updates: &SetupUpdates) -> Result<SetupConfig, SetupError> {
105    let mut out = existing.clone();
106    if let Some(v) = &updates.router {
107        out.router = v.clone();
108    }
109    if let Some(v) = &updates.router_api_key {
110        validate_router_api_key(v)?;
111        out.router_api_key = v.clone();
112    }
113    if let Some(v) = &updates.site_url {
114        out.site_url = v.clone();
115    }
116    if let Some(v) = &updates.upgrade_url {
117        out.upgrade_url = v.clone();
118    }
119    if let Some(v) = &updates.compute {
120        out.compute = v.clone();
121    }
122    if let Some(v) = &updates.hf_token {
123        out.hf_token = v.clone();
124    }
125    if let Some(v) = &updates.modelscope_api_token {
126        out.modelscope_api_token = v.clone();
127    }
128    match out.compute.as_str() {
129        "auto" | "cpu" | "cuda" => {}
130        other => return Err(SetupError::InvalidCompute(other.into())),
131    }
132    validate_router_api_key(&out.router_api_key)?;
133    Ok(fill_setup_urls(out))
134}
aria_engine/setup.rs

aria_engine/
setup.rs
