feat: add vertex ai sample#2623
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! 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! This pull request introduces a new, locally deployable Vertex AI Agent sample designed as an Enterprise AI Assistant, featuring static authentication and Google Chat integration. Concurrently, it refactors the existing Highlights
Changelog
Using Gemini Code AssistThe 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 by creating a comment using either
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 Limitations & Feedback Gemini Code Assist 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. 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. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new local version of the Enterprise AI agent sample, which uses a static access token for authentication. It also includes some refactoring of the existing agent. My review focuses on the new local agent implementation. I've identified a critical issue with the authentication mechanism which will cause the agent to fail after a short time. I've also found high-severity issues related to dependency management, where dependencies are unpinned and defined inconsistently across pyproject.toml and requirements.txt. My suggestions aim to make the new sample more robust and the project's dependencies more maintainable.
| # Access token for authentication | ||
| ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN") | ||
| if not ACCESS_TOKEN: | ||
| raise ValueError("ACCESS_TOKEN environment variable must be set") |
There was a problem hiding this comment.
The current implementation reads a static ACCESS_TOKEN from an environment variable at module load time. Access tokens from gcloud auth application-default print-access-token are short-lived (typically 1 hour). For a long-running service started with adk web, this will cause authentication failures once the token expires, requiring a manual restart of the service.
A more robust approach is to use Application Default Credentials (ADC), which can automatically refresh tokens. This avoids the need for the ACCESS_TOKEN environment variable and makes the agent more resilient.
Here's how you can refactor the code to use ADC:
-
Remove
ACCESS_TOKENlogic: Delete lines 29-32 and theACCESS_TOKENenvironment variable dependency. -
Initialize ADC: Get credentials that can be refreshed.
from google.adk.tools import ToolContext from google.auth.transport.requests import Request SCOPES = [ "https://www.googleapis.com/auth/cloud-platform", "https://www.googleapis.com/auth/chat.spaces", "https://www.googleapis.com/auth/chat.messages", ] creds, _ = google.auth.default(scopes=SCOPES)
-
Create a header provider for
McpToolset: This will provide a fresh token for each API call.def auth_header_provider(_: ToolContext) -> dict[str, str]: """Provides auth headers, refreshing token if necessary.""" creds.refresh(Request()) return {"Authorization": f"Bearer {creds.token}"}
-
Update
McpToolsetandsend_direct_message:def send_direct_message(email: str, message: str) -> dict: """Sends a Google Chat Direct Message (DM) to a specific user by email address.""" chat_client = chat_v1.ChatServiceClient(credentials=creds) # ... rest of the function is the same vertexai_mcp = McpToolset( connection_params=StreamableHTTPConnectionParams( url="https://discoveryengine.googleapis.com/mcp", timeout=VERTEXAI_SEARCH_TIMEOUT, sse_read_timeout=VERTEXAI_SEARCH_TIMEOUT, ), header_provider=auth_header_provider, tool_filter=['search'] )
This approach also requires updating README.md to remove instructions about setting ACCESS_TOKEN.
| The agent requires a valid OAuth access token to authenticate with Google APIs (Vertex AI Search, Google Chat). | ||
| To set the `ACCESS_TOKEN` environment variable with a valid token, you must authenticate using a **Desktop app OAuth client**. | ||
|
|
||
| 1. Download your Desktop app OAuth client JSON file (e.g., `client_secret.json`) in the root directory. | ||
| 2. Authenticate using `gcloud` with the client ID and required scopes: | ||
|
|
||
| ```bash | ||
| gcloud auth application-default login \ | ||
| --client-id-file=client_secret.json \ | ||
| --scopes=https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/chat.spaces,https://www.googleapis.com/auth/chat.messages | ||
| ``` | ||
|
|
||
| 3. Generate the access token and set the environment variable: | ||
|
|
||
| ```bash | ||
| export ACCESS_TOKEN=$(gcloud auth application-default print-access-token) | ||
| ``` |
There was a problem hiding this comment.
The authentication mechanism described here relies on a short-lived ACCESS_TOKEN. This will cause the agent to fail after about an hour. It's better to rely on Application Default Credentials (ADC) directly in the code, which can handle token refreshes automatically. If the authentication logic in agent.py is updated to use ADC as suggested in another comment, these instructions should be changed to only require the user to run gcloud auth application-default login ... (step 2), and remove all mentions of ACCESS_TOKEN (steps 1, 3).
| google-adk (>=1.25.1,<2.0.0) | ||
| google-cloud-aiplatform[adk,agent_engines] (>=1.126.1,<2.0.0) | ||
| google-genai (>=1.9.0,<2.0.0) | ||
| pydantic (>=2.10.6,<3.0.0) | ||
| absl-py (>=2.2.1,<3.0.0) | ||
| google-cloud-discoveryengine (>=0.13.12,<0.14.0) | ||
| google-apps-chat (>=0.6.0,<0.7.0) |
There was a problem hiding this comment.
This requirements.txt file seems to be maintained separately from pyproject.toml, and they are inconsistent. For example, pyproject.toml lists python-dotenv which is missing here, and this file lists absl-py as a main dependency while it's a 'deployment' dependency in pyproject.toml. Having two competing dependency definition files increases maintenance overhead and can lead to confusion and errors. It's recommended to use pyproject.toml as the single source of truth and generate requirements.txt from it if needed (e.g., using poetry export).
| "google-cloud-aiplatform[adk,agent-engines]>=1.126.1", | ||
| "google-genai>=1.9.0", | ||
| "pydantic>=2.10.6", |
There was a problem hiding this comment.
Some dependencies have open-ended version constraints (e.g., >=1.126.1). This can lead to non-reproducible builds and unexpected breakages if a new major version with breaking changes is released. It's recommended to specify an upper bound for all dependencies to ensure stability. The requirements.txt file in the same directory already uses more restrictive version ranges (e.g., <2.0.0). The constraints in pyproject.toml should be at least as restrictive.
| "google-cloud-aiplatform[adk,agent-engines]>=1.126.1", | |
| "google-genai>=1.9.0", | |
| "pydantic>=2.10.6", | |
| "google-cloud-aiplatform[adk,agent-engines] (>=1.126.1,<2.0.0)", | |
| "google-genai (>=1.9.0,<2.0.0)", | |
| "pydantic (>=2.10.6,<3.0.0)", |
1 participant
No description provided.