Skip to content

Commit 4312cd2

Browse files
realarkrealark
committed
Add Instrumentation for Google GenAI Client
1 parent 21d2582 commit 4312cd2

8 files changed

Lines changed: 1041 additions & 0 deletions

File tree

‎build.gradle‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,10 @@ dependencies {
8282
// Anthropic Instrumentation
8383
compileOnly "com.anthropic:anthropic-java:2.8.1"
8484
testImplementation "com.anthropic:anthropic-java:2.8.1"
85+
86+
// Google GenAI Instrumentation
87+
compileOnly "com.google.genai:google-genai:1.20.0"
88+
testImplementation "com.google.genai:google-genai:1.20.0"
8589
}
8690

8791
/**

‎examples/build.gradle‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,14 @@ dependencies {
2727
implementation 'com.openai:openai-java:2.8.1'
2828
// to run anthropic examples
2929
implementation "com.anthropic:anthropic-java:2.8.1"
30+
// to run gemini examples
31+
implementation 'com.google.genai:google-genai:1.20.0'
32+
// spring ai examples
33+
implementation 'org.springframework.ai:spring-ai-google-genai:1.1.0'
34+
// spring boot for SpringAIExample (exclude logback, use slf4j-simple like other examples)
35+
implementation('org.springframework.boot:spring-boot-starter:3.4.1') {
36+
exclude group: 'org.springframework.boot', module: 'spring-boot-starter-logging'
37+
}
3038
}
3139

3240
application {
@@ -105,3 +113,31 @@ task runPromptFetching(type: JavaExec) {
105113
suspend = false
106114
}
107115
}
116+
117+
task runGeminiInstrumentation(type: JavaExec) {
118+
group = 'Braintrust SDK Examples'
119+
description = 'Run the Gemini instrumentation example. NOTE: this requires GOOGLE_API_KEY or GEMINI_API_KEY to be exported and will make a small call to google, using your tokens'
120+
classpath = sourceSets.main.runtimeClasspath
121+
mainClass = 'dev.braintrust.examples.GeminiInstrumentationExample'
122+
systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel
123+
debugOptions {
124+
enabled = true
125+
port = 5566
126+
server = true
127+
suspend = false
128+
}
129+
}
130+
131+
task runSpringAI(type: JavaExec) {
132+
group = 'Braintrust SDK Examples'
133+
description = 'Run the Spring Boot + Spring AI + Gemini example.'
134+
classpath = sourceSets.main.runtimeClasspath
135+
mainClass = 'dev.braintrust.examples.SpringAIExample'
136+
systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel
137+
debugOptions {
138+
enabled = true
139+
port = 5566
140+
server = true
141+
suspend = false
142+
}
143+
}

‎examples/src/main/java/dev/braintrust/examples/GeminiInstrumentationExample.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package dev.braintrust.examples;
2+
3+
import com.google.genai.Client;
4+
import com.google.genai.types.GenerateContentConfig;
5+
import com.google.genai.types.GenerateContentResponse;
6+
import dev.braintrust.Braintrust;
7+
import dev.braintrust.instrumentation.genai.BraintrustGenAI;
8+
import io.opentelemetry.api.OpenTelemetry;
9+
10+
/** Basic OTel + Gemini instrumentation example */
11+
public class GeminiInstrumentationExample {
12+
public static void main(String[] args) throws Exception {
13+
if (null == System.getenv("GOOGLE_API_KEY") && null == System.getenv("GEMINI_API_KEY")) {
14+
System.err.println(
15+
"\n"
16+
+ "WARNING: Neither GOOGLE_API_KEY nor GEMINI_API_KEY found. This"
17+
+ " example will likely fail.\n"
18+
+ "Set either: export GOOGLE_API_KEY='your-key' (recommended) or export"
19+
+ " GEMINI_API_KEY='your-key'\n");
20+
}
21+
22+
Braintrust braintrust = Braintrust.get();
23+
OpenTelemetry openTelemetry = braintrust.openTelemetryCreate();
24+
// CLAUDE: don't change the type of geminiClient -- sdk users must use the google genai
25+
// client in their signature, not our instrumented client.
26+
Client geminiClient = BraintrustGenAI.wrap(openTelemetry, new Client.Builder());
27+
28+
var tracer = openTelemetry.getTracer("my-instrumentation");
29+
var rootSpan = tracer.spanBuilder("gemini-java-instrumentation-example").startSpan();
30+
try (var ignored = rootSpan.makeCurrent()) {
31+
generateContentExample(geminiClient);
32+
// generateContentStreamingExample(client);
33+
} finally {
34+
rootSpan.end();
35+
}
36+
37+
var url =
38+
braintrust.projectUri()
39+
+ "/logs?r=%s&s=%s"
40+
.formatted(
41+
rootSpan.getSpanContext().getTraceId(),
42+
rootSpan.getSpanContext().getSpanId());
43+
44+
System.out.println(
45+
"\n\n Example complete! View your data in Braintrust: %s\n".formatted(url));
46+
}
47+
48+
private static void generateContentExample(Client client) {
49+
var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build();
50+
51+
var response =
52+
client.models.generateContent(
53+
"gemini-2.0-flash-lite", "What is the third planet from the sun?", config);
54+
55+
System.out.println("\n~~~ GENERATE CONTENT RESPONSE: %s\n".formatted(response.text()));
56+
}
57+
58+
private static void generateContentStreamingExample(Client client) {
59+
var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build();
60+
61+
System.out.println("\n~~~ STREAMING RESPONSE:");
62+
var stream =
63+
client.models.generateContentStream(
64+
"gemini-2.0-flash-exp",
65+
"Who was the first president of the United States?",
66+
config);
67+
68+
for (GenerateContentResponse chunk : stream) {
69+
String text = chunk.text();
70+
if (text != null && !text.isEmpty()) {
71+
System.out.print(text);
72+
}
73+
}
74+
System.out.println("\n");
75+
}
76+
}

