From 660b8887b07cf88beb74d1cf259c9ad79eebbf57 Mon Sep 17 00:00:00 2001
From: Waleed Latif <walif6@gmail.com>
Date: Fri, 7 Feb 2025 00:19:36 -0800
Subject: [PATCH] Updated gemini model to 2.0 flash

---
 blocks/consts.ts          |  2 +-
 providers/google/index.ts | 75 +++++----------------------------------
 tools/google/chat.ts      |  4 +--
 3 files changed, 12 insertions(+), 69 deletions(-)

diff --git a/blocks/consts.ts b/blocks/consts.ts
index d6e9c9a872..0d50d3d6b3 100644
--- a/blocks/consts.ts
+++ b/blocks/consts.ts
@@ -5,7 +5,7 @@ export const MODEL_TOOLS = {
   'deepseek-v3': 'deepseek_chat',
   'deepseek-r1': 'deepseek_reasoner',
   'claude-3-5-sonnet-20241022': 'anthropic_chat',
-  'gemini-pro': 'google_chat',
+  'gemini-2.0-flash': 'google_chat',
   'grok-2-latest': 'xai_chat',
 } as const
 
diff --git a/providers/google/index.ts b/providers/google/index.ts
index 381eee8b36..de78d36809 100644
--- a/providers/google/index.ts
+++ b/providers/google/index.ts
@@ -6,10 +6,11 @@ export const googleProvider: ProviderConfig = {
   name: 'Google',
   description: "Google's Gemini models",
   version: '1.0.0',
-  models: ['gemini-pro'],
-  defaultModel: 'gemini-pro',
+  models: ['gemini-2.0-flash-001'],
+  defaultModel: 'gemini-2.0-flash-001',
 
-  baseUrl: 'https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent',
+  baseUrl:
+    'https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-001:generateContent',
   headers: (apiKey: string) => ({
     'Content-Type': 'application/json',
     'x-goog-api-key': apiKey,
@@ -114,68 +115,7 @@ export const googleProvider: ProviderConfig = {
 
   transformRequest: (request: ProviderRequest, functions?: any) => {
     // Combine system prompt and context into a single message if both exist
-    let initialMessage = request.systemPrompt
-
-    // Add response format for structured output if specified
-    let generationConfig: any = {
-      temperature: request.temperature || 0.7,
-      maxOutputTokens: request.maxTokens || 1024,
-    }
-
-    if (request.responseFormat) {
-      // Convert our standard JSON Schema format to Google's expected format
-      const googleSchema = {
-        type: 'OBJECT',
-        properties: request.responseFormat.fields.reduce(
-          (acc, field) => ({
-            ...acc,
-            [field.name]: {
-              type:
-                field.type === 'array'
-                  ? 'ARRAY'
-                  : field.type === 'object'
-                    ? 'OBJECT'
-                    : field.type === 'number'
-                      ? 'NUMBER'
-                      : field.type === 'boolean'
-                        ? 'BOOLEAN'
-                        : 'STRING',
-              description: field.description || '',
-              ...(field.type === 'array' && {
-                items: {
-                  type: 'STRING',
-                },
-              }),
-              ...(field.type === 'object' && {
-                properties: {
-                  // Add a flexible string field to satisfy Google's validation
-                  _any: {
-                    type: 'STRING',
-                    description: 'Any additional metadata',
-                  },
-                },
-              }),
-            },
-          }),
-          {}
-        ),
-      }
-
-      // Add schema to generation config
-      generationConfig = {
-        ...generationConfig,
-        response_mime_type: 'application/json',
-        response_schema: googleSchema,
-      }
-
-      // Add clear instructions in the initial message
-      initialMessage = `${initialMessage}\n\nPlease provide your response as a valid JSON object following the specified schema. The response should only contain the JSON data without any additional text or explanations.`
-    }
-
-    // Add context if provided
-    if (request.context) {
-      initialMessage += `\n\n${request.context}`
-    }
+    const initialMessage = request.systemPrompt + (request.context ? `\n\n${request.context}` : '')
 
     const messages = [
       { role: 'user', parts: [{ text: initialMessage }] },
@@ -234,7 +174,10 @@ export const googleProvider: ProviderConfig = {
     return {
       contents: messages,
       tools: functions ? [{ functionDeclarations: functions.functionDeclarations }] : undefined,
-      generationConfig,
+      generationConfig: {
+        temperature: request.temperature || 0.7,
+        maxOutputTokens: request.maxTokens || 1024,
+      },
     }
   },
 
diff --git a/tools/google/chat.ts b/tools/google/chat.ts
index 972ff8d40c..408af2c069 100644
--- a/tools/google/chat.ts
+++ b/tools/google/chat.ts
@@ -44,7 +44,7 @@ export const chatTool: ToolConfig<ChatParams, ChatResponse> = {
     },
     model: {
       type: 'string',
-      default: 'gemini-pro',
+      default: 'gemini-2.0-flash-001',
       description: 'Model to use',
     },
     temperature: {
@@ -55,7 +55,7 @@ export const chatTool: ToolConfig<ChatParams, ChatResponse> = {
   },
 
   request: {
-    url: 'https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent',
+    url: 'https://generativelanguage.googleapis.com/v1/models/gemini-2.0-flash-001:generateContent',
     method: 'POST',
     headers: (params) => ({
       'Content-Type': 'application/json',