Files
optimclaw/src/llm/models.rs
T
8638895879 feat(gemini_oauth): full Gemini CLI OAuth integration with Cloud Code API (#1356)
* feat: integrate Gemini CLI OAuth with Cloud Code API

- Add gemini_oauth.rs: full OAuth flow with PKCE, token refresh,
  and Cloud Code project discovery (loadCodeAssist + onboardUser)
- Route preview/gemini-3 models through cloudcode-pa.googleapis.com
  with proper project ID injection in request payload
- Trigger OAuth login during onboarding wizard (not first chat message)
- Support manual redirect URL paste as fallback (tokio::select race)
- Parse 429 rate-limit errors with retry_after from Google response
- Add static model list: gemini-1.5/2.0/2.5/3.0/3.1 variants
- Add GeminiOauthConfig with default credentials path (~/.gemini/)

* feat(gemini): implement function calling, generationConfig, and update models

- Implement function calling support (functionDeclarations, functionResponse)
- Add functionCall SSE parsing and empty stream retry support
- Add generationConfig (temperature, maxOutputTokens)
- Add thinkingConfig for Gemini 3 and thinking models
- Add toolConfig (functionCallingConfig.mode)
- Fix .expect() panics with .ok_or_else()
- Restrict oauth credentials file permissions to 0600
- Update docs and FEATURE_PARITY.md
- Update wizard to current Gemini 3.1 and 2.5 models

* fix: address code review issues in gemini-cli OAuth integration

- Add cache_read_input_tokens/cache_creation_input_tokens fields (value 0)
- Implement manual Debug for OAuthCredential to redact tokens
- Fix hardcoded /tmp: use GeminiOauthConfig::default_credentials_path()
- Replace emoji output with plain text markers
- Propagate Client::builder() errors instead of silent fallback
- Use tokio::fs for all file I/O in CredentialManager (was std::fs)
- Use if let Some(ref pid) to avoid consuming credential.project_id
- Extract uses_cloud_code_api() helper; route by major version (gemini-2+)
- Concatenate multiple system messages into systemInstruction
- Include functionCall parts in assistant message conversion
- Add 401 retry loop with allow_retry flag for auth failures
- Remove biased from tokio::select! in OAuth callback handler
- Remove hardcoded context_length 1M; vary by model family
- Change GOOG_API_CLIENT from Node.js spoof to gl-rust/1.0.0
- Implement list_models() with static model list
- Move create_gemini_oauth_provider() before test module (clippy)
- Fix 9 additional clippy warnings (collapsible_if, map_or, needless_borrow)
- Run cargo fmt

* Add dedicated regression tests for Gemini OAuth fixes

* style: fix formatting in Gemini OAuth regression tests

* feat(gemini-oauth): implement code review v3 refinements

- Add force_refresh() for 401 retry (bypass timestamp check)
- Standardize Gemini model list across docs, wizard, and provider
- Restore gemini-3 check for thinkingConfig
- Redact sensitive tokens in GoogleTokenRefreshResponse Debug output
- Use dynamic version for GOOG_API_CLIENT
- Improve model_metadata() context length heuristics
- Use strip_prefix("data:") for safer SSE parsing
- Skip re-auth in wizard if keeping existing provider

* feat(gemini_oauth): full Cloud Code API integration with project discovery

- Register gemini_oauth as a dedicated backend in config/llm.rs (skip
  registry fallback, preserve backend name, suppress unknown-backend warning)
- Fix app.rs credential guard to exclude backends with dedicated configs
  (gemini_oauth, bedrock) from the provider.is_none() check
- Auto-discover Cloud Code project_id via loadCodeAssist when credentials
  lack it (e.g. created by the original Gemini CLI)
- Persist discovered project_id to credentials file for subsequent runs
- Add safety settings (BLOCK_NONE), gated behind GEMINI_SAFETY_BLOCK_NONE env
- Add thinkingConfig: budget-based for Gemini 2.5, level-based for Gemini 3.x
  (without includeThoughts to avoid empty responses from reasoning.rs stripping)
- Add thought signature injection for Gemini 3.x preview APIs
- Add history curation to filter invalid model outputs before re-sending
- Add extended generationConfig env vars (topP, topK, seed, penalties,
  responseMimeType, responseJsonSchema, cachedContent)
- Add custom headers support via GEMINI_CLI_CUSTOM_HEADERS
- Add API key auth mode (GEMINI_API_KEY + GEMINI_API_KEY_AUTH_MECHANISM)
- Add SSE metadata extraction (modelVersion, credits, promptFeedback,
  groundingMetadata, citationMetadata, cachedContentTokenCount)
- Add countTokens API support
- Add new models to wizard (gemini-3.1-pro-preview-customtools,
  gemini-3-pro-preview, gemini-3.1-flash-lite-preview)
- Update docs/LLM_PROVIDERS.md with new models and routing rules
- Rewrite regression tests with comprehensive coverage (23 unit tests pass)

* fix: CI violations — add safety comment on expect, fix fmt

- Add '// safety: hardcoded literal' to regex .expect() to satisfy
  the no-panic-in-prod CI check
- Fix cargo fmt whitespace in collapsible if-let chain

* fix: address PR review feedback from gemini-code-assist

- Fix parse_custom_headers to preserve commas in values by splitting
  only on commas followed by a header-name:colon pattern (manual scan
  instead of simple split(','))
- Use matches! macro for backend exclusion check in app.rs
- Merge SSE metadata extraction into single pass (was iterating twice)
- Replace fragile substring-based context_length with explicit match
  on known Gemini model IDs via gemini_context_length()
- Add missing models to regression test (8 models, not 5)

* fix: address Copilot PR review feedback

- Fix empty text part for assistant messages with tool calls
  (curate_contents could drop entire model turn)
- Propagate cache_read/creation_input_tokens in complete_with_tools
- Log warning on save_credential failure instead of silently ignoring
- Fix doc comment to mention underscore in header name pattern
- Handle gemini-oauth (hyphen variant) in setup wizard display
- Fix docs: thinkingConfig uses thinkingBudget/thinkingLevel, not
  includeThoughts

* fix: add missing allow_always field after staging merge

* fix(gemini_oauth): align header parser doc with implementation [skip-regression-check]

Update parse_custom_headers doc comments to include underscore in the
header-name character class, matching the actual implementation.
Also fix formatting from merge.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(gemini_oauth): curate_contents per-part filtering and dead code removal

Fix curate_contents to filter invalid parts individually instead of
dropping entire model turn sequences. Previously a single empty text
part would discard all consecutive model turns including valid
functionCall parts, breaking the tool-call flow.

Also remove unused MID_STREAM_* constants.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* style(gemini_oauth): rustfmt formatting [skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(llm): support smart routing cheap model for gemini_oauth backend

Add explicit gemini_oauth handling in create_cheap_provider_for_backend()
to create a GeminiOauthProvider with the cheap model swapped in. Without
this, setting LLM_CHEAP_MODEL with gemini_oauth backend would fail with
a confusing "no registry provider config available" error.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* docs: add Gemini OAuth env vars to .env.example [skip-regression-check]

Document GEMINI_MODEL, GEMINI_CREDENTIALS_PATH, GEMINI_API_KEY, and
all extended generation config env vars in the example config file.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: [email protected] <[email protected]>
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 22:41:44 -07:00

354 lines
10 KiB
Rust

//! Model discovery and fetching for multiple LLM providers.
/// Fetch models from the Anthropic API.
///
/// Returns `(model_id, display_label)` pairs. Falls back to static defaults on error.
pub(crate) async fn fetch_anthropic_models(cached_key: Option<&str>) -> Vec<(String, String)> {
let static_defaults = vec![
(
"claude-opus-4-6".into(),
"Claude Opus 4.6 (latest flagship)".into(),
),
("claude-sonnet-4-6".into(), "Claude Sonnet 4.6".into()),
("claude-opus-4-5".into(), "Claude Opus 4.5".into()),
("claude-sonnet-4-5".into(), "Claude Sonnet 4.5".into()),
("claude-haiku-4-5".into(), "Claude Haiku 4.5 (fast)".into()),
];
let api_key = cached_key
.map(String::from)
.or_else(|| std::env::var("ANTHROPIC_API_KEY").ok())
.filter(|k| !k.is_empty() && k != crate::config::OAUTH_PLACEHOLDER);
// Fall back to OAuth token if no API key
let oauth_token = if api_key.is_none() {
crate::config::helpers::optional_env("ANTHROPIC_OAUTH_TOKEN")
.ok()
.flatten()
.filter(|t| !t.is_empty())
} else {
None
};
let (key_or_token, is_oauth) = match (api_key, oauth_token) {
(Some(k), _) => (k, false),
(None, Some(t)) => (t, true),
(None, None) => return static_defaults,
};
let client = reqwest::Client::new();
let mut request = client
.get("https://api.anthropic.com/v1/models")
.header("anthropic-version", "2023-06-01")
.timeout(std::time::Duration::from_secs(5));
if is_oauth {
request = request
.bearer_auth(&key_or_token)
.header("anthropic-beta", "oauth-2025-04-20");
} else {
request = request.header("x-api-key", &key_or_token);
}
let resp = match request.send().await {
Ok(r) if r.status().is_success() => r,
_ => return static_defaults,
};
#[derive(serde::Deserialize)]
struct ModelEntry {
id: String,
}
#[derive(serde::Deserialize)]
struct ModelsResponse {
data: Vec<ModelEntry>,
}
match resp.json::<ModelsResponse>().await {
Ok(body) => {
let mut models: Vec<(String, String)> = body
.data
.into_iter()
.filter(|m| !m.id.contains("embedding") && !m.id.contains("audio"))
.map(|m| {
let label = m.id.clone();
(m.id, label)
})
.collect();
if models.is_empty() {
return static_defaults;
}
models.sort_by(|a, b| a.0.cmp(&b.0));
models
}
Err(_) => static_defaults,
}
}
/// Fetch models from the OpenAI API.
///
/// Returns `(model_id, display_label)` pairs. Falls back to static defaults on error.
pub(crate) async fn fetch_openai_models(cached_key: Option<&str>) -> Vec<(String, String)> {
let static_defaults = vec![
(
"gpt-5.3-codex".into(),
"GPT-5.3 Codex (latest flagship)".into(),
),
("gpt-5.2-codex".into(), "GPT-5.2 Codex".into()),
("gpt-5.2".into(), "GPT-5.2".into()),
(
"gpt-5.1-codex-mini".into(),
"GPT-5.1 Codex Mini (fast)".into(),
),
("gpt-5".into(), "GPT-5".into()),
("gpt-5-mini".into(), "GPT-5 Mini".into()),
("gpt-4.1".into(), "GPT-4.1".into()),
("gpt-4.1-mini".into(), "GPT-4.1 Mini".into()),
("o4-mini".into(), "o4-mini (fast reasoning)".into()),
("o3".into(), "o3 (reasoning)".into()),
];
let api_key = cached_key
.map(String::from)
.or_else(|| std::env::var("OPENAI_API_KEY").ok())
.filter(|k| !k.is_empty());
let api_key = match api_key {
Some(k) => k,
None => return static_defaults,
};
let client = reqwest::Client::new();
let resp = match client
.get("https://api.openai.com/v1/models")
.bearer_auth(&api_key)
.timeout(std::time::Duration::from_secs(5))
.send()
.await
{
Ok(r) if r.status().is_success() => r,
_ => return static_defaults,
};
#[derive(serde::Deserialize)]
struct ModelEntry {
id: String,
}
#[derive(serde::Deserialize)]
struct ModelsResponse {
data: Vec<ModelEntry>,
}
match resp.json::<ModelsResponse>().await {
Ok(body) => {
let mut models: Vec<(String, String)> = body
.data
.into_iter()
.filter(|m| is_openai_chat_model(&m.id))
.map(|m| {
let label = m.id.clone();
(m.id, label)
})
.collect();
if models.is_empty() {
return static_defaults;
}
sort_openai_models(&mut models);
models
}
Err(_) => static_defaults,
}
}
pub(crate) fn is_openai_chat_model(model_id: &str) -> bool {
let id = model_id.to_ascii_lowercase();
let is_chat_family = id.starts_with("gpt-")
|| id.starts_with("chatgpt-")
|| id.starts_with("o1")
|| id.starts_with("o3")
|| id.starts_with("o4")
|| id.starts_with("o5");
let is_non_chat_variant = id.contains("realtime")
|| id.contains("audio")
|| id.contains("transcribe")
|| id.contains("tts")
|| id.contains("embedding")
|| id.contains("moderation")
|| id.contains("image");
is_chat_family && !is_non_chat_variant
}
pub(crate) fn openai_model_priority(model_id: &str) -> usize {
let id = model_id.to_ascii_lowercase();
const EXACT_PRIORITY: &[&str] = &[
"gpt-5.3-codex",
"gpt-5.2-codex",
"gpt-5.2",
"gpt-5.1-codex-mini",
"gpt-5",
"gpt-5-mini",
"gpt-5-nano",
"o4-mini",
"o3",
"o1",
"gpt-4.1",
"gpt-4.1-mini",
"gpt-4o",
"gpt-4o-mini",
];
if let Some(pos) = EXACT_PRIORITY.iter().position(|m| id == *m) {
return pos;
}
const PREFIX_PRIORITY: &[&str] = &[
"gpt-5.", "gpt-5-", "o3-", "o4-", "o1-", "gpt-4.1-", "gpt-4o-", "gpt-3.5-", "chatgpt-",
];
if let Some(pos) = PREFIX_PRIORITY
.iter()
.position(|prefix| id.starts_with(prefix))
{
return EXACT_PRIORITY.len() + pos;
}
EXACT_PRIORITY.len() + PREFIX_PRIORITY.len() + 1
}
pub(crate) fn sort_openai_models(models: &mut [(String, String)]) {
models.sort_by(|a, b| {
openai_model_priority(&a.0)
.cmp(&openai_model_priority(&b.0))
.then_with(|| a.0.cmp(&b.0))
});
}
/// Fetch installed models from a local Ollama instance.
///
/// Returns `(model_name, display_label)` pairs. Falls back to static defaults on error.
pub(crate) async fn fetch_ollama_models(base_url: &str) -> Vec<(String, String)> {
let static_defaults = vec![
("llama3".into(), "llama3".into()),
("mistral".into(), "mistral".into()),
("codellama".into(), "codellama".into()),
];
let url = format!("{}/api/tags", base_url.trim_end_matches('/'));
let client = reqwest::Client::new();
let resp = match client
.get(&url)
.timeout(std::time::Duration::from_secs(5))
.send()
.await
{
Ok(r) if r.status().is_success() => r,
Ok(_) => return static_defaults,
Err(_) => {
tracing::warn!(
"Could not connect to Ollama at {base_url}. Is it running? Using static defaults."
);
return static_defaults;
}
};
#[derive(serde::Deserialize)]
struct ModelEntry {
name: String,
}
#[derive(serde::Deserialize)]
struct TagsResponse {
models: Vec<ModelEntry>,
}
match resp.json::<TagsResponse>().await {
Ok(body) => {
let models: Vec<(String, String)> = body
.models
.into_iter()
.map(|m| {
let label = m.name.clone();
(m.name, label)
})
.collect();
if models.is_empty() {
return static_defaults;
}
models
}
Err(_) => static_defaults,
}
}
/// Fetch models from a generic OpenAI-compatible /v1/models endpoint.
///
/// Used for registry providers like Groq, NVIDIA NIM, etc.
pub(crate) async fn fetch_openai_compatible_models(
base_url: &str,
cached_key: Option<&str>,
) -> Vec<(String, String)> {
if base_url.is_empty() {
return vec![];
}
let url = format!("{}/models", base_url.trim_end_matches('/'));
let client = reqwest::Client::new();
let mut req = client.get(&url).timeout(std::time::Duration::from_secs(5));
if let Some(key) = cached_key {
req = req.bearer_auth(key);
}
let resp = match req.send().await {
Ok(r) if r.status().is_success() => r,
_ => return vec![],
};
#[derive(serde::Deserialize)]
struct Model {
id: String,
}
#[derive(serde::Deserialize)]
struct ModelsResponse {
data: Vec<Model>,
}
match resp.json::<ModelsResponse>().await {
Ok(body) => body
.data
.into_iter()
.map(|m| {
let label = m.id.clone();
(m.id, label)
})
.collect(),
Err(_) => vec![],
}
}
/// Build the `LlmConfig` used by `fetch_nearai_models` to list available models.
///
/// Uses [`NearAiConfig::for_model_discovery()`] to construct a minimal NEAR AI
/// config, then wraps it in an `LlmConfig` with session config for auth.
pub(crate) fn build_nearai_model_fetch_config() -> crate::config::LlmConfig {
let auth_base_url = crate::config::helpers::env_or_override("NEARAI_AUTH_URL")
.unwrap_or_else(|| "https://private.near.ai".to_string());
crate::config::LlmConfig {
backend: "nearai".to_string(),
session: crate::llm::session::SessionConfig {
auth_base_url,
session_path: crate::config::llm::default_session_path(),
},
nearai: crate::config::NearAiConfig::for_model_discovery(),
provider: None,
bedrock: None,
gemini_oauth: None,
request_timeout_secs: 120,
cheap_model: None,
smart_routing_cascade: false,
openai_codex: None,
}
}