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
28 changes: 12 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 @@ -95,27 +95,21 @@ 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
Map<String, String> callerHeaders = customHeaders != null ? customHeaders : ImmutableMap.of();
HttpOptions.Builder httpOptionsBuilder =
HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(trackingHeaders);
HttpOptions.builder()
.baseUrl(effectiveProxyUrl)
// Chat completions clients send no language tag, so they get both tags; the genai SDK
// adds its own gl-java tag.
.headers(
this.apiType == ApiType.CHAT_COMPLETIONS
? TrackingHeaders.mergeInto(callerHeaders)
: TrackingHeaders.mergeIntoGenAiSdkHeaders(callerHeaders));
String apiVersion = identifyApiVersion(modelName);
if (!apiVersion.isEmpty()) {
httpOptionsBuilder.apiVersion(apiVersion);
}
if (customHeaders != null) {
httpOptionsBuilder.headers(
ImmutableMap.<String, String>builder()
.putAll(trackingHeaders)
.putAll(customHeaders)
.buildOrThrow());
}
this.httpOptions = httpOptionsBuilder.build();

if (this.apiType == ApiType.CHAT_COMPLETIONS) {
Expand Down Expand Up @@ -289,7 +283,9 @@ public Builder proxyUrl(String proxyUrl) {
}

/**
* Sets a dictionary of headers to be sent with the request.
* Sets a dictionary of headers to be sent with the request. A {@code user-agent} or {@code
* x-goog-api-client} value, matched case-insensitively, is kept and extended with the ADK
* tracking tags it lacks.
*
* @param customHeaders the custom headers.
* @return this builder.
Expand Down
38 changes: 35 additions & 3 deletions core/src/main/java/com/google/adk/models/Claude.java
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
import static java.nio.charset.StandardCharsets.UTF_8;

import com.anthropic.client.AnthropicClient;
import com.anthropic.core.Properties;
import com.anthropic.core.http.Headers;
import com.anthropic.models.messages.ContentBlock;
import com.anthropic.models.messages.ContentBlockParam;
import com.anthropic.models.messages.Message;
Expand Down Expand Up @@ -89,22 +91,52 @@ public class Claude extends BaseLlm {
private final AnthropicClient anthropicClient;

/**
* Constructs a new Claude instance.
* Constructs a new Claude instance. Requests sent through {@code anthropicClient} carry ADK's
* usage-tracking headers, added to the ones the client already sends.
*
* @param modelName The name of the Claude model to use (e.g., "claude-3-opus-20240229").
* @param anthropicClient The Anthropic API client instance.
*/
public Claude(String modelName, AnthropicClient anthropicClient) {
super(modelName);
this.anthropicClient = anthropicClient;
this.anthropicClient = withTrackingHeaders(anthropicClient);
}

public Claude(String modelName, AnthropicClient anthropicClient, int maxTokens) {
super(modelName);
this.anthropicClient = anthropicClient;
this.anthropicClient = withTrackingHeaders(anthropicClient);
this.maxTokens = maxTokens;
}

/**
* Returns a copy of {@code client} that shares its connections and sends ADK's tracking headers,
* extending any the client already sets. A null client, or a test double that returns no copy, is
* returned unchanged.
*/
private static @Nullable AnthropicClient withTrackingHeaders(@Nullable AnthropicClient client) {
if (client == null) {
return null;
}
AnthropicClient tracked =
client.withOptions(
options -> {
Headers configured = options.build().headers();
Map<String, String> current = new HashMap<>();
// Start from the SDK's default User-Agent, which it adds only when none is set.
current.put(
"user-agent",
client.getClass().getSimpleName() + "/Java " + Properties.getPackageVersion());
for (String name : TrackingHeaders.HEADERS.keySet()) {
List<String> values = configured.values(name);
if (!values.isEmpty()) {
current.put(name, String.join(" ", values));
}
}
TrackingHeaders.mergeInto(current).forEach(options::replaceHeaders);
});
return tracked != null ? tracked : client;
}

@Override
public Flowable<LlmResponse> generateContent(LlmRequest llmRequest, boolean stream) {
// TODO: Switch to streaming API.
Expand Down
45 changes: 39 additions & 6 deletions core/src/test/java/com/google/adk/models/ApigeeLlmTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -331,15 +331,23 @@ public void validateModelString_rejectsOpenAiWithVersion() {
}

@Test
public void build_withCustomHeadersOverlappingTrackingHeaders_throwsException() {
ImmutableMap<String, String> overlappingHeaders = ImmutableMap.of("user-agent", "custom-agent");
ApigeeLlm.Builder builder =
public void build_chatCompletions_addsTheTrackingTagsCustomHeadersLack() {
ImmutableMap<String, String> overlappingHeaders =
ImmutableMap.of(
"User-Agent", "custom-agent", "x-goog-api-client", "lib/1 google-adk/0.0.1");
ApigeeLlm llm =
ApigeeLlm.builder()
.modelName("apigee/openai/gpt-4o")
.proxyUrl(PROXY_URL)
.customHeaders(overlappingHeaders);
IllegalArgumentException e = assertThrows(IllegalArgumentException.class, builder::build);
assertThat(e).hasMessageThat().contains("Multiple entries with same key: user-agent=");
.customHeaders(overlappingHeaders)
.build();

assertThat(llm.getHttpOptions().headers().get())
.containsExactly(
"user-agent",
"custom-agent " + TrackingHeaders.VALUE,
"x-goog-api-client",
"lib/1 google-adk/0.0.1 gl-java/" + JAVA_VERSION.value());
}

@Test
Expand Down Expand Up @@ -428,6 +436,31 @@ public void generateContent_chatCompletions_sendsTrackingHeaders() throws Except
}
}

@Test
public void generateContent_genAi_extendsACustomUserAgentWithTheAdkTag() throws Exception {
Headers headers =
headersSentToProxy(
"apigee/gemini/v1beta/gemini-test",
ImmutableMap.of("User-Agent", "custom-agent"),
GENERATE_CONTENT_JSON);

assertThat(getOnlyElement(headers.get("user-agent")))
.endsWith(" custom-agent " + TrackingHeaders.ADK_TAG);
assertGenAiTrackingHeaders(headers);
}

@Test
public void generateContent_chatCompletions_extendsACustomUserAgent() throws Exception {
Headers headers =
headersSentToProxy(
"apigee/openai/gpt-4o",
ImmutableMap.of("User-Agent", "custom-agent"),
CHAT_COMPLETION_JSON);

assertThat(headers.get("user-agent")).containsExactly("custom-agent " + TrackingHeaders.VALUE);
assertThat(headers.get("x-goog-api-client")).containsExactly(TrackingHeaders.VALUE);
}

/**
* Sends one request through an {@link ApigeeLlm} for {@code modelName} to a local proxy that
* answers with {@code responseJson}, and returns the request headers the proxy received.
Expand Down
88 changes: 88 additions & 0 deletions core/src/test/java/com/google/adk/models/ClaudeTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

package com.google.adk.models;

import static com.google.common.collect.Iterables.getOnlyElement;
import static com.google.common.truth.Truth.assertThat;
import static java.nio.charset.StandardCharsets.UTF_8;
import static org.junit.Assert.assertEquals;
Expand All @@ -26,7 +27,9 @@
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.anthropic.backends.Backend;
import com.anthropic.client.AnthropicClient;
import com.anthropic.client.okhttp.AnthropicOkHttpClient;
import com.anthropic.core.JsonValue;
import com.anthropic.models.messages.ContentBlock;
import com.anthropic.models.messages.ContentBlockParam;
Expand All @@ -46,17 +49,24 @@
import com.fasterxml.jackson.core.type.TypeReference;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.net.InetAddresses;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.Part;
import com.google.genai.types.Schema;
import com.sun.net.httpserver.Headers;
import com.sun.net.httpserver.HttpServer;
import java.io.OutputStream;
import java.lang.reflect.Method;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
Expand Down Expand Up @@ -704,4 +714,82 @@ public void generateContent_userTurnOfOnlyEmptyText_isSentAsPlaceholder() {
assertThat(blocks).hasSize(1);
assertThat(blocks.get(0).asText().text()).isEqualTo(".");
}

@Test
public void generateContent_extendsTheSdkUserAgentWithAdkTrackingHeaders() throws Exception {
Headers headers = headersSentThrough(AnthropicOkHttpClient.builder());

String userAgent = getOnlyElement(headers.get("user-agent"));
assertThat(userAgent).startsWith("AnthropicClientImpl/Java ");
assertThat(userAgent).endsWith(" " + TrackingHeaders.VALUE);
assertThat(headers.get("x-goog-api-client")).containsExactly(TrackingHeaders.VALUE);
}

@Test
public void generateContent_extendsTheCallerUserAgentWithAdkTrackingHeaders() throws Exception {
Headers headers =
headersSentThrough(AnthropicOkHttpClient.builder().putHeader("User-Agent", "my-app/1.0"));

assertThat(headers.get("user-agent")).containsExactly("my-app/1.0 " + TrackingHeaders.VALUE);
assertThat(headers.get("x-goog-api-client")).containsExactly(TrackingHeaders.VALUE);
}

@Test
public void constructor_nullClient_doesNotThrow() {
assertThat(new Claude("claude-test", (AnthropicClient) null).model()).isEqualTo("claude-test");
}

private static final String MESSAGE_JSON =
"{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-test\","
+ "\"content\":[{\"type\":\"text\",\"text\":\"pong\"}],\"stop_reason\":\"end_turn\","
+ "\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}";

/**
* Sends one request through Claude using a client built from {@code clientBuilder} for a local
* server, and returns the headers the server received.
*/
private static Headers headersSentThrough(AnthropicOkHttpClient.Builder clientBuilder)
throws Exception {
AtomicReference<Headers> requestHeaders = new AtomicReference<>();
HttpServer server =
HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0);
server.createContext(
"/",
exchange -> {
requestHeaders.set(exchange.getRequestHeaders());
byte[] body = MESSAGE_JSON.getBytes(UTF_8);
exchange.getResponseHeaders().set("Content-Type", "application/json");
exchange.sendResponseHeaders(200, body.length);
try (OutputStream out = exchange.getResponseBody()) {
out.write(body);
}
});
server.start();
try {
String baseUrl =
"http://"
+ InetAddresses.toUriString(server.getAddress().getAddress())
+ ":"
+ server.getAddress().getPort();
Backend backend =
new Backend() {
@Override
public String baseUrl() {
return baseUrl;
}

@Override
public void close() {}
};
Claude model = new Claude("claude-test", clientBuilder.backend(backend).build());

LlmResponse response =
model.generateContent(request(userText("Ping")), /* stream= */ false).blockingFirst();

assertThat(response.content().get().parts().get().get(0).text()).hasValue("pong");
return requestHeaders.get();
} finally {
server.stop(0);
}
}
}
Loading