From 94a9ac83d304d865ce40c22b012fff8e5f86bfcb Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Wed, 11 Dec 2024 23:01:06 -0500
Subject: [PATCH] feat: add Gemini 2.0 Flash Experimental support to Google AI
 and Vertex AI plugins (#1499)

---
 js/plugins/googleai/src/gemini.ts | 16 ++++++++++++++++
 js/plugins/googleai/src/index.ts  | 11 ++++++++++-
 js/plugins/vertexai/src/gemini.ts | 16 ++++++++++++++++
 js/plugins/vertexai/src/index.ts  |  2 ++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts
index 125a62f1e..67b1932d6 100644
--- a/js/plugins/googleai/src/gemini.ts
+++ b/js/plugins/googleai/src/gemini.ts
@@ -163,6 +163,21 @@ export const gemini15Flash8b = modelRef({
   configSchema: GeminiConfigSchema,
 });
 
+export const gemini20FlashExp = modelRef({
+  name: 'googleai/gemini-2.0-flash-exp',
+  info: {
+    label: 'Google AI - Gemini 2.0 Flash (Experimental)',
+    versions: [],
+    supports: {
+      multiturn: true,
+      media: true,
+      tools: true,
+      systemRole: true,
+    },
+  },
+  configSchema: GeminiConfigSchema,
+});
+
 export const SUPPORTED_V1_MODELS = {
   'gemini-1.0-pro': gemini10Pro,
 };
@@ -171,6 +186,7 @@ export const SUPPORTED_V15_MODELS = {
   'gemini-1.5-pro': gemini15Pro,
   'gemini-1.5-flash': gemini15Flash,
   'gemini-1.5-flash-8b': gemini15Flash8b,
+  'gemini-2.0-flash-exp': gemini20FlashExp,
 };
 
 export const SUPPORTED_GEMINI_MODELS: Record<
diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts
index abb731426..be1de65f1 100644
--- a/js/plugins/googleai/src/index.ts
+++ b/js/plugins/googleai/src/index.ts
@@ -27,9 +27,18 @@ import {
   defineGoogleAIModel,
   gemini10Pro,
   gemini15Flash,
+  gemini15Flash8b,
   gemini15Pro,
+  gemini20FlashExp,
 } from './gemini.js';
-export { gemini10Pro, gemini15Flash, gemini15Pro, textEmbeddingGecko001 };
+export {
+  gemini10Pro,
+  gemini15Flash,
+  gemini15Flash8b,
+  gemini15Pro,
+  gemini20FlashExp,
+  textEmbeddingGecko001,
+};
 
 export interface PluginOptions {
   apiKey?: string;
diff --git a/js/plugins/vertexai/src/gemini.ts b/js/plugins/vertexai/src/gemini.ts
index a9e5f2529..4c45d7db9 100644
--- a/js/plugins/vertexai/src/gemini.ts
+++ b/js/plugins/vertexai/src/gemini.ts
@@ -131,6 +131,21 @@ export const gemini15Flash = modelRef({
   configSchema: GeminiConfigSchema,
 });
 
+export const gemini20FlashExp = modelRef({
+  name: 'vertexai/gemini-2.0-flash-exp',
+  info: {
+    label: 'Vertex AI - Gemini 2.0 Flash (Experimental)',
+    versions: [],
+    supports: {
+      multiturn: true,
+      media: true,
+      tools: true,
+      systemRole: true,
+    },
+  },
+  configSchema: GeminiConfigSchema,
+});
+
 export const SUPPORTED_V1_MODELS = {
   'gemini-1.0-pro': gemini10Pro,
 };
@@ -138,6 +153,7 @@ export const SUPPORTED_V1_MODELS = {
 export const SUPPORTED_V15_MODELS = {
   'gemini-1.5-pro': gemini15Pro,
   'gemini-1.5-flash': gemini15Flash,
+  'gemini-2.0-flash-exp': gemini20FlashExp,
 };
 
 export const SUPPORTED_GEMINI_MODELS = {
diff --git a/js/plugins/vertexai/src/index.ts b/js/plugins/vertexai/src/index.ts
index 688be3f82..6648ca686 100644
--- a/js/plugins/vertexai/src/index.ts
+++ b/js/plugins/vertexai/src/index.ts
@@ -32,6 +32,7 @@ import {
   gemini10Pro,
   gemini15Flash,
   gemini15Pro,
+  gemini20FlashExp,
 } from './gemini.js';
 import {
   SUPPORTED_IMAGEN_MODELS,
@@ -45,6 +46,7 @@ export {
   gemini10Pro,
   gemini15Flash,
   gemini15Pro,
+  gemini20FlashExp,
   imagen2,
   imagen3,
   imagen3Fast,