mirror of
https://github.com/CyberSecurityUP/NeuroSploit.git
synced 2026-07-17 08:17:18 +02:00
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:<deployment>`. - #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) <noreply@anthropic.com>
This commit is contained in:
@@ -17,7 +17,15 @@ ANTHROPIC_API_KEY=
|
|||||||
OPENAI_API_KEY=
|
OPENAI_API_KEY=
|
||||||
|
|
||||||
# gemini: https://aistudio.google.com/app/apikey
|
# 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=
|
GEMINI_API_KEY=
|
||||||
|
#GOOGLE_API_KEY=
|
||||||
|
|
||||||
|
# azure: Azure OpenAI (OpenAI-compatible). Use `--model azure:<deployment>`
|
||||||
|
# (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: https://console.x.ai/
|
||||||
XAI_API_KEY=
|
XAI_API_KEY=
|
||||||
|
|||||||
@@ -55,6 +55,15 @@ and severity-calibrated).
|
|||||||
neurosploit whitebox https://github.com/digininja/DVWA \
|
neurosploit whitebox https://github.com/digininja/DVWA \
|
||||||
--subscription --model anthropic:claude-opus-4-8 -v
|
--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:<deployment>`
|
||||||
|
(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
|
## Notes
|
||||||
|
|
||||||
|
|||||||
@@ -49,6 +49,12 @@ pub fn providers() -> Vec<Provider> {
|
|||||||
models: vec!["gpt-4o", "claude-3-7-sonnet", "gemini/gemini-2.5-pro"] },
|
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",
|
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"] },
|
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://<resource>.openai.azure.com), optionally AZURE_OPENAI_API_VERSION
|
||||||
|
// (default 2024-10-21), and use `azure:<your-deployment-name>` 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",
|
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"] },
|
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<Provider> {
|
|||||||
providers().into_iter().find(|p| p.key == key)
|
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.
|
/// A `provider:model` selection.
|
||||||
#[derive(Clone, Debug)]
|
#[derive(Clone, Debug)]
|
||||||
pub struct ModelRef {
|
pub struct ModelRef {
|
||||||
@@ -97,17 +114,32 @@ impl ChatClient {
|
|||||||
pub async fn chat(&self, m: &ModelRef, system: &str, user: &str) -> Result<String> {
|
pub async fn chat(&self, m: &ModelRef, system: &str, user: &str) -> Result<String> {
|
||||||
let p = provider_for(&m.provider)
|
let p = provider_for(&m.provider)
|
||||||
.ok_or_else(|| anyhow!("unknown provider '{}'", 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" {
|
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, …).
|
// Azure OpenAI uses a per-resource endpoint + deployment + api-version,
|
||||||
let base = match p.key {
|
// and authenticates with an `api-key` header instead of Bearer.
|
||||||
"litellm" => std::env::var("LITELLM_BASE_URL").unwrap_or_else(|_| p.base_url.to_string()),
|
let azure = p.key == "azure";
|
||||||
"ollama" => std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| p.base_url.to_string()),
|
let url = if azure {
|
||||||
_ => p.base_url.to_string(),
|
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://<resource>.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:<deployment>`).
|
||||||
|
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!({
|
let body = serde_json::json!({
|
||||||
"model": m.model,
|
"model": m.model,
|
||||||
"max_tokens": 4096,
|
"max_tokens": 4096,
|
||||||
@@ -119,7 +151,7 @@ impl ChatClient {
|
|||||||
});
|
});
|
||||||
let mut req = self.http.post(&url).json(&body);
|
let mut req = self.http.post(&url).json(&body);
|
||||||
if !key.is_empty() {
|
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 resp = req.send().await?;
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
|
|||||||
Reference in New Issue
Block a user