vrr/zed
2
0
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2026-10-19 02:54:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

google_ai: Load Gemini models from Google's model catalog

Browse source 
Google's model catalog now drives the list of Gemini models instead of a
hardcoded enum, so new models appear without a Zed release. This replaces
the `Gemini4Argon` variant added in #65565 with the same behavior: Argon
offers Low/Medium/High thinking, defaults to Medium, and sends Low when a
request asks for no thinking. Facts the catalog doesn't report (thinking
budgets vs. levels, Minimal support, default levels, excluded media/alias
models) live in one documented block of rules in `google_ai`.

- Fetch models from `v1beta/models`, which unlike `v1` includes previews.
- Only offer `gemini-<version>-…` models, excluding media and `-latest` aliases.
- Send the lowest accepted thinking level when a request asks for no
  thinking but the model can't disable it.
- Pick the default model in release order.
- Don't report usage until the prompt token count is known; Gemini streams
  empty `usageMetadata` until the final chunk, which reset the context
  indicator to zero.

Release Notes:

- Improved the Google AI provider to load available Gemini models from Google, so new models appear without a Zed update
This commit is contained in:
John Tur 2026-10-14 00:13:41 -04:00
parent 76b8ab814f
commit 5427a1f251
9 changed files with 524 additions and 421 deletions
Download patch file Download diff file Expand all files Collapse all files

2
Cargo.lock generated
View file

@ -7338,10 +7338,8 @@ dependencies = [
"http_client",
"language_model_core",
"log",
"schemars 1.0.4",
"serde",
"serde_json",
"strum 0.28.0",
]
[[package]]

5
crates/google_ai/Cargo.toml
View file

@ -12,16 +12,11 @@ workspace = true
path = "src/google_ai.rs"
doctest = false
[features]
schemars = ["dep:schemars"]
[dependencies]
anyhow.workspace = true
futures.workspace = true
http_client.workspace = true
language_model_core.workspace = true
log.workspace = true
schemars = { workspace = true, optional = true }
serde.workspace = true
serde_json.workspace = true
strum.workspace = true

120
crates/google_ai/src/completion.rs
View file

