feat: Gemini latest GA released models (run-llama#898)

Co-authored-by: Alex Yang <himself65@outlook.com>
himself65 · Jun 3, 2024 · 54230f0 · 54230f0
1 parent 3d484da
commit 54230f0
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/.changeset/tame-chefs-tickle.md b/.changeset/tame-chefs-tickle.md
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: Gemini GA release models
diff --git a/packages/core/src/llm/gemini/base.ts b/packages/core/src/llm/gemini/base.ts
@@ -38,13 +38,17 @@ export const GEMINI_MODEL_INFO_MAP: Record<GEMINI_MODEL, GeminiModelInfo> = {
   // multi-modal/multi turn
   [GEMINI_MODEL.GEMINI_PRO_1_5_PRO_PREVIEW]: { contextWindow: 10 ** 6 },
   [GEMINI_MODEL.GEMINI_PRO_1_5_FLASH_PREVIEW]: { contextWindow: 10 ** 6 },
+  [GEMINI_MODEL.GEMINI_PRO_1_5]: { contextWindow: 2 * 10 ** 6 },
+  [GEMINI_MODEL.GEMINI_PRO_1_5_FLASH]: { contextWindow: 10 ** 6 },
 };
 
 const SUPPORT_TOOL_CALL_MODELS: GEMINI_MODEL[] = [
   GEMINI_MODEL.GEMINI_PRO,
   GEMINI_MODEL.GEMINI_PRO_VISION,
   GEMINI_MODEL.GEMINI_PRO_1_5_PRO_PREVIEW,
   GEMINI_MODEL.GEMINI_PRO_1_5_FLASH_PREVIEW,
+  GEMINI_MODEL.GEMINI_PRO_1_5,
+  GEMINI_MODEL.GEMINI_PRO_1_5_FLASH,
 ];
 
 const DEFAULT_GEMINI_PARAMS = {

diff --git a/packages/core/src/llm/gemini/types.ts b/packages/core/src/llm/gemini/types.ts
@@ -53,6 +53,8 @@ export enum GEMINI_MODEL {
   GEMINI_PRO_LATEST = "gemini-1.5-pro-latest",
   GEMINI_PRO_1_5_PRO_PREVIEW = "gemini-1.5-pro-preview-0514",
   GEMINI_PRO_1_5_FLASH_PREVIEW = "gemini-1.5-flash-preview-0514",
+  GEMINI_PRO_1_5 = "gemini-1.5-pro-001",
+  GEMINI_PRO_1_5_FLASH = "gemini-1.5-flash-001",
 }
 
 export interface GeminiModelInfo {