Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions solutions/enterprise-ai-agent/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
# Enterprise AI Agent (built as Gemini Enterprise BYO Agent)

**Note:** This project is part of an official Google Codelab ([link pending](#)).

This sample contains a specialized Gemini 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. **Dynamic Authentication (`ToolContext`):**
When deployed as a Bring-Your-Own (BYO) model via Gemini Enterprise, the session state dynamically passes an authentication token (e.g., `enterprise-ai_12345`). This agent intercepts the `ToolContext` state and extracts the token at runtime using regex pattern matching (`^enterprise-ai_\d+$`) 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 `ToolContext` used for Vertex AI.

## Deployment

This agent is designed exclusively to be deployed as a backend for a Gemini Enterprise (GE) Bring-Your-Own (BYO) Agent. It **will not** work successfully if tested locally via standard ADK run commands because it relies entirely on the external GE gateway to dynamically inject OAuth tokens into the `ToolContext` at runtime.

Deploy this agent directly to Vertex AI Agent Engines using the ADK CLI:

```bash
adk deploy agent_engine \
--project=your-gcp-project-id \
--region=us-central1 \
--display_name="Enterprise AI" \
enterprise_ai
```

[Register](https://docs.cloud.google.com/gemini/enterprise/docs/register-and-manage-an-adk-agent) the deployed agent in the Gemini Enterprise UI as a BYO agent.
15 changes: 15 additions & 0 deletions solutions/enterprise-ai-agent/enterprise_ai/__init__.py
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
129 changes: 129 additions & 0 deletions solutions/enterprise-ai-agent/enterprise_ai/agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
# 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 re
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 ToolContext, FunctionTool
from google.apps import chat_v1
from google.oauth2.credentials import Credentials

MODEL = "gemini-2.5-flash"

# Gemini Enterprise authentication injects a bearer token into the ToolContext state.
# The key pattern is "GE_AUTH_NAME_<random_digits>".
# We dynamically parse this token to authenticate our MCP and API calls.
GE_AUTH_NAME = "enterprise-ai"

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 _get_access_token_from_context(tool_context: ToolContext) -> str:
"""Helper method to dynamically parse the intercepted bearer token from the context state."""
escaped_name = re.escape(GE_AUTH_NAME)
pattern = re.compile(fr"^{escaped_name}_\d+$")
# Handle ADK varying state object types (Raw Dict vs ADK State)
state_dict = tool_context.state.to_dict() if hasattr(tool_context.state, 'to_dict') else tool_context.state
matching_keys = [k for k in state_dict.keys() if pattern.match(k)]
if matching_keys:
return state_dict.get(matching_keys[0])
raise Exception(f"No bearer token found in ToolContext state matching pattern {pattern.pattern}")

def auth_header_provider(tool_context: ToolContext) -> dict[str, str]:
token = _get_access_token_from_context(tool_context)
return {"Authorization": f"Bearer {token}"}

def send_direct_message(email: str, message: str, tool_context: ToolContext) -> dict:
"""Sends a Google Chat Direct Message (DM) to a specific user by email address."""
chat_client = chat_v1.ChatServiceClient(
credentials=Credentials(token=_get_access_token_from_context(tool_context))
)

# 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
),
tool_filter=['search'],
# The auth_header_provider dynamically injects the bearer token from the ToolContext
# into the MCP call for authentication.
header_provider=auth_header_provider
)

# 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)]
)
21 changes: 21 additions & 0 deletions solutions/enterprise-ai-agent/enterprise_ai/requirements.txt
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)
Loading
Loading