feat: [vertexai] add GenerateContentConfig to sendMessage method (#10436)

copybara-service[bot] · jaycee-li · web-flow · commit 25d00c7c4f1a · 2024-03-05T12:43:26.000-08:00
PiperOrigin-RevId: 610386082

Co-authored-by: Jaycee Li &lt;jayceeli@google.com&gt;
diff --git a/java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/generativeai/ChatSession.java b/java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/generativeai/ChatSession.java
@@ -58,6 +58,21 @@ public ResponseStream<GenerateContentResponse> sendMessageStream(String text) th
     return sendMessageStream(text, null, null);
   }
 
+  /**
+   * Sends a message to the model and returns a stream of responses.
+   *
+   * @param text the message to be sent.
+   * @param config a {@link GenerateContentConfig} that contains all the configs for sending message
+   *     in a chat session.
+   * @return an iterable in which each element is a GenerateContentResponse. Can be converted to
+   *     stream by stream() method.
+   */
+  @BetaApi
+  public ResponseStream<GenerateContentResponse> sendMessageStream(
+      String text, GenerateContentConfig config) throws IOException {
+    return sendMessageStream(ContentMaker.fromString(text), config);
+  }
+
   /**
    * Sends a message to the model and returns a stream of responses.
    *
@@ -123,6 +138,27 @@ public ResponseStream<GenerateContentResponse> sendMessageStream(Content content
     return sendMessageStream(content, null, null);
   }
 
+  /**
+   * Sends a message to the model and returns a stream of responses.
+   *
+   * @param content the content to be sent.
+   * @param config a {@link GenerateContentConfig} that contains all the configs for sending message
+   *     in a chat session.
+   * @return an iterable in which each element is a GenerateContentResponse. Can be converted to
+   *     stream by stream() method.
+   */
+  @BetaApi
+  public ResponseStream<GenerateContentResponse> sendMessageStream(
+      Content content, GenerateContentConfig config) throws IOException {
+    checkLastResponseAndEditHistory();
+    history.add(content);
+    ResponseStream<GenerateContentResponse> respStream =
+        model.generateContentStream(history, config);
+    currentResponseStream = respStream;
+    currentResponse = null;
+    return respStream;
+  }
+
   /**
    * Sends a message to the model and returns a stream of responses.
    *
@@ -187,6 +223,20 @@ public GenerateContentResponse sendMessage(String text) throws IOException {
     return sendMessage(text, null, null);
   }
 
+  /**
+   * Sends a message to the model and returns a response.
+   *
+   * @param text the message to be sent.
+   * @param config a {@link GenerateContentConfig} that contains all the configs for sending message
+   *     in a chat session.
+   * @return a response.
+   */
+  @BetaApi
+  public GenerateContentResponse sendMessage(String text, GenerateContentConfig config)
+      throws IOException {
+    return sendMessage(ContentMaker.fromString(text), config);
+  }
+
   /**
    * Sends a message to the model and returns a response.
    *
@@ -246,6 +296,25 @@ public GenerateContentResponse sendMessage(Content content) throws IOException {
     return sendMessage(content, null, null);
   }
 
+  /**
+   * Sends a message to the model and returns a response.
+   *
+   * @param content the content to be sent.
+   * @param config a {@link GenerateContentConfig} that contains all the configs for sending message
+   *     in a chat session.
+   * @return a response.
+   */
+  @BetaApi
+  public GenerateContentResponse sendMessage(Content content, GenerateContentConfig config)
+      throws IOException {
+    checkLastResponseAndEditHistory();
+    history.add(content);
+    GenerateContentResponse response = model.generateContent(history, config);
+    currentResponse = response;
+    currentResponseStream = null;
+    return response;
+  }
+
   /**
    * Sends a message to the model and returns a response.
    *
@@ -303,7 +372,7 @@ private void removeLastContent() {
    *
    * @throws IllegalStateException if the response stream is not finished.
    */
-  private void checkLastResponseAndEditHistory() throws IllegalStateException {
+  private void checkLastResponseAndEditHistory() {
     if (currentResponseStream == null && currentResponse == null) {
       return;
     } else if (currentResponseStream != null && !currentResponseStream.isConsumed()) {
