Skip to main content

agent_engine/runtime/openai/catalog/
nvidia.rs

1//! NVIDIA NIM `/models` parser, context inference, and reasoning inference.
2
3use serde::Deserialize;
4
5use super::{CatalogModel, CatalogProviderKind, CatalogSource, ReasoningSupport};
6
7#[derive(Debug, Deserialize)]
8struct NvidiaModelsResponse {
9    data: Vec<NvidiaModelItem>,
10}
11
12#[derive(Debug, Deserialize)]
13struct NvidiaModelItem {
14    id: String,
15    #[serde(default)]
16    owned_by: Option<String>,
17}
18
19fn infer_nvidia_context_tokens(model_id: &str) -> Option<u64> {
20    let id = model_id.to_ascii_lowercase();
21    if id.contains("kimi-k2-thinking") {
22        Some(256_000)
23    } else if id.contains("nemotron-ultra")
24        || id.contains("nemotron-super")
25        || id.contains("qwen3-next")
26        || id.contains("deepseek-v3.1")
27        || id.contains("llama-3.1-405b")
28        || id.contains("llama-3.2-90b-vision")
29        || id.contains("mistral-large-2")
30    {
31        Some(128_000)
32    } else {
33        None
34    }
35}
36
37pub fn infer_nvidia_reasoning(model_id: &str) -> ReasoningSupport {
38    let id = model_id.to_ascii_lowercase();
39    if id.contains("thinking")
40        || id.contains("cosmos-reason")
41        || id.contains("nemotron-ultra")
42        || id.contains("nemotron-super")
43        || id.contains("mistral-nemotron")
44        || id.contains("magistral")
45        || id.contains("deepseek-v4")
46        || id.contains("kimi")
47        || id.contains("reasoning")
48    {
49        ReasoningSupport::NvidiaInlineThinking
50    } else {
51        ReasoningSupport::None
52    }
53}
54
55pub fn parse_nvidia_catalog_models(body: &str) -> Result<Vec<CatalogModel>, serde_json::Error> {
56    let resp: NvidiaModelsResponse = serde_json::from_str(body)?;
57    let mut seen = std::collections::BTreeSet::new();
58    Ok(resp
59        .data
60        .into_iter()
61        .filter_map(|item| {
62            let id = item.id;
63            if id.trim().is_empty() || !seen.insert(id.clone()) {
64                return None;
65            }
66            let mut m = CatalogModel::new("nvidia", "NVIDIA NIM", id)?;
67            m.provider_kind = CatalogProviderKind::NvidiaNim;
68            m.label = item
69                .owned_by
70                .as_ref()
71                .map(|owner| format!("{} — {}", m.id, owner));
72            m.context_tokens = infer_nvidia_context_tokens(&m.id);
73            m.reasoning = infer_nvidia_reasoning(&m.id);
74            m.source = if m.context_tokens.is_some() || m.reasoning != ReasoningSupport::None {
75                CatalogSource::Inferred
76            } else {
77                CatalogSource::Live
78            };
79            Some(m)
80        })
81        .collect())
82}