Instrumentation for Gemini via Google GenAIClient library

build.gradle

@@ -82,6 +82,10 @@ dependencies {
   // Anthropic Instrumentation
   compileOnly "com.anthropic:anthropic-java:2.8.1"
   testImplementation "com.anthropic:anthropic-java:2.8.1"
+
+  // Google GenAI Instrumentation
+  compileOnly "com.google.genai:google-genai:1.20.0"
+  testImplementation "com.google.genai:google-genai:1.20.0"
 }
 
 /**

examples/build.gradle

@@ -27,6 +27,14 @@ dependencies {
   implementation 'com.openai:openai-java:2.8.1'
   // to run anthropic examples
   implementation "com.anthropic:anthropic-java:2.8.1"
+  // to run gemini examples
+  implementation 'com.google.genai:google-genai:1.20.0'
+  // spring ai examples
+  implementation 'org.springframework.ai:spring-ai-google-genai:1.1.0'
+  // spring boot for SpringAIExample (exclude logback, use slf4j-simple like other examples)
+  implementation('org.springframework.boot:spring-boot-starter:3.4.1') {
+    exclude group: 'org.springframework.boot', module: 'spring-boot-starter-logging'
+  }
 }
 
 application {
@@ -105,3 +113,31 @@ task runPromptFetching(type: JavaExec) {
     suspend = false
   }
 }
+
+task runGeminiInstrumentation(type: JavaExec) {
+  group = 'Braintrust SDK Examples'
+  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'
+  classpath = sourceSets.main.runtimeClasspath
+  mainClass = 'dev.braintrust.examples.GeminiInstrumentationExample'
+  systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel
+  debugOptions {
+    enabled = true
+    port = 5566
+    server = true
+    suspend = false
+  }
+}
+
+task runSpringAI(type: JavaExec) {
+  group = 'Braintrust SDK Examples'
+  description = 'Run the Spring Boot + Spring AI + Gemini example.'
+  classpath = sourceSets.main.runtimeClasspath
+  mainClass = 'dev.braintrust.examples.SpringAIExample'
+  systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel
+  debugOptions {
+    enabled = true
+    port = 5566
+    server = true
+    suspend = false
+  }
+}

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

@@ -0,0 +1,76 @@
+package dev.braintrust.examples;
+
+import com.google.genai.Client;
+import com.google.genai.types.GenerateContentConfig;
+import com.google.genai.types.GenerateContentResponse;
+import dev.braintrust.Braintrust;
+import dev.braintrust.instrumentation.genai.BraintrustGenAI;
+import io.opentelemetry.api.OpenTelemetry;
+
+/** Basic OTel + Gemini instrumentation example */
+public class GeminiInstrumentationExample {
+    public static void main(String[] args) throws Exception {
+        if (null == System.getenv("GOOGLE_API_KEY") && null == System.getenv("GEMINI_API_KEY")) {
+            System.err.println(
+                    "\n"
+                            + "WARNING: Neither GOOGLE_API_KEY nor GEMINI_API_KEY found. This"
+                            + " example will likely fail.\n"
+                            + "Set either: export GOOGLE_API_KEY='your-key' (recommended) or export"
+                            + " GEMINI_API_KEY='your-key'\n");
+        }
+
+        Braintrust braintrust = Braintrust.get();
+        OpenTelemetry openTelemetry = braintrust.openTelemetryCreate();
+        // CLAUDE: don't change the type of geminiClient -- sdk users must use the google genai
+        // client in their signature, not our instrumented client.
+        Client geminiClient = BraintrustGenAI.wrap(openTelemetry, new Client.Builder());
+
+        var tracer = openTelemetry.getTracer("my-instrumentation");
+        var rootSpan = tracer.spanBuilder("gemini-java-instrumentation-example").startSpan();
+        try (var ignored = rootSpan.makeCurrent()) {
+            generateContentExample(geminiClient);
+            // generateContentStreamingExample(client);
+        } finally {
+            rootSpan.end();
+        }
+
+        var url =
+                braintrust.projectUri()
+                        + "/logs?r=%s&s=%s"
+                                .formatted(
+                                        rootSpan.getSpanContext().getTraceId(),
+                                        rootSpan.getSpanContext().getSpanId());
+
+        System.out.println(
+                "\n\n  Example complete! View your data in Braintrust: %s\n".formatted(url));
+    }
+
+    private static void generateContentExample(Client client) {
+        var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build();
+
+        var response =
+                client.models.generateContent(
+                        "gemini-2.0-flash-lite", "What is the third planet from the sun?", config);
+
+        System.out.println("\n~~~ GENERATE CONTENT RESPONSE: %s\n".formatted(response.text()));
+    }
+
+    private static void generateContentStreamingExample(Client client) {
+        var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build();
+
+        System.out.println("\n~~~ STREAMING RESPONSE:");
+        var stream =
+                client.models.generateContentStream(
+                        "gemini-2.0-flash-exp",
+                        "Who was the first president of the United States?",
+                        config);
+
+        for (GenerateContentResponse chunk : stream) {
+            String text = chunk.text();
+            if (text != null && !text.isEmpty()) {
+                System.out.print(text);
+            }
+        }
+        System.out.println("\n");
+    }
+}

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

