language_models: Use `JsonSchemaSubset` for Gemini models in OpenRouter (#33477)

Umesh Yadav created 5 months ago

Closes #33466

Release Notes:

- N/A

Change summary

crates/language_models/src/provider/open_router.rs | 13 +++++++++++--
1 file changed, 11 insertions(+), 2 deletions(-)

Detailed changes

crates/language_models/src/provider/open_router.rs

@@ -11,8 +11,8 @@ use language_model::{
     AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
     LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,
     LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,
-    LanguageModelToolChoice, LanguageModelToolResultContent, LanguageModelToolUse, MessageContent,
-    RateLimiter, Role, StopReason, TokenUsage,
+    LanguageModelToolChoice, LanguageModelToolResultContent, LanguageModelToolSchemaFormat,
+    LanguageModelToolUse, MessageContent, RateLimiter, Role, StopReason, TokenUsage,
 };
 use open_router::{
     Model, ModelMode as OpenRouterModelMode, ResponseStreamEvent, list_models, stream_completion,
@@ -374,6 +374,15 @@ impl LanguageModel for OpenRouterLanguageModel {
         self.model.supports_tool_calls()
     }
 
+    fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
+        let model_id = self.model.id().trim().to_lowercase();
+        if model_id.contains("gemini") {
+            LanguageModelToolSchemaFormat::JsonSchemaSubset
+        } else {
+            LanguageModelToolSchemaFormat::JsonSchema
+        }
+    }
+
     fn telemetry_id(&self) -> String {
         format!("openrouter/{}", self.model.id())
     }