feat: allow user to pass "models/model-ID" to instantiate model

PiperOrigin-RevId: 592906210

Author: yyyu-google

src/index.ts

@@ -270,6 +270,7 @@ export class GenerativeModel {
   safety_settings?: SafetySetting[];
   private _vertex_instance: VertexAI_Internal;
   private _use_non_stream = false;
+  private publisherModelEndpoint: string;
 
   /**
    * @constructor
@@ -288,6 +289,11 @@ export class GenerativeModel {
     this.model = model;
     this.generation_config = generation_config;
     this.safety_settings = safety_settings;
+    if (model.startsWith("models/")) {
+      this.publisherModelEndpoint = `publishers/google/${this.model}`;
+    } else {
+      this.publisherModelEndpoint = `publishers/google/models/${this.model}`;
+    }
   }
 
   /**
@@ -314,8 +320,6 @@ export class GenerativeModel {
       return Promise.resolve(result);
     }
 
-    const publisherModelEndpoint = `publishers/google/models/${this.model}`;
-
     const generateContentRequest: GenerateContentRequest = {
       contents: request.contents,
       generation_config: request.generation_config ?? this.generation_config,
@@ -327,7 +331,7 @@ export class GenerativeModel {
       response = await postRequest({
         region: this._vertex_instance.location,
         project: this._vertex_instance.project,
-        resourcePath: publisherModelEndpoint,
+        resourcePath: this.publisherModelEndpoint,
         resourceMethod: constants.GENERATE_CONTENT_METHOD,
         token: await this._vertex_instance.token,
         data: generateContentRequest,
@@ -361,8 +365,6 @@ export class GenerativeModel {
           validateGenerationConfig(request.generation_config);
     }
 
-    const publisherModelEndpoint = `publishers/google/models/${this.model}`;
-
     const generateContentRequest: GenerateContentRequest = {
       contents: request.contents,
       generation_config: request.generation_config ?? this.generation_config,
@@ -374,7 +376,7 @@ export class GenerativeModel {
       response = await postRequest({
         region: this._vertex_instance.location,
         project: this._vertex_instance.project,
-        resourcePath: publisherModelEndpoint,
+        resourcePath: this.publisherModelEndpoint,
         resourceMethod: constants.STREAMING_GENERATE_CONTENT_METHOD,
         token: await this._vertex_instance.token,
         data: generateContentRequest,
@@ -405,7 +407,7 @@ export class GenerativeModel {
       response = await postRequest({
         region: this._vertex_instance.location,
         project: this._vertex_instance.project,
-        resourcePath: `publishers/google/models/${this.model}`,
+        resourcePath: this.publisherModelEndpoint,
         resourceMethod: 'countTokens',
         token: await this._vertex_instance.token,
         data: request,

system_test/end_to_end_sample_test.ts

@@ -62,14 +62,22 @@ const generativeTextModel = vertex_ai.preview.getGenerativeModel({
     max_output_tokens: 256,
   },
 });
-
+const generativeTextModelWithPrefix = vertex_ai.preview.getGenerativeModel({
+  model: 'models/gemini-pro',
+  generation_config: {
+    max_output_tokens: 256,
+  },
+});
 const textModelNoOutputLimit = vertex_ai.preview.getGenerativeModel({
   model: 'gemini-pro',
 });
 
 const generativeVisionModel = vertex_ai.preview.getGenerativeModel({
   model: 'gemini-pro-vision',
 });
+const generativeVisionModelWithPrefix = vertex_ai.preview.getGenerativeModel({
+  model: 'models/gemini-pro-vision',
+});
 
 // TODO (b/316599049): update tests to use jasmine expect syntax:
 // expect(...).toBeInstanceOf(...)
@@ -92,7 +100,7 @@ describe('generateContentStream', () => {
     const aggregatedResp = await streamingResp.response;
     assert(
       aggregatedResp.candidates[0],
-      `sys test failure on generateContentStream for aggregated response: ${aggregatedResp.candidates[0]}`
+      `sys test failure on generateContentStream for aggregated response: ${aggregatedResp}`
     );
   });
   it('should not return a invalid unicode', async () => {
@@ -213,3 +221,46 @@ describe('countTokens', () => {
     );
   });
 });
+
+describe('generateContentStream using models/model-id', () => {
+  beforeEach(() => {
+    jasmine.DEFAULT_TIMEOUT_INTERVAL = 10000;
+  });
+
+  it('should should return a stream and aggregated response when passed text', async () => {
+    const streamingResp =
+      await generativeTextModelWithPrefix.generateContentStream(TEXT_REQUEST);
+
+    for await (const item of streamingResp.stream) {
+      assert(
+        item.candidates[0],
+        `sys test failure on generateContentStream using models/gemini-pro, for item ${item}`
+      );
+    }
+
+    const aggregatedResp = await streamingResp.response;
+    assert(
+      aggregatedResp.candidates[0],
+      `sys test failure on generateContentStream using models/gemini-pro for aggregated response: ${aggregatedResp}`
+    );
+  });
+  
+  it('should should return a stream and aggregated response when passed multipart base64 content when using models/gemini-pro-vision', async () => {
+    const streamingResp = await generativeVisionModelWithPrefix.generateContentStream(
+      MULTI_PART_BASE64_REQUEST
+    );
+
+    for await (const item of streamingResp.stream) {
+      assert(
+        item.candidates[0],
+        `sys test failure on generateContentStream using models/gemini-pro-vision, for item ${item}`
+      );
+    }
+
+    const aggregatedResp = await streamingResp.response;
+    assert(
+      aggregatedResp.candidates[0],
+      `sys test failure on generateContentStream using models/gemini-pro-visionfor aggregated response: ${aggregatedResp}`
+    );
+  });
+});