ares_llm/
nvidia_catalog.rs

1//! NVIDIA Model Catalog Cache
2//!
3//! Fetches the live model catalog from `https://integrate.api.nvidia.com/v1/models`,
4//! caches it in memory, and supports periodic background refresh.
5
6use arc_swap::ArcSwap;
7use parking_lot::RwLock;
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11use tracing::{info, warn};
12
13/// Configuration for the NVIDIA provider and catalog fetch.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct NvidiaConfig {
16    /// Environment variable holding the NVIDIA API key (default: `NVIDIA_API_KEY`).
17    #[serde(default = "default_api_key_env")]
18    pub api_key_env: String,
19
20    /// Base URL for NVIDIA NIM API calls (default: `https://integrate.api.nvidia.com/v1`).
21    #[serde(default = "default_api_base")]
22    pub api_base: String,
23
24    /// URL to fetch the model catalog (default: `https://integrate.api.nvidia.com/v1/models`).
25    #[serde(default = "default_models_url")]
26    pub models_url: String,
27
28    /// Background refresh interval in seconds. `0` disables background refresh.
29    #[serde(default = "default_catalog_refresh_seconds")]
30    pub catalog_refresh_seconds: u64,
31
32    /// Fallback model id used when the catalog is empty or fetch fails.
33    #[serde(default = "default_default_model")]
34    pub default_model: String,
35}
36
37impl Default for NvidiaConfig {
38    fn default() -> Self {
39        Self {
40            api_key_env: default_api_key_env(),
41            api_base: default_api_base(),
42            models_url: default_models_url(),
43            catalog_refresh_seconds: default_catalog_refresh_seconds(),
44            default_model: default_default_model(),
45        }
46    }
47}
48
49fn default_api_key_env() -> String {
50    "NVIDIA_API_KEY".to_string()
51}
52
53fn default_api_base() -> String {
54    "https://integrate.api.nvidia.com/v1".to_string()
55}
56
57fn default_models_url() -> String {
58    "https://integrate.api.nvidia.com/v1/models".to_string()
59}
60
61fn default_catalog_refresh_seconds() -> u64 {
62    3600
63}
64
65fn default_default_model() -> String {
66    "meta/llama-3.3-70b-instruct".to_string()
67}
68
69/// A single entry from the NVIDIA catalog endpoint.
70#[derive(Debug, Clone, Serialize, Deserialize)]
71pub struct CatalogEntry {
72    /// Full model id, e.g. `meta/llama-3.3-70b-instruct`.
73    pub id: String,
74
75    /// Organization that owns the model, e.g. `meta`.
76    pub owned_by: String,
77
78    /// Unix timestamp when the model was created.
79    pub created: i64,
80
81    /// Derived quality score (0-100).
82    #[serde(skip)]
83    pub quality_score: u8,
84}
85
86/// In-memory cache of the NVIDIA model catalog.
87///
88/// `cfg` is wrapped in `Arc<ArcSwap<...>>` so the admin can hot-swap the
89/// `api_key_env`, `api_base`, and `models_url` fields at runtime without
90/// restarting the service. The actual API key is NOT cached here — it is
91/// resolved at refresh time from either the env var named in `cfg.api_key_env`
92/// or the fleet provider secrets override.
93pub struct NvidiaCatalogCache {
94    inner: ArcSwap<Vec<CatalogEntry>>,
95    last_fetch: RwLock<Option<Instant>>,
96    last_error: RwLock<Option<String>>,
97    cfg: Arc<ArcSwap<NvidiaConfig>>,
98}
99
100/// Errors that can occur during catalog refresh.
101#[derive(Debug, thiserror::Error)]
102pub enum NvidiaCatalogError {
103    #[error("HTTP request failed: {0}")]
104    Http(#[from] reqwest::Error),
105    #[error("JSON parse failed: {0}")]
106    Json(#[from] serde_json::Error),
107    #[error("API key not found in environment variable {0}")]
108    MissingApiKey(String),
109    #[error("NVIDIA API returned HTTP {status}: {body}")]
110    BadStatus { status: u16, body: String },
111}
112
113/// NVIDIA API response shape.
114#[derive(Debug, Deserialize)]
115struct NvidiaModelsResponse {
116    #[serde(default)]
117    data: Vec<NvidiaModelItem>,
118}
119
120#[derive(Debug, Deserialize)]
121struct NvidiaModelItem {
122    id: String,
123    #[serde(default)]
124    owned_by: String,
125    #[serde(default)]
126    created: i64,
127}
128
129impl NvidiaCatalogCache {
130    /// Create a new empty cache from configuration.
131    pub fn new(cfg: NvidiaConfig) -> Self {
132        Self {
133            inner: ArcSwap::from_pointee(Vec::new()),
134            last_fetch: RwLock::new(None),
135            last_error: RwLock::new(None),
136            cfg: Arc::new(ArcSwap::from_pointee(cfg)),
137        }
138    }
139
140    /// Build from a pre-constructed `Arc<ArcSwap<NvidiaConfig>>`. The
141    /// caller's wrapper is shared with the registry/admin endpoint so a
142    /// hot-swap is visible to all readers (including `refresh`).
143    pub fn from_arcswap(cfg: Arc<ArcSwap<NvidiaConfig>>) -> Self {
144        Self {
145            inner: ArcSwap::from_pointee(Vec::new()),
146            last_fetch: RwLock::new(None),
147            last_error: RwLock::new(None),
148            cfg,
149        }
150    }
151
152    /// Fetch the catalog from NVIDIA and update the cache.
153    ///
154    /// Returns the number of chat models that were stored.
155    pub async fn refresh(&self) -> Result<usize, NvidiaCatalogError> {
156        // Snapshot the current config. The reference is short-lived; if an
157        // admin hot-swaps mid-refresh, the next refresh sees the new value.
158        let cfg_snapshot = self.cfg.load_full();
159        let cfg_ref: &NvidiaConfig = cfg_snapshot.as_ref();
160
161        let api_key = std::env::var(&cfg_ref.api_key_env).map_err(|_| {
162            NvidiaCatalogError::MissingApiKey(cfg_ref.api_key_env.clone())
163        })?;
164
165        let client = reqwest::Client::new();
166        let resp = client
167            .get(&cfg_ref.models_url)
168            .header("Authorization", format!("Bearer {}", api_key))
169            .header("Accept", "application/json")
170            .timeout(Duration::from_secs(10))
171            .send()
172            .await?;
173
174        let status = resp.status();
175        if !status.is_success() {
176            let body = resp.text().await.unwrap_or_default();
177            return Err(NvidiaCatalogError::BadStatus {
178                status: status.as_u16(),
179                body,
180            });
181        }
182
183        let parsed: NvidiaModelsResponse = resp.json().await?;
184
185        let mut entries: Vec<CatalogEntry> = parsed
186            .data
187            .into_iter()
188            .filter(|item| is_chat_model(&item.id))
189            .map(|item| CatalogEntry {
190                quality_score: quality_score_for(&item.id),
191                id: item.id,
192                owned_by: if item.owned_by.is_empty() {
193                    "nvidia".to_string()
194                } else {
195                    item.owned_by
196                },
197                created: item.created,
198            })
199            .collect();
200
201        // Stable sort by quality score descending
202        entries.sort_by_key(|e| std::cmp::Reverse(e.quality_score));
203
204        let count = entries.len();
205        self.inner.store(Arc::new(entries));
206        *self.last_fetch.write() = Some(Instant::now());
207        *self.last_error.write() = None;
208
209        info!("NVIDIA catalog refreshed with {} chat models", count);
210        Ok(count)
211    }
212
213    /// Atomically replace the cached `NvidiaConfig` (e.g. after an admin
214    /// updates `api_key_env`, `api_base`, or `models_url`). Subsequent
215    /// `refresh()` calls use the new config.
216    pub fn update_config(&self, new_cfg: NvidiaConfig) {
217        self.cfg.store(Arc::new(new_cfg));
218    }
219
220    /// Borrow a read handle to the current `NvidiaConfig` for callers that
221    /// need to introspect the live values.
222    pub fn config_handle(&self) -> Arc<ArcSwap<NvidiaConfig>> {
223        Arc::clone(&self.cfg)
224    }
225
226    /// Snapshot the current refresh interval (seconds). Returns 0 if disabled.
227    pub fn refresh_seconds(&self) -> u64 {
228        self.cfg.load().catalog_refresh_seconds
229    }
230
231    /// Return a snapshot of the currently cached entries.
232    pub fn snapshot(&self) -> Vec<CatalogEntry> {
233        self.inner.load_full().as_ref().clone()
234    }
235
236    /// Return the age of the last successful fetch, if any.
237    pub fn last_fetch_age(&self) -> Option<Duration> {
238        self.last_fetch.read().map(|t| t.elapsed())
239    }
240
241    /// Return the last error message, if any.
242    pub fn last_error(&self) -> Option<String> {
243        self.last_error.read().clone()
244    }
245
246    /// Spawn a background Tokio task that refreshes the catalog periodically.
247    ///
248    /// Does nothing if `catalog_refresh_seconds` is `0`.
249    pub fn start_background_refresh(self: Arc<Self>) {
250        let seconds = self.cfg.load().catalog_refresh_seconds;
251        if seconds == 0 {
252            return;
253        }
254
255        tokio::spawn(async move {
256            loop {
257                // Re-read the interval on each tick so an admin hot-swap of
258                // `catalog_refresh_seconds` takes effect on the next loop.
259                let secs = self.cfg.load().catalog_refresh_seconds;
260                if secs == 0 {
261                    return;
262                }
263                tokio::time::sleep(Duration::from_secs(secs)).await;
264                if let Err(e) = self.refresh().await {
265                    warn!("NVIDIA catalog background refresh failed: {}", e);
266                    *self.last_error.write() = Some(e.to_string());
267                }
268            }
269        });
270    }
271}
272
273/// Filter out non-chat models by id substring.
274fn is_chat_model(id: &str) -> bool {
275    let lower = id.to_lowercase();
276    let denied = [
277        "embed", "rerank", "retriev", "parse", "reward",
278        "safety", "guard", "detect", "asr", "tts",
279        "kosmos", "vila", "vision-encoder",
280    ];
281    !denied.iter().any(|d| lower.contains(d))
282}
283
284/// Compute a quality score (0-100) from a model id.
285fn quality_score_for(id: &str) -> u8 {
286    let lower = id.to_lowercase();
287    let mut score: u8 = 75;
288
289    // Vendor-based adjustments
290    if lower.contains("qwen") {
291        score = score.saturating_add(12);
292    } else if lower.contains("llama-3.3") || lower.contains("llama-3.1") {
293        score = score.saturating_add(10);
294    } else if lower.contains("mistral") || lower.contains("codestral") {
295        score = score.saturating_add(8);
296    } else if lower.contains("gemma-3") || lower.contains("nemotron") {
297        score = score.saturating_add(6);
298    } else if lower.contains("glm") || lower.contains("phi") || lower.contains("step") || lower.contains("granite") {
299        score = score.saturating_add(3);
300    }
301
302    // Size-based adjustments
303    if lower.contains("405b") {
304        score = score.saturating_add(10);
305    } else if lower.contains("70b") {
306        score = score.saturating_add(8);
307    } else if lower.contains("32b") {
308        score = score.saturating_add(5);
309    } else if lower.contains("14b") {
310        score = score.saturating_add(3);
311    } else if lower.contains("8b") {
312        score = score.saturating_add(2);
313    }
314
315    score.min(100)
316}
ares_llm/nvidia_catalog.rs

ares_llm/
nvidia_catalog.rs
