Skip to content

Commit 90e1e3e

Browse files
seanzhougooglecopybara-github
authored andcommitted
chore: Add sample agent for testing oatuh2 client credentials grant type
PiperOrigin-RevId: 815863131
1 parent 4bb089d commit 90e1e3e

5 files changed

Lines changed: 794 additions & 0 deletions

File tree

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
# OAuth2 Client Credentials Weather Agent
2+
3+
This sample demonstrates OAuth2 client credentials flow with ADK's `AuthenticatedFunctionTool` using a practical weather assistant agent.
4+
5+
## Overview
6+
7+
The OAuth2 client credentials grant type is used for server-to-server authentication where no user interaction is required. This demo shows:
8+
9+
- How to configure OAuth2 client credentials in ADK
10+
- Using `AuthenticatedFunctionTool` for automatic token management
11+
- Transparent authentication in a practical weather assistant
12+
- Testing the OAuth2 client credentials implementation
13+
14+
## Architecture
15+
16+
```
17+
[WeatherAssistant] -> [AuthenticatedFunctionTool] -> [OAuth2CredentialExchanger] -> [OAuth2 Server] -> [Weather API]
18+
```
19+
20+
1. **WeatherAssistant** calls weather tool when user asks for weather data
21+
2. **AuthenticatedFunctionTool** automatically handles OAuth2 flow
22+
3. **OAuth2CredentialExchanger** exchanges client credentials for access token
23+
4. **Authenticated requests** are made to weather API
24+
25+
## Files
26+
27+
### `agent.py` - WeatherAssistant Agent
28+
29+
Weather assistant agent that demonstrates OAuth2 client credentials flow transparently:
30+
31+
- **OAuth2 Configuration**: Client credentials setup with token URL and scopes
32+
- **Weather Tool**: Single `get_weather_data` tool for fetching weather information
33+
- **Agent Definition**: ADK LLM agent focused on providing weather information
34+
35+
**Key Features:**
36+
- Automatic token exchange using client ID and secret
37+
- Bearer token authentication
38+
- Transparent OAuth2 handling (invisible to the model)
39+
- Practical use case demonstrating machine-to-machine authentication
40+
41+
### `main.py` - CLI Interface
42+
43+
Command-line interface for running the WeatherAssistant agent:
44+
45+
```bash
46+
# Ask for weather
47+
python contributing/samples/oauth2_client_credentials/main.py "What's the weather in Tokyo?"
48+
```
49+
50+
**Requirements:**
51+
- LLM API key (Google AI or Vertex AI)
52+
- OAuth2 test server running
53+
54+
### `oauth2_test_server.py` - Local OAuth2 Server
55+
56+
Mock OAuth2 server for testing the client credentials flow:
57+
58+
```bash
59+
python contributing/samples/oauth2_client_credentials/oauth2_test_server.py
60+
```
61+
62+
**Features:**
63+
- OIDC discovery endpoint (`/.well-known/openid_configuration`)
64+
- Client credentials token exchange (`/token`)
65+
- Protected weather API (`/api/weather`)
66+
- Supports both `authorization_code` and `client_credentials` grant types
67+
- Test credentials: `client_id="test_client"`, `client_secret="test_secret"`
68+
69+
**Endpoints:**
70+
- `GET /.well-known/openid_configuration` - OIDC discovery
71+
- `POST /token` - Token exchange
72+
- `GET /api/weather` - Weather API (requires Bearer token)
73+
- `GET /` - Server info
74+
75+
## Quick Start
76+
77+
1. **Start the OAuth2 server:**
78+
```bash
79+
python contributing/samples/oauth2_client_credentials/oauth2_test_server.py &
80+
```
81+
2. Create a `.env` file in the project root with your API credentials:
82+
83+
```bash
84+
# Choose Model Backend: 0 -> ML Dev, 1 -> Vertex
85+
GOOGLE_GENAI_USE_VERTEXAI=1
86+
87+
# ML Dev backend config
88+
GOOGLE_API_KEY=your_google_api_key_here
89+
90+
# Vertex backend config
91+
GOOGLE_CLOUD_PROJECT=your_project_id
92+
GOOGLE_CLOUD_LOCATION=us-central1
93+
```
94+
95+
3. **Run the agent:**
96+
```bash
97+
# Ask for weather
98+
python contributing/samples/oauth2_client_credentials/main.py "What's the weather in Tokyo?"
99+
```
100+
101+
3. **Interactive demo (use ADK commands):**
102+
```bash
103+
# Interactive CLI
104+
adk run contributing/samples/oauth2_client_credentials
105+
106+
# Interactive web UI
107+
adk web contributing/samples
108+
```
109+
110+
## OAuth2 Configuration
111+
112+
The agent uses these OAuth2 settings (configured in `agent.py`):
113+
114+
```python
115+
flows = OAuthFlows(
116+
clientCredentials=OAuthFlowClientCredentials(
117+
tokenUrl="http://localhost:8000/token",
118+
scopes={
119+
"read": "Read access to weather data",
120+
"write": "Write access for data updates",
121+
"admin": "Administrative access",
122+
},
123+
)
124+
)
125+
126+
raw_credential = AuthCredential(
127+
auth_type=AuthCredentialTypes.OAUTH2,
128+
oauth2=OAuth2Auth(
129+
client_id="test_client",
130+
client_secret="test_secret",
131+
),
132+
)
133+
```
134+
135+
## Authentication Flow
136+
137+
1. **Weather Request**: User asks WeatherAssistant for weather information
138+
2. **Tool Invocation**: Agent calls `get_weather_data` authenticated function tool
139+
3. **Credential Loading**: CredentialManager loads OAuth2 configuration
140+
4. **Token Exchange**: OAuth2CredentialExchanger uses client credentials to get access token
141+
5. **Request Enhancement**: AuthenticatedFunctionTool adds `Authorization: Bearer <token>` header
142+
6. **API Call**: Weather API accessed with valid token
143+
7. **Response**: Weather data returned to user
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright 2025 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+
from . import agent
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
# Copyright 2025 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+
"""Weather Assistant Agent.
16+
17+
This agent provides weather information for cities worldwide.
18+
It demonstrates OAuth2 client credentials flow transparently
19+
through AuthenticatedFunctionTool usage.
20+
"""
21+
22+
from fastapi.openapi.models import OAuth2
23+
from fastapi.openapi.models import OAuthFlowClientCredentials
24+
from fastapi.openapi.models import OAuthFlows
25+
from google.adk.agents.llm_agent import Agent
26+
from google.adk.auth.auth_credential import AuthCredential
27+
from google.adk.auth.auth_credential import AuthCredentialTypes
28+
from google.adk.auth.auth_credential import OAuth2Auth
29+
from google.adk.auth.auth_tool import AuthConfig
30+
from google.adk.tools.authenticated_function_tool import AuthenticatedFunctionTool
31+
import requests
32+
33+
34+
# OAuth2 configuration for weather API access
35+
def create_auth_config() -> AuthConfig:
36+
"""Create OAuth2 auth configuration for weather API."""
37+
38+
# Define OAuth2 scheme with client credentials flow
39+
flows = OAuthFlows(
40+
clientCredentials=OAuthFlowClientCredentials(
41+
tokenUrl="http://localhost:8080/token",
42+
scopes={
43+
"read": "Read access to weather data",
44+
"write": "Write access for data updates",
45+
"admin": "Administrative access",
46+
},
47+
)
48+
)
49+
auth_scheme = OAuth2(flows=flows)
50+
51+
# Create credential with client ID and secret
52+
raw_credential = AuthCredential(
53+
auth_type=AuthCredentialTypes.OAUTH2,
54+
oauth2=OAuth2Auth(
55+
client_id="test_client",
56+
client_secret="test_secret",
57+
),
58+
)
59+
60+
return AuthConfig(
61+
auth_scheme=auth_scheme,
62+
raw_auth_credential=raw_credential,
63+
credential_key="weather_api_client",
64+
)
65+
66+
67+
def get_weather_data(city: str = "San Francisco", credential=None) -> str:
68+
"""Get current weather data for a specified city.
69+
70+
Args:
71+
city: City name to get weather for
72+
credential: API credential (automatically injected by AuthenticatedFunctionTool)
73+
74+
Returns:
75+
Current weather information for the city.
76+
"""
77+
78+
try:
79+
# Use the credential to make authenticated requests to weather API
80+
headers = {}
81+
if credential and credential.oauth2 and credential.oauth2.access_token:
82+
headers["Authorization"] = f"Bearer {credential.oauth2.access_token}"
83+
84+
# Call weather API endpoint
85+
params = {"city": city, "units": "metric"}
86+
response = requests.get(
87+
"http://localhost:8080/api/weather",
88+
headers=headers,
89+
params=params,
90+
timeout=10,
91+
)
92+
93+
if response.status_code == 200:
94+
data = response.json()
95+
result = f"🌤️ Weather for {city}:\n"
96+
result += f"Temperature: {data.get('temperature', 'N/A')}°C\n"
97+
result += f"Condition: {data.get('condition', 'N/A')}\n"
98+
result += f"Humidity: {data.get('humidity', 'N/A')}%\n"
99+
result += f"Wind Speed: {data.get('wind_speed', 'N/A')} km/h\n"
100+
result += f"Last Updated: {data.get('timestamp', 'N/A')}\n"
101+
return result
102+
else:
103+
return (
104+
f"❌ Failed to get weather data: {response.status_code} -"
105+
f" {response.text}"
106+
)
107+
108+
except Exception as e:
109+
return f"❌ Error getting weather data: {str(e)}"
110+
111+
112+
# Create the weather assistant agent
113+
root_agent = Agent(
114+
name="WeatherAssistant",
115+
description=(
116+
"Weather assistant that provides current weather information for cities"
117+
" worldwide."
118+
),
119+
model="gemini-2.5-pro",
120+
instruction=(
121+
"You are a helpful Weather Assistant that provides current weather"
122+
" information for any city worldwide.\n\nWhen users ask for weather:\n•"
123+
" Ask for the city name if not provided\n• Provide temperature in"
124+
" Celsius\n• Include helpful details like humidity, wind speed, and"
125+
" conditions\n• Be friendly and conversational about the weather\n\nIf"
126+
" there are any issues getting weather data, apologize and suggest"
127+
" trying again or checking for a different city name."
128+
),
129+
tools=[
130+
AuthenticatedFunctionTool(
131+
func=get_weather_data, auth_config=create_auth_config()
132+
),
133+
],
134+
)

0 commit comments

Comments
 (0)