@@ -0,0 +1,115 @@
+package dev.braintrust.examples;
+
+import com.google.genai.Client;
+import dev.braintrust.Braintrust;
+import dev.braintrust.config.BraintrustConfig;
+import dev.braintrust.instrumentation.genai.BraintrustGenAI;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Scope;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.google.genai.GoogleGenAiChatModel;
+import org.springframework.ai.google.genai.GoogleGenAiChatOptions;
+import org.springframework.boot.CommandLineRunner;
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.autoconfigure.http.client.HttpClientAutoConfiguration;
+import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
+import org.springframework.context.annotation.Bean;
+
+/** Spring Boot application demonstrating Braintrust + Spring AI integration */
+@SpringBootApplication(
+        // NOTE: these excludes are specific to the Braintrust examples project to play nice with
+        // other examples' classpaths. Excludes are not required for production spring apps
+        exclude = {HttpClientAutoConfiguration.class, RestClientAutoConfiguration.class})
+public class SpringAIExample {
+
+    public static void main(String[] args) {
+        SpringApplication.run(SpringAIExample.class, args);
+    }
+
+    @Bean
+    public CommandLineRunner run(ChatModel chatModel, Tracer tracer, Braintrust braintrust) {
+        return args -> {
+            Span rootSpan = tracer.spanBuilder("spring-ai-example").startSpan();
+            try (Scope scope = rootSpan.makeCurrent()) {
+                System.out.println("\n=== Running Spring Boot Example ===\n");
+
+                // Make a simple chat call
+                var prompt = new Prompt("what's the name of the most popular java DI framework?");
+                var response = chatModel.call(prompt);
+
+                System.out.println(
+                        "~~~ SPRING AI CHAT RESPONSE: %s\n"
+                                .formatted(response.getResult().getOutput().getText()));
+            } finally {
+                rootSpan.end();
+            }
+
+            var url =
+                    braintrust.projectUri()
+                            + "/logs?r=%s&s=%s"
+                                    .formatted(
+                                            rootSpan.getSpanContext().getTraceId(),
+                                            rootSpan.getSpanContext().getSpanId());
+
+            System.out.println(
+                    "\n  Example complete! View your data in Braintrust: %s\n".formatted(url));
+        };
+    }
+
+    @Bean
+    public Braintrust braintrust() {
+        return Braintrust.get(BraintrustConfig.fromEnvironment());
+    }
+
+    @Bean
+    public OpenTelemetry openTelemetry(Braintrust braintrust) {
+        return braintrust.openTelemetryCreate();
+    }
+
+    @Bean
+    public Tracer tracer(OpenTelemetry openTelemetry) {
+        return openTelemetry.getTracer("spring-ai-instrumentation");
+    }
+
+    @Bean
+    public String aiProvider() {
+        // return "openai";
+        // return "anthropic";
+        return "google";
+    }
+
+    @Bean
+    public ChatModel chatModel(String aiProvider, OpenTelemetry openTelemetry) {
+        return switch (aiProvider) {
+            case "openai", "anthropic" -> {
+                throw new RuntimeException("TODO: " + aiProvider);
+            }
+            case "google" -> {
+                if (null == System.getenv("GOOGLE_API_KEY")
+                        && null == System.getenv("GEMINI_API_KEY")) {
+                    System.err.println(
+                            "\n"
+                                + "WARNING: Neither GOOGLE_API_KEY nor GEMINI_API_KEY found. This"
+                                + " example will likely fail.\n"
+                                + "Set either: export GOOGLE_API_KEY='your-key' (recommended) or"
+                                + " export GEMINI_API_KEY='your-key'\n");
+                }
+                Client genAIClient = BraintrustGenAI.wrap(openTelemetry, new Client.Builder());
+                yield GoogleGenAiChatModel.builder()
+                        .genAiClient(genAIClient)
+                        .defaultOptions(
+                                GoogleGenAiChatOptions.builder()
+                                        .model("gemini-2.0-flash-lite")
+                                        .temperature(0.0)
+                                        .maxOutputTokens(50)
+                                        .build())
+                        .build();
+            }
+            default -> throw new RuntimeException("unsupported provider: " + aiProvider);
+        };
+    }
+}