Add support for GCS FileData in inference

firebase · May 6, 2024 · 2c1c31b · 2c1c31b
1 parent 070e0cc
commit 2c1c31b
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 2 deletions.
diff --git a/packages/vertexai/src/requests/request-helpers.test.ts b/packages/vertexai/src/requests/request-helpers.test.ts
@@ -170,6 +170,33 @@ describe('request formatting methods', () => {
         ],
         systemInstruction: { role: 'system', parts: [{ text: 'be excited' }] }
       });
-    });
+    }),
+      it('formats fileData as part if provided as part', () => {
+        const result = formatGenerateContentInput([
+          'What is this?',
+          {
+            fileData: {
+              mimeType: 'image/jpeg',
+              fileUri: 'gs://sample.appspot.com/image.jpeg'
+            }
+          }
+        ]);
+        expect(result).to.be.deep.equal({
+          contents: [
+            {
+              role: 'user',
+              parts: [
+                {
+                  fileData: {
+                    mimeType: 'image/jpeg',
+                    fileUri: 'gs://sample.appspot.com/image.jpeg'
+                  }
+                },
+                { text: 'What is this?' }
+              ]
+            }
+          ]
+        });
+      });
   });
 });
diff --git a/packages/vertexai/src/types/content.ts b/packages/vertexai/src/types/content.ts
@@ -34,7 +34,8 @@ export type Part =
   | TextPart
   | InlineDataPart
   | FunctionCallPart
-  | FunctionResponsePart;
+  | FunctionResponsePart
+  | FileDataPart;
 
 /**
  * Content part interface if the part represents a text string.
@@ -101,6 +102,18 @@ export interface FunctionResponsePart {
   functionResponse: FunctionResponse;
 }
 
+/**
+ * Content part interface if the part represents {@link FileData}
+ * @public
+ */
+export interface FileDataPart {
+  text?: never;
+  inlineData?: never;
+  functionCall?: never;
+  functionResponse?: never;
+  fileData: FileData;
+}
+
 /**
  * A predicted [FunctionCall] returned from the model
  * that contains a string representing the [FunctionDeclaration.name]
@@ -137,3 +150,12 @@ export interface GenerativeContentBlob {
    */
   data: string;
 }
+
+/**
+ * Data pointing to a file uploaded on Google Cloud Storage.
+ * @public
+ */
+export interface FileData {
+  mimeType: string;
+  fileUri: string;
+}