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

) PiperOrigin-RevId: 610386082 Co-authored-by: Jaycee Li <[email protected]>
googleapis · Mar 5, 2024 · 25d00c7 · 25d00c7
1 parent d3110ae
commit 25d00c7
Showing 1 changed file with 70 additions and 1 deletion.
diff --git a/...ogle-cloud-vertexai/src/main/java/com/google/cloud/vertexai/generativeai/ChatSession.java b/...ogle-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()) {