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
14 changes: 14 additions & 0 deletions core/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,20 @@
</environmentVariables>
</configuration>
</execution>
<execution>
<id>gemini-tracking-headers</id>
<goals>
<goal>test</goal>
</goals>
<configuration>
<test>GeminiTrackingHeadersTest#generateContent_adkBuiltRegistryClient_sendsTrackingHeaders, GeminiTrackingHeadersTest#generateContent_adkBuiltVertexCredentialsClient_sendsTrackingHeaders</test>
<!-- these tests read the API key from the environment -->
<environmentVariables>
<GOOGLE_API_KEY>api-key</GOOGLE_API_KEY>
<GOOGLE_GENAI_USE_VERTEXAI>false</GOOGLE_GENAI_USE_VERTEXAI>
</environmentVariables>
</configuration>
</execution>
</executions>
</plugin>
</plugins>
Expand Down
27 changes: 11 additions & 16 deletions core/src/main/java/com/google/adk/models/ApigeeLlm.java
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,8 @@
*/
package com.google.adk.models;

import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;
import static com.google.common.base.Strings.isNullOrEmpty;

import com.google.adk.Version;
import com.google.adk.models.chat.ChatCompletionsClient;
import com.google.adk.models.chat.ChatCompletionsHttpClient;
import com.google.common.annotations.VisibleForTesting;
Expand Down Expand Up @@ -46,17 +44,6 @@ public class ApigeeLlm extends BaseLlm {
private static final String GOOGLE_GENAI_USE_VERTEXAI_ENV_VARIABLE_NAME =
"GOOGLE_GENAI_USE_VERTEXAI";
private static final String APIGEE_PROXY_URL_ENV_VARIABLE_NAME = "APIGEE_PROXY_URL";
private static final ImmutableMap<String, String> TRACKING_HEADERS;

static {
final String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
final String languageLabel = "gl-java/" + JAVA_VERSION.value();
final String versionHeaderValue = String.format("%s %s", frameworkLabel, languageLabel);
TRACKING_HEADERS =
ImmutableMap.of(
"x-goog-api-client", versionHeaderValue,
"user-agent", versionHeaderValue);
}

/** Defines the type of API to be used by the Apigee proxy. */
public enum ApiType {
Expand Down Expand Up @@ -108,17 +95,24 @@ private ApigeeLlm(
+ " APIGEE_PROXY_URL.");
}

// Chat completions clients send no language tag, so they get both tags; the genai SDK adds its
// own gl-java tag.
ImmutableMap<String, String> trackingHeaders =
this.apiType == ApiType.CHAT_COMPLETIONS
? TrackingHeaders.HEADERS
: TrackingHeaders.GENAI_SDK_HEADERS;

// Build the Client
HttpOptions.Builder httpOptionsBuilder =
HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(TRACKING_HEADERS);
HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(trackingHeaders);
String apiVersion = identifyApiVersion(modelName);
if (!apiVersion.isEmpty()) {
httpOptionsBuilder.apiVersion(apiVersion);
}
if (customHeaders != null) {
httpOptionsBuilder.headers(
ImmutableMap.<String, String>builder()
.putAll(TRACKING_HEADERS)
.putAll(trackingHeaders)
.putAll(customHeaders)
.buildOrThrow());
}
Expand All @@ -134,7 +128,8 @@ private ApigeeLlm(
apiClientBuilder.vertexAI(true);
}
this.apiClient = apiClientBuilder.build();
this.geminiDelegate = new Gemini(modelName, apiClient);
this.geminiDelegate =
new Gemini(modelName, apiClient, /* clientSendsTrackingHeaders= */ true);
this.chatCompletionsClient = null;
}

Expand Down
102 changes: 58 additions & 44 deletions core/src/main/java/com/google/adk/models/Gemini.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,6 @@

package com.google.adk.models;

import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;

import com.google.adk.Version;
import com.google.adk.internal.http.HttpClientFactory;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -60,7 +57,6 @@
public class Gemini extends BaseLlm {

private static final Logger logger = LoggerFactory.getLogger(Gemini.class);
private static final ImmutableMap<String, String> TRACKING_HEADERS;

private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executorService) {
OkHttpClient client =
Expand All @@ -75,61 +71,53 @@ private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executor
.build();
}

