diff --git a/crates/codegen/kigi-models/src/lib.rs b/crates/codegen/kigi-models/src/lib.rs index 6eb9323..c2b17fe 100644 --- a/crates/codegen/kigi-models/src/lib.rs +++ b/crates/codegen/kigi-models/src/lib.rs @@ -531,6 +531,38 @@ const CEREBRAS_SPEC: PlatformSpec = PlatformSpec { restrict_to_enriched: false, }; +/// Base-URL override for NVIDIA NIM (dev/test escape hatch). +pub const NVIDIA_BASE_URL_ENV: &str = "KIGI_NVIDIA_BASE_URL"; + +const NVIDIA_SPEC: PlatformSpec = PlatformSpec { + id: "nvidia", + display_name: "NVIDIA NIM", + base_url: BaseUrlSource::EnvOr { + env: NVIDIA_BASE_URL_ENV, + default: "https://integrate.api.nvidia.com/v1", + }, + uses_oauth: false, + allowed_model_prefixes: None, + api_key_envs: &["NVIDIA_API_KEY"], + vendor: "NVIDIA", + console_host: Some("build.nvidia.com"), + login_label: Some("NVIDIA NIM (API key)"), + models_dev_id: Some("nvidia"), + wire_serves_metadata: false, + wire_api: PlatformWireApi::ChatCompletions, + listing: ListingDialect::OpenAi, + // NIM exposes raw vLLM behavior; stream_options support varies per model + // and some 4xx on it, so strip it (StrictOpenAi) to keep streaming + // working across the fleet. + chat_compat: PlatformChatCompat::StrictOpenAi, + key_header: PlatformKeyHeader::Bearer, + key_validation_path: None, + strip_listing_id_prefix: None, + // The listing mixes chat/embedding/rerank/vision/image models; keep + // tool-calling enrichment-known chat models only. + restrict_to_enriched: true, +}; + /// The platform registry. Platforms are compiled-in spec rows; there is no /// dynamic provider registration (PRD F2). #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] @@ -561,12 +593,14 @@ pub enum PlatformId { Together, /// Cerebras platform API (API key, OpenAI-compatible ChatCompletions). Cerebras, + /// NVIDIA NIM platform API (API key, OpenAI-compatible ChatCompletions). + Nvidia, } impl PlatformId { /// All platforms, in catalog precedence order: the subscription channel /// first so "default model = first list item" favors it when present. - pub const ALL: [PlatformId; 13] = [ + pub const ALL: [PlatformId; 14] = [ Self::KimiCode, Self::MoonshotCn, Self::MoonshotAi, @@ -580,6 +614,7 @@ impl PlatformId { Self::OpenRouter, Self::Together, Self::Cerebras, + Self::Nvidia, ]; /// The registry row backing this platform (single source of per-platform @@ -599,6 +634,7 @@ impl PlatformId { Self::OpenRouter => &OPENROUTER_SPEC, Self::Together => &TOGETHER_SPEC, Self::Cerebras => &CEREBRAS_SPEC, + Self::Nvidia => &NVIDIA_SPEC, } } @@ -1350,9 +1386,10 @@ mod tests { PlatformId::OpenRouter => 10, PlatformId::Together => 11, PlatformId::Cerebras => 12, + PlatformId::Nvidia => 13, } } - const VARIANT_COUNT: usize = 13; // update together with `ordinal` + const VARIANT_COUNT: usize = 14; // update together with `ordinal` let mut seen: Vec = PlatformId::ALL.iter().map(|&p| ordinal(p)).collect(); seen.sort_unstable(); seen.dedup(); diff --git a/crates/codegen/kigi-shell/src/agent/auth_method.rs b/crates/codegen/kigi-shell/src/agent/auth_method.rs index f0c2111..d2e60e8 100644 --- a/crates/codegen/kigi-shell/src/agent/auth_method.rs +++ b/crates/codegen/kigi-shell/src/agent/auth_method.rs @@ -605,7 +605,8 @@ mod tests { "google", "openrouter", "together", - "cerebras" + "cerebras", + "nvidia" ] ); assert_eq!(default_id(&built), Some(XAI_API_KEY_METHOD_ID)); @@ -641,7 +642,8 @@ mod tests { "google", "openrouter", "together", - "cerebras" + "cerebras", + "nvidia" ] ); assert_eq!(default_id(&built), Some(CACHED_TOKEN_AUTH_METHOD_ID)); @@ -670,7 +672,8 @@ mod tests { "google", "openrouter", "together", - "cerebras" + "cerebras", + "nvidia" ] ); assert_eq!(default_id(&built), Some(CACHED_TOKEN_AUTH_METHOD_ID)); @@ -702,7 +705,8 @@ mod tests { "google", "openrouter", "together", - "cerebras" + "cerebras", + "nvidia" ] ); assert_eq!(default_id(&built), None); diff --git a/crates/codegen/kigi-shell/src/agent/models_fetch.rs b/crates/codegen/kigi-shell/src/agent/models_fetch.rs index 2ceff13..e23af4c 100644 --- a/crates/codegen/kigi-shell/src/agent/models_fetch.rs +++ b/crates/codegen/kigi-shell/src/agent/models_fetch.rs @@ -1805,6 +1805,98 @@ mod tests { ); } + /// NVIDIA-cycle e2e: polluted listing (embedding/image models) restricted + /// to tool-calling enrichment-known chat models; slashed org/model ids; + /// StrictOpenAi dialect (strips stream_options — some NIM models reject it). + #[tokio::test(flavor = "multi_thread")] + #[serial_test::serial] + async fn nvidia_restricts_slashed_ids_and_maps_strict_dialect() { + let platform_server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/models")) + .and(wiremock::matchers::header( + "Authorization", + "Bearer nvapi-1", + )) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json( + serde_json::json!({ "data": [ + { "id": "meta/llama-3.3-70b-instruct", "object": "model" }, + { "id": "baai/bge-m3", "object": "model" } + ]}), + )) + .expect(1) + .mount(&platform_server) + .await; + let modelsdev_server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/api.json")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json( + serde_json::json!({ "nvidia": { "models": { + "meta/llama-3.3-70b-instruct": { + "limit": {"context": 128000, "output": 32768}, + "tool_call": true + }, + "baai/bge-m3": { "limit": {"context": 8192} } + }}}), + )) + .expect(1) + .mount(&modelsdev_server) + .await; + let cache_dir = tempfile::tempdir().unwrap(); + let _base = kigi_test_support::EnvGuard::set( + kigi_models::NVIDIA_BASE_URL_ENV, + platform_server.uri(), + ); + let _mdev = kigi_test_support::EnvGuard::set( + crate::agent::enrichment_fetch::MODELS_DEV_URL_ENV, + format!("{}/api.json", modelsdev_server.uri()), + ); + let _mdev_cache = kigi_test_support::EnvGuard::set( + crate::agent::enrichment_fetch::MODELS_DEV_CACHE_DIR_ENV, + cache_dir.path(), + ); + let endpoints = crate::agent::config::EndpointsConfig::default(); + let keys = crate::agent::models::PlatformApiKeys::test_single( + kigi_models::PlatformId::Nvidia, + "nvapi-1", + ); + let result = tokio::task::spawn_blocking(move || { + fetch_platform_models_blocking(&endpoints, None, &keys) + }) + .await + .unwrap() + .expect("fetch must succeed"); + assert_eq!( + result + .models + .iter() + .map(|m| m.id.as_deref().unwrap_or_default()) + .collect::>(), + vec!["nvidia/meta/llama-3.3-70b-instruct"], + "bge-m3 embedding (not tool-calling) dropped; slashed org/model id kept" + ); + let entry = &result.models[0]; + assert_eq!(entry.context_window.get(), 128_000); + assert_eq!(entry.max_completion_tokens, Some(32_768)); + assert_eq!( + entry.model, "meta/llama-3.3-70b-instruct", + "the native slashed id rides the wire" + ); + let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry); + let creds = crate::agent::config::ResolvedCredentials { + api_key: Some("nvapi-1".into()), + base_url: entry.base_url.clone(), + auth_type: kigi_chat_state::AuthType::ApiKey, + auth_scheme: Default::default(), + }; + let cfg = crate::agent::config::sampling_config_for_model(&model_entry, creds, None); + assert_eq!( + cfg.chat_compat, + kigi_sampling_types::ChatCompat::StrictOpenAi, + "NVIDIA uses StrictOpenAi (strips stream_options for NIM models that reject it)" + ); + } + #[test] fn get_env_keys_parses_strings_and_rejects_non_strings() { use crate::agent::config::EnvKeys; diff --git a/crates/codegen/kigi-tui/src/app/app_view.rs b/crates/codegen/kigi-tui/src/app/app_view.rs index e37e28d..7f6716b 100644 --- a/crates/codegen/kigi-tui/src/app/app_view.rs +++ b/crates/codegen/kigi-tui/src/app/app_view.rs @@ -6894,7 +6894,7 @@ pub(crate) mod tests { #[test] fn pending_menu_items_lists_interactive_methods_plus_quit() { let items = pending_menu_items(&fresh_user_auth_methods(), None); - assert_eq!(items.len(), 14, "13 login rows + Quit, got {items:?}"); + assert_eq!(items.len(), 15, "14 login rows + Quit, got {items:?}"); assert!( matches!(&items[0], PendingMenuItem::Login { label } if label == "Kimi Code (OAuth)"), "row 0 must be the OAuth login, got {:?}", @@ -6985,7 +6985,14 @@ pub(crate) mod tests { label: "Cerebras (API key)".into(), } ); - assert_eq!(items[13], PendingMenuItem::Quit); + assert_eq!( + items[13], + PendingMenuItem::ApiKey { + target: PlatformLogin(kigi_shell::models::PlatformId::Nvidia), + label: "NVIDIA NIM (API key)".into(), + } + ); + assert_eq!(items[14], PendingMenuItem::Quit); // The non-interactive methods must never appear as rows. let byok = kigi_shell::agent::auth_method::build_auth_methods( kigi_shell::agent::auth_method::AuthMethodsBuildInputs { @@ -6996,7 +7003,7 @@ pub(crate) mod tests { ); assert_eq!( pending_menu_items(&byok.methods, None).len(), - 14, + 15, "xai.api_key / cached_token must not add rows" ); }