@ -209,9 +209,7 @@ fn thinking_config_for_request(
model_id: &str,
mode: GoogleModelMode,
) -> Option<ThinkingConfig> {
let supports_thinking =
matches!(mode, GoogleModelMode::Thinking { .. }) || is_google_thinking_model(model_id);
if !supports_thinking {
if !matches!(mode, GoogleModelMode::Thinking { .. }) {
return None;
}
@ -231,9 +229,9 @@ fn thinking_config_for_request(
{
config.thinking_budget = Some(budget_tokens);
}
} else if let Some(thinking_level) = disabled_thinking_level(model_id) {
} else if let Some(thinking_level) = crate::lowest_thinking_level(model_id) {
config.thinking_level = Some(thinking_level);
} else if supports_thinking_budget_disable(model_id) {
} else if crate::supports_disabling_thinking(model_id) {
config.thinking_budget = Some(0);
}
@ -248,41 +246,6 @@ impl ThinkingConfig {
}
}
fn is_google_thinking_model(model_id: &str) -> bool {
model_id.starts_with("gemini-2.5-")
|| model_id.starts_with("gemini-3")
|| model_id.starts_with("gemini-4")
}
fn disabled_thinking_level(model_id: &str) -> Option<ThinkingLevel> {
match model_id {
// Gemini 3.7 and 3.8 Flash and Gemini 4 Argon reject `MINIMAL` with a
// validation error, so `LOW` is the lowest level available to them.
model_id
if model_id.starts_with("gemini-3.7-flash")
|| model_id.starts_with("gemini-3.8-flash")
|| (model_id.starts_with("gemini-3") && model_id.contains("-pro"))
|| (model_id.starts_with("gemini-4") && model_id.contains("-argon")) =>
{
Some(ThinkingLevel::Low)
}
model_id if model_id.starts_with("gemini-3") => Some(ThinkingLevel::Minimal),
_ => None,
}
}
fn supports_thinking_budget_disable(model_id: &str) -> bool {
matches!(
model_id,
"gemini-2.5-flash"
| "gemini-2.5-flash-lite"
| "gemini-2.5-flash-preview-latest"
| "gemini-2.5-flash-preview-04-17"
| "gemini-2.5-flash-preview-05-20"
| "gemini-2.5-flash-lite-preview-06-17"
)
}
pub struct GoogleEventMapper {
usage: UsageMetadata,
stop_emitted: bool,
@ -330,9 +293,13 @@ impl GoogleEventMapper {
let mut stop_reason = None;
if let Some(usage_metadata) = event.usage_metadata {
update_usage(&mut self.usage, &usage_metadata);
events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(
convert_usage(&self.usage),
)))
// Streamed chunks carry an empty `usageMetadata` until the final one.
// Reporting usage before the prompt count is known would show zero context.
if self.usage.prompt_token_count.is_some() {
events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(
convert_usage(&self.usage),
)))
}
}
if let Some(prompt_feedback) = event.prompt_feedback
@ -803,6 +770,48 @@ mod tests {
));
}
#[test]
fn usage_is_reported_only_once_prompt_tokens_are_known() {
let mut mapper = GoogleEventMapper::new();
let chunks = [
json!({
"candidates": [{"content": {"parts": [{"text": "Let's check"}], "role": "model"}}],
"usageMetadata": {}
}),
json!({
"candidates": [{
"content": {"parts": [{"text": ""}], "role": "model"},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": 33732,
"cachedContentTokenCount": 30436,
"candidatesTokenCount": 107,
"totalTokenCount": 33839
}
}),
];
let usage_updates: Vec<TokenUsage> = chunks
.into_iter()
.flat_map(|chunk| mapper.map_event(serde_json::from_value(chunk).unwrap()))
.filter_map(|event| match event {
Ok(LanguageModelCompletionEvent::UsageUpdate(usage)) => Some(usage),
_ => None,
})
.collect();
assert_eq!(
usage_updates,
vec![TokenUsage {
input_tokens: 3296,
output_tokens: 107,
cache_read_input_tokens: 30436,
cache_creation_input_tokens: 0,
}]
);
}
#[test]
fn signed_non_thought_text_part_preserves_signature() {
let part: Part = serde_json::from_value(json!({
@ -885,31 +894,6 @@ mod tests {
));
}
#[test]
fn test_disabled_thinking_level_per_model() {
assert_eq!(
disabled_thinking_level("gemini-4-argon"),
Some(ThinkingLevel::Low)
);
assert_eq!(
disabled_thinking_level("gemini-3.8-flash"),
Some(ThinkingLevel::Low)
);
assert_eq!(
disabled_thinking_level("gemini-3.7-flash"),
Some(ThinkingLevel::Low)
);
assert_eq!(
disabled_thinking_level("gemini-3.1-pro-preview"),
Some(ThinkingLevel::Low)
);
assert_eq!(
disabled_thinking_level("gemini-3.6-flash"),
Some(ThinkingLevel::Minimal)
);
assert_eq!(disabled_thinking_level("gemini-2.5-flash"), None);
}
#[test]
fn test_function_call_with_signature_creates_tool_use_with_signature() {
let mut mapper = GoogleEventMapper::new();

543
crates/google_ai/src/google_ai.rs
View file

@ -4,7 +4,7 @@ use anyhow::{Result, anyhow, bail};
use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
use http_client::{
AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt,
StatusCode, http::HeaderMap,
StatusCode, Url, http::HeaderMap,
};
pub use language_model_core::{GOOGLE_AI_API_URL as API_URL, ModelMode as GoogleModelMode};
use language_model_core::{
@ -13,6 +13,174 @@ use language_model_core::{
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub mod completion;
// Google's model catalog reports token limits and whether a model thinks, but not the facts
// below. Update these lists when Google releases or retires models.
/// Catalog entries Zed doesn't offer: models that produce media or serve specialized APIs,
/// and `-latest` aliases, which silently retarget to models already listed by version.
///
/// Matched by ID segment so that new versions of these families stay excluded.
fn is_excluded_model(model_id: &str) -> bool {
model_id.split('-').any(|segment| {
matches!(
segment,
"latest"
| "image"
| "banana"
| "tts"
| "audio"
| "live"
| "transcribe"
| "customtools"
| "computer"
| "robotics"
| "research"
| "omni"
| "embedding"
)
})
}
/// Models that budget thinking in tokens rather than accepting thinking levels.
fn uses_thinking_budget(model_id: &str) -> bool {
matches!(
model_id,
"gemini-2.5-pro" | "gemini-2.5-flash" | "gemini-2.5-flash-lite"
)
}
/// Thinking-budget models whose thinking can be turned off with a zero budget.
pub(crate) fn supports_disabling_thinking(model_id: &str) -> bool {
matches!(model_id, "gemini-2.5-flash" | "gemini-2.5-flash-lite")
}
/// Level-based models that accept `MINIMAL`; the others reject it and start at `LOW`.
fn supports_minimal_thinking(model_id: &str) -> bool {
matches!(
model_id,
"gemini-3-flash-preview"
| "gemini-3.5-flash"
| "gemini-3.5-flash-lite"
| "gemini-3.6-flash"
)
}
/// The thinking level Zed selects for a level-based model until the user picks another.
fn default_thinking_level(model_id: &str) -> ThinkingLevel {
match model_id {
"gemini-3.5-flash-lite" => ThinkingLevel::Minimal,
"gemini-3-flash-preview" | "gemini-3.1-pro-preview" => ThinkingLevel::High,
_ => ThinkingLevel::Medium,
}
}
/// The least thinking a model accepts, sent when a request asks for none but the model can't
/// turn thinking off. `None` for thinking-budget models.
pub(crate) fn lowest_thinking_level(model_id: &str) -> Option<ThinkingLevel> {
if uses_thinking_budget(model_id) {
None
} else if supports_minimal_thinking(model_id) {
Some(ThinkingLevel::Minimal)
} else {
Some(ThinkingLevel::Low)
}
}
/// Whether the ID has a version after `gemini-`, as in `gemini-3.1-pro-preview`.
///
/// Every Gemini text model is versioned. Unversioned IDs such as `gemini-flash-latest` or
/// `gemini-nano-banana-2.1` are aliases or other model families, so Zed doesn't offer them.
fn is_versioned_model(model_id: &str) -> bool {
model_id
.strip_prefix("gemini-")
.and_then(|rest| rest.split_once('-'))
.is_some_and(|(version, _)| {
version
.split('.')
.all(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit()))
})
}
/// Fetches the models available to the API key, following every catalog page.
///
/// Entries that aren't general-purpose text models are skipped; see [`Model::from_listed`].
///
/// See <https://ai.google.dev/api/models#method:-models.list>.
///
/// # Errors
/// Returns a completion error if a request fails or a response cannot be decoded.
pub async fn list_models(
client: &dyn HttpClient,
api_url: &str,
api_key: &str,
extra_headers: &CustomHeaders,
) -> Result<Vec<Model>, LanguageModelCompletionError> {
let mut models = Vec::new();
let mut page_token: Option<String> = None;
loop {
// `v1` omits preview models such as `gemini-3.1-pro-preview`, which completions
// reach through `v1beta`, so list them from the same API version.
let mut url = Url::parse(&format!("{}/v1beta/models", api_url.trim_end_matches('/')))
.map_err(|error| LanguageModelCompletionError::Other(error.into()))?;
url.query_pairs_mut().append_pair("pageSize", "1000");
if let Some(token) = &page_token {
url.query_pairs_mut().append_pair("pageToken", token);
}
let request = HttpRequest::builder()
.method(Method::GET)
.uri(url.as_str())
.header("x-goog-api-key", api_key.trim())
.header("Accept", "application/json")
.extra_headers(extra_headers)
.body(AsyncBody::empty())
.map_err(|error| LanguageModelCompletionError::BuildRequestBody {
provider: GOOGLE_PROVIDER_NAME,
error,
})?;
let mut response =
client
.send(request)
.await
.map_err(|error| LanguageModelCompletionError::HttpSend {
provider: GOOGLE_PROVIDER_NAME,
host: url.host_str().unwrap_or(api_url).to_owned(),
error,
})?;
let status = response.status();
let headers = response.headers().clone();
let mut body = String::new();
response
.body_mut()
.read_to_string(&mut body)
.await
.map_err(|error| LanguageModelCompletionError::ApiReadResponseError {
provider: GOOGLE_PROVIDER_NAME,
error,
})?;
if !status.is_success() {
return Err(completion_error_from_response(status, &headers, body));
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct Page {
#[serde(default)]
models: Vec<ListModelEntry>,
next_page_token: Option<String>,
}
let page: Page = serde_json::from_str(&body).map_err(|error| {
LanguageModelCompletionError::DeserializeResponse {
provider: GOOGLE_PROVIDER_NAME,
error,
}
})?;
models.extend(page.models.into_iter().filter_map(Model::from_listed));
page_token = page.next_page_token.filter(|token| !token.is_empty());
if page_token.is_none() {
return Ok(models);
}
}
}
pub async fn stream_generate_content(
client: &dyn HttpClient,
api_url: &str,
@ -631,324 +799,123 @@ impl<'de> Deserialize<'de> for ModelName {
}
}
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[derive(Clone, Default, Debug, Deserialize, Serialize, PartialEq, Eq, strum::EnumIter)]
pub enum Model {
#[serde(rename = "gemini-3.5-flash-lite")]
#[default]
Gemini35FlashLite,
#[serde(rename = "gemini-3-flash-preview")]
Gemini3Flash,
#[serde(rename = "gemini-3.5-flash")]
Gemini35Flash,
#[serde(rename = "gemini-3.6-flash")]
Gemini36Flash,
#[serde(rename = "gemini-3.7-flash")]
Gemini37Flash,
#[serde(rename = "gemini-3.8-flash")]
Gemini38Flash,
#[serde(rename = "gemini-3.1-pro-preview")]
Gemini31Pro,
#[serde(rename = "gemini-4-argon")]
Gemini4Argon,
#[serde(rename = "custom")]
Custom {
name: String,
/// The name displayed in the UI, such as in the agent panel model dropdown menu.
display_name: Option<String>,
max_tokens: u64,
#[serde(default)]
mode: GoogleModelMode,
},
/// A Gemini model and the request settings Zed uses for it.
#[derive(Clone, Debug, PartialEq)]
pub struct Model {
pub id: String,
pub display_name: String,
pub max_input_tokens: u64,
pub max_output_tokens: Option<u64>,
pub mode: GoogleModelMode,
/// Empty for models that budget thinking by tokens instead of by level.
pub supported_thinking_levels: Vec<ThinkingLevel>,
pub default_thinking_level: Option<ThinkingLevel>,
}
impl Model {
pub fn default_fast() -> Self {
Self::Gemini35FlashLite
}
pub fn id(&self) -> &str {
match self {
Self::Gemini35FlashLite => "gemini-3.5-flash-lite",
Self::Gemini3Flash => "gemini-3-flash-preview",
Self::Gemini35Flash => "gemini-3.5-flash",
Self::Gemini36Flash => "gemini-3.6-flash",
Self::Gemini37Flash => "gemini-3.7-flash",
Self::Gemini38Flash => "gemini-3.8-flash",
Self::Gemini31Pro => "gemini-3.1-pro-preview",
Self::Gemini4Argon => "gemini-4-argon",
Self::Custom { name, .. } => name,
/// Builds a model from a catalog entry, or `None` if Zed doesn't offer it.
fn from_listed(entry: ListModelEntry) -> Option<Self> {
let id = entry.name.strip_prefix("models/")?;
let generates_content = entry
.supported_generation_methods
.iter()
.any(|method| method == "generateContent");
if !is_versioned_model(id) || is_excluded_model(id) || !generates_content {
return None;
}
}
pub fn request_id(&self) -> &str {
match self {
Self::Gemini35FlashLite => "gemini-3.5-flash-lite",
Self::Gemini3Flash => "gemini-3-flash-preview",
Self::Gemini35Flash => "gemini-3.5-flash",
Self::Gemini36Flash => "gemini-3.6-flash",
Self::Gemini37Flash => "gemini-3.7-flash",
Self::Gemini38Flash => "gemini-3.8-flash",
Self::Gemini31Pro => "gemini-3.1-pro-preview",
Self::Gemini4Argon => "gemini-4-argon",
Self::Custom { name, .. } => name,
}
}
let max_input_tokens = entry.input_token_limit.filter(|limit| *limit > 0)?;
let max_output_tokens = entry.output_token_limit.filter(|limit| *limit > 0)?;
pub fn display_name(&self) -> &str {
match self {
Self::Gemini35FlashLite => "Gemini 3.5 Flash-Lite",
Self::Gemini3Flash => "Gemini 3 Flash",
Self::Gemini35Flash => "Gemini 3.5 Flash",
Self::Gemini36Flash => "Gemini 3.6 Flash",
Self::Gemini37Flash => "Gemini 3.7 Flash",
Self::Gemini38Flash => "Gemini 3.8 Flash",
Self::Gemini31Pro => "Gemini 3.1 Pro",
Self::Gemini4Argon => "Gemini 4 Argon",
Self::Custom {
name, display_name, ..
} => display_name.as_ref().unwrap_or(name),
}
}
let (mode, supported_thinking_levels, default_thinking_level) =
match (entry.thinking, lowest_thinking_level(id)) {
(false, _) => (GoogleModelMode::Default, Vec::new(), None),
(true, None) => (
GoogleModelMode::Thinking {
budget_tokens: None,
},
Vec::new(),
None,
),
(true, Some(lowest)) => (
GoogleModelMode::Thinking {
budget_tokens: None,
},
[
ThinkingLevel::Minimal,
ThinkingLevel::Low,
ThinkingLevel::Medium,
ThinkingLevel::High,
]
.into_iter()
.skip_while(|level| *level != lowest)
.collect(),
Some(default_thinking_level(id)),
),
};
pub fn max_token_count(&self) -> u64 {
match self {
Self::Gemini35FlashLite
| Self::Gemini3Flash
| Self::Gemini35Flash
| Self::Gemini36Flash
| Self::Gemini37Flash
| Self::Gemini38Flash
| Self::Gemini31Pro
| Self::Gemini4Argon => 1_048_576,
Self::Custom { max_tokens, .. } => *max_tokens,
}
}
pub fn max_output_tokens(&self) -> Option<u64> {
match self {
Model::Gemini35FlashLite
| Model::Gemini3Flash
| Model::Gemini35Flash
| Model::Gemini36Flash
| Model::Gemini37Flash
| Model::Gemini38Flash
| Model::Gemini31Pro
| Model::Gemini4Argon => Some(65_536),
Model::Custom { .. } => None,
}
}
pub fn supports_tools(&self) -> bool {
true
}
pub fn supports_images(&self) -> bool {
true
Some(Self {
id: id.to_string(),
display_name: entry.display_name.unwrap_or_else(|| id.to_string()),
max_input_tokens,
max_output_tokens: Some(max_output_tokens),
mode,
supported_thinking_levels,
default_thinking_level,
})
}
pub fn supports_thinking(&self) -> bool {
matches!(
self,
Self::Gemini35FlashLite
| Self::Gemini3Flash
| Self::Gemini35Flash
| Self::Gemini36Flash
| Self::Gemini37Flash
| Self::Gemini38Flash
| Self::Gemini31Pro
| Self::Gemini4Argon
| Self::Custom {
mode: GoogleModelMode::Thinking { .. },
..
}
)
matches!(self.mode, GoogleModelMode::Thinking { .. })
}
pub fn supported_thinking_levels(&self) -> &'static [ThinkingLevel] {
match self {
Self::Gemini35FlashLite
| Self::Gemini3Flash
| Self::Gemini35Flash
| Self::Gemini36Flash => &[
ThinkingLevel::Minimal,
ThinkingLevel::Low,
ThinkingLevel::Medium,
ThinkingLevel::High,
],
Self::Gemini37Flash | Self::Gemini38Flash | Self::Gemini31Pro | Self::Gemini4Argon => {
&[
ThinkingLevel::Low,
ThinkingLevel::Medium,
ThinkingLevel::High,
]
}
_ => &[],
}
}
pub fn default_thinking_level(&self) -> Option<ThinkingLevel> {
match self {
Self::Gemini35FlashLite => Some(ThinkingLevel::Minimal),
Self::Gemini3Flash => Some(ThinkingLevel::High),
Self::Gemini35Flash => Some(ThinkingLevel::Medium),
Self::Gemini36Flash => Some(ThinkingLevel::Medium),
Self::Gemini37Flash => Some(ThinkingLevel::Medium),
Self::Gemini38Flash => Some(ThinkingLevel::Medium),
Self::Gemini31Pro => Some(ThinkingLevel::High),
Self::Gemini4Argon => Some(ThinkingLevel::Medium),
_ => None,
}
}
pub fn mode(&self) -> GoogleModelMode {
match self {
Self::Gemini35FlashLite
| Self::Gemini3Flash
| Self::Gemini35Flash
| Self::Gemini36Flash
| Self::Gemini37Flash
| Self::Gemini38Flash
| Self::Gemini31Pro
| Self::Gemini4Argon => GoogleModelMode::Thinking {
budget_tokens: None,
},
Self::Custom { mode, .. } => *mode,
}
/// Whether the model accepts requests that turn off reasoning entirely.
pub fn supports_disabling_thinking(&self) -> bool {
supports_disabling_thinking(&self.id)
}
}
impl std::fmt::Display for Model {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.id())
}
/// An entry returned by the `models.list` endpoint.
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct ListModelEntry {
name: String,
display_name: Option<String>,
input_token_limit: Option<u64>,
output_token_limit: Option<u64>,
#[serde(default)]
supported_generation_methods: Vec<String>,
#[serde(default)]
thinking: bool,
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use strum::IntoEnumIterator;
#[test]
fn test_gemini_3_6_flash_model_metadata() {
let model = Model::Gemini36Flash;
assert_eq!(model.id(), "gemini-3.6-flash");
assert_eq!(model.request_id(), "gemini-3.6-flash");
assert_eq!(model.display_name(), "Gemini 3.6 Flash");
let serialized = serde_json::to_value(&model).unwrap();
assert_eq!(serialized, json!("gemini-3.6-flash"));
let deserialized: Model = serde_json::from_value(json!("gemini-3.6-flash")).unwrap();
assert_eq!(deserialized, Model::Gemini36Flash);
}
#[test]
fn test_gemini_3_5_flash_lite_model_metadata() {
let model = Model::Gemini35FlashLite;
assert_eq!(model.id(), "gemini-3.5-flash-lite");
assert_eq!(model.request_id(), "gemini-3.5-flash-lite");
assert_eq!(model.display_name(), "Gemini 3.5 Flash-Lite");
fn test_lowest_thinking_level_per_model() {
assert_eq!(
model.supported_thinking_levels(),
&[
ThinkingLevel::Minimal,
ThinkingLevel::Low,
ThinkingLevel::Medium,
ThinkingLevel::High
]
lowest_thinking_level("gemini-3.6-flash"),
Some(ThinkingLevel::Minimal)
);
assert_eq!(model.default_thinking_level(), Some(ThinkingLevel::Minimal));
let serialized = serde_json::to_value(&model).unwrap();
assert_eq!(serialized, json!("gemini-3.5-flash-lite"));
let deserialized: Model = serde_json::from_value(json!("gemini-3.5-flash-lite")).unwrap();
assert_eq!(deserialized, Model::Gemini35FlashLite);
}
#[test]
fn test_gemini_3_7_flash_model_metadata() {
let model = Model::Gemini37Flash;
assert_eq!(model.id(), "gemini-3.7-flash");
assert_eq!(model.request_id(), "gemini-3.7-flash");
assert_eq!(model.display_name(), "Gemini 3.7 Flash");
assert_eq!(
model.supported_thinking_levels(),
&[
ThinkingLevel::Low,
ThinkingLevel::Medium,
ThinkingLevel::High
]
lowest_thinking_level("gemini-3.8-flash"),
Some(ThinkingLevel::Low)
);
let serialized = serde_json::to_value(&model).unwrap();
assert_eq!(serialized, json!("gemini-3.7-flash"));
let deserialized: Model = serde_json::from_value(json!("gemini-3.7-flash")).unwrap();
assert_eq!(deserialized, Model::Gemini37Flash);
}
#[test]
fn test_gemini_3_8_flash_model_metadata() {
let model = Model::Gemini38Flash;
assert_eq!(model.id(), "gemini-3.8-flash");
assert_eq!(model.request_id(), "gemini-3.8-flash");
assert_eq!(model.display_name(), "Gemini 3.8 Flash");
assert_eq!(model.max_token_count(), 1_048_576);
assert_eq!(model.max_output_tokens(), Some(65_536));
assert!(model.supports_thinking());
assert_eq!(
model.supported_thinking_levels(),
&[
ThinkingLevel::Low,
ThinkingLevel::Medium,
ThinkingLevel::High
]
lowest_thinking_level("gemini-4-argon"),
Some(ThinkingLevel::Low)
);
assert_eq!(model.default_thinking_level(), Some(ThinkingLevel::Medium));
let serialized = serde_json::to_value(&model).unwrap();
assert_eq!(serialized, json!("gemini-3.8-flash"));
let deserialized: Model = serde_json::from_value(json!("gemini-3.8-flash")).unwrap();
assert_eq!(deserialized, Model::Gemini38Flash);
assert_eq!(lowest_thinking_level("gemini-2.5-flash"), None);
}
#[test]
fn test_gemini_4_argon_model_metadata() {
let model = Model::Gemini4Argon;
assert_eq!(model.id(), "gemini-4-argon");
assert_eq!(model.request_id(), "gemini-4-argon");
assert_eq!(model.display_name(), "Gemini 4 Argon");
assert_eq!(model.max_token_count(), 1_048_576);
assert_eq!(model.max_output_tokens(), Some(65_536));
assert!(model.supports_thinking());
assert_eq!(
model.supported_thinking_levels(),
&[
ThinkingLevel::Low,
ThinkingLevel::Medium,
ThinkingLevel::High
]
);
assert_eq!(model.default_thinking_level(), Some(ThinkingLevel::Medium));
let serialized = serde_json::to_value(&model).unwrap();
assert_eq!(serialized, json!("gemini-4-argon"));
let deserialized: Model = serde_json::from_value(json!("gemini-4-argon")).unwrap();
assert_eq!(deserialized, Model::Gemini4Argon);
}
#[test]
fn test_max_output_tokens_leave_room_for_input() {
// The agent reserves `max_output_tokens` from the shared context window,
// so an output limit equal to the window leaves no room for input.
for model in Model::iter() {
if let Some(max_output_tokens) = model.max_output_tokens() {
assert!(
max_output_tokens < model.max_token_count(),
"{} reserves its whole context window for output",
model.id()
);
}
}
fn test_is_versioned_model() {
assert!(is_versioned_model("gemini-3.1-pro-preview"));
assert!(is_versioned_model("gemini-4-argon"));
assert!(!is_versioned_model("gemini-flash-latest"));
assert!(!is_versioned_model("gemini-nano-banana-2.1"));
}
#[test]

2
crates/language_models/Cargo.toml
View file

@ -39,7 +39,7 @@ extension.workspace = true
extension_host.workspace = true
fs.workspace = true
futures.workspace = true
google_ai = { workspace = true, features = ["schemars"] }
google_ai.workspace = true
gpui.workspace = true
gpui_tokio.workspace = true
http_client.workspace = true

197
crates/language_models/src/provider/google.rs
View file

@ -2,8 +2,8 @@ use anyhow::Result;
use collections::BTreeMap;
use credentials_provider::CredentialsProvider;
use futures::{FutureExt, StreamExt, future::BoxFuture};
use google_ai::GenerateContentResponse;
pub use google_ai::completion::{GoogleEventMapper, into_google};
use google_ai::{GenerateContentResponse, ThinkingLevel};
use gpui::{App, AppContext, AsyncApp, Context, Entity, SharedString, Task};
use http_client::{CustomHeaders, HttpClient};
use language_model::{
@ -21,8 +21,8 @@ use serde::{Deserialize, Serialize};
pub use settings::GoogleAvailableModel as AvailableModel;
use settings::{Settings, SettingsStore};
use std::sync::{Arc, LazyLock};
use strum::IntoEnumIterator;
use ui::IconName;
use util::ResultExt as _;
use language_model::ApiKeyState;
@ -56,6 +56,9 @@ pub struct GoogleLanguageModelProvider {
pub struct State {
api_key_state: ApiKeyState,
credentials_provider: Arc<dyn CredentialsProvider>,
http_client: Arc<dyn HttpClient>,
fetched_models: Vec<google_ai::Model>,
fetch_models_task: Option<Task<()>>,
}
const GEMINI_API_KEY_VAR_NAME: &str = "GEMINI_API_KEY";
@ -74,24 +77,76 @@ impl State {
fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {
let credentials_provider = self.credentials_provider.clone();
let api_url = GoogleLanguageModelProvider::api_url(cx);
self.api_key_state.store(
let should_fetch_models = api_key.is_some();
let task = self.api_key_state.store(
api_url,
api_key,
|this| &mut this.api_key_state,
credentials_provider,
cx,
)
);
self.clear_fetched_models();
cx.spawn(async move |this, cx| {
let result = task.await;
if result.is_ok() && should_fetch_models {
this.update(cx, |this, cx| this.restart_fetch_models_task(cx))
.ok();
}
result
})
}
fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
let credentials_provider = self.credentials_provider.clone();
let api_url = GoogleLanguageModelProvider::api_url(cx);
self.api_key_state.load_if_needed(
let task = self.api_key_state.load_if_needed(
api_url,
|this| &mut this.api_key_state,
credentials_provider,
cx,
)
);
cx.spawn(async move |this, cx| {
let result = task.await;
if result.is_ok() {
this.update(cx, |this, cx| this.restart_fetch_models_task(cx))
.ok();
}
result
})
}
/// Drops models fetched with an outdated key or endpoint, cancelling any fetch in flight.
fn clear_fetched_models(&mut self) {
self.fetched_models.clear();
self.fetch_models_task = None;
}
fn restart_fetch_models_task(&mut self, cx: &mut Context<Self>) {
let api_url = GoogleLanguageModelProvider::api_url(cx);
let Some(api_key) = self.api_key_state.key(&api_url) else {
self.fetch_models_task = None;
return;
};
let http_client = self.http_client.clone();
let extra_headers = GoogleLanguageModelProvider::settings(cx)
.custom_headers
.clone();
self.fetch_models_task = Some(cx.spawn(async move |this, cx| {
let Some(models) =
google_ai::list_models(http_client.as_ref(), &api_url, &api_key, &extra_headers)
.await
.log_err()
else {
return;
};
this.update(cx, |this, cx| {
this.fetched_models = models;
cx.notify();
})
.ok();
}));
}
}
@ -102,21 +157,37 @@ impl GoogleLanguageModelProvider {
cx: &mut App,
) -> Self {
let state = cx.new(|cx| {
cx.observe_global::<SettingsStore>(|this: &mut State, cx| {
let credentials_provider = this.credentials_provider.clone();
let api_url = Self::api_url(cx);
this.api_key_state.handle_url_change(
api_url,
|this| &mut this.api_key_state,
credentials_provider,
cx,
);
cx.notify();
cx.observe_global::<SettingsStore>({
let mut last_api_url = Self::api_url(cx);
let mut last_custom_headers = Self::settings(cx).custom_headers.clone();
move |this: &mut State, cx| {
let credentials_provider = this.credentials_provider.clone();
let api_url = Self::api_url(cx);
let custom_headers = Self::settings(cx).custom_headers.clone();
let catalog_changed =
api_url != last_api_url || custom_headers != last_custom_headers;
last_api_url = api_url.clone();
last_custom_headers = custom_headers;
this.api_key_state.handle_url_change(
api_url,
|this| &mut this.api_key_state,
credentials_provider,
cx,
);
if catalog_changed {
this.clear_fetched_models();
this.authenticate(cx).detach();
}
cx.notify();
}
})
.detach();
State {
api_key_state: ApiKeyState::new(Self::api_url(cx), (*API_KEY_ENV_VAR).clone()),
credentials_provider,
http_client: http_client.clone(),
fetched_models: Vec::new(),
fetch_models_task: None,
}
});
@ -127,32 +198,36 @@ impl GoogleLanguageModelProvider {
}
}
/// Every model this provider offers, keyed by id: the built-in models,
/// with settings entries added or overriding built-in ones.
/// Every model this provider offers, keyed by id: the fetched catalog,
/// with settings entries added or overriding fetched ones.
fn google_models(&self, cx: &App) -> BTreeMap<String, google_ai::Model> {
let mut models = BTreeMap::default();
for model in google_ai::Model::iter() {
if !matches!(model, google_ai::Model::Custom { .. }) {
models.insert(model.id().to_string(), model);
}
// Empty until authentication succeeds and the first catalog fetch completes.
for model in &self.state.read(cx).fetched_models {
models.insert(model.id.clone(), model.clone());
}
for model in &GoogleLanguageModelProvider::settings(cx).available_models {
models.insert(
model.name.clone(),
google_ai::Model::Custom {
name: model.name.clone(),
display_name: model.display_name.clone(),
max_tokens: model.max_tokens,
mode: model.mode.unwrap_or_default(),
},
);
for available in &GoogleLanguageModelProvider::settings(cx).available_models {
let model = available_model_to_google_model(available);
models.insert(model.id.clone(), model);
}
models
}
/// The newest offered Flash-Lite model, else the newest Flash model.
fn preferred_model(&self, cx: &App) -> Option<LanguageModel> {
let models = self.google_models(cx);
["-flash-lite", "-flash"].iter().find_map(|family| {
models
.values()
.filter(|model| model.id.contains(family))
.max_by(|a, b| a.id.cmp(&b.id))
.map(language_model)
})
}
/// The current configuration of `model`, if this provider still offers it.
fn config(
&self,
@ -239,15 +314,11 @@ impl LanguageModelProvider for GoogleLanguageModelProvider {
}
fn default_model(&self, cx: &App) -> Option<LanguageModel> {
self.google_models(cx)
.get(google_ai::Model::default().id())
.map(language_model)
self.preferred_model(cx)
}
fn default_fast_model(&self, cx: &App) -> Option<LanguageModel> {
self.google_models(cx)
.get(google_ai::Model::default_fast().id())
.map(language_model)
self.preferred_model(cx)
}
fn provided_models(&self, cx: &App) -> Vec<LanguageModel> {
@ -297,9 +368,19 @@ impl LanguageModelClient for GoogleLanguageModelProvider {
let mut request = request;
if request.max_output_tokens.is_some() {
request.max_output_tokens =
request.effective_max_output_tokens(config.max_output_tokens());
request.effective_max_output_tokens(config.max_output_tokens);
}
let request = match into_google(request, config.request_id().to_string(), config.mode()) {
// A thread keeps its effort across model switches, and Gemini rejects levels a model
// doesn't accept, including any level for models that budget thinking by tokens.
let effort_supported = request
.thinking_effort
.as_deref()
.and_then(ThinkingLevel::from_effort)
.is_some_and(|level| config.supported_thinking_levels.contains(&level));
if !effort_supported {
request.thinking_effort = None;
}
let request = match into_google(request, config.id.clone(), config.mode) {
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
@ -313,29 +394,45 @@ impl LanguageModelClient for GoogleLanguageModelProvider {
}
fn language_model(model: &google_ai::Model) -> LanguageModel {
let default_level = model.default_thinking_level();
LanguageModel {
supports_tools: model.supports_tools(),
supports_images: model.supports_images(),
supports_tools: true,
supports_images: true,
supports_thinking: model.supports_thinking(),
supports_disabling_thinking: model.supports_disabling_thinking(),
supported_effort_levels: model
.supported_thinking_levels()
.supported_thinking_levels
.iter()
.map(|level| LanguageModelEffortLevel {
name: level.name().into(),
value: level.value().into(),
is_default: Some(*level) == default_level,
is_default: Some(*level) == model.default_thinking_level,
})
.collect(),
tool_choice_support: LanguageModelToolChoiceSupport::ALL,
max_output_tokens: model.max_output_tokens(),
max_output_tokens: model.max_output_tokens,
..LanguageModel::new(
LanguageModelId::from(model.id().to_string()),
LanguageModelName::from(model.display_name().to_string()),
LanguageModelId::from(model.id.clone()),
LanguageModelName::from(model.display_name.clone()),
PROVIDER_ID,
PROVIDER_NAME,
format!("google/{}", model.request_id()),
model.max_token_count(),
format!("google/{}", model.id),
model.max_input_tokens,
)
}
}
/// Converts a settings-defined `available_models` entry into a `google_ai::Model`.
fn available_model_to_google_model(available: &AvailableModel) -> google_ai::Model {
google_ai::Model {
id: available.name.clone(),
display_name: available
.display_name
.clone()
.unwrap_or_else(|| available.name.clone()),
max_input_tokens: available.max_tokens,
max_output_tokens: None,
mode: available.mode.unwrap_or_default(),
supported_thinking_levels: Vec::new(),
default_thinking_level: None,
}
}

2
crates/language_models/src/provider/opencode.rs
View file

@ -871,7 +871,7 @@ impl LanguageModelClient for OpenCodeLanguageModelProvider {
request.max_output_tokens =
request.effective_max_output_tokens(config.model.max_output_tokens());
}
let mode = if config.supports_thinking && request.thinking_allowed {
let mode = if config.supports_thinking {
google_ai::GoogleModelMode::Thinking {
budget_tokens: None,
}

2
crates/language_models_cloud/Cargo.toml
View file

@ -18,7 +18,7 @@ anyhow.workspace = true
cloud_api_client.workspace = true
cloud_llm_client.workspace = true
futures.workspace = true
google_ai = { workspace = true, features = ["schemars"] }
google_ai.workspace = true
gpui.workspace = true
http_client.workspace = true
language_model.workspace = true

72
crates/language_models_cloud/src/language_models_cloud.rs
View file

@ -1161,11 +1161,17 @@ impl<TP: CloudLlmTokenProvider + 'static> CloudModelProvider<TP> {
cloud_llm_client::LanguageModelProvider::Google => {
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
let request =
match into_google(request, config.id.to_string(), GoogleModelMode::Default) {
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
let mode = if config.supports_thinking {
GoogleModelMode::Thinking {
budget_tokens: None,
}
} else {
GoogleModelMode::Default
};
let request = match into_google(request, config.id.to_string(), mode) {
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
let auth_context = token_provider.auth_context(cx);
let future = request_limiter.stream(async move {
let PerformLlmCompletionResponse {
@ -2390,6 +2396,62 @@ mod tests {
}
}
#[gpui::test]
async fn hosted_google_requests_summaries_using_catalog_thinking_support(
cx: &mut gpui::TestAppContext,
) {
let mut config = cloud_test_config();
config.provider = cloud_llm_client::LanguageModelProvider::Google;
config.id = cloud_llm_client::LanguageModelId("future-google-model".into());
config.supports_thinking = true;
let (provider, model) = bind_test_model(
FakeHttpClient::create(|mut request| async move {
let mut body = String::new();
request.body_mut().read_to_string(&mut body).await?;
let body: serde_json::Value = serde_json::from_str(&body)?;
assert_eq!(
body["provider_request"]["generationConfig"]["thinkingConfig"],
json!({"thinkingLevel": "MEDIUM", "includeThoughts": true})
);
let completed = CompletionEvent::Event(json!({
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Summary", "thought": true}, {"text": "Answer"}]
},
"finishReason": "STOP"
}]
}));
let ended = CompletionEvent::<serde_json::Value>::Status(
CompletionRequestStatus::StreamEnded,
);
Ok(Response::builder()
.status(200)
.header(SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, "true")
.body(AsyncBody::from(format!(
"{}\n{}\n",
serde_json::to_string(&completed)?,
serde_json::to_string(&ended)?,
)))?)
}),
config,
);
let mut request = compact_test_request();
request.thinking_allowed = true;
request.thinking_effort = Some("medium".into());
let mut stream = cx
.update(|cx| provider.stream_completion(&model, request, cx))
.await
.unwrap();
let mut summaries = Vec::new();
while let Some(event) = stream.next().await {
if let LanguageModelCompletionEvent::Thinking { text, .. } = event.unwrap() {
summaries.push(text);
}
}
assert_eq!(summaries, ["Summary"]);
}
fn cloud_google_test_model(
events: Vec<serde_json::Value>,
) -> (CloudModelProvider<TestTokenProvider>, LanguageModel) {