Skip to content

Commit

Permalink
feat: [vertexai] add GenerateContentConfig to sendMessage method (#10436
Browse files Browse the repository at this point in the history
)

PiperOrigin-RevId: 610386082

Co-authored-by: Jaycee Li <jayceeli@google.com>
  • Loading branch information
copybara-service[bot] and jaycee-li authored Mar 5, 2024
1 parent d3110ae commit 25d00c7
Showing 1 changed file with 70 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down Expand Up @@ -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.
*
Expand Down Expand Up @@ -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.
*
Expand Down Expand Up @@ -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.
*
Expand Down Expand Up @@ -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()) {
Expand Down

0 comments on commit 25d00c7

Please sign in to comment.