From 19833f013284300e1d4fb8de9005f19a98592f54 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Wed, 17 Dec 2025 15:49:28 -0500 Subject: [PATCH] More minor Google AI fixes --- .../wit/since_v0.8.0/llm-provider.wit | 3 +++ .../src/wasm_host/llm_provider.rs | 3 +++ extensions/google-ai/src/google_ai.rs | 20 ++++++++++++------- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/crates/extension_api/wit/since_v0.8.0/llm-provider.wit b/crates/extension_api/wit/since_v0.8.0/llm-provider.wit index 0d5dea54c388d7e31b2b6d6e71faa90bb12c74c1..a46c267b14468897207bc8c34df6476777841f2d 100644 --- a/crates/extension_api/wit/since_v0.8.0/llm-provider.wit +++ b/crates/extension_api/wit/since_v0.8.0/llm-provider.wit @@ -33,6 +33,9 @@ interface llm-provider { enum tool-input-format { /// Standard JSON Schema format. json-schema, + /// A subset of JSON Schema supported by Google AI. + /// See https://ai.google.dev/api/caching#Schema + json-schema-subset, /// Simplified schema format for certain providers. simplified, } diff --git a/crates/extension_host/src/wasm_host/llm_provider.rs b/crates/extension_host/src/wasm_host/llm_provider.rs index 236a63d353cd558b6f9c639501dcd2d3ecaa5c63..ca3fd780a83ee704ad6c080984965d4f0aa3f222 100644 --- a/crates/extension_host/src/wasm_host/llm_provider.rs +++ b/crates/extension_host/src/wasm_host/llm_provider.rs @@ -1637,6 +1637,9 @@ impl LanguageModel for ExtensionLanguageModel { fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { match self.model_info.capabilities.tool_input_format { LlmToolInputFormat::JsonSchema => LanguageModelToolSchemaFormat::JsonSchema, + LlmToolInputFormat::JsonSchemaSubset => { + LanguageModelToolSchemaFormat::JsonSchemaSubset + } LlmToolInputFormat::Simplified => LanguageModelToolSchemaFormat::JsonSchema, } } diff --git a/extensions/google-ai/src/google_ai.rs b/extensions/google-ai/src/google_ai.rs index dfa65004c0f65a80f5f7448aebfb5630a1485f0e..808e15c6576802ba2ced79db775cd264284e1498 100644 --- a/extensions/google-ai/src/google_ai.rs +++ b/extensions/google-ai/src/google_ai.rs @@ -1,4 +1,7 @@ use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +static TOOL_CALL_COUNTER: AtomicU64 = AtomicU64::new(0); use serde::{Deserialize, Deserializer, Serialize, Serializer}; use zed_extension_api::{ @@ -273,7 +276,7 @@ fn get_default_models() -> Vec { supports_tool_choice_any: true, supports_tool_choice_none: true, supports_thinking: true, - tool_input_format: LlmToolInputFormat::JsonSchema, + tool_input_format: LlmToolInputFormat::JsonSchemaSubset, }, is_default: false, is_default_fast: true, @@ -290,7 +293,7 @@ fn get_default_models() -> Vec { supports_tool_choice_any: true, supports_tool_choice_none: true, supports_thinking: true, - tool_input_format: LlmToolInputFormat::JsonSchema, + tool_input_format: LlmToolInputFormat::JsonSchemaSubset, }, is_default: true, is_default_fast: false, @@ -307,7 +310,7 @@ fn get_default_models() -> Vec { supports_tool_choice_any: true, supports_tool_choice_none: true, supports_thinking: true, - tool_input_format: LlmToolInputFormat::JsonSchema, + tool_input_format: LlmToolInputFormat::JsonSchemaSubset, }, is_default: false, is_default_fast: false, @@ -324,7 +327,7 @@ fn get_default_models() -> Vec { supports_tool_choice_any: true, supports_tool_choice_none: true, supports_thinking: true, - tool_input_format: LlmToolInputFormat::JsonSchema, + tool_input_format: LlmToolInputFormat::JsonSchemaSubset, }, is_default: false, is_default_fast: false, @@ -341,7 +344,7 @@ fn get_default_models() -> Vec { supports_tool_choice_any: true, supports_tool_choice_none: true, supports_thinking: true, - tool_input_format: LlmToolInputFormat::JsonSchema, + tool_input_format: LlmToolInputFormat::JsonSchemaSubset, }, is_default: false, is_default_fast: false, @@ -405,7 +408,7 @@ fn get_models() -> Vec { supports_tool_choice_any: true, supports_tool_choice_none: true, supports_thinking: custom_model.thinking_budget.is_some(), - tool_input_format: LlmToolInputFormat::JsonSchema, + tool_input_format: LlmToolInputFormat::JsonSchemaSubset, }, is_default: false, is_default_fast: false, @@ -517,8 +520,11 @@ fn stream_generate_content_next( // Normalize empty string signatures to None let thought_signature = fc_part.thought_signature.filter(|s| !s.is_empty()); + // Generate unique tool use ID like hardcoded implementation + let next_tool_id = TOOL_CALL_COUNTER.fetch_add(1, Ordering::SeqCst); + let tool_use_id = format!("{}-{}", fc_part.function_call.name, next_tool_id); return Ok(Some(LlmCompletionEvent::ToolUse(LlmToolUse { - id: fc_part.function_call.name.clone(), + id: tool_use_id, name: fc_part.function_call.name, input: serde_json::to_string(&fc_part.function_call.args) .unwrap_or_default(),