private static Client buildApiKeyClient(
String apiKey, @Nullable ExecutorService executorService) {
// Every client ADK builds starts here, so each sends ADK's tracking tag.
private static Client.Builder adkClientBuilder(@Nullable ExecutorService executorService) {
return Client.builder()
.apiKey(apiKey)
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.httpOptions(HttpOptions.builder().headers(TrackingHeaders.GENAI_SDK_HEADERS).build())
.clientOptions(
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build())
.build();
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build());
}

private static Client buildApiKeyClient(
String apiKey, @Nullable ExecutorService executorService) {
return adkClientBuilder(executorService).apiKey(apiKey).build();
}

private static Client buildVertexClient(
VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) {
Client.Builder apiClientBuilder =
Client.builder()
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.clientOptions(
ClientOptions.builder()
.customHttpClient(prepareHttpClient(executorService))
.build());
Client.Builder apiClientBuilder = adkClientBuilder(executorService);
vertexCredentials.project().ifPresent(apiClientBuilder::project);
vertexCredentials.location().ifPresent(apiClientBuilder::location);
vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
return apiClientBuilder.build();
}

private static Client buildDefaultClient(@Nullable ExecutorService executorService) {
return Client.builder()
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.clientOptions(
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build())
.build();
}

static {
String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
String languageLabel = "gl-java/" + JAVA_VERSION.value();
String versionHeaderValue = String.format("%s %s", frameworkLabel, languageLabel);

TRACKING_HEADERS =
ImmutableMap.of(
"x-goog-api-client", versionHeaderValue,
"user-agent", versionHeaderValue);
return adkClientBuilder(executorService).build();
}

private final Client apiClient;
// True when apiClient already sends ADK's tracking tag.
private final boolean clientSendsTrackingHeaders;

/**
* Constructs a new Gemini instance.
* Constructs a new Gemini instance. Each generate request sent through {@code apiClient} carries
* ADK's tracking tag, which the genai SDK appends to the client's tracking headers. If the
* client's own headers already carry an ADK tag, requests send two because the SDK does not
* expose a client's headers.
*
* @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
* @param apiClient The genai {@link com.google.genai.Client} instance for making API calls.
*/
public Gemini(String modelName, Client apiClient) {
this(modelName, apiClient, /* clientSendsTrackingHeaders= */ false);
}

Gemini(String modelName, Client apiClient, boolean clientSendsTrackingHeaders) {
super(modelName);
this.apiClient = Objects.requireNonNull(apiClient, "apiClient cannot be null");
this.clientSendsTrackingHeaders = clientSendsTrackingHeaders;
}

/**
Expand All @@ -139,9 +127,10 @@ public Gemini(String modelName, Client apiClient) {
* @param apiKey The Google Gemini API key.
*/
public Gemini(String modelName, String apiKey) {
super(modelName);
Objects.requireNonNull(apiKey, "apiKey cannot be null");
this.apiClient = buildApiKeyClient(apiKey, null);
this(
modelName,
buildApiKeyClient(Objects.requireNonNull(apiKey, "apiKey cannot be null"), null),
/* clientSendsTrackingHeaders= */ true);
}

/**
Expand All @@ -151,9 +140,11 @@ public Gemini(String modelName, String apiKey) {
* @param vertexCredentials The Vertex AI credentials to access the Gemini model.
*/
public Gemini(String modelName, VertexCredentials vertexCredentials) {
super(modelName);
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null");
this.apiClient = buildVertexClient(vertexCredentials, null);
this(
modelName,
buildVertexClient(
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null"), null),
/* clientSendsTrackingHeaders= */ true);
}

/**
Expand Down Expand Up @@ -256,13 +247,16 @@ public Gemini build() {

if (apiClient != null) {
return new Gemini(modelName, apiClient);
} else if (apiKey != null) {
return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService));
}
Client client;
if (apiKey != null) {
client = buildApiKeyClient(apiKey, httpExecutorService);
} else if (vertexCredentials != null) {
return new Gemini(modelName, buildVertexClient(vertexCredentials, httpExecutorService));
client = buildVertexClient(vertexCredentials, httpExecutorService);
} else {
return new Gemini(modelName, buildDefaultClient(httpExecutorService));
client = buildDefaultClient(httpExecutorService);
}
return new Gemini(modelName, client, /* clientSendsTrackingHeaders= */ true);
}
}

