Skip to content

Commit d85932f

Browse files
google-genai-botcopybara-github
authored andcommitted
chore: Add X-Goog-API-Client header to GDA API requests
PiperOrigin-RevId: 869372484
1 parent c6b1c74 commit d85932f

4 files changed

Lines changed: 26 additions & 2 deletions

File tree

src/google/adk/tools/bigquery/data_insights_tool.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
from . import client
2626
from .config import BigQueryToolConfig
2727

28+
_GDA_CLIENT_ID = "GOOGLE_ADK"
29+
2830

2931
def ask_data_insights(
3032
project_id: str,
@@ -129,6 +131,7 @@ def ask_data_insights(
129131
headers = {
130132
"Authorization": f"Bearer {credentials.token}",
131133
"Content-Type": "application/json",
134+
"X-Goog-API-Client": _GDA_CLIENT_ID,
132135
}
133136
ca_url = f"https://geminidataanalytics.googleapis.com/v1alpha/projects/{project_id}/locations/{location}:chat"
134137

@@ -149,7 +152,7 @@ def ask_data_insights(
149152
"systemInstruction": instructions,
150153
"options": {"chart": {"image": {"noImage": {}}}},
151154
},
152-
"clientIdEnum": "GOOGLE_ADK",
155+
"clientIdEnum": _GDA_CLIENT_ID,
153156
}
154157

155158
resp = _get_stream(

src/google/adk/tools/data_agent/data_agent_tool.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from .config import DataAgentToolConfig
2424

2525
BASE_URL = "https://geminidataanalytics.googleapis.com/v1beta"
26+
_GDA_CLIENT_ID = "GOOGLE_ADK"
2627

2728

2829
def _get_http_headers(
@@ -41,6 +42,7 @@ def _get_http_headers(
4142
return {
4243
"Authorization": f"Bearer {credentials.token}",
4344
"Content-Type": "application/json",
45+
"X-Goog-API-Client": _GDA_CLIENT_ID,
4446
}
4547

4648

@@ -294,7 +296,7 @@ def ask_data_agent(
294296
"dataAgentContext": {
295297
"dataAgent": data_agent_name,
296298
},
297-
"clientIdEnum": "GOOGLE_ADK",
299+
"clientIdEnum": _GDA_CLIENT_ID,
298300
}
299301
resp = _get_stream(
300302
chat_url,

tests/unittests/tools/bigquery/test_bigquery_data_insights_tool.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,12 @@ def test_ask_data_insights_success(mock_get_stream):
8989
assert result["response"] == "Final formatted string from stream"
9090
mock_get_stream.assert_called_once()
9191

92+
# Verify that the correct headers and client ID were passed to _get_stream
93+
args, _ = mock_get_stream.call_args
94+
headers = args[2]
95+
assert headers["X-Goog-API-Client"] == "GOOGLE_ADK"
96+
assert headers["Authorization"] == "Bearer fake-token"
97+
9298

9399
@mock.patch.object(data_insights_tool, "_get_stream")
94100
def test_ask_data_insights_handles_exception(mock_get_stream):

tests/unittests/tools/data_agent/test_data_agent_tool.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,3 +196,16 @@ def test_get_stream_from_file(mock_post, case_file_path):
196196

197197
# 6. Assert that the final list of dicts matches the expected output
198198
assert result == expected_final_list
199+
200+
201+
def test_get_http_headers_includes_client_id():
202+
"""Tests _get_http_headers includes the correct GDA client ID."""
203+
mock_creds = mock.Mock()
204+
mock_creds.token = "fake-token"
205+
206+
# pylint: disable=protected-access
207+
headers = data_agent_tool._get_http_headers(mock_creds)
208+
209+
assert headers["X-Goog-API-Client"] == "GOOGLE_ADK"
210+
assert headers["Content-Type"] == "application/json"
211+
assert headers["Authorization"] == "Bearer fake-token"

0 commit comments

Comments
 (0)