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)
162            .map_err(|_| NvidiaCatalogError::MissingApiKey(cfg_ref.api_key_env.clone()))?;
163
164        // Client-level timeouts are belt-and-braces: the request below also
165        // carries its own total timeout, but future fetch paths may not.
166        let client = reqwest::ClientBuilder::new()
167            .connect_timeout(Duration::from_secs(10))
168            .timeout(Duration::from_secs(10))
169            .build()
170            .expect("failed to build reqwest client");
171        let resp = client
172            .get(&cfg_ref.models_url)
173            .header("Authorization", format!("Bearer {}", api_key))
174            .header("Accept", "application/json")
175            .timeout(Duration::from_secs(10))
176            .send()
177            .await?;
178
179        let status = resp.status();
180        if !status.is_success() {
181            let body = resp.text().await.unwrap_or_default();
182            return Err(NvidiaCatalogError::BadStatus {
183                status: status.as_u16(),
184                body,
185            });
186        }
187
188        let parsed: NvidiaModelsResponse = resp.json().await?;
189
190        let mut entries: Vec<CatalogEntry> = parsed
191            .data
192            .into_iter()
193            .filter(|item| is_chat_model(&item.id))
194            .map(|item| CatalogEntry {
195                quality_score: quality_score_for(&item.id),
196                id: item.id,
197                owned_by: if item.owned_by.is_empty() {
198                    "nvidia".to_string()
199                } else {
200                    item.owned_by
201                },
202                created: item.created,
203            })
204            .collect();
205
206        // Stable sort by quality score descending
207        entries.sort_by_key(|e| std::cmp::Reverse(e.quality_score));
208
209        let count = entries.len();
210        self.inner.store(Arc::new(entries));
211        *self.last_fetch.write() = Some(Instant::now());
212        *self.last_error.write() = None;
213
214        info!("NVIDIA catalog refreshed with {} chat models", count);
215        Ok(count)
216    }
217
218    /// Atomically replace the cached `NvidiaConfig` (e.g. after an admin
219    /// updates `api_key_env`, `api_base`, or `models_url`). Subsequent
220    /// `refresh()` calls use the new config.
221    pub fn update_config(&self, new_cfg: NvidiaConfig) {
222        self.cfg.store(Arc::new(new_cfg));
223    }
224
225    /// Borrow a read handle to the current `NvidiaConfig` for callers that
226    /// need to introspect the live values.
227    pub fn config_handle(&self) -> Arc<ArcSwap<NvidiaConfig>> {
228        Arc::clone(&self.cfg)
229    }
230
231    /// Snapshot the current refresh interval (seconds). Returns 0 if disabled.
232    pub fn refresh_seconds(&self) -> u64 {
233        self.cfg.load().catalog_refresh_seconds
234    }
235
236    /// Return a snapshot of the currently cached entries.
237    pub fn snapshot(&self) -> Vec<CatalogEntry> {
238        self.inner.load_full().as_ref().clone()
239    }
240
241    /// Return the age of the last successful fetch, if any.
242    pub fn last_fetch_age(&self) -> Option<Duration> {
243        self.last_fetch.read().map(|t| t.elapsed())
244    }
245
246    /// Return the last error message, if any.
247    pub fn last_error(&self) -> Option<String> {
248        self.last_error.read().clone()
249    }
250
251    /// Spawn a background Tokio task that refreshes the catalog periodically.
252    ///
253    /// Does nothing if `catalog_refresh_seconds` is `0`.
254    pub fn start_background_refresh(self: Arc<Self>) {
255        let seconds = self.cfg.load().catalog_refresh_seconds;
256        if seconds == 0 {
257            return;
258        }
259
260        tokio::spawn(async move {
261            loop {
262                // Re-read the interval on each tick so an admin hot-swap of
263                // `catalog_refresh_seconds` takes effect on the next loop.
264                let secs = self.cfg.load().catalog_refresh_seconds;
265                if secs == 0 {
266                    return;
267                }
268                tokio::time::sleep(Duration::from_secs(secs)).await;
269                if let Err(e) = self.refresh().await {
270                    warn!("NVIDIA catalog background refresh failed: {}", e);
271                    *self.last_error.write() = Some(e.to_string());
272                }
273            }
274        });
275    }
276}
277
278/// Filter out non-chat models by id substring.
279fn is_chat_model(id: &str) -> bool {
280    let lower = id.to_lowercase();
281    let denied = [
282        "embed",
283        "rerank",
284        "retriev",
285        "parse",
286        "reward",
287        "safety",
288        "guard",
289        "detect",
290        "asr",
291        "tts",
292        "kosmos",
293        "vila",
294        "vision-encoder",
295    ];
296    !denied.iter().any(|d| lower.contains(d))
297}
298
299/// Compute a quality score (0-100) from a model id.
300fn quality_score_for(id: &str) -> u8 {
301    let lower = id.to_lowercase();
302    let mut score: u8 = 75;
303
304    // Vendor-based adjustments
305    if lower.contains("qwen") {
306        score = score.saturating_add(12);
307    } else if lower.contains("llama-3.3") || lower.contains("llama-3.1") {
308        score = score.saturating_add(10);
309    } else if lower.contains("mistral") || lower.contains("codestral") {
310        score = score.saturating_add(8);
311    } else if lower.contains("gemma-3") || lower.contains("nemotron") {
312        score = score.saturating_add(6);
313    } else if lower.contains("glm")
314        || lower.contains("phi")
315        || lower.contains("step")
316        || lower.contains("granite")
317    {
318        score = score.saturating_add(3);
319    }
320
321    // Size-based adjustments
322    if lower.contains("405b") {
323        score = score.saturating_add(10);
324    } else if lower.contains("70b") {
325        score = score.saturating_add(8);
326    } else if lower.contains("32b") {
327        score = score.saturating_add(5);
328    } else if lower.contains("14b") {
329        score = score.saturating_add(3);
330    } else if lower.contains("8b") {
331        score = score.saturating_add(2);
332    }
333
334    score.min(100)
335}
ares_llm/nvidia_catalog.rs

ares_llm/
nvidia_catalog.rs
