-
Notifications
You must be signed in to change notification settings - Fork 920
feat: add vertex ai sample #2623
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,47 @@ | ||
| # Enterprise AI Agent (local) | ||
|
|
||
| **Note:** This project is part of the official Google Codelabs [Integrate Vertex AI Agents with Google Workspace](https://codelabs.developers.google.com/vertexai-agents-gws). | ||
|
|
||
| This sample contains a specialized Enterprise Agent built using the Google Agent Development Kit (ADK). This agent acts as an Enterprise AI Assistant by querying user's data corpus using the Vertex AI Search MCP toolset and sending Chat messages to DM spaces using a custom Function tool & Google Chat API. | ||
|
|
||
| ## Key Features | ||
|
|
||
| 1. **Dynamic Vertex AI Serving Configs:** | ||
| The agent automatically discovers your project's `default_collection` engine and dynamically binds its queries to the `default_serving_config`. | ||
|
|
||
| 2. **Static Authentication (`ACCESS_TOKEN`):** | ||
| The client (e.g. ADK Web) passes an authentication token in the `ACCESS_TOKEN` environment variable. This agent extracts the token at runtime to securely execute calls using a Bearer token. | ||
|
|
||
| 3. **Graceful Timeouts:** | ||
| The `McpToolset` streaming components have been intentionally configured with an explicit 15-second `timeout` and `sse_read_timeout` to prevent the agent from hanging infinitely on backend network issues. | ||
|
|
||
| 4. **Google Chat Integration:** | ||
| The agent natively includes a `send_direct_message` tool powered by the `google-apps-chat` SDK. This allows the AI to immediately send direct messages to users inside Google Chat. It seamlessly reuses the same authentication token extracted from the `ACCESS_TOKEN` environment variable. | ||
|
|
||
| ## Deployment | ||
|
|
||
| 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) | ||
| ``` | ||
|
|
||
| 4. Optionally, you can set the `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` environment variables (defaults to current gcloud project and `us-central1`). | ||
|
|
||
| 5. Deploy the agent locally using the ADK `web` command: | ||
|
|
||
| ```bash | ||
| adk web | ||
| ``` | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,15 @@ | ||
| # Copyright 2026 Google LLC | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from . import agent |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,112 @@ | ||
| # Copyright 2026 Google LLC | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| import os | ||
| import google.auth | ||
| from dotenv import load_dotenv | ||
| load_dotenv() | ||
|
|
||
| from google.cloud import discoveryengine_v1 | ||
| from google.adk.agents.llm_agent import LlmAgent | ||
| from google.adk.tools.mcp_tool.mcp_toolset import McpToolset, StreamableHTTPConnectionParams | ||
| from google.adk.tools import FunctionTool | ||
| from google.apps import chat_v1 | ||
| from google.oauth2.credentials import Credentials | ||
|
|
||
| MODEL = "gemini-2.5-flash" | ||
|
|
||
| # Access token for authentication | ||
| ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN") | ||
| if not ACCESS_TOKEN: | ||
| raise ValueError("ACCESS_TOKEN environment variable must be set") | ||
|
Comment on lines
+29
to
+32
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The current implementation reads a static A more robust approach is to use Application Default Credentials (ADC), which can automatically refresh tokens. This avoids the need for the Here's how you can refactor the code to use ADC:
This approach also requires updating |
||
|
|
||
| VERTEXAI_SEARCH_TIMEOUT = 15.0 | ||
|
|
||
| def get_project_id(): | ||
| """Fetches the consumer project ID from the environment natively.""" | ||
| _, project = google.auth.default() | ||
| if project: | ||
| return project | ||
| raise Exception(f"Failed to resolve GCP Project ID from environment.") | ||
|
|
||
| def find_serving_config_path(): | ||
| """Dynamically finds the default serving config in the engine.""" | ||
| project_id = get_project_id() | ||
| engines = discoveryengine_v1.EngineServiceClient().list_engines( | ||
| parent=f"projects/{project_id}/locations/global/collections/default_collection" | ||
| ) | ||
| for engine in engines: | ||
| # engine.name natively contains the numeric Project Number | ||
| return f"{engine.name}/servingConfigs/default_serving_config" | ||
| raise Exception(f"No Discovery Engines found in project {project_id}") | ||
|
|
||
| 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=Credentials(token=ACCESS_TOKEN) | ||
| ) | ||
|
|
||
| # 1. Setup the DM space or find existing one | ||
| person = chat_v1.User( | ||
| name=f"users/{email}", | ||
| type_=chat_v1.User.Type.HUMAN | ||
| ) | ||
| membership = chat_v1.Membership(member=person) | ||
| space_req = chat_v1.Space(space_type=chat_v1.Space.SpaceType.DIRECT_MESSAGE) | ||
| setup_request = chat_v1.SetUpSpaceRequest( | ||
| space=space_req, | ||
| memberships=[membership] | ||
| ) | ||
| space_response = chat_client.set_up_space(request=setup_request) | ||
| space_name = space_response.name | ||
|
|
||
| # 2. Send the message | ||
| msg = chat_v1.Message(text=message) | ||
| message_request = chat_v1.CreateMessageRequest( | ||
| parent=space_name, | ||
| message=msg | ||
| ) | ||
| message_response = chat_client.create_message(request=message_request) | ||
|
|
||
| return {"status": "success", "message_id": message_response.name, "space": space_name} | ||
|
|
||
| vertexai_mcp = McpToolset( | ||
| connection_params=StreamableHTTPConnectionParams( | ||
| url="https://discoveryengine.googleapis.com/mcp", | ||
| timeout=VERTEXAI_SEARCH_TIMEOUT, | ||
| sse_read_timeout=VERTEXAI_SEARCH_TIMEOUT, | ||
| headers={"Authorization": f"Bearer {ACCESS_TOKEN}"} | ||
| ), | ||
| tool_filter=['search'] | ||
| ) | ||
|
|
||
| # Answer nicely the following user queries: | ||
| # - Please find my meetings for today, I need their titles and links | ||
| # - What is the latest Drive file I created? | ||
| # - What is the latest Gmail message I received? | ||
| # - Please send the following message to someone@example.com: Hello, this is a test message. | ||
|
|
||
| root_agent = LlmAgent( | ||
| model=MODEL, | ||
| name='enterprise_ai', | ||
| instruction=f""" | ||
| You are a helpful assistant that always uses the Vertex AI MCP search tool to answer the user's message, unless the user asks you to send a message to someone. | ||
| If the user asks you to send a message to someone, use the send_direct_message tool to send the message. | ||
| You MUST unconditionally use the Vertex AI MCP search tool to find answer, even if you believe you already know the answer or believe the Vertex AI MCP search tool does not contain the data. | ||
| The Vertex AI MCP search tool accesses the user's data through datastores including Google Drive, Google Calendar, and Gmail. | ||
| Only use the Vertex AI MCP search tool with servingConfig and query parameters, do not use any other parameters. | ||
| Always use the servingConfig {find_serving_config_path()} while using the Vertex AI MCP search tool. | ||
| """, | ||
| tools=[vertexai_mcp, FunctionTool(send_direct_message)] | ||
| ) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| # Copyright 2026 Google LLC | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| 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) | ||
|
Comment on lines
+15
to
+21
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This |
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 inagent.pyis updated to use ADC as suggested in another comment, these instructions should be changed to only require the user to rungcloud auth application-default login ...(step 2), and remove all mentions ofACCESS_TOKEN(steps 1, 3).