From 54230f0477f37b30d77d2c37b21eeb99c581eddd Mon Sep 17 00:00:00 2001
From: Parham Saidi <parham@parha.me>
Date: Mon, 3 Jun 2024 20:41:15 +0200
Subject: [PATCH] feat: Gemini latest GA released models (#898)

Co-authored-by: Alex Yang <himself65@outlook.com>
---
 .changeset/tame-chefs-tickle.md       | 5 +++++
 packages/core/src/llm/gemini/base.ts  | 4 ++++
 packages/core/src/llm/gemini/types.ts | 2 ++
 3 files changed, 11 insertions(+)
 create mode 100644 .changeset/tame-chefs-tickle.md

diff --git a/.changeset/tame-chefs-tickle.md b/.changeset/tame-chefs-tickle.md
new file mode 100644
index 0000000000..2b3ba211e6
--- /dev/null
+++ 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
index 5cc072bca8..65491fd98b 100644
--- a/packages/core/src/llm/gemini/base.ts
+++ b/packages/core/src/llm/gemini/base.ts
@@ -38,6 +38,8 @@ 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[] = [
@@ -45,6 +47,8 @@ const SUPPORT_TOOL_CALL_MODELS: GEMINI_MODEL[] = [
   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
index 1ef3a881a4..998ec12ea0 100644
--- 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 {