Merge pull request #70 from FormulaMonks/chore/update-gemini-supported-models

Chore/update gemini supported models

Author: theoremonemehdi

examples/basic/src/vertex.ts

@@ -8,7 +8,7 @@ const vertexAdapter = new KurtVertexAI({
     project: process.env.VERTEX_AI_PROJECT ?? "my-project",
     location: process.env.VERTEX_AI_LOCATION ?? "us-central1",
   }),
-  model: "gemini-1.0-pro", // or any other supported model
+  model: "gemini-1.5-pro", // or any other supported model
 })
 
 const kurt = new Kurt(vertexAdapter)

packages/kurt-vertex-ai/spec/isSupportedModel.spec.ts

@@ -1,20 +1,20 @@
 import { expect, test } from "@jest/globals"
-import {
-  KurtVertexAI,
-  type KurtVertexAISupportedModel,
-} from "../src/KurtVertexAI"
+import { KurtVertexAI, type KurtVertexAISupportedModel } from "../src"
 
 test("KurtVertexAI.isSupportedModel", () => {
   // For updating this test, the current list of models can be found at:
-  // https://platform.openai.com/docs/models/overview
+  // https://ai.google.dev/gemini-api/docs/models/gemini
   expect(KurtVertexAI.isSupportedModel("gemini-1.5-pro")).toBe(true)
-  expect(KurtVertexAI.isSupportedModel("gemini-1.0-pro")).toBe(true)
+  expect(KurtVertexAI.isSupportedModel("gemini-1.5-flash")).toBe(true)
+  expect(KurtVertexAI.isSupportedModel("gemini-1.5-flash-8b")).toBe(true)
+  expect(KurtVertexAI.isSupportedModel("gemini-2.0-flash")).toBe(true)
+  expect(KurtVertexAI.isSupportedModel("gemini-1.0-pro")).toBe(false)
   expect(KurtVertexAI.isSupportedModel("gemini-1.0-pro-vision")).toBe(false)
 
   expect(KurtVertexAI.isSupportedModel("bogus")).toBe(false)
 
   // The below code proves that the function works as a type guard.
-  const modelName = "gemini-1.0-pro"
+  const modelName = "gemini-1.5-pro"
   if (KurtVertexAI.isSupportedModel(modelName)) {
     const modelNameGood: KurtVertexAISupportedModel = modelName
   }

packages/kurt-vertex-ai/src/KurtVertexAI.ts

@@ -34,7 +34,12 @@ import type {
 import { ZodError } from "zod"
 
 // These models support function calling.
-const COMPATIBLE_MODELS = ["gemini-1.0-pro", "gemini-1.5-pro"] as const
+const COMPATIBLE_MODELS = [
+  "gemini-1.5-pro",
+  "gemini-1.5-flash",
+  "gemini-1.5-flash-8b",
+  "gemini-2.0-flash",
+] as const
 
 export type KurtVertexAISupportedModel = (typeof COMPATIBLE_MODELS)[number]