feature: added latest gemini pro models (run-llama#876)

himself65 · May 23, 2024 · 94543de · 94543de
1 parent b963782
commit 94543de
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 19 deletions.
diff --git a/.changeset/shiny-dogs-argue.md b/.changeset/shiny-dogs-argue.md
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Added the latest preview gemini models and multi modal images taken into account
diff --git a/packages/core/src/llm/gemini/base.ts b/packages/core/src/llm/gemini/base.ts
@@ -35,11 +35,16 @@ export const GEMINI_MODEL_INFO_MAP: Record<GEMINI_MODEL, GeminiModelInfo> = {
   [GEMINI_MODEL.GEMINI_PRO]: { contextWindow: 30720 },
   [GEMINI_MODEL.GEMINI_PRO_VISION]: { contextWindow: 12288 },
   [GEMINI_MODEL.GEMINI_PRO_LATEST]: { contextWindow: 10 ** 6 },
+  // 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 },
 };
 
 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,
 ];
 
 const DEFAULT_GEMINI_PARAMS = {

diff --git a/packages/core/src/llm/gemini/types.ts b/packages/core/src/llm/gemini/types.ts
@@ -51,6 +51,8 @@ export enum GEMINI_MODEL {
   GEMINI_PRO = "gemini-pro",
   GEMINI_PRO_VISION = "gemini-pro-vision",
   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",
 }
 
 export interface GeminiModelInfo {

diff --git a/packages/core/src/llm/gemini/utils.ts b/packages/core/src/llm/gemini/utils.ts
@@ -99,7 +99,13 @@ export const cleanParts = (
 ): GeminiMessageContent => {
   return {
     ...message,
-    parts: message.parts.filter((part) => part.text?.trim()),
+    parts: message.parts.filter(
+      (part) =>
+        part.text?.trim() ||
+        part.inlineData ||
+        part.fileData ||
+        part.functionCall,
+    ),
   };
 };
 
@@ -147,24 +153,28 @@ export class GeminiHelper {
   public static mergeNeighboringSameRoleMessages(
     messages: GeminiMessageContent[],
   ): GeminiMessageContent[] {
-    return messages.reduce(
-      (
-        result: GeminiMessageContent[],
-        current: GeminiMessageContent,
-        index: number,
-      ) => {
-        if (index > 0 && messages[index - 1].role === current.role) {
-          result[result.length - 1].parts = [
-            ...result[result.length - 1].parts,
-            ...current.parts,
-          ];
-        } else {
-          result.push(current);
-        }
-        return result;
-      },
-      [],
-    );
+    return messages
+      .map(cleanParts)
+      .filter((message) => message.parts.length)
+      .reduce(
+        (
+          result: GeminiMessageContent[],
+          current: GeminiMessageContent,
+          index: number,
+          original: GeminiMessageContent[],
+        ) => {
+          if (index > 0 && original[index - 1].role === current.role) {
+            result[result.length - 1].parts = [
+              ...result[result.length - 1].parts,
+              ...current.parts,
+            ];
+          } else {
+            result.push(current);
+          }
+          return result;
+        },
+        [],
+      );
   }
 
   public static messageContentToGeminiParts(content: MessageContent): Part[] {