|
| 1 | +# Copyright 2026 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +import os |
| 16 | +import google.auth |
| 17 | +from dotenv import load_dotenv |
| 18 | +load_dotenv() |
| 19 | + |
| 20 | +from google.cloud import discoveryengine_v1 |
| 21 | +from google.adk.agents.llm_agent import LlmAgent |
| 22 | +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset, StreamableHTTPConnectionParams |
| 23 | +from google.adk.tools import FunctionTool |
| 24 | +from google.apps import chat_v1 |
| 25 | +from google.oauth2.credentials import Credentials |
| 26 | + |
| 27 | +MODEL = "gemini-2.5-flash" |
| 28 | + |
| 29 | +# Access token for authentication |
| 30 | +ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN") |
| 31 | +if not ACCESS_TOKEN: |
| 32 | + raise ValueError("ACCESS_TOKEN environment variable must be set") |
| 33 | + |
| 34 | +VERTEXAI_SEARCH_TIMEOUT = 15.0 |
| 35 | + |
| 36 | +def get_project_id(): |
| 37 | + """Fetches the consumer project ID from the environment natively.""" |
| 38 | + _, project = google.auth.default() |
| 39 | + if project: |
| 40 | + return project |
| 41 | + raise Exception(f"Failed to resolve GCP Project ID from environment.") |
| 42 | + |
| 43 | +def find_serving_config_path(): |
| 44 | + """Dynamically finds the default serving config in the engine.""" |
| 45 | + project_id = get_project_id() |
| 46 | + engines = discoveryengine_v1.EngineServiceClient().list_engines( |
| 47 | + parent=f"projects/{project_id}/locations/global/collections/default_collection" |
| 48 | + ) |
| 49 | + for engine in engines: |
| 50 | + # engine.name natively contains the numeric Project Number |
| 51 | + return f"{engine.name}/servingConfigs/default_serving_config" |
| 52 | + raise Exception(f"No Discovery Engines found in project {project_id}") |
| 53 | + |
| 54 | +def send_direct_message(email: str, message: str) -> dict: |
| 55 | + """Sends a Google Chat Direct Message (DM) to a specific user by email address.""" |
| 56 | + chat_client = chat_v1.ChatServiceClient( |
| 57 | + credentials=Credentials(token=ACCESS_TOKEN) |
| 58 | + ) |
| 59 | + |
| 60 | + # 1. Setup the DM space or find existing one |
| 61 | + person = chat_v1.User( |
| 62 | + name=f"users/{email}", |
| 63 | + type_=chat_v1.User.Type.HUMAN |
| 64 | + ) |
| 65 | + membership = chat_v1.Membership(member=person) |
| 66 | + space_req = chat_v1.Space(space_type=chat_v1.Space.SpaceType.DIRECT_MESSAGE) |
| 67 | + setup_request = chat_v1.SetUpSpaceRequest( |
| 68 | + space=space_req, |
| 69 | + memberships=[membership] |
| 70 | + ) |
| 71 | + space_response = chat_client.set_up_space(request=setup_request) |
| 72 | + space_name = space_response.name |
| 73 | + |
| 74 | + # 2. Send the message |
| 75 | + msg = chat_v1.Message(text=message) |
| 76 | + message_request = chat_v1.CreateMessageRequest( |
| 77 | + parent=space_name, |
| 78 | + message=msg |
| 79 | + ) |
| 80 | + message_response = chat_client.create_message(request=message_request) |
| 81 | + |
| 82 | + return {"status": "success", "message_id": message_response.name, "space": space_name} |
| 83 | + |
| 84 | +vertexai_mcp = McpToolset( |
| 85 | + connection_params=StreamableHTTPConnectionParams( |
| 86 | + url="https://discoveryengine.googleapis.com/mcp", |
| 87 | + timeout=VERTEXAI_SEARCH_TIMEOUT, |
| 88 | + sse_read_timeout=VERTEXAI_SEARCH_TIMEOUT, |
| 89 | + headers={"Authorization": f"Bearer {ACCESS_TOKEN}"} |
| 90 | + ), |
| 91 | + tool_filter=['search'] |
| 92 | +) |
| 93 | + |
| 94 | +# Answer nicely the following user queries: |
| 95 | +# - Please find my meetings for today, I need their titles and links |
| 96 | +# - What is the latest Drive file I created? |
| 97 | +# - What is the latest Gmail message I received? |
| 98 | +# - Please send the following message to someone@example.com: Hello, this is a test message. |
| 99 | + |
| 100 | +root_agent = LlmAgent( |
| 101 | + model=MODEL, |
| 102 | + name='enterprise_ai', |
| 103 | + instruction=f""" |
| 104 | + 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. |
| 105 | + If the user asks you to send a message to someone, use the send_direct_message tool to send the message. |
| 106 | + 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. |
| 107 | + The Vertex AI MCP search tool accesses the user's data through datastores including Google Drive, Google Calendar, and Gmail. |
| 108 | + Only use the Vertex AI MCP search tool with servingConfig and query parameters, do not use any other parameters. |
| 109 | + Always use the servingConfig {find_serving_config_path()} while using the Vertex AI MCP search tool. |
| 110 | + """, |
| 111 | + tools=[vertexai_mcp, FunctionTool(send_direct_message)] |
| 112 | +) |
0 commit comments