From d957429c0980ba4e05bfb374d7fbab0372d5f303 Mon Sep 17 00:00:00 2001 From: CyberSecurityUP Date: Fri, 26 Jun 2026 14:17:25 -0300 Subject: [PATCH] feat(models): add Azure OpenAI provider + GOOGLE_API_KEY alias for Gemini Resolves the only two open issues that still apply to the Rust build: - #21 Azure OpenAI: new `azure` provider (OpenAI-compatible). Endpoint comes from AZURE_OPENAI_ENDPOINT, api-version from AZURE_OPENAI_API_VERSION (default 2024-10-21); the model name is the Azure deployment; auth uses the `api-key` header instead of Bearer. Use `--model azure:`. - #25 Gemini key confusion: GEMINI_API_KEY now also accepts GOOGLE_API_KEY (Google's standard env var) as an alias; local providers (ollama/litellm) require no key. .env.example documents both. Kept under the v3.5.2 line (additive provider support). Co-Authored-By: Claude Opus 4.8 (1M context) --- .env.example | 8 ++++ RELEASE.md | 9 ++++ neurosploit-rs/crates/harness/src/models.rs | 50 +++++++++++++++++---- 3 files changed, 58 insertions(+), 9 deletions(-) diff --git a/.env.example b/.env.example index 5b654bb..fd8774e 100755 --- a/.env.example +++ b/.env.example @@ -17,7 +17,15 @@ ANTHROPIC_API_KEY= OPENAI_API_KEY= # gemini: https://aistudio.google.com/app/apikey +# (GOOGLE_API_KEY is also accepted as an alias if GEMINI_API_KEY is unset) GEMINI_API_KEY= +#GOOGLE_API_KEY= + +# azure: Azure OpenAI (OpenAI-compatible). Use `--model azure:` +# (the model name is your Azure *deployment* name). +#AZURE_OPENAI_API_KEY= +#AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com +#AZURE_OPENAI_API_VERSION=2024-10-21 # xai: https://console.x.ai/ XAI_API_KEY= diff --git a/RELEASE.md b/RELEASE.md index 1486db9..4370c4f 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -55,6 +55,15 @@ and severity-calibrated). neurosploit whitebox https://github.com/digininja/DVWA \ --subscription --model anthropic:claude-opus-4-8 -v ``` +- **Azure OpenAI provider** (resolves #21). OpenAI-compatible: set + `AZURE_OPENAI_ENDPOINT` (+ optional `AZURE_OPENAI_API_VERSION`, default + `2024-10-21`) and `AZURE_OPENAI_API_KEY`, then `--model azure:` + (the model name is your Azure *deployment* name; auth via the `api-key` + header). +- **`GOOGLE_API_KEY` alias for Gemini** (resolves #25 confusion). Gemini's API + path reads `GEMINI_API_KEY`, and now also accepts `GOOGLE_API_KEY` (Google's + standard env var) when the former is unset. Local providers (ollama/litellm) + still need **no** key at all. ## Notes diff --git a/neurosploit-rs/crates/harness/src/models.rs b/neurosploit-rs/crates/harness/src/models.rs index 6161e18..7741550 100644 --- a/neurosploit-rs/crates/harness/src/models.rs +++ b/neurosploit-rs/crates/harness/src/models.rs @@ -49,6 +49,12 @@ pub fn providers() -> Vec { models: vec!["gpt-4o", "claude-3-7-sonnet", "gemini/gemini-2.5-pro"] }, Provider { key: "openrouter", label: "OpenRouter", base_url: "https://openrouter.ai/api/v1", env_key: "OPENROUTER_API_KEY", kind: "api", models: vec!["anthropic/claude-opus-4-8", "qwen/qwen-2.5-coder-32b-instruct", "deepseek/deepseek-r1", "meta-llama/llama-3.3-70b-instruct"] }, + // Azure OpenAI (OpenAI-compatible). Set AZURE_OPENAI_ENDPOINT (e.g. + // https://.openai.azure.com), optionally AZURE_OPENAI_API_VERSION + // (default 2024-10-21), and use `azure:` as the model. + // base_url is resolved from the endpoint at call time; auth uses an api-key header. + Provider { key: "azure", label: "Azure OpenAI", base_url: "", env_key: "AZURE_OPENAI_API_KEY", kind: "api", + models: vec!["gpt-4o", "gpt-4o-mini", "gpt-5.1", "o4-mini"] }, Provider { key: "ollama", label: "Ollama (local)", base_url: "http://localhost:11434/v1", env_key: "OLLAMA_API_KEY", kind: "api", models: vec!["qwen2.5-coder:32b", "qwq:32b", "deepseek-r1:32b", "llama3.3:70b"] }, ] @@ -58,6 +64,17 @@ pub fn provider_for(key: &str) -> Option { providers().into_iter().find(|p| p.key == key) } +/// Resolve a provider's API key from the environment, honoring common aliases. +/// For Gemini we also accept `GOOGLE_API_KEY` (Google's standard env var name) +/// when `GEMINI_API_KEY` is unset. +fn resolve_key(p: &Provider) -> String { + let mut k = std::env::var(p.env_key).unwrap_or_default(); + if k.is_empty() && p.key == "gemini" { + k = std::env::var("GOOGLE_API_KEY").unwrap_or_default(); + } + k +} + /// A `provider:model` selection. #[derive(Clone, Debug)] pub struct ModelRef { @@ -97,17 +114,32 @@ impl ChatClient { pub async fn chat(&self, m: &ModelRef, system: &str, user: &str) -> Result { let p = provider_for(&m.provider) .ok_or_else(|| anyhow!("unknown provider '{}'", m.provider))?; - let key = std::env::var(p.env_key).unwrap_or_default(); + let key = resolve_key(&p); if key.is_empty() && p.key != "ollama" && p.key != "litellm" { - return Err(anyhow!("no API key ({}) for provider '{}'", p.env_key, p.key)); + let hint = if p.key == "gemini" { format!("{} (or GOOGLE_API_KEY)", p.env_key) } else { p.env_key.to_string() }; + return Err(anyhow!("no API key ({}) for provider '{}'", hint, p.key)); } - // Allow an env base-URL override (LiteLLM gateway, self-hosted proxies, …). - let base = match p.key { - "litellm" => std::env::var("LITELLM_BASE_URL").unwrap_or_else(|_| p.base_url.to_string()), - "ollama" => std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| p.base_url.to_string()), - _ => p.base_url.to_string(), + // Azure OpenAI uses a per-resource endpoint + deployment + api-version, + // and authenticates with an `api-key` header instead of Bearer. + let azure = p.key == "azure"; + let url = if azure { + let endpoint = std::env::var("AZURE_OPENAI_ENDPOINT").unwrap_or_default(); + if endpoint.is_empty() { + return Err(anyhow!("set AZURE_OPENAI_ENDPOINT (e.g. https://.openai.azure.com) for the azure provider")); + } + let ver = std::env::var("AZURE_OPENAI_API_VERSION").unwrap_or_else(|_| "2024-10-21".to_string()); + // `model` is the Azure DEPLOYMENT name (use `azure:`). + format!("{}/openai/deployments/{}/chat/completions?api-version={}", + endpoint.trim_end_matches('/'), m.model, ver) + } else { + // Allow an env base-URL override (LiteLLM gateway, self-hosted proxies, …). + let base = match p.key { + "litellm" => std::env::var("LITELLM_BASE_URL").unwrap_or_else(|_| p.base_url.to_string()), + "ollama" => std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| p.base_url.to_string()), + _ => p.base_url.to_string(), + }; + format!("{}/chat/completions", base.trim_end_matches('/')) }; - let url = format!("{}/chat/completions", base.trim_end_matches('/')); let body = serde_json::json!({ "model": m.model, "max_tokens": 4096, @@ -119,7 +151,7 @@ impl ChatClient { }); let mut req = self.http.post(&url).json(&body); if !key.is_empty() { - req = req.bearer_auth(&key); + if azure { req = req.header("api-key", &key); } else { req = req.bearer_auth(&key); } } let resp = req.send().await?; let status = resp.status();