agent_engine/runtime/openai/catalog/
nvidia.rs1use 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}