‎examples/src/main/java/dev/braintrust/examples/SpringAIExample.java‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package dev.braintrust.examples;
2+
3+
import com.google.genai.Client;
4+
import dev.braintrust.Braintrust;
5+
import dev.braintrust.config.BraintrustConfig;
6+
import dev.braintrust.instrumentation.genai.BraintrustGenAI;
7+
import io.opentelemetry.api.OpenTelemetry;
8+
import io.opentelemetry.api.trace.Span;
9+
import io.opentelemetry.api.trace.Tracer;
10+
import io.opentelemetry.context.Scope;
11+
import org.springframework.ai.chat.model.ChatModel;
12+
import org.springframework.ai.chat.prompt.Prompt;
13+
import org.springframework.ai.google.genai.GoogleGenAiChatModel;
14+
import org.springframework.ai.google.genai.GoogleGenAiChatOptions;
15+
import org.springframework.boot.CommandLineRunner;
16+
import org.springframework.boot.SpringApplication;
17+
import org.springframework.boot.autoconfigure.SpringBootApplication;
18+
import org.springframework.boot.autoconfigure.http.client.HttpClientAutoConfiguration;
19+
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
20+
import org.springframework.context.annotation.Bean;
21+
22+
/** Spring Boot application demonstrating Braintrust + Spring AI integration */
23+
@SpringBootApplication(
24+
// NOTE: these excludes are specific to the Braintrust examples project to play nice with
25+
// other examples' classpaths. Excludes are not required for production spring apps
26+
exclude = {HttpClientAutoConfiguration.class, RestClientAutoConfiguration.class})
27+
public class SpringAIExample {
28+
29+
public static void main(String[] args) {
30+
SpringApplication.run(SpringAIExample.class, args);
31+
}
32+
33+
@Bean
34+
public CommandLineRunner run(ChatModel chatModel, Tracer tracer, Braintrust braintrust) {
35+
return args -> {
36+
Span rootSpan = tracer.spanBuilder("spring-ai-example").startSpan();
37+
try (Scope scope = rootSpan.makeCurrent()) {
38+
System.out.println("\n=== Running Spring Boot Example ===\n");
39+
40+
// Make a simple chat call
41+
var prompt = new Prompt("what's the name of the most popular java DI framework?");
42+
var response = chatModel.call(prompt);
43+
44+
System.out.println(
45+
"~~~ SPRING AI CHAT RESPONSE: %s\n"
46+
.formatted(response.getResult().getOutput().getText()));
47+
} finally {
48+
rootSpan.end();
49+
}
50+
51+
var url =
52+
braintrust.projectUri()
53+
+ "/logs?r=%s&s=%s"
54+
.formatted(
55+
rootSpan.getSpanContext().getTraceId(),
56+
rootSpan.getSpanContext().getSpanId());
57+
58+
System.out.println(
59+
"\n Example complete! View your data in Braintrust: %s\n".formatted(url));
60+
};
61+
}
62+
63+
@Bean
64+
public Braintrust braintrust() {
65+
return Braintrust.get(BraintrustConfig.fromEnvironment());
66+
}
67+
68+
@Bean
69+
public OpenTelemetry openTelemetry(Braintrust braintrust) {
70+
return braintrust.openTelemetryCreate();
71+
}
72+
73+
@Bean
74+
public Tracer tracer(OpenTelemetry openTelemetry) {
75+
return openTelemetry.getTracer("spring-ai-instrumentation");
76+
}
77+
78+
@Bean
79+
public String aiProvider() {
80+
// return "openai";
81+
// return "anthropic";
82+
return "google";
83+
}
84+
85+
@Bean
86+
public ChatModel chatModel(String aiProvider, OpenTelemetry openTelemetry) {
87+
return switch (aiProvider) {
88+
case "openai", "anthropic" -> {
89+
throw new RuntimeException("TODO: " + aiProvider);
90+
}
91+
case "google" -> {
92+
if (null == System.getenv("GOOGLE_API_KEY")
93+
&& null == System.getenv("GEMINI_API_KEY")) {
94+
System.err.println(
95+
"\n"
96+
+ "WARNING: Neither GOOGLE_API_KEY nor GEMINI_API_KEY found. This"
97+
+ " example will likely fail.\n"
98+
+ "Set either: export GOOGLE_API_KEY='your-key' (recommended) or"
99+
+ " export GEMINI_API_KEY='your-key'\n");
100+
}
101+
Client genAIClient = BraintrustGenAI.wrap(openTelemetry, new Client.Builder());
102+
yield GoogleGenAiChatModel.builder()
103+
.genAiClient(genAIClient)
104+
.defaultOptions(
105+
GoogleGenAiChatOptions.builder()
106+
.model("gemini-2.0-flash-lite")
107+
.temperature(0.0)
108+
.maxOutputTokens(50)
109+
.build())
110+
.build();
111+
}
112+
default -> throw new RuntimeException("unsupported provider: " + aiProvider);
113+
};
114+
}
115+
}

0 commit comments

Comments
 (0)