Expand Down Expand Up @@ -292,6 +286,9 @@ private Flowable<LlmResponse> sendRequest(
logger.trace("Request Contents: {}", llmRequest.contents());
logger.trace("Request Config: {}", config);

if (!clientSendsTrackingHeaders) {
config = withTrackingHeaders(config);
}
if (stream) {
logger.debug("Sending streaming generateContent request to model {}", effectiveModelName);
CompletableFuture<ResponseStream<GenerateContentResponse>> streamFuture =
Expand All @@ -316,6 +313,23 @@ private Flowable<LlmResponse> sendRequest(
}
}

/**
* Returns {@code config} with ADK's tracking tag merged into its request headers, which the genai
* SDK appends to the client's own.
*/
private static GenerateContentConfig withTrackingHeaders(@Nullable GenerateContentConfig config) {
GenerateContentConfig base = config != null ? config : GenerateContentConfig.builder().build();
HttpOptions httpOptions = base.httpOptions().orElseGet(() -> HttpOptions.builder().build());
return base.toBuilder()
.httpOptions(
httpOptions.toBuilder()
.headers(
TrackingHeaders.mergeIntoGenAiSdkHeaders(
httpOptions.headers().orElse(ImmutableMap.of())))
.build())
.build();
}

private static LlmResponse withCacheMetadata(
LlmResponse response, @Nullable CacheMetadata cacheMetadata) {
return cacheMetadata == null
Expand Down
109 changes: 109 additions & 0 deletions core/src/main/java/com/google/adk/models/TrackingHeaders.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
/*
* 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.
*/

package com.google.adk.models;

import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static java.util.stream.Collectors.joining;

import com.google.adk.Version;
import com.google.common.base.Ascii;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

/**
* Headers that let Google APIs attribute model calls to ADK, in the format all ADK SDKs use. A
* header value is a space-separated list of {@code name/version} tags.
*/
final class TrackingHeaders {

private static final Splitter TAG_SPLITTER = Splitter.on(' ').omitEmptyStrings();

/** ADK's own tag. */
static final String ADK_TAG = "google-adk/" + Version.JAVA_ADK_VERSION;

/** ADK's tag followed by the Java version tag, for clients that send no language tag. */
static final String VALUE = ADK_TAG + " gl-java/" + JAVA_VERSION.value();

static final ImmutableMap<String, String> HEADERS =
ImmutableMap.of("x-goog-api-client", VALUE, "user-agent", VALUE);

/**
* Headers for a genai SDK client, which puts its own {@code gl-java} tag first. Both carry only
* {@link #ADK_TAG}.
*/
static final ImmutableMap<String, String> GENAI_SDK_HEADERS =
ImmutableMap.of("x-goog-api-client", ADK_TAG, "user-agent", ADK_TAG);

/**
* Like {@link #mergeInto}, but adds only {@link #ADK_TAG}, for headers that the genai SDK appends
* to its own label.
*/
static ImmutableMap<String, String> mergeIntoGenAiSdkHeaders(Map<String, String> headers) {
return mergeInto(headers, ADK_TAG);
}

/**
* Returns {@code headers} with each tracking header, matched case-insensitively, followed by
* every tag of {@link #VALUE} whose name it lacks.
*/
static ImmutableMap<String, String> mergeInto(Map<String, String> headers) {
return mergeInto(headers, VALUE);
}

private static ImmutableMap<String, String> mergeInto(
Map<String, String> headers, String adkTags) {
Map<String, String> merged = new LinkedHashMap<>();
Map<String, String> callerValues = new LinkedHashMap<>();
headers.forEach(
(name, value) -> {
String trackingName = Ascii.toLowerCase(name);
if (HEADERS.containsKey(trackingName)) {
callerValues.merge(trackingName, value, (first, second) -> first + " " + second);
} else {
merged.put(name, value);
}
});
for (String name : HEADERS.keySet()) {
merged.put(name, extend(callerValues.getOrDefault(name, ""), adkTags));
}
return ImmutableMap.copyOf(merged);
}

/** Returns {@code callerValue} followed by each tag of {@code adkTags} whose name it lacks. */
private static String extend(String callerValue, String adkTags) {
List<String> tags = TAG_SPLITTER.splitToList(callerValue);
ImmutableSet<String> names =
tags.stream().map(TrackingHeaders::tagName).collect(toImmutableSet());
return Stream.concat(
tags.stream(),
TAG_SPLITTER.splitToStream(adkTags).filter(tag -> !names.contains(tagName(tag))))
.collect(joining(" "));
}

private static String tagName(String tag) {
int slash = tag.indexOf('/');
return slash < 0 ? tag : tag.substring(0, slash);
}

private TrackingHeaders() {}
}
Loading
Loading