[VertexAI] Support cancellation in GenerateContent (#1239)

* [VertexAI] Support cancellation in GenerateContent

* Comment cleanup

Author: a-maurice

vertexai/src/Chat.cs

@@ -14,10 +14,11 @@
  * limitations under the License.
  */
 
-using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Linq;
+using System.Runtime.CompilerServices;
+using System.Threading;
 using System.Threading.Tasks;
 using Firebase.VertexAI.Internal;
 
@@ -60,72 +61,78 @@ internal static Chat InternalCreateChat(GenerativeModel model, IEnumerable<Model
   /// Sends a message using the existing history of this chat as context. If successful, the message
   /// and response will be added to the history. If unsuccessful, history will remain unchanged.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The model's response if no error occurred.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public Task<GenerateContentResponse> SendMessageAsync(
-      params ModelContent[] content) {
-    return SendMessageAsync((IEnumerable<ModelContent>)content);
+      ModelContent content, CancellationToken cancellationToken = default) {
+    return SendMessageAsync(new[] { content }, cancellationToken);
   }
   /// <summary>
   /// Sends a message using the existing history of this chat as context. If successful, the message
   /// and response will be added to the history. If unsuccessful, history will remain unchanged.
   /// </summary>
   /// <param name="text">The text given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The model's response if no error occurred.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public Task<GenerateContentResponse> SendMessageAsync(
-      string text) {
-    return SendMessageAsync(new ModelContent[] { ModelContent.Text(text) });
+      string text, CancellationToken cancellationToken = default) {
+    return SendMessageAsync(new ModelContent[] { ModelContent.Text(text) }, cancellationToken);
   }
   /// <summary>
   /// Sends a message using the existing history of this chat as context. If successful, the message
   /// and response will be added to the history. If unsuccessful, history will remain unchanged.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The model's response if no error occurred.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public Task<GenerateContentResponse> SendMessageAsync(
-      IEnumerable<ModelContent> content) {
-    return SendMessageAsyncInternal(content);
+      IEnumerable<ModelContent> content, CancellationToken cancellationToken = default) {
+    return SendMessageAsyncInternal(content, cancellationToken);
   }
 
   /// <summary>
   /// Sends a message using the existing history of this chat as context. If successful, the message
   /// and response will be added to the history. If unsuccessful, history will remain unchanged.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>A stream of generated content responses from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public IAsyncEnumerable<GenerateContentResponse> SendMessageStreamAsync(
-      params ModelContent[] content) {
-    return SendMessageStreamAsync((IEnumerable<ModelContent>)content);
+      ModelContent content, CancellationToken cancellationToken = default) {
+    return SendMessageStreamAsync(new[] { content }, cancellationToken);
   }
   /// <summary>
   /// Sends a message using the existing history of this chat as context. If successful, the message
   /// and response will be added to the history. If unsuccessful, history will remain unchanged.
   /// </summary>
   /// <param name="text">The text given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>A stream of generated content responses from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public IAsyncEnumerable<GenerateContentResponse> SendMessageStreamAsync(
-      string text) {
-    return SendMessageStreamAsync(new ModelContent[] { ModelContent.Text(text) });
+      string text, CancellationToken cancellationToken = default) {
+    return SendMessageStreamAsync(new ModelContent[] { ModelContent.Text(text) }, cancellationToken);
   }
   /// <summary>
   /// Sends a message using the existing history of this chat as context. If successful, the message
   /// and response will be added to the history. If unsuccessful, history will remain unchanged.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>A stream of generated content responses from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public IAsyncEnumerable<GenerateContentResponse> SendMessageStreamAsync(
