From 2c1c31b88da4859d01310599b59f00242fbab327 Mon Sep 17 00:00:00 2001
From: Daniel La Rocque <dlarocque@google.com>
Date: Mon, 6 May 2024 13:34:16 -0400
Subject: [PATCH] Add support for GCS FileData in inference

---
 .../src/requests/request-helpers.test.ts      | 29 ++++++++++++++++++-
 packages/vertexai/src/types/content.ts        | 24 ++++++++++++++-
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/packages/vertexai/src/requests/request-helpers.test.ts b/packages/vertexai/src/requests/request-helpers.test.ts
index 41278ee657d..d0dfbe63fc5 100644
--- 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
index c626ce75fc9..168eeda0ee9 100644
--- 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;
+}