Skip to content

Commit 03dd84e

Browse files
google-genai-botcopybara-github
authored andcommitted
feat:Add telemetry headers
The change adds custom telemetry headers (x-goog-api-client, user-agent) to BigQuery and BigQuery Storage API calls made by the plugin, including the ADK and Java versions. PiperOrigin-RevId: 901368290
1 parent 589328e commit 03dd84e

3 files changed

Lines changed: 23 additions & 5 deletions

File tree

core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.adk.plugins.agentanalytics;
1818

1919
import static com.google.adk.plugins.agentanalytics.BigQueryUtils.createAnalyticsViews;
20+
import static com.google.adk.plugins.agentanalytics.BigQueryUtils.getVersionHeaderValue;
2021
import static com.google.adk.plugins.agentanalytics.BigQueryUtils.maybeUpgradeSchema;
2122
import static com.google.adk.plugins.agentanalytics.JsonFormatter.convertToJsonNode;
2223
import static com.google.adk.plugins.agentanalytics.JsonFormatter.smartTruncate;
@@ -40,6 +41,7 @@
4041
import com.google.adk.tools.ToolContext;
4142
import com.google.adk.tools.mcp.AbstractMcpTool;
4243
import com.google.adk.utils.AgentEnums.AgentOrigin;
44+
import com.google.api.gax.rpc.FixedHeaderProvider;
4345
import com.google.auth.oauth2.GoogleCredentials;
4446
import com.google.cloud.bigquery.BigQuery;
4547
import com.google.cloud.bigquery.BigQueryException;
@@ -110,6 +112,8 @@ public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQue
110112

111113
private static BigQuery createBigQuery(BigQueryLoggerConfig config) throws IOException {
112114
BigQueryOptions.Builder builder = BigQueryOptions.newBuilder();
115+
builder.setHeaderProvider(
116+
FixedHeaderProvider.create(ImmutableMap.of("user-agent", getVersionHeaderValue())));
113117
if (config.credentials() != null) {
114118
builder.setCredentials(config.credentials());
115119
} else {

core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,12 @@
1818

1919
import static com.google.adk.plugins.agentanalytics.BigQuerySchema.SCHEMA_VERSION;
2020
import static com.google.adk.plugins.agentanalytics.BigQuerySchema.SCHEMA_VERSION_LABEL_KEY;
21+
import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;
2122
import static com.google.common.collect.ImmutableList.toImmutableList;
2223
import static com.google.common.collect.ImmutableMap.toImmutableMap;
2324
import static java.util.stream.Collectors.toCollection;
2425

26+
import com.google.adk.Version;
2527
import com.google.cloud.bigquery.BigQuery;
2628
import com.google.cloud.bigquery.BigQueryException;
2729
import com.google.cloud.bigquery.Field;
@@ -136,6 +138,13 @@ final class BigQueryUtils {
136138
"JSON_QUERY(content, '$.args') AS tool_args"))
137139
.buildOrThrow();
138140

141+
private static final String FRAMEWORK_PREFIX = "google-adk-bq-logger-java";
142+
143+
/** Returns the telemetry header value. */
144+
static String getVersionHeaderValue() {
145+
return FRAMEWORK_PREFIX + "/" + Version.JAVA_ADK_VERSION;
146+
}
147+
139148
/** Creates and/or replaces the analytics views in BigQuery. */
140149
static void createAnalyticsViews(BigQuery bigQuery, BigQueryLoggerConfig config) {
141150
for (Map.Entry<String, ImmutableList<String>> entry : EVENT_VIEW_DEFS.entrySet()) {

core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,19 @@
11
package com.google.adk.plugins.agentanalytics;
22

3+
import static com.google.adk.plugins.agentanalytics.BigQueryUtils.getVersionHeaderValue;
34
import static java.util.concurrent.TimeUnit.MILLISECONDS;
45

56
import com.google.api.gax.core.FixedCredentialsProvider;
67
import com.google.api.gax.retrying.RetrySettings;
8+
import com.google.api.gax.rpc.FixedHeaderProvider;
79
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
810
import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
911
import com.google.cloud.bigquery.storage.v1.StreamWriter;
1012
import com.google.common.annotations.VisibleForTesting;
1113
import com.google.common.base.VerifyException;
1214
import com.google.common.cache.Cache;
1315
import com.google.common.cache.CacheBuilder;
16+
import com.google.common.collect.ImmutableMap;
1417
import java.io.IOException;
1518
import java.util.Collection;
1619
import java.util.concurrent.ConcurrentHashMap;
@@ -67,13 +70,14 @@ void markProcessed(String invocationId) {
6770
}
6871

6972
protected BigQueryWriteClient createWriteClient(BigQueryLoggerConfig config) throws IOException {
73+
BigQueryWriteSettings.Builder settingsBuilder =
74+
BigQueryWriteSettings.newBuilder()
75+
.setHeaderProvider(
76+
FixedHeaderProvider.create(ImmutableMap.of("user-agent", getVersionHeaderValue())));
7077
if (config.credentials() != null) {
71-
return BigQueryWriteClient.create(
72-
BigQueryWriteSettings.newBuilder()
73-
.setCredentialsProvider(FixedCredentialsProvider.create(config.credentials()))
74-
.build());
78+
settingsBuilder.setCredentialsProvider(FixedCredentialsProvider.create(config.credentials()));
7579
}
76-
return BigQueryWriteClient.create();
80+
return BigQueryWriteClient.create(settingsBuilder.build());
7781
}
7882

7983
protected StreamWriter createWriter() {
@@ -89,6 +93,7 @@ protected StreamWriter createWriter() {
8993
String streamName = getStreamName(config);
9094
try {
9195
return StreamWriter.newBuilder(streamName, writeClient)
96+
.setTraceId(BigQueryUtils.getVersionHeaderValue())
9297
.setRetrySettings(retrySettings)
9398
.setWriterSchema(BigQuerySchema.getArrowSchema())
9499
.build();

0 commit comments

Comments
 (0)