-      IEnumerable<ModelContent> content) {
-    return SendMessageStreamAsyncInternal(content);
+      IEnumerable<ModelContent> content, CancellationToken cancellationToken = default) {
+    return SendMessageStreamAsyncInternal(content, cancellationToken);
   }
 
   private async Task<GenerateContentResponse> SendMessageAsyncInternal(
-      IEnumerable<ModelContent> requestContent) {
+      IEnumerable<ModelContent> requestContent, CancellationToken cancellationToken = default) {
     // Make sure that the requests are set to to role "user".
     List<ModelContent> fixedRequests = requestContent.Select(VertexAIExtensions.ConvertToUser).ToList();
     // Set up the context to send in the request
@@ -134,7 +141,7 @@ private async Task<GenerateContentResponse> SendMessageAsyncInternal(
 
     // Note: GenerateContentAsync can throw exceptions if there was a problem, but
     // we allow it to just be passed back to the user.
-    GenerateContentResponse response = await generativeModel.GenerateContentAsync(fullRequest);
+    GenerateContentResponse response = await generativeModel.GenerateContentAsync(fullRequest, cancellationToken);
 
     // Only after getting a valid response, add both to the history for later.
     // But either way pass the response along to the user.
@@ -149,7 +156,8 @@ private async Task<GenerateContentResponse> SendMessageAsyncInternal(
   }
 
   private async IAsyncEnumerable<GenerateContentResponse> SendMessageStreamAsyncInternal(
-      IEnumerable<ModelContent> requestContent) {
+      IEnumerable<ModelContent> requestContent,
+      [EnumeratorCancellation] CancellationToken cancellationToken = default) {
     // Make sure that the requests are set to to role "user".
     List<ModelContent> fixedRequests = requestContent.Select(VertexAIExtensions.ConvertToUser).ToList();
     // Set up the context to send in the request
@@ -161,7 +169,7 @@ private async IAsyncEnumerable<GenerateContentResponse> SendMessageStreamAsyncIn
     // Note: GenerateContentStreamAsync can throw exceptions if there was a problem, but
     // we allow it to just be passed back to the user.
     await foreach (GenerateContentResponse response in
-        generativeModel.GenerateContentStreamAsync(fullRequest)) {
+        generativeModel.GenerateContentStreamAsync(fullRequest, cancellationToken)) {
       // If the response had a problem, we still want to pass it along to the user for context,
       // but we don't want to save the history anymore.
       if (response.Candidates.Any()) {

vertexai/src/GenerativeModel.cs

@@ -19,7 +19,9 @@
 using System.IO;
 using System.Linq;
 using System.Net.Http;
+using System.Runtime.CompilerServices;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Google.MiniJSON;
 using Firebase.VertexAI.Internal;
@@ -81,94 +83,102 @@ internal GenerativeModel(FirebaseApp firebaseApp,
   /// <summary>
   /// Generates new content from input `ModelContent` given to the model as a prompt.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The generated content response from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public Task<GenerateContentResponse> GenerateContentAsync(
-      params ModelContent[] content) {
-    return GenerateContentAsync((IEnumerable<ModelContent>)content);
+      ModelContent content, CancellationToken cancellationToken = default) {
+    return GenerateContentAsync(new[] { content }, cancellationToken);
   }
   /// <summary>
   /// Generates new content from input text given to the model as a prompt.
   /// </summary>
   /// <param name="text">The text given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The generated content response from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public Task<GenerateContentResponse> GenerateContentAsync(
-      string text) {
-    return GenerateContentAsync(new ModelContent[] { ModelContent.Text(text) });
+      string text, CancellationToken cancellationToken = default) {
+    return GenerateContentAsync(new[] { ModelContent.Text(text) }, cancellationToken);
   }
   /// <summary>
   /// Generates new content from input `ModelContent` given to the model as a prompt.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The generated content response from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public Task<GenerateContentResponse> GenerateContentAsync(
-      IEnumerable<ModelContent> content) {
-    return GenerateContentAsyncInternal(content);
+      IEnumerable<ModelContent> content, CancellationToken cancellationToken = default) {
+    return GenerateContentAsyncInternal(content, cancellationToken);
   }
 
   /// <summary>
   /// Generates new content as a stream from input `ModelContent` given to the model as a prompt.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>A stream of generated content responses from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsync(
-      params ModelContent[] content) {
-    return GenerateContentStreamAsync((IEnumerable<ModelContent>)content);
+      ModelContent content, CancellationToken cancellationToken = default) {
+    return GenerateContentStreamAsync(new[] { content }, cancellationToken);
   }
   /// <summary>
   /// Generates new content as a stream from input text given to the model as a prompt.
   /// </summary>
   /// <param name="text">The text given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>A stream of generated content responses from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsync(
-      string text) {
-    return GenerateContentStreamAsync(new ModelContent[] { ModelContent.Text(text) });
+      string text, CancellationToken cancellationToken = default) {
+    return GenerateContentStreamAsync(new[] { ModelContent.Text(text) }, cancellationToken);
   }
   /// <summary>
   /// Generates new content as a stream from input `ModelContent` given to the model as a prompt.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>A stream of generated content responses from the model.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during content generation.</exception>
   public IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsync(
-      IEnumerable<ModelContent> content) {
-    return GenerateContentStreamAsyncInternal(content);
+      IEnumerable<ModelContent> content, CancellationToken cancellationToken = default) {
+    return GenerateContentStreamAsyncInternal(content, cancellationToken);
   }
 
   /// <summary>
   /// Counts the number of tokens in a prompt using the model's tokenizer.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
   /// <returns>The `CountTokensResponse` of running the model's tokenizer on the input.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during the request.</exception>
   public Task<CountTokensResponse> CountTokensAsync(
-      params ModelContent[] content) {
-    return CountTokensAsync((IEnumerable<ModelContent>)content);
+      ModelContent content, CancellationToken cancellationToken = default) {
+    return CountTokensAsync(new[] { content }, cancellationToken);
   }
   /// <summary>
   /// Counts the number of tokens in a prompt using the model's tokenizer.
   /// </summary>
   /// <param name="text">The text input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The `CountTokensResponse` of running the model's tokenizer on the input.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during the request.</exception>
   public Task<CountTokensResponse> CountTokensAsync(
-      string text) {
-    return CountTokensAsync(new ModelContent[] { ModelContent.Text(text) });
+      string text, CancellationToken cancellationToken = default) {
+    return CountTokensAsync(new[] { ModelContent.Text(text) }, cancellationToken);
   }
   /// <summary>
   /// Counts the number of tokens in a prompt using the model's tokenizer.
   /// </summary>
-  /// <param name="content">The input(s) given to the model as a prompt.</param>
+  /// <param name="content">The input given to the model as a prompt.</param>
+  /// <param name="cancellationToken">An optional token to cancel the operation.</param>
   /// <returns>The `CountTokensResponse` of running the model's tokenizer on the input.</returns>
   /// <exception cref="VertexAIException">Thrown when an error occurs during the request.</exception>
   public Task<CountTokensResponse> CountTokensAsync(
-      IEnumerable<ModelContent> content) {
-    return CountTokensAsyncInternal(content);
+      IEnumerable<ModelContent> content, CancellationToken cancellationToken = default) {
+    return CountTokensAsyncInternal(content, cancellationToken);
   }
 
   /// <summary>
@@ -188,7 +198,8 @@ public Chat StartChat(IEnumerable<ModelContent> history) {
 #endregion
 
   private async Task<GenerateContentResponse> GenerateContentAsyncInternal(
-      IEnumerable<ModelContent> content) {
+      IEnumerable<ModelContent> content,
+      CancellationToken cancellationToken) {
     HttpRequestMessage request = new(HttpMethod.Post, GetURL() + ":generateContent");
 
     // Set the request headers
@@ -204,7 +215,7 @@ private async Task<GenerateContentResponse> GenerateContentAsyncInternal(
 
     HttpResponseMessage response;
     try {
-      response = await _httpClient.SendAsync(request);
+      response = await _httpClient.SendAsync(request, cancellationToken);
       response.EnsureSuccessStatusCode();
     } catch (TaskCanceledException e) when (e.InnerException is TimeoutException) {
       throw new VertexAIRequestTimeoutException("Request timed out.", e);
@@ -223,7 +234,8 @@ private async Task<GenerateContentResponse> GenerateContentAsyncInternal(
   }
 
   private async IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsyncInternal(
-      IEnumerable<ModelContent> content) {
+      IEnumerable<ModelContent> content,
+      [EnumeratorCancellation] CancellationToken cancellationToken) {
     HttpRequestMessage request = new(HttpMethod.Post, GetURL() + ":streamGenerateContent?alt=sse");
 
     // Set the request headers
@@ -239,7 +251,7 @@ private async IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsy
 
     HttpResponseMessage response;
     try {
-      response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead);
+      response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
       response.EnsureSuccessStatusCode();
     } catch (TaskCanceledException e) when (e.InnerException is TimeoutException) {
       throw new VertexAIRequestTimeoutException("Request timed out.", e);
@@ -266,7 +278,8 @@ private async IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsy
   }
 
   private async Task<CountTokensResponse> CountTokensAsyncInternal(
-      IEnumerable<ModelContent> content) {
+      IEnumerable<ModelContent> content,
+      CancellationToken cancellationToken) {
     HttpRequestMessage request = new(HttpMethod.Post, GetURL() + ":countTokens");
 
     // Set the request headers
@@ -282,7 +295,7 @@ private async Task<CountTokensResponse> CountTokensAsyncInternal(
 
     HttpResponseMessage response;
     try {
-      response = await _httpClient.SendAsync(request);
+      response = await _httpClient.SendAsync(request, cancellationToken);
       response.EnsureSuccessStatusCode();
     } catch (TaskCanceledException e) when (e.InnerException is TimeoutException) {
       throw new VertexAIRequestTimeoutException("Request timed out.", e);