Visitar URL original
feat: [vertexai] add GenerateContentConfig to generateContent method … · googleapis/google-cloud-java@903abf3 · GitHub
Skip to content

Commit 903abf3

Browse files
copybara-service[bot]jaycee-li
copybara-service[bot]
and
jaycee-li
authored
feat: [vertexai] add GenerateContentConfig to generateContent method (#10425)
PiperOrigin-RevId: 609363710 Co-authored-by: Jaycee Li <jayceeli@google.com>
1 parent ec9dd00 commit 903abf3

2 files changed

Lines changed: 97 additions & 0 deletions

File tree

‎java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/generativeai/GenerativeModel.java‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,22 @@ public GenerateContentResponse generateContent(String text) throws IOException {
409409
return generateContent(text, null, null);
410410
}
411411

412+
/**
413+
* Generates content from generative model given a text and configs.
414+
*
415+
* @param text a text message to send to the generative model
416+
* @param config a {@link GenerateContentConfig} that contains all the configs in making a
417+
* generate content api call
418+
* @return a {@link com.google.cloud.vertexai.api.GenerateContentResponse} instance that contains
419+
* response contents and other metadata
420+
* @throws IOException if an I/O error occurs while making the API call
421+
*/
422+
@BetaApi
423+
public GenerateContentResponse generateContent(String text, GenerateContentConfig config)
424+
throws IOException {
425+
return generateContent(ContentMaker.fromString(text), config);
426+
}
427+
412428
/**
413429
* Generate content from generative model given a text and generation config.
414430
*
@@ -511,6 +527,41 @@ public GenerateContentResponse generateContent(
511527
return generateContent(contents, null, safetySettings);
512528
}
513529

530+
/**
531+
* Generates content from generative model given a list of contents and configs.
532+
*
533+
* @param contents a list of {@link com.google.cloud.vertexai.api.Content} to send to the
534+
* generative model
535+
* @param config a {@link GenerateContentConfig} that contains all the configs in making a
536+
* generate content api call
537+
* @return a {@link com.google.cloud.vertexai.api.GenerateContentResponse} instance that contains
538+
* response contents and other metadata
539+
* @throws IOException if an I/O error occurs while making the API call
540+
*/
541+
@BetaApi
542+
public GenerateContentResponse generateContent(
543+
List<Content> contents, GenerateContentConfig config) throws IOException {
544+
GenerateContentRequest.Builder requestBuilder =
545+
GenerateContentRequest.newBuilder().addAllContents(contents);
546+
if (config.getGenerationConfig() != null) {
547+
requestBuilder.setGenerationConfig(config.getGenerationConfig());
548+
} else if (this.generationConfig != null) {
549+
requestBuilder.setGenerationConfig(this.generationConfig);
550+
}
551+
if (config.getSafetySettings().isEmpty() == false) {
552+
requestBuilder.addAllSafetySettings(config.getSafetySettings());
553+
} else if (this.safetySettings != null) {
554+
requestBuilder.addAllSafetySettings(this.safetySettings);
555+
}
556+
if (config.getTools().isEmpty() == false) {
557+
requestBuilder.addAllTools(config.getTools());
558+
} else if (this.tools != null) {
559+
requestBuilder.addAllTools(this.tools);
560+
}
561+
562+
return generateContent(requestBuilder);
563+
}
564+
514565
/**
515566
* Generate content from generative model given a list of contents, generation config, and safety
516567
* settings.
@@ -581,6 +632,22 @@ public GenerateContentResponse generateContent(Content content) throws IOExcepti
581632
return generateContent(content, null, null);
582633
}
583634

635+
/**
636+
* Generates content from generative model given a single content and configs.
637+
*
638+
* @param content a {@link com.google.cloud.vertexai.api.Content} to send to the generative model
639+
* @param config a {@link GenerateContentConfig} that contains all the configs in making a
640+
* generate content api call
641+
* @return a {@link com.google.cloud.vertexai.api.GenerateContentResponse} instance that contains
642+
* response contents and other metadata
643+
* @throws IOException if an I/O error occurs while making the API call
644+
*/
645+
@BetaApi
646+
public GenerateContentResponse generateContent(Content content, GenerateContentConfig config)
647+
throws IOException {
648+
return generateContent(Arrays.asList(content), config);
649+
}
650+
584651
/**
585652
* Generate content from this model given a single content and generation config.
586653
*

‎java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/generativeai/GenerativeModelTest.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -452,6 +452,36 @@ public void testGenerateContentwithDefaultTools() throws Exception {
452452
assertThat(request.getValue().getTools(0)).isEqualTo(TOOL);
453453
}
454454

455+
@Test
456+
public void testGenerateContentwithGenerateContentConfig() throws Exception {
457+
model = new GenerativeModel(MODEL_NAME, vertexAi);
458+
GenerateContentConfig config =
459+
GenerateContentConfig.newBuilder()
460+
.setGenerationConfig(GENERATION_CONFIG)
461+
.setSafetySettings(safetySettings)
462+
.setTools(tools)
463+
.build();
464+
465+
Field field = VertexAI.class.getDeclaredField("predictionServiceClient");
466+
field.setAccessible(true);
467+
field.set(vertexAi, mockPredictionServiceClient);
468+
469+
when(mockPredictionServiceClient.generateContentCallable()).thenReturn(mockUnaryCallable);
470+
when(mockUnaryCallable.call(any(GenerateContentRequest.class)))
471+
.thenReturn(mockGenerateContentResponse);
472+
473+
GenerateContentResponse unused = model.generateContent(TEXT, config);
474+
475+
ArgumentCaptor<GenerateContentRequest> request =
476+
ArgumentCaptor.forClass(GenerateContentRequest.class);
477+
verify(mockUnaryCallable).call(request.capture());
478+
479+
assertThat(request.getValue().getContents(0).getParts(0).getText()).isEqualTo(TEXT);
480+
assertThat(request.getValue().getGenerationConfig()).isEqualTo(GENERATION_CONFIG);
481+
assertThat(request.getValue().getSafetySettings(0)).isEqualTo(SAFETY_SETTING);
482+
assertThat(request.getValue().getTools(0)).isEqualTo(TOOL);
483+
}
484+
455485
@Test
456486
public void testGenerateContentStreamwithText() throws Exception {
457487
model = new GenerativeModel(MODEL_NAME, vertexAi);

0 commit comments

Comments
 (0)