[VertexAI] Support cancellation in GenerateContent

[a-maurice]

Description

Provide details of the change, and generalize the change in the PR title above.

Support providing a CancellationToken to the various GenerateContent calls. Removes the version with params, instead supporting a single input, or a collection of inputs, along with the optional cancellation token.

---

Testing

Describe how you've tested these changes.

Running the tests locally.

---

Type of Change

Place an x the applicable box:

• Bug fix. Add the issue # below if applicable.
• New feature. A non-breaking change which adds functionality.
• Other, such as a build process or documentation change.

---

[gemini-code-assist]

Hello @a-maurice, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

This pull request focuses on adding cancellation support to the GenerateContent methods within the VertexAI library. The primary change involves introducing a CancellationToken parameter to the GenerateContentAsync and GenerateContentStreamAsync methods in both the Chat.cs and GenerativeModel.cs files. Additionally, the original methods that accepted a variable number of ModelContent parameters (params ModelContent[] content) have been replaced with single ModelContent content and IEnumerable<ModelContent> content overloads, both of which now accept an optional CancellationToken. This enhancement allows users to cancel long-running content generation tasks, improving the responsiveness and control of their applications. The test app was also updated to use the new signatures.

Highlights

• Cancellation Support: Introduces CancellationToken to GenerateContentAsync and GenerateContentStreamAsync methods, allowing users to cancel long-running operations.
• Method Signature Changes: Replaces params ModelContent[] content overloads with ModelContent content and IEnumerable<ModelContent> content overloads, both accepting an optional CancellationToken.
• Internal Method Updates: Updates internal methods to accept and pass the CancellationToken to the underlying HTTP client calls.
• Test App Updates: Updates the test app to use the new method signatures.

Changelog

Click here to see the changelog

• vertexai/src/Chat.cs
  • Added using System.Runtime.CompilerServices; and using System.Threading; to support cancellation tokens.
  • Modified SendMessageAsync and SendMessageStreamAsync methods to accept a CancellationToken.
  • Removed the params ModelContent[] content overloads of SendMessageAsync and SendMessageStreamAsync.
  • Updated internal methods SendMessageAsyncInternal and SendMessageStreamAsyncInternal to accept and pass the CancellationToken.
• vertexai/src/GenerativeModel.cs
  • Added using System.Runtime.CompilerServices; and using System.Threading; to support cancellation tokens.
  • Modified GenerateContentAsync, GenerateContentStreamAsync, and CountTokensAsync methods to accept a CancellationToken.
  • Removed the params ModelContent[] content overloads of GenerateContentAsync, GenerateContentStreamAsync, and CountTokensAsync.
  • Updated internal methods GenerateContentAsyncInternal, GenerateContentStreamAsyncInternal, and CountTokensAsyncInternal to accept and pass the CancellationToken.
  • The SendAsync calls in the internal methods now include the cancellation token.
• vertexai/testapp/Assets/Firebase/Sample/VertexAI/UIHandlerAutomated.cs
  • Updated calls to model.GenerateContentAsync to use the new ModelContent[] { ... } syntax.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

---

A token's call,
A task finds its release,
Control's gentle hand.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces cancellation token support to the GenerateContent calls, enhancing the flexibility and responsiveness of the VertexAI SDK. The changes involve modifying the method signatures to accept an optional CancellationToken parameter and passing it down to the underlying API calls. Overall, the changes seem well-structured and address the intended functionality.

Summary of Findings

• Inconsistent documentation of 'content' parameter: The documentation for the content parameter in Chat.cs and GenerativeModel.cs is inconsistent. Some methods describe it as 'The input given to the model as a prompt,' while others say 'The input(s) given to the model as a prompt.' This inconsistency should be resolved for clarity.
• Missing cancellation token in SendMessageAsyncInternal in Chat.cs: The SendMessageAsyncInternal method in Chat.cs does not pass the cancellationToken to the generativeModel.GenerateContentAsync method, which could lead to un-cancellable operations.

Merge Readiness

The pull request introduces a valuable feature by adding cancellation token support. However, there are a couple of issues that need to be addressed before merging. Specifically, the documentation for the content parameter should be consistent across all methods, and the cancellationToken should be passed to the generativeModel.GenerateContentAsync method in SendMessageAsyncInternal in Chat.cs. I am unable to approve this pull request, and recommend that it not be merged until these issues are addressed (at a minimum), and that others review and approve this code before merging.