fix: enable passing only a string to generateContent and generateContentStream

PiperOrigin-RevId: 596067266

Author: sararob

src/index.ts

@@ -347,8 +347,14 @@ export class GenerativeModel {
    * @return The GenerateContentResponse object with the response candidates.
    */
   async generateContent(
-    request: GenerateContentRequest
+    request: GenerateContentRequest | string
   ): Promise<GenerateContentResult> {
+    request = formatContentRequest(
+      request,
+      this.generation_config,
+      this.safety_settings
+    );
+
     validateGcsInput(request.contents);
 
     if (request.generation_config) {
@@ -396,8 +402,13 @@ export class GenerativeModel {
    * @return {Promise<StreamGenerateContentResult>} Promise of {@link StreamGenerateContentResult}
    */
   async generateContentStream(
-    request: GenerateContentRequest
+    request: GenerateContentRequest | string
   ): Promise<StreamGenerateContentResult> {
+    request = formatContentRequest(
+      request,
+      this.generation_config,
+      this.safety_settings
+    );
     validateGcsInput(request.contents);
 
     if (request.generation_config) {
@@ -531,3 +542,19 @@ function validateGenerationConfig(
   }
   return generation_config;
 }
+
+function formatContentRequest(
+  request: GenerateContentRequest | string,
+  generation_config?: GenerationConfig,
+  safety_settings?: SafetySetting[]
+): GenerateContentRequest {
+  if (typeof request === 'string') {
+    return {
+      contents: [{role: constants.USER_ROLE, parts: [{text: request}]}],
+      generation_config: generation_config,
+      safety_settings: safety_settings,
+    };
+  } else {
+    return request;
+  }
+}

test/index_test.ts

@@ -41,16 +41,17 @@ import {constants} from '../src/util';
 
 const PROJECT = 'test_project';
 const LOCATION = 'test_location';
+const TEST_CHAT_MESSSAGE_TEXT = 'How are you doing today?';
 const TEST_USER_CHAT_MESSAGE = [
-  {role: constants.USER_ROLE, parts: [{text: 'How are you doing today?'}]},
+  {role: constants.USER_ROLE, parts: [{text: TEST_CHAT_MESSSAGE_TEXT}]},
 ];
 const TEST_TOKEN = 'testtoken';
 
 const TEST_USER_CHAT_MESSAGE_WITH_GCS_FILE = [
   {
     role: constants.USER_ROLE,
     parts: [
-      {text: 'How are you doing today?'},
+      {text: TEST_CHAT_MESSSAGE_TEXT},
       {
         file_data: {
           file_uri: 'gs://test_bucket/test_image.jpeg',
@@ -65,7 +66,7 @@ const TEST_USER_CHAT_MESSAGE_WITH_INVALID_GCS_FILE = [
   {
     role: constants.USER_ROLE,
     parts: [
-      {text: 'How are you doing today?'},
+      {text: TEST_CHAT_MESSSAGE_TEXT},
       {file_data: {file_uri: 'test_image.jpeg', mime_type: 'image/jpeg'}},
     ],
   },
@@ -234,6 +235,16 @@ describe('VertexAI', () => {
       const resp = await model.generateContent(req);
       expect(resp).toEqual(expectedResult);
     });
+    it('returns a GenerateContentResponse when passed a string', async () => {
+      const expectedResult: GenerateContentResult = {
+        response: TEST_MODEL_RESPONSE,
+      };
+      spyOn(StreamFunctions, 'processStream').and.returnValue(
+        expectedStreamResult
+      );
+      const resp = await model.generateContent(TEST_CHAT_MESSSAGE_TEXT);
+      expect(resp).toEqual(expectedResult);
+    });
   });
 
   describe('generateContent', () => {
@@ -450,6 +461,15 @@ describe('VertexAI', () => {
       const resp = await model.generateContentStream(req);
       expect(resp).toEqual(expectedResult);
     });
+    it('returns a GenerateContentResponse when passed a string', async () => {
+      const expectedResult: StreamGenerateContentResult = {
+        response: Promise.resolve(TEST_MODEL_RESPONSE),
+        stream: testGenerator(),
+      };
+      spyOn(StreamFunctions, 'processStream').and.returnValue(expectedResult);
+      const resp = await model.generateContentStream(TEST_CHAT_MESSSAGE_TEXT);
+      expect(resp).toEqual(expectedResult);
+    });
   });
 
   describe('generateContentStream', () => {