diff --git a/braintrust-java-agent/build.gradle b/braintrust-java-agent/build.gradle index abb5926d..74e15cf9 100644 --- a/braintrust-java-agent/build.gradle +++ b/braintrust-java-agent/build.gradle @@ -44,11 +44,11 @@ subprojects { } } -// Add muzzle generation task and muzzle check plugin to all instrumentation subprojects. +// Add muzzle generation task and muzzle check plugin to smoke-test test-instrumentation. +// (Instrumentation subprojects under :braintrust-sdk are configured in braintrust-sdk/build.gradle.) subprojects { subproject -> subproject.afterEvaluate { - if (subproject.path.startsWith(':braintrust-java-agent:instrumentation:') - || subproject.path == ':braintrust-java-agent:smoke-test:test-instrumentation') { + if (subproject.path == ':braintrust-java-agent:smoke-test:test-instrumentation') { def instrumentationApi = project(':braintrust-java-agent:instrumenter') // Bootstrap classes (e.g. BraintrustBridge) are on the bootstrap classpath at runtime. diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchain.java b/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchain.java deleted file mode 100644 index bb9d767a..00000000 --- a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchain.java +++ /dev/null @@ -1,150 +0,0 @@ -package dev.braintrust.instrumentation.langchain.v1_8_0; - -import dev.langchain4j.model.openai.OpenAiChatModel; -import dev.langchain4j.model.openai.OpenAiStreamingChatModel; -import dev.langchain4j.service.AiServiceContext; -import dev.langchain4j.service.AiServices; -import dev.langchain4j.service.tool.ToolExecutor; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Tracer; -import java.util.Map; -import lombok.extern.slf4j.Slf4j; - -/** Braintrust LangChain4j client instrumentation for the auto-instrumentation agent. */ -@Slf4j -public final class BraintrustLangchain { - - private static final String INSTRUMENTATION_NAME = "braintrust-langchain4j"; - - /** - * Wrap an already-built OpenAiChatModel by replacing its internal HTTP client with a tracing - * wrapper. - */ - public static OpenAiChatModel wrapChatModel( - OpenTelemetry openTelemetry, OpenAiChatModel model) { - try { - Object internalClient = getPrivateField(model, "client"); - dev.langchain4j.http.client.HttpClient httpClient = - getPrivateField(internalClient, "httpClient"); - - if (httpClient instanceof WrappedHttpClient) { - log.debug("model already instrumented, skipping"); - return model; - } - - dev.langchain4j.http.client.HttpClient wrappedHttpClient = - new WrappedHttpClient(openTelemetry, httpClient, new Options("openai")); - setPrivateField(internalClient, "httpClient", wrappedHttpClient); - return model; - } catch (Exception e) { - log.warn("failed to instrument OpenAiChatModel", e); - return model; - } - } - - /** - * Wrap an already-built OpenAiStreamingChatModel by replacing its internal HTTP client with a - * tracing wrapper. - */ - public static OpenAiStreamingChatModel wrapStreamingChatModel( - OpenTelemetry openTelemetry, OpenAiStreamingChatModel model) { - try { - Object internalClient = getPrivateField(model, "client"); - dev.langchain4j.http.client.HttpClient httpClient = - getPrivateField(internalClient, "httpClient"); - - if (httpClient instanceof WrappedHttpClient) { - log.debug("model already instrumented, skipping"); - return model; - } - - dev.langchain4j.http.client.HttpClient wrappedHttpClient = - new WrappedHttpClient(openTelemetry, httpClient, new Options("openai")); - setPrivateField(internalClient, "httpClient", wrappedHttpClient); - return model; - } catch (Exception e) { - log.warn("failed to instrument OpenAiStreamingChatModel", e); - return model; - } - } - - /** - * Wrap an already-built AiService with TracingProxy and TracingToolExecutors. Called from - * AiServices.build() advice — the service is already built, so we wrap it with tracing proxy - * and instrument tool executors. - */ - @SuppressWarnings("unchecked") - public static T wrapAiService( - OpenTelemetry openTelemetry, AiServices aiServices, Object builtService) { - try { - AiServiceContext context = getPrivateField(aiServices, "context"); - Tracer tracer = openTelemetry.getTracer(INSTRUMENTATION_NAME); - - // Wrap tool executors with tracing - if (context.toolService != null) { - for (Map.Entry entry : - context.toolService.toolExecutors().entrySet()) { - String toolName = entry.getKey(); - ToolExecutor original = entry.getValue(); - if (!(original instanceof TracingToolExecutor)) { - entry.setValue(new TracingToolExecutor(original, toolName, tracer)); - } - } - - // Link spans across concurrent tool calls - var underlyingExecutor = context.toolService.executor(); - if (underlyingExecutor != null) { - // Replace the executor with one that passes OTel context - try { - setPrivateField( - context.toolService, - "executor", - new OtelContextPassingExecutor(underlyingExecutor)); - } catch (Exception e) { - log.debug("Could not replace tool executor for context propagation", e); - } - } - } - - // Wrap service with tracing proxy - Class serviceInterface = (Class) context.aiServiceClass; - return TracingProxy.create(serviceInterface, (T) builtService, tracer); - } catch (Exception e) { - log.warn("failed to apply langchain AI services instrumentation", e); - return (T) builtService; - } - } - - public record Options(String providerName) {} - - @SuppressWarnings("unchecked") - static T getPrivateField(Object obj, String fieldName) throws ReflectiveOperationException { - Class clazz = obj.getClass(); - while (clazz != null) { - try { - java.lang.reflect.Field field = clazz.getDeclaredField(fieldName); - field.setAccessible(true); - return (T) field.get(obj); - } catch (NoSuchFieldException e) { - clazz = clazz.getSuperclass(); - } - } - throw new NoSuchFieldException(fieldName); - } - - static void setPrivateField(Object obj, String fieldName, Object value) - throws ReflectiveOperationException { - Class clazz = obj.getClass(); - while (clazz != null) { - try { - java.lang.reflect.Field field = clazz.getDeclaredField(fieldName); - field.setAccessible(true); - field.set(obj, value); - return; - } catch (NoSuchFieldException e) { - clazz = clazz.getSuperclass(); - } - } - throw new NoSuchFieldException(fieldName); - } -} diff --git a/braintrust-java-agent/internal/build.gradle b/braintrust-java-agent/internal/build.gradle index 98ddb67e..d176c784 100644 --- a/braintrust-java-agent/internal/build.gradle +++ b/braintrust-java-agent/internal/build.gradle @@ -48,15 +48,18 @@ dependencies { testRuntimeOnly 'org.junit.platform:junit-platform-launcher' } -// Auto-discover all instrumentation subprojects and add them as implementation dependencies. -// Any subproject under :braintrust-java-agent:instrumentation is automatically bundled -// into the shadow JAR, and its META-INF/services entries are merged. -project(':braintrust-java-agent').subprojects.each { subproject -> - if (subproject.path.startsWith(':braintrust-java-agent:instrumentation:') - || subproject.path == ':braintrust-java-agent:smoke-test:test-instrumentation') { +// Auto-discover instrumentation subprojects from braintrust-sdk and add them as implementation +// dependencies. These get bundled into the shadow JAR with merged META-INF/services entries. +project(':braintrust-sdk').subprojects.each { subproject -> + if (subproject.path.startsWith(':braintrust-sdk:instrumentation:')) { dependencies.add('implementation', subproject) } } +// Also include test-instrumentation for smoke tests +// TODO: don't bundle this instrumentation in the agent jar. Instead create a flag to look for instrumentation on the system classpath +if (findProject(':braintrust-java-agent:smoke-test:test-instrumentation') != null) { + dependencies.add('implementation', project(':braintrust-java-agent:smoke-test:test-instrumentation')) +} test { useJUnitPlatform() diff --git a/braintrust-sdk/build.gradle b/braintrust-sdk/build.gradle index 4ce6f825..88279421 100644 --- a/braintrust-sdk/build.gradle +++ b/braintrust-sdk/build.gradle @@ -23,7 +23,86 @@ repositories { mavenCentral() } -def langchainVersion = '1.9.1' +// Shared config for instrumentation subprojects +subprojects { + apply plugin: 'java' + apply plugin: 'dev.braintrust.muzzle' + + java { + toolchain { + languageVersion = JavaLanguageVersion.of(17) + vendor = JvmVendorSpec.ADOPTIUM + } + } + + tasks.withType(JavaCompile).configureEach { + options.release = 17 + } + + repositories { + mavenCentral() + } +} + +// Add muzzle generation task and muzzle check plugin to all instrumentation subprojects. +subprojects { subproject -> + subproject.afterEvaluate { + if (subproject.path.startsWith(':braintrust-sdk:instrumentation:')) { + def instrumentationApi = project(':braintrust-java-agent:instrumenter') + + // Bootstrap classes (e.g. BraintrustBridge) are on the bootstrap classpath at runtime. + // Make them available at compile time so instrumentation modules can reference them. + dependencies.add('compileOnly', project(':braintrust-java-agent:bootstrap')) + + // --- $Muzzle side-class generation (compile-time) --- + + // Configuration for the muzzle generator classpath + configurations.maybeCreate('muzzleGenerator') + dependencies.add('muzzleGenerator', instrumentationApi) + dependencies.add('muzzleGenerator', 'net.bytebuddy:byte-buddy:1.17.5') + + task generateMuzzle(type: JavaExec) { + dependsOn compileJava, instrumentationApi.tasks.named('compileJava') + description = 'Generates $Muzzle side-classes for InstrumentationModule subclasses' + group = 'build' + + // Classpath: compiled instrumentation classes + instrumenter + bytebuddy + deps + classpath = files( + sourceSets.main.output.classesDirs, + configurations.muzzleGenerator, + configurations.compileClasspath + ) + mainClass = 'dev.braintrust.instrumentation.muzzle.MuzzleGenerator' + args = [sourceSets.main.java.classesDirectory.get().asFile.absolutePath] + } + + // Run muzzle generation after compileJava, before jar + tasks.named('classes').configure { dependsOn generateMuzzle } + + // Run muzzle checks as part of the check lifecycle + tasks.named('check').configure { dependsOn 'muzzle' } + } + } +} + +// Configuration for embedding instrumentation subproject classes into the SDK JAR +configurations { + // Non-transitive: we only want the instrumentation subproject JARs themselves, + // not their transitive dependencies (which include braintrust-sdk itself). + embed { + transitive = false + } +} + +// Auto-discover instrumentation subprojects and add them to the embed configuration. +// This must be done after project evaluation so the subprojects are resolved. +afterEvaluate { + project.subprojects.each { sub -> + if (sub.path.startsWith(':braintrust-sdk:instrumentation:')) { + dependencies.add('embed', sub) + } + } +} dependencies { api "io.opentelemetry:opentelemetry-api:${otelVersion}" @@ -40,7 +119,6 @@ dependencies { implementation "org.slf4j:slf4j-api:${slf4jVersion}" - implementation 'org.apache.commons:commons-lang3:3.14.0' implementation 'com.google.code.findbugs:jsr305:3.0.2' // for @Nullable annotations implementation "com.github.spullara.mustache.java:compiler:0.9.14" @@ -53,27 +131,44 @@ dependencies { testImplementation "org.junit.jupiter:junit-jupiter-params:${junitVersion}" testImplementation 'org.wiremock:wiremock:3.13.1' - // OAI instrumentation - compileOnly 'com.openai:openai-java:2.8.1' - testImplementation 'com.openai:openai-java:2.8.1' - - // 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" - - // LangChain4j Instrumentation - compileOnly "dev.langchain4j:langchain4j:${langchainVersion}" - compileOnly "dev.langchain4j:langchain4j-http-client:${langchainVersion}" - compileOnly "dev.langchain4j:langchain4j-open-ai:${langchainVersion}" - testImplementation "dev.langchain4j:langchain4j:${langchainVersion}" - testImplementation "dev.langchain4j:langchain4j-http-client:${langchainVersion}" - testImplementation "dev.langchain4j:langchain4j-open-ai:${langchainVersion}" + } +// Merge META-INF/services files from embedded instrumentation JARs. +// This ensures ServiceLoader can discover InstrumentationModule implementations. +task mergeEmbedServices { + dependsOn configurations.embed + def mergedDir = layout.buildDirectory.dir("merged-services") + outputs.dir mergedDir + + doLast { + def servicesMap = [:] + configurations.embed.each { file -> + if (file.isFile()) { + zipTree(file).matching { include 'META-INF/services/**' }.visit { details -> + if (!details.isDirectory()) { + def serviceName = details.relativePath.toString() + if (!servicesMap.containsKey(serviceName)) { + servicesMap[serviceName] = new LinkedHashSet() + } + servicesMap[serviceName].addAll(details.file.readLines().findAll { it.trim() }) + } + } + } + } + def outDir = mergedDir.get().asFile + outDir.deleteDir() + servicesMap.each { path, lines -> + def outFile = new File(outDir, path) + outFile.parentFile.mkdirs() + outFile.text = lines.join('\n') + '\n' + } + } +} + +jar.dependsOn mergeEmbedServices +jar.from(layout.buildDirectory.dir("merged-services")) + // Generate braintrust.properties at build time with smart versioning task generateBraintrustProperties { description = 'Generate braintrust.properties with smart git-based versioning' @@ -129,6 +224,19 @@ jar { 'Main-Class': 'dev.braintrust.SDKMain' ) } + + // Embed instrumentation subproject classes into the SDK JAR. + // Use dependsOn to ensure instrumentation jars are built first. + dependsOn configurations.embed + + from({ + configurations.embed.collect { it.isDirectory() ? it : zipTree(it) } + }) { + // Exclude META-INF from embedded jars (we handle service files separately below) + exclude 'META-INF/**' + } + + duplicatesStrategy = DuplicatesStrategy.EXCLUDE } import com.vanniktech.maven.publish.JavadocJar diff --git a/braintrust-java-agent/instrumentation/anthropic_2_2_0/build.gradle b/braintrust-sdk/instrumentation/anthropic_2_2_0/build.gradle similarity index 96% rename from braintrust-java-agent/instrumentation/anthropic_2_2_0/build.gradle rename to braintrust-sdk/instrumentation/anthropic_2_2_0/build.gradle index 155d7c1d..6a9e8752 100644 --- a/braintrust-java-agent/instrumentation/anthropic_2_2_0/build.gradle +++ b/braintrust-sdk/instrumentation/anthropic_2_2_0/build.gradle @@ -10,7 +10,7 @@ muzzle { } dependencies { - implementation project(':braintrust-java-agent:instrumenter') + compileOnly project(':braintrust-java-agent:instrumenter') implementation "io.opentelemetry:opentelemetry-api:${otelVersion}" implementation 'com.google.code.findbugs:jsr305:3.0.2' // for @Nullable annotations implementation "org.slf4j:slf4j-api:${slf4jVersion}" diff --git a/braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropic.java b/braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropic.java new file mode 100644 index 00000000..88e9fac8 --- /dev/null +++ b/braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropic.java @@ -0,0 +1,14 @@ +package dev.braintrust.instrumentation.anthropic; + +import com.anthropic.client.AnthropicClient; +import io.opentelemetry.api.OpenTelemetry; + +/** Braintrust Anthropic client instrumentation. */ +public final class BraintrustAnthropic { + + /** Instrument Anthropic client with Braintrust traces. */ + public static AnthropicClient wrap(OpenTelemetry openTelemetry, AnthropicClient client) { + return dev.braintrust.instrumentation.anthropic.v2_2_0.BraintrustAnthropic.wrap( + openTelemetry, client); + } +} diff --git a/braintrust-java-agent/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropic.java b/braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropic.java similarity index 100% rename from braintrust-java-agent/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropic.java rename to braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropic.java diff --git a/braintrust-java-agent/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/TracingHttpClient.java b/braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/TracingHttpClient.java similarity index 100% rename from braintrust-java-agent/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/TracingHttpClient.java rename to braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/TracingHttpClient.java diff --git a/braintrust-java-agent/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/auto/AnthropicInstrumentationModule.java b/braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/auto/AnthropicInstrumentationModule.java similarity index 100% rename from braintrust-java-agent/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/auto/AnthropicInstrumentationModule.java rename to braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/auto/AnthropicInstrumentationModule.java diff --git a/braintrust-java-agent/instrumentation/anthropic_2_2_0/src/test/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropicTest.java b/braintrust-sdk/instrumentation/anthropic_2_2_0/src/test/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropicTest.java similarity index 100% rename from braintrust-java-agent/instrumentation/anthropic_2_2_0/src/test/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropicTest.java rename to braintrust-sdk/instrumentation/anthropic_2_2_0/src/test/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropicTest.java diff --git a/braintrust-java-agent/instrumentation/genai_1_18_0/build.gradle b/braintrust-sdk/instrumentation/genai_1_18_0/build.gradle similarity index 96% rename from braintrust-java-agent/instrumentation/genai_1_18_0/build.gradle rename to braintrust-sdk/instrumentation/genai_1_18_0/build.gradle index 6e4cf633..ccc9f0c8 100644 --- a/braintrust-java-agent/instrumentation/genai_1_18_0/build.gradle +++ b/braintrust-sdk/instrumentation/genai_1_18_0/build.gradle @@ -10,7 +10,7 @@ muzzle { } dependencies { - implementation project(':braintrust-java-agent:instrumenter') + compileOnly project(':braintrust-java-agent:instrumenter') implementation "io.opentelemetry:opentelemetry-api:${otelVersion}" implementation 'com.google.code.findbugs:jsr305:3.0.2' // for @Nullable annotations implementation "org.slf4j:slf4j-api:${slf4jVersion}" diff --git a/braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustApiClient.java b/braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustApiClient.java similarity index 100% rename from braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustApiClient.java rename to braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustApiClient.java diff --git a/braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustInstrumentation.java b/braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustInstrumentation.java similarity index 100% rename from braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustInstrumentation.java rename to braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustInstrumentation.java diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/genai/BraintrustGenAI.java b/braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/BraintrustGenAI.java similarity index 65% rename from braintrust-sdk/src/main/java/dev/braintrust/instrumentation/genai/BraintrustGenAI.java rename to braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/BraintrustGenAI.java index 3de28008..3825e39b 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/genai/BraintrustGenAI.java +++ b/braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/BraintrustGenAI.java @@ -1,12 +1,9 @@ package dev.braintrust.instrumentation.genai; -import com.google.genai.BraintrustInstrumentation; import com.google.genai.Client; import io.opentelemetry.api.OpenTelemetry; -import lombok.extern.slf4j.Slf4j; /** Braintrust Google GenAI client instrumentation. */ -@Slf4j public class BraintrustGenAI { /** * Instrument Google GenAI Client with Braintrust traces. @@ -19,11 +16,7 @@ public class BraintrustGenAI { * @return an instrumented Gemini client */ public static Client wrap(OpenTelemetry openTelemetry, Client.Builder genAIClientBuilder) { - try { - return BraintrustInstrumentation.wrapClient(genAIClientBuilder.build(), openTelemetry); - } catch (Throwable t) { - log.error("failed to instrument gemini client", t); - return genAIClientBuilder.build(); - } + return dev.braintrust.instrumentation.genai.v1_18_0.BraintrustGenAI.wrap( + openTelemetry, genAIClientBuilder.build()); } } diff --git a/braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAI.java b/braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAI.java similarity index 100% rename from braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAI.java rename to braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAI.java diff --git a/braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/auto/GenAIInstrumentationModule.java b/braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/auto/GenAIInstrumentationModule.java similarity index 100% rename from braintrust-java-agent/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/auto/GenAIInstrumentationModule.java rename to braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/dev/braintrust/instrumentation/genai/v1_18_0/auto/GenAIInstrumentationModule.java diff --git a/braintrust-java-agent/instrumentation/genai_1_18_0/src/test/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAITest.java b/braintrust-sdk/instrumentation/genai_1_18_0/src/test/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAITest.java similarity index 100% rename from braintrust-java-agent/instrumentation/genai_1_18_0/src/test/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAITest.java rename to braintrust-sdk/instrumentation/genai_1_18_0/src/test/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAITest.java diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/build.gradle b/braintrust-sdk/instrumentation/langchain_1_8_0/build.gradle similarity index 97% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/build.gradle rename to braintrust-sdk/instrumentation/langchain_1_8_0/build.gradle index d32ee262..db4cf87c 100644 --- a/braintrust-java-agent/instrumentation/langchain_1_8_0/build.gradle +++ b/braintrust-sdk/instrumentation/langchain_1_8_0/build.gradle @@ -16,7 +16,7 @@ muzzle { } dependencies { - implementation project(':braintrust-java-agent:instrumenter') + compileOnly project(':braintrust-java-agent:instrumenter') implementation "io.opentelemetry:opentelemetry-api:${otelVersion}" implementation 'com.google.code.findbugs:jsr305:3.0.2' // for @Nullable annotations implementation "org.slf4j:slf4j-api:${slf4jVersion}" diff --git a/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/BraintrustLangchain.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/BraintrustLangchain.java new file mode 100644 index 00000000..a996ec40 --- /dev/null +++ b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/BraintrustLangchain.java @@ -0,0 +1,37 @@ +package dev.braintrust.instrumentation.langchain; + +import dev.langchain4j.model.openai.OpenAiChatModel; +import dev.langchain4j.model.openai.OpenAiStreamingChatModel; +import dev.langchain4j.service.AiServices; +import io.opentelemetry.api.OpenTelemetry; + +/** Braintrust LangChain4j client instrumentation. */ +public final class BraintrustLangchain { + + /** Instrument a LangChain4j AiServices builder with Braintrust traces. */ + @SuppressWarnings("unchecked") + public static T wrap(OpenTelemetry openTelemetry, AiServices aiServices) { + return dev.braintrust.instrumentation.langchain.v1_8_0.BraintrustLangchain.wrap( + openTelemetry, aiServices); + } + + /** Instrument langchain openai chat model with braintrust traces. */ + public static OpenAiChatModel wrap( + OpenTelemetry otel, OpenAiChatModel.OpenAiChatModelBuilder builder) { + return dev.braintrust.instrumentation.langchain.v1_8_0.BraintrustLangchain.wrap( + otel, builder); + } + + /** Instrument langchain openai streaming chat model with braintrust traces. */ + public static OpenAiStreamingChatModel wrap( + OpenTelemetry otel, OpenAiStreamingChatModel.OpenAiStreamingChatModelBuilder builder) { + return dev.braintrust.instrumentation.langchain.v1_8_0.BraintrustLangchain.wrap( + otel, builder.build()); + } + + public static OpenAiStreamingChatModel wrap( + OpenTelemetry otel, OpenAiStreamingChatModel model) { + return dev.braintrust.instrumentation.langchain.v1_8_0.BraintrustLangchain.wrap( + otel, model); + } +} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/BraintrustLangchain.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchain.java similarity index 93% rename from braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/BraintrustLangchain.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchain.java index 6947bd49..33d68841 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/BraintrustLangchain.java +++ b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchain.java @@ -1,4 +1,4 @@ -package dev.braintrust.instrumentation.langchain; +package dev.braintrust.instrumentation.langchain.v1_8_0; import dev.langchain4j.model.openai.OpenAiChatModel; import dev.langchain4j.model.openai.OpenAiStreamingChatModel; @@ -15,9 +15,16 @@ public final class BraintrustLangchain { private static final String INSTRUMENTATION_NAME = "braintrust-langchain4j"; + private static final ThreadLocal AI_SERVICES_RECURSION_GUARD = + ThreadLocal.withInitial(() -> false); @SuppressWarnings("unchecked") public static T wrap(OpenTelemetry openTelemetry, AiServices aiServices) { + if (AI_SERVICES_RECURSION_GUARD.get()) { + // already wrapped + return null; + } + AI_SERVICES_RECURSION_GUARD.set(true); try { AiServiceContext context = getPrivateField(aiServices, "context"); Tracer tracer = openTelemetry.getTracer(INSTRUMENTATION_NAME); @@ -72,6 +79,8 @@ public static T wrap(OpenTelemetry openTelemetry, AiServices aiServices) } catch (Exception e) { log.warn("failed to apply langchain AI services instrumentation", e); return aiServices.build(); + } finally { + AI_SERVICES_RECURSION_GUARD.set(false); } } @@ -81,7 +90,7 @@ public static OpenAiChatModel wrap( return wrap(otel, builder.build()); } - private static OpenAiChatModel wrap(OpenTelemetry otel, OpenAiChatModel model) { + public static OpenAiChatModel wrap(OpenTelemetry otel, OpenAiChatModel model) { try { // Get the internal OpenAiClient from the chat model Object internalClient = getPrivateField(model, "client"); diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/OtelContextPassingExecutor.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/OtelContextPassingExecutor.java similarity index 100% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/OtelContextPassingExecutor.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/OtelContextPassingExecutor.java diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingProxy.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingProxy.java similarity index 100% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingProxy.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingProxy.java diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutor.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutor.java similarity index 100% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutor.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutor.java diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClient.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClient.java similarity index 100% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClient.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClient.java diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClientBuilder.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClientBuilder.java similarity index 100% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClientBuilder.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClientBuilder.java diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/auto/LangchainInstrumentationModule.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/auto/LangchainInstrumentationModule.java similarity index 90% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/auto/LangchainInstrumentationModule.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/auto/LangchainInstrumentationModule.java index d422dc90..4e478fd0 100644 --- a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/auto/LangchainInstrumentationModule.java +++ b/braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/auto/LangchainInstrumentationModule.java @@ -71,9 +71,10 @@ private static class OpenAiChatModelBuilderAdvice { @Advice.OnMethodExit public static void build( @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) - OpenAiChatModel returnedModel) { + Object returnedModel) { returnedModel = - BraintrustLangchain.wrapChatModel(GlobalOpenTelemetry.get(), returnedModel); + BraintrustLangchain.wrap( + GlobalOpenTelemetry.get(), (OpenAiChatModel) returnedModel); } } @@ -102,14 +103,14 @@ private static class OpenAiStreamingChatModelBuilderAdvice { @Advice.OnMethodExit public static void build( @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) - OpenAiStreamingChatModel returnedModel) { + Object returnedModel) { returnedModel = - BraintrustLangchain.wrapStreamingChatModel( - GlobalOpenTelemetry.get(), returnedModel); + BraintrustLangchain.wrap( + GlobalOpenTelemetry.get(), (OpenAiStreamingChatModel) returnedModel); } } - // ------------------------------------------------------------------------- + // ------------------------------------------------------------------------ - // Intercept AiServices.build() to wrap with TracingProxy + TracingToolExecutor // ------------------------------------------------------------------------- @@ -136,9 +137,10 @@ public static void build( @Advice.This AiServices aiServices, @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object returnedService) { - returnedService = - BraintrustLangchain.wrapAiService( - GlobalOpenTelemetry.get(), aiServices, returnedService); + var wrapped = BraintrustLangchain.wrap(GlobalOpenTelemetry.get(), aiServices); + if (wrapped != null) { + returnedService = wrapped; + } } } } diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchainTest.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchainTest.java similarity index 100% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchainTest.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/BraintrustLangchainTest.java diff --git a/braintrust-java-agent/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutorTest.java b/braintrust-sdk/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutorTest.java similarity index 100% rename from braintrust-java-agent/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutorTest.java rename to braintrust-sdk/instrumentation/langchain_1_8_0/src/test/java/dev/braintrust/instrumentation/langchain/v1_8_0/TracingToolExecutorTest.java diff --git a/braintrust-java-agent/instrumentation/openai_2_8_0/build.gradle b/braintrust-sdk/instrumentation/openai_2_8_0/build.gradle similarity index 95% rename from braintrust-java-agent/instrumentation/openai_2_8_0/build.gradle rename to braintrust-sdk/instrumentation/openai_2_8_0/build.gradle index 2dde3eb6..6c6981e7 100644 --- a/braintrust-java-agent/instrumentation/openai_2_8_0/build.gradle +++ b/braintrust-sdk/instrumentation/openai_2_8_0/build.gradle @@ -7,7 +7,7 @@ muzzle { } dependencies { - implementation project(':braintrust-java-agent:instrumenter') + compileOnly project(':braintrust-java-agent:instrumenter') implementation "io.opentelemetry:opentelemetry-api:${otelVersion}" implementation 'com.google.code.findbugs:jsr305:3.0.2' // for @Nullable annotations implementation "org.slf4j:slf4j-api:${slf4jVersion}" diff --git a/braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java b/braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java new file mode 100644 index 00000000..92d1677a --- /dev/null +++ b/braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java @@ -0,0 +1,21 @@ +package dev.braintrust.instrumentation.openai; + +import com.openai.client.OpenAIClient; +import com.openai.models.chat.completions.ChatCompletionCreateParams; +import dev.braintrust.prompt.BraintrustPrompt; +import io.opentelemetry.api.OpenTelemetry; +import java.util.Map; + +public class BraintrustOpenAI { + /** Instrument openai client with braintrust traces */ + public static OpenAIClient wrapOpenAI(OpenTelemetry openTelemetry, OpenAIClient openAIClient) { + return dev.braintrust.instrumentation.openai.v2_8_0.BraintrustOpenAI.wrapOpenAI( + openTelemetry, openAIClient); + } + + public static ChatCompletionCreateParams buildChatCompletionsPrompt( + BraintrustPrompt prompt, Map parameters) { + return dev.braintrust.instrumentation.openai.v2_8_0.BraintrustOpenAI + .buildChatCompletionsPrompt(prompt, parameters); + } +} diff --git a/braintrust-java-agent/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAI.java b/braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAI.java similarity index 100% rename from braintrust-java-agent/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAI.java rename to braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAI.java diff --git a/braintrust-java-agent/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/TracingHttpClient.java b/braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/TracingHttpClient.java similarity index 100% rename from braintrust-java-agent/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/TracingHttpClient.java rename to braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/TracingHttpClient.java diff --git a/braintrust-java-agent/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/auto/OpenAIInstrumentationModule.java b/braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/auto/OpenAIInstrumentationModule.java similarity index 100% rename from braintrust-java-agent/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/auto/OpenAIInstrumentationModule.java rename to braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/auto/OpenAIInstrumentationModule.java diff --git a/braintrust-java-agent/instrumentation/openai_2_8_0/src/test/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAITest.java b/braintrust-sdk/instrumentation/openai_2_8_0/src/test/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAITest.java similarity index 100% rename from braintrust-java-agent/instrumentation/openai_2_8_0/src/test/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAITest.java rename to braintrust-sdk/instrumentation/openai_2_8_0/src/test/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAITest.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/build.gradle b/braintrust-sdk/instrumentation/springai_1_0_0/build.gradle similarity index 98% rename from braintrust-java-agent/instrumentation/springai_1_0_0/build.gradle rename to braintrust-sdk/instrumentation/springai_1_0_0/build.gradle index 06fc820f..541c6178 100644 --- a/braintrust-java-agent/instrumentation/springai_1_0_0/build.gradle +++ b/braintrust-sdk/instrumentation/springai_1_0_0/build.gradle @@ -20,7 +20,7 @@ muzzle { } dependencies { - implementation project(':braintrust-java-agent:instrumenter') + compileOnly project(':braintrust-java-agent:instrumenter') implementation "io.opentelemetry:opentelemetry-api:${otelVersion}" implementation 'com.google.code.findbugs:jsr305:3.0.2' implementation "org.slf4j:slf4j-api:${slf4jVersion}" diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/AnthropicBuilderWrapper.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/AnthropicBuilderWrapper.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/AnthropicBuilderWrapper.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/AnthropicBuilderWrapper.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustObservationHandler.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustObservationHandler.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustObservationHandler.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustObservationHandler.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAI.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAI.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAI.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAI.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/OpenAIBuilderWrapper.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/OpenAIBuilderWrapper.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/OpenAIBuilderWrapper.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/OpenAIBuilderWrapper.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/TriConsumer.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/TriConsumer.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/TriConsumer.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/TriConsumer.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIAnthropicInstrumentationModule.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIAnthropicInstrumentationModule.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIAnthropicInstrumentationModule.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIAnthropicInstrumentationModule.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIOpenAIInstrumentationModule.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIOpenAIInstrumentationModule.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIOpenAIInstrumentationModule.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/auto/SpringAIOpenAIInstrumentationModule.java diff --git a/braintrust-java-agent/instrumentation/springai_1_0_0/src/test/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAITest.java b/braintrust-sdk/instrumentation/springai_1_0_0/src/test/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAITest.java similarity index 100% rename from braintrust-java-agent/instrumentation/springai_1_0_0/src/test/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAITest.java rename to braintrust-sdk/instrumentation/springai_1_0_0/src/test/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAITest.java diff --git a/braintrust-sdk/src/main/java/com/google/genai/BraintrustApiClient.java b/braintrust-sdk/src/main/java/com/google/genai/BraintrustApiClient.java deleted file mode 100644 index 35f4441b..00000000 --- a/braintrust-sdk/src/main/java/com/google/genai/BraintrustApiClient.java +++ /dev/null @@ -1,411 +0,0 @@ -package com.google.genai; - -import static dev.braintrust.json.BraintrustJsonMapper.fromJson; -import static dev.braintrust.json.BraintrustJsonMapper.toJson; - -import com.google.genai.types.HttpOptions; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; -import io.opentelemetry.context.Scope; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; -import javax.annotation.Nullable; -import lombok.SneakyThrows; -import lombok.extern.slf4j.Slf4j; -import okhttp3.Headers; -import okhttp3.MediaType; -import okhttp3.ResponseBody; - -/** - * Instrumented wrapper for ApiClient that adds OpenTelemetry spans. - * - *

This class lives in com.google.genai package to access package-private ApiClient class. - */ -@Slf4j -class BraintrustApiClient extends ApiClient { - private final ApiClient delegate; - private final Tracer tracer; - - public BraintrustApiClient(ApiClient delegate, OpenTelemetry openTelemetry) { - // We must call super(), but we'll override all methods to delegate - // Pass the delegate's config to minimize differences - super( - delegate.apiKey != null ? delegate.apiKey : Optional.empty(), - delegate.project != null ? delegate.project : Optional.empty(), - delegate.location != null ? delegate.location : Optional.empty(), - delegate.credentials != null ? delegate.credentials : Optional.empty(), - delegate.httpOptions != null ? Optional.of(delegate.httpOptions) : Optional.empty(), - delegate.clientOptions != null ? delegate.clientOptions : Optional.empty()); - this.delegate = delegate; - this.tracer = openTelemetry.getTracer("io.opentelemetry.gemini-java-1.20"); - } - - private void tagSpan( - Span span, - @Nullable String genAIEndpoint, - @Nullable String requestMethod, - @Nullable String requestBody, - @Nullable String responseBody) { - try { - Map metadata = new java.util.HashMap<>(); - metadata.put("provider", "gemini"); - - // Parse request - if (requestBody != null) { - var requestJson = fromJson(requestBody, Map.class); - - // Extract metadata fields - for (String field : - List.of( - "model", - "systemInstruction", - "tools", - "toolConfig", - "safetySettings", - "cachedContent")) { - if (requestJson.containsKey(field)) { - metadata.put(field, requestJson.get(field)); - } - } - - // Extract generationConfig fields into metadata - if (requestJson.get("generationConfig") instanceof Map) { - var genConfig = (Map) requestJson.get("generationConfig"); - for (String field : - List.of( - "temperature", - "topP", - "topK", - "candidateCount", - "maxOutputTokens", - "stopSequences", - "responseMimeType", - "responseSchema")) { - if (genConfig.containsKey(field)) { - metadata.put(field, genConfig.get(field)); - } - } - } - - // Build input_json - Map inputJson = new java.util.HashMap<>(); - String model = getModel(genAIEndpoint); - if (requestJson.containsKey("model")) { - inputJson.put("model", requestJson.get("model")); - } else if (model != null) { - inputJson.put("model", model); - } - if (requestJson.containsKey("contents")) { - inputJson.put("contents", requestJson.get("contents")); - } - if (requestJson.containsKey("generationConfig")) { - inputJson.put("config", requestJson.get("generationConfig")); - } - - span.setAttribute("braintrust.input_json", toJson(inputJson)); - } - - // Parse response - if (responseBody != null) { - var responseJson = fromJson(responseBody, Map.class); - - // Extract model version from response - if (responseJson.containsKey("modelVersion")) { - metadata.put("model", responseJson.get("modelVersion")); - } - - // Set full response as output_json - span.setAttribute("braintrust.output_json", toJson(responseJson)); - - // Parse usage metadata for metrics - if (responseJson.get("usageMetadata") instanceof Map) { - var usage = (Map) responseJson.get("usageMetadata"); - Map metrics = new java.util.HashMap<>(); - - if (usage.containsKey("promptTokenCount")) { - metrics.put("prompt_tokens", (Number) usage.get("promptTokenCount")); - } - if (usage.containsKey("candidatesTokenCount")) { - metrics.put( - "completion_tokens", (Number) usage.get("candidatesTokenCount")); - } - if (usage.containsKey("totalTokenCount")) { - metrics.put("tokens", (Number) usage.get("totalTokenCount")); - } - if (usage.containsKey("cachedContentTokenCount")) { - metrics.put( - "prompt_cached_tokens", - (Number) usage.get("cachedContentTokenCount")); - } - - span.setAttribute("braintrust.metrics", toJson(metrics)); - } - } - - // Set metadata - span.setAttribute("braintrust.metadata", toJson(metadata)); - - // Set span_attributes to mark as LLM span - span.setAttribute("braintrust.span_attributes", toJson(Map.of("type", "llm"))); - - } catch (Throwable t) { - log.warn("failed to tag gemini span", t); - } - } - - // Override accessor methods to delegate to original client - @Override - public boolean vertexAI() { - return delegate.vertexAI(); - } - - @Override - public String project() { - return delegate.project(); - } - - @Override - public String location() { - return delegate.location(); - } - - @Override - public String apiKey() { - return delegate.apiKey(); - } - - @Override - @SneakyThrows - public ApiResponse request( - String requestMethod, - String genAIUrl, - String requestBody, - Optional options) { - Span span = - tracer.spanBuilder(getOperation(genAIUrl)).setSpanKind(SpanKind.CLIENT).startSpan(); - try (Scope scope = span.makeCurrent()) { - ApiResponse response = delegate.request(requestMethod, genAIUrl, requestBody, options); - BufferedApiResponse bufferedResponse = new BufferedApiResponse(response); - span.setStatus(StatusCode.OK); - tagSpan(span, genAIUrl, requestMethod, requestBody, bufferedResponse.getBodyAsString()); - return bufferedResponse; - } catch (Throwable t) { - span.setStatus(StatusCode.ERROR, t.getMessage()); - span.recordException(t); - throw t; - } finally { - span.end(); - } - } - - @Override - @SneakyThrows - public ApiResponse request( - String requestMethod, - String genAIUrl, - byte[] requestBodyBytes, - Optional options) { - Span span = - tracer.spanBuilder(getOperation(genAIUrl)).setSpanKind(SpanKind.CLIENT).startSpan(); - try (Scope scope = span.makeCurrent()) { - ApiResponse response = - delegate.request(requestMethod, genAIUrl, requestBodyBytes, options); - BufferedApiResponse bufferedResponse = new BufferedApiResponse(response); - span.setStatus(StatusCode.OK); - tagSpan( - span, - genAIUrl, - requestMethod, - new String(requestBodyBytes), - bufferedResponse.getBodyAsString()); - return bufferedResponse; - } catch (Throwable t) { - span.setStatus(StatusCode.ERROR, t.getMessage()); - span.recordException(t); - throw t; - } finally { - span.end(); - } - } - - @Override - public CompletableFuture asyncRequest( - String method, String url, String body, Optional options) { - Span span = tracer.spanBuilder(getOperation(url)).setSpanKind(SpanKind.CLIENT).startSpan(); - Context context = Context.current().with(span); - - return delegate.asyncRequest(method, url, body, options) - .handle( - (response, throwable) -> { - try (Scope scope = context.makeCurrent()) { - if (throwable != null) { - span.setStatus(StatusCode.ERROR, throwable.getMessage()); - span.recordException(throwable); - throw new RuntimeException(throwable); - } - - try { - // Buffer the response so we can read it for instrumentation - BufferedApiResponse bufferedResponse = - new BufferedApiResponse(response); - span.setStatus(StatusCode.OK); - tagSpan( - span, - url, - method, - body, - bufferedResponse.getBodyAsString()); - return (ApiResponse) bufferedResponse; - } catch (Exception e) { - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.recordException(e); - throw new RuntimeException(e); - } - } finally { - span.end(); - } - }); - } - - @Override - public CompletableFuture asyncRequest( - String method, String url, byte[] body, Optional options) { - Span span = tracer.spanBuilder(getOperation(url)).setSpanKind(SpanKind.CLIENT).startSpan(); - Context context = Context.current().with(span); - - return delegate.asyncRequest(method, url, body, options) - .handle( - (response, throwable) -> { - try (Scope scope = context.makeCurrent()) { - if (throwable != null) { - span.setStatus(StatusCode.ERROR, throwable.getMessage()); - span.recordException(throwable); - throw new RuntimeException(throwable); - } - - try { - // Buffer the response so we can read it for instrumentation - BufferedApiResponse bufferedResponse = - new BufferedApiResponse(response); - span.setStatus(StatusCode.OK); - tagSpan( - span, - url, - method, - new String(body), - bufferedResponse.getBodyAsString()); - return (ApiResponse) bufferedResponse; - } catch (Exception e) { - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.recordException(e); - throw new RuntimeException(e); - } - } finally { - span.end(); - } - }); - } - - private static String getModel(String genAIEndpoint) { - try { - // endpoint has model and request type. Example: - // models/gemini-2.0-flash-lite:generateContent - var segments = genAIEndpoint.split("/"); - var lastSegment = segments[segments.length - 1].split(":"); - return lastSegment[0]; - } catch (Exception e) { - log.debug("unable to determine model name", e); - return "gemini"; - } - } - - private static String getOperation(String genAIEndpoint) { - try { - // endpoint has model and request type. Example: - // models/gemini-2.0-flash-lite:generateContent - var segments = genAIEndpoint.split("/"); - var lastSegment = segments[segments.length - 1].split(":"); - return toSnakeCase(lastSegment[1]); - } catch (Exception e) { - log.debug("unable to determine operation name", e); - return "gemini.api.call"; - } - } - - /** convert a camelCaseString to a snake_case_string */ - private static String toSnakeCase(String camelCase) { - if (camelCase == null || camelCase.isEmpty()) return camelCase; - - StringBuilder sb = new StringBuilder(camelCase.length() + 5); - - for (int i = 0; i < camelCase.length(); i++) { - char c = camelCase.charAt(i); - if (Character.isUpperCase(c)) { - if (i > 0) sb.append('_'); - sb.append(Character.toLowerCase(c)); - } else { - sb.append(c); - } - } - - return sb.toString(); - } - - /** - * Wrapper for ApiResponse that buffers the response body so it can be read multiple times. - * - *

This allows us to capture the response body for instrumentation while still allowing the - * delegate to read it. - */ - static class BufferedApiResponse extends ApiResponse { - private final ApiResponse delegate; - private final byte[] bufferedBody; - - public BufferedApiResponse(ApiResponse delegate) throws Exception { - this.delegate = delegate; - // Read the body once and buffer it - ResponseBody body = delegate.getBody(); - this.bufferedBody = body != null ? body.bytes() : null; - } - - @Override - public ResponseBody getBody() { - if (bufferedBody == null) { - return null; - } - // Create a new ResponseBody from the buffered bytes - // Get the original content type if available - MediaType contentType = null; - try { - ResponseBody originalBody = delegate.getBody(); - if (originalBody != null) { - contentType = originalBody.contentType(); - } - } catch (Exception e) { - // Ignore, use null content type - } - return ResponseBody.create(bufferedBody, contentType); - } - - @Override - public Headers getHeaders() { - return delegate.getHeaders(); - } - - @Override - public void close() { - delegate.close(); - } - - /** Get the buffered body as a string for instrumentation. */ - public String getBodyAsString() { - return bufferedBody != null ? new String(bufferedBody) : null; - } - } -} diff --git a/braintrust-sdk/src/main/java/com/google/genai/BraintrustInstrumentation.java b/braintrust-sdk/src/main/java/com/google/genai/BraintrustInstrumentation.java deleted file mode 100644 index 0b3f462e..00000000 --- a/braintrust-sdk/src/main/java/com/google/genai/BraintrustInstrumentation.java +++ /dev/null @@ -1,93 +0,0 @@ -package com.google.genai; - -import io.opentelemetry.api.OpenTelemetry; -import java.lang.reflect.Field; -import java.lang.reflect.Modifier; -import lombok.extern.slf4j.Slf4j; - -/** - * Helper class for instrumenting Gemini Client by replacing its internal ApiClient. - * - *

This class lives in com.google.genai package to access package-private ApiClient class. - */ -@Slf4j -public class BraintrustInstrumentation { - /** - * Wraps a Client's internal ApiClient with an instrumented version. - * - * @param client the client to instrument - * @param openTelemetry the OpenTelemetry instance - * @return the same client instance, but with instrumented ApiClient - */ - public static Client wrapClient(Client client, OpenTelemetry openTelemetry) throws Exception { - // Get the apiClient field from Client - Field clientApiClientField = Client.class.getDeclaredField("apiClient"); - clientApiClientField.setAccessible(true); - ApiClient originalApiClient = (ApiClient) clientApiClientField.get(client); - - // Create instrumented wrapper - BraintrustApiClient instrumentedApiClient = - new BraintrustApiClient(originalApiClient, openTelemetry); - - // Replace apiClient in Client - setFinalField(client, clientApiClientField, instrumentedApiClient); - - // Replace apiClient in all Client service fields - replaceApiClientInService(client.models, instrumentedApiClient); - replaceApiClientInService(client.batches, instrumentedApiClient); - replaceApiClientInService(client.caches, instrumentedApiClient); - replaceApiClientInService(client.operations, instrumentedApiClient); - replaceApiClientInService(client.chats, instrumentedApiClient); - replaceApiClientInService(client.files, instrumentedApiClient); - replaceApiClientInService(client.tunings, instrumentedApiClient); - - // Replace apiClient in all Client.async service fields - if (client.async != null) { - replaceApiClientInService(client.async.models, instrumentedApiClient); - replaceApiClientInService(client.async.batches, instrumentedApiClient); - replaceApiClientInService(client.async.caches, instrumentedApiClient); - replaceApiClientInService(client.async.operations, instrumentedApiClient); - replaceApiClientInService(client.async.chats, instrumentedApiClient); - replaceApiClientInService(client.async.files, instrumentedApiClient); - replaceApiClientInService(client.async.tunings, instrumentedApiClient); - } - - log.debug("Successfully instrumented Gemini client"); - return client; - } - - /** Replaces the apiClient field in a service object (Models, Batches, etc). */ - private static void replaceApiClientInService(Object service, ApiClient instrumentedApiClient) - throws Exception { - if (service == null) { - return; - } - try { - Field apiClientField = service.getClass().getDeclaredField("apiClient"); - apiClientField.setAccessible(true); - setFinalField(service, apiClientField, instrumentedApiClient); - } catch (NoSuchFieldException e) { - // Some services might not have an apiClient field - log.info("No apiClient field found in " + service.getClass().getSimpleName()); - } - } - - /** - * Sets a final field using reflection. - * - *

This works by making the field accessible and, on older Java versions, removing the final - * modifier. - */ - private static void setFinalField(Object target, Field field, Object value) throws Exception { - field.setAccessible(true); - // Try to remove final modifier - try { - Field modifiersField = Field.class.getDeclaredField("modifiers"); - modifiersField.setAccessible(true); - modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL); - } catch (NoSuchFieldException e) { - // ignore - } - field.set(target, value); - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java b/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java index ca363556..0d41a2c6 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java @@ -4,7 +4,6 @@ import java.util.*; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import org.jspecify.annotations.NonNull; /** A dataset loaded externally from Braintrust using paginated API fetches */ public class DatasetBrainstoreImpl implements Dataset { @@ -70,7 +69,7 @@ private class BrainstoreCursor implements Cursor> { private boolean closed; private final @Nonnull String cursorVersion; - BrainstoreCursor(@NonNull String cursorVersion) { + BrainstoreCursor(@Nonnull String cursorVersion) { this.currentBatch = new ArrayList<>(); this.currentIndex = 0; this.cursor = null; diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java index fa497302..91e567aa 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java @@ -10,9 +10,9 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; import lombok.SneakyThrows; -import org.jspecify.annotations.NonNull; -import org.jspecify.annotations.Nullable; public class InstrumentationSemConv { public static final String PROVIDER_NAME_OPENAI = "openai"; @@ -27,10 +27,10 @@ public class InstrumentationSemConv { @SneakyThrows public static void tagLLMSpanRequest( Span span, - @NonNull String providerName, - @NonNull String baseUrl, - @NonNull List pathSegments, - @NonNull String method, + @Nonnull String providerName, + @Nonnull String baseUrl, + @Nonnull List pathSegments, + @Nonnull String method, @Nullable String requestBody) { switch (providerName) { case PROVIDER_NAME_OPENAI -> @@ -46,15 +46,15 @@ public static void tagLLMSpanRequest( } public static void tagLLMSpanResponse( - Span span, @NonNull String providerName, @NonNull String responseBody) { + Span span, @Nonnull String providerName, @Nonnull String responseBody) { tagLLMSpanResponse(span, providerName, responseBody, null); } @SneakyThrows public static void tagLLMSpanResponse( Span span, - @NonNull String providerName, - @NonNull String responseBody, + @Nonnull String providerName, + @Nonnull String responseBody, @Nullable Long timeToFirstTokenNanoseconds) { switch (providerName) { case PROVIDER_NAME_OPENAI -> @@ -65,7 +65,7 @@ public static void tagLLMSpanResponse( } } - public static void tagLLMSpanResponse(Span span, @NonNull Throwable responseError) { + public static void tagLLMSpanResponse(Span span, @Nonnull Throwable responseError) { span.setStatus(StatusCode.ERROR, responseError.getMessage()); span.recordException(responseError); } diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropic.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropic.java deleted file mode 100644 index 1d980b20..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropic.java +++ /dev/null @@ -1,133 +0,0 @@ -package dev.braintrust.instrumentation.anthropic; - -import com.anthropic.client.AnthropicClient; -import com.anthropic.core.ClientOptions; -import com.anthropic.core.http.HttpClient; -import io.opentelemetry.api.OpenTelemetry; -import java.lang.reflect.Field; -import java.lang.reflect.Modifier; -import java.util.function.BiConsumer; -import java.util.function.Consumer; -import kotlin.Lazy; -import lombok.extern.slf4j.Slf4j; - -/** Braintrust Anthropic client instrumentation. */ -@Slf4j -public final class BraintrustAnthropic { - - /** Instrument Anthropic client with Braintrust traces. */ - public static AnthropicClient wrap(OpenTelemetry openTelemetry, AnthropicClient client) { - try { - instrumentHttpClient(openTelemetry, client); - return client; - } catch (Exception e) { - log.error("failed to apply anthropic instrumentation", e); - return client; - } - } - - private static void instrumentHttpClient(OpenTelemetry openTelemetry, AnthropicClient client) { - forAllFields( - client, - fieldName -> { - try { - var field = getField(client, fieldName); - if (field instanceof ClientOptions clientOptions) { - instrumentClientOptions( - openTelemetry, clientOptions, "originalHttpClient"); - instrumentClientOptions(openTelemetry, clientOptions, "httpClient"); - } else if (field instanceof Lazy lazyField) { - var resolved = lazyField.getValue(); - forAllFieldsOfType( - resolved, - ClientOptions.class, - (clientOptions, subfieldName) -> - instrumentClientOptions( - openTelemetry, clientOptions, subfieldName)); - } else { - forAllFieldsOfType( - field, - ClientOptions.class, - (clientOptions, subfieldName) -> - instrumentClientOptions( - openTelemetry, clientOptions, subfieldName)); - } - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } - }); - } - - private static void instrumentClientOptions( - OpenTelemetry openTelemetry, ClientOptions clientOptions, String fieldName) { - try { - HttpClient httpClient = getField(clientOptions, fieldName); - if (!(httpClient instanceof TracingHttpClient)) { - setPrivateField( - clientOptions, fieldName, new TracingHttpClient(openTelemetry, httpClient)); - } - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } - } - - private static void forAllFields(Object object, Consumer consumer) { - if (object == null || consumer == null) return; - Class clazz = object.getClass(); - while (clazz != null && clazz != Object.class) { - for (Field field : clazz.getDeclaredFields()) { - if (field.isSynthetic()) continue; - if (Modifier.isStatic(field.getModifiers())) continue; - consumer.accept(field.getName()); - } - clazz = clazz.getSuperclass(); - } - } - - private static void forAllFieldsOfType( - Object object, Class targetClazz, BiConsumer consumer) { - forAllFields( - object, - fieldName -> { - try { - if (targetClazz.isAssignableFrom(object.getClass())) { - consumer.accept(getField(object, fieldName), fieldName); - } - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } - }); - } - - @SuppressWarnings("unchecked") - private static T getField(Object obj, String fieldName) - throws ReflectiveOperationException { - Class clazz = obj.getClass(); - while (clazz != null) { - try { - Field field = clazz.getDeclaredField(fieldName); - field.setAccessible(true); - return (T) field.get(obj); - } catch (NoSuchFieldException e) { - clazz = clazz.getSuperclass(); - } - } - throw new NoSuchFieldException(fieldName); - } - - private static void setPrivateField(Object obj, String fieldName, Object value) - throws ReflectiveOperationException { - Class clazz = obj.getClass(); - while (clazz != null) { - try { - Field field = clazz.getDeclaredField(fieldName); - field.setAccessible(true); - field.set(obj, value); - return; - } catch (NoSuchFieldException e) { - clazz = clazz.getSuperclass(); - } - } - throw new NoSuchFieldException(fieldName); - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/anthropic/TracingHttpClient.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/anthropic/TracingHttpClient.java deleted file mode 100644 index ba561c28..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/anthropic/TracingHttpClient.java +++ /dev/null @@ -1,370 +0,0 @@ -package dev.braintrust.instrumentation.anthropic; - -import com.anthropic.core.RequestOptions; -import com.anthropic.core.http.HttpClient; -import com.anthropic.core.http.HttpRequest; -import com.anthropic.core.http.HttpRequestBody; -import com.anthropic.core.http.HttpResponse; -import com.anthropic.helpers.MessageAccumulator; -import com.anthropic.models.messages.RawMessageStreamEvent; -import dev.braintrust.instrumentation.InstrumentationSemConv; -import dev.braintrust.json.BraintrustJsonMapper; -import dev.braintrust.trace.BraintrustTracing; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.Tracer; -import java.io.BufferedReader; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.InputStream; -import java.io.InputStreamReader; -import java.io.OutputStream; -import java.nio.charset.StandardCharsets; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; -import lombok.extern.slf4j.Slf4j; -import org.jspecify.annotations.NonNull; - -@Slf4j -public class TracingHttpClient implements HttpClient { - private final Tracer tracer; - private final HttpClient underlying; - - public TracingHttpClient(OpenTelemetry openTelemetry, HttpClient underlying) { - this.tracer = BraintrustTracing.getTracer(openTelemetry); - this.underlying = underlying; - } - - @Override - public void close() { - underlying.close(); - } - - @Override - public @NonNull HttpResponse execute( - @NonNull HttpRequest httpRequest, @NonNull RequestOptions requestOptions) { - var span = tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME).startSpan(); - try (var ignored = span.makeCurrent()) { - var bufferedRequest = bufferRequestBody(httpRequest); - - String inputJson = - bufferedRequest.body() != null - ? readBodyAsString(bufferedRequest.body()) - : null; - - InstrumentationSemConv.tagLLMSpanRequest( - span, - InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC, - bufferedRequest.baseUrl() != null ? bufferedRequest.baseUrl() : "", - bufferedRequest.pathSegments(), - bufferedRequest.method().name(), - inputJson); - - var response = underlying.execute(bufferedRequest, requestOptions); - return new TeeingStreamHttpResponse(response, span); - } catch (Exception e) { - InstrumentationSemConv.tagLLMSpanResponse(span, e); - span.end(); - throw e; - } - } - - @Override - public @NonNull CompletableFuture executeAsync( - @NonNull HttpRequest httpRequest, @NonNull RequestOptions requestOptions) { - var span = tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME).startSpan(); - try { - var bufferedRequest = bufferRequestBody(httpRequest); - String inputJson = - bufferedRequest.body() != null - ? readBodyAsString(bufferedRequest.body()) - : null; - InstrumentationSemConv.tagLLMSpanRequest( - span, - InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC, - bufferedRequest.baseUrl() != null ? bufferedRequest.baseUrl() : "", - bufferedRequest.pathSegments(), - bufferedRequest.method().name(), - inputJson); - return underlying - .executeAsync(bufferedRequest, requestOptions) - .thenApply( - response -> (HttpResponse) new TeeingStreamHttpResponse(response, span)) - .whenComplete( - (response, t) -> { - if (t != null) { - // this means the future itself failed - InstrumentationSemConv.tagLLMSpanResponse(span, t); - span.end(); - } - }); - } catch (Exception e) { - InstrumentationSemConv.tagLLMSpanResponse(span, e); - span.end(); - throw e; - } - } - - // ------------------------------------------------------------------------- - // Request buffering — identical pattern to OpenAI TracingHttpClient - // ------------------------------------------------------------------------- - - private static HttpRequest bufferRequestBody(HttpRequest request) { - HttpRequestBody originalBody = request.body(); - if (originalBody == null) { - return request; - } - var baos = new ByteArrayOutputStream(); - originalBody.writeTo(baos); - byte[] bytes = baos.toByteArray(); - String contentType = originalBody.contentType(); - - HttpRequestBody bufferedBody = - new HttpRequestBody() { - @Override - public void writeTo(OutputStream outputStream) { - try { - outputStream.write(bytes); - } catch (java.io.IOException e) { - throw new RuntimeException(e); - } - } - - @Override - public String contentType() { - return contentType; - } - - @Override - public long contentLength() { - return bytes.length; - } - - @Override - public boolean repeatable() { - return true; - } - - @Override - public void close() {} - }; - - return request.toBuilder().body(bufferedBody).build(); - } - - private static String readBodyAsString(HttpRequestBody body) { - var baos = new ByteArrayOutputStream((int) Math.max(body.contentLength(), 0)); - body.writeTo(baos); - return baos.toString(StandardCharsets.UTF_8); - } - - // ------------------------------------------------------------------------- - // Response tee — identical pattern to OpenAI TracingHttpClient - // ------------------------------------------------------------------------- - - /** - * Tees the response body so bytes are accumulated as the caller reads, then on close tags the - * span by auto-detecting SSE vs plain JSON from the first non-empty line. - */ - private static final class TeeingStreamHttpResponse implements HttpResponse { - private final HttpResponse delegate; - private final Span span; - private final long spanStartNanos = System.nanoTime(); - private final AtomicLong timeToFirstTokenNanos = new AtomicLong(); - private final ByteArrayOutputStream teeBuffer = new ByteArrayOutputStream(); - private final InputStream teeStream; - - TeeingStreamHttpResponse(HttpResponse delegate, Span span) { - this.delegate = delegate; - this.span = span; - this.teeStream = - new TeeInputStream( - delegate.body(), teeBuffer, this::onFirstByte, this::onStreamClosed); - } - - private void onFirstByte() { - timeToFirstTokenNanos.set(System.nanoTime() - spanStartNanos); - } - - private void onStreamClosed() { - try { - byte[] bytes; - synchronized (teeBuffer) { - bytes = teeBuffer.toByteArray(); - } - tagSpanFromBuffer(span, bytes, timeToFirstTokenNanos.get()); - } finally { - span.end(); - } - } - - @Override - public int statusCode() { - return delegate.statusCode(); - } - - @Override - public com.anthropic.core.http.Headers headers() { - return delegate.headers(); - } - - @Override - public InputStream body() { - return teeStream; - } - - @Override - public void close() { - try { - teeStream.close(); - } catch (java.io.IOException ignored) { - } - delegate.close(); - } - } - - private static final class TeeInputStream extends InputStream { - private final InputStream source; - private final OutputStream sink; - private final Runnable onFirstByte; - private final Runnable onClose; - private final AtomicBoolean firstByteSeen = new AtomicBoolean(false); - private final AtomicBoolean closed = new AtomicBoolean(false); - - TeeInputStream( - InputStream source, OutputStream sink, Runnable onFirstByte, Runnable onClose) { - this.source = source; - this.sink = sink; - this.onFirstByte = onFirstByte; - this.onClose = onClose; - } - - @Override - public int read() throws java.io.IOException { - int b = source.read(); - if (b == -1) { - notifyClosed(); - } else { - notifyFirstByte(); - sink.write(b); - } - return b; - } - - @Override - public int read(byte[] buf, int off, int len) throws java.io.IOException { - int n = source.read(buf, off, len); - if (n == -1) { - notifyClosed(); - } else { - notifyFirstByte(); - sink.write(buf, off, n); - } - return n; - } - - @Override - public void close() throws java.io.IOException { - notifyClosed(); - source.close(); - } - - private void notifyFirstByte() { - if (!firstByteSeen.compareAndExchange(false, true)) { - onFirstByte.run(); - } - } - - private void notifyClosed() { - if (!closed.compareAndExchange(false, true)) { - onClose.run(); - } - } - } - - // ------------------------------------------------------------------------- - // Span tagging from buffered bytes - // ------------------------------------------------------------------------- - - private static void tagSpanFromBuffer(Span span, byte[] bytes, Long timeToFirstTokenNanos) { - if (bytes.length == 0) return; - try { - String firstLine = firstNonEmptyLine(bytes); - // Anthropic SSE starts with "event: message_start\ndata: ..." so we detect - // either prefix. OpenAI SSE starts directly with "data:". - boolean isSse = - firstLine != null - && (firstLine.startsWith("data:") || firstLine.startsWith("event:")); - if (isSse) { - tagSpanFromSseBytes(span, bytes, timeToFirstTokenNanos); - } else { - // Non-streaming: plain Message JSON — pass it whole, no time_to_first_token - InstrumentationSemConv.tagLLMSpanResponse( - span, - InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC, - new String(bytes, StandardCharsets.UTF_8)); - } - } catch (Exception e) { - log.error("Could not tag span from Anthropic response buffer", e); - } - } - - private static String firstNonEmptyLine(byte[] bytes) { - int start = 0; - for (int i = 0; i <= bytes.length; i++) { - if (i == bytes.length || bytes[i] == '\n') { - String line = new String(bytes, start, i - start, StandardCharsets.UTF_8).strip(); - if (!line.isEmpty()) return line; - start = i + 1; - } - } - return null; - } - - /** - * Anthropic SSE wire format has named events: - * - * 

-     * event: message_start
-     * data: {"type":"message_start","message":{...}}
-     *
-     * event: content_block_delta
-     * data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}
-     *
- * - * We only need the {@code data:} lines — the event name is redundant with the {@code type} - * field inside the JSON. Feed each data payload to {@link MessageAccumulator} and serialize the - * assembled {@link com.anthropic.models.messages.Message} for the span. - */ - private static void tagSpanFromSseBytes( - Span span, byte[] sseBytes, Long timeToFirstTokenNanos) { - try { - var mapper = BraintrustJsonMapper.get(); - var reader = - new BufferedReader( - new InputStreamReader( - new ByteArrayInputStream(sseBytes), StandardCharsets.UTF_8)); - var accumulator = MessageAccumulator.create(); - String line; - while ((line = reader.readLine()) != null) { - if (!line.startsWith("data:")) continue; - String data = line.substring("data:".length()).strip(); - if (data.isEmpty()) continue; - try { - accumulator.accumulate(mapper.readValue(data, RawMessageStreamEvent.class)); - } catch (Exception ignored) { - // skip unrecognized event types (e.g. ping) - } - } - String assembledMessageJson = BraintrustJsonMapper.toJson(accumulator.message()); - InstrumentationSemConv.tagLLMSpanResponse( - span, - InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC, - assembledMessageJson, - timeToFirstTokenNanos); - } catch (Exception e) { - log.error("Could not parse Anthropic SSE buffer to tag streaming span output", e); - } - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/OtelContextPassingExecutor.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/OtelContextPassingExecutor.java deleted file mode 100644 index cdf54d92..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/OtelContextPassingExecutor.java +++ /dev/null @@ -1,31 +0,0 @@ -package dev.braintrust.instrumentation.langchain; - -import io.opentelemetry.context.Context; -import java.util.concurrent.Executor; -import org.jspecify.annotations.NonNull; - -/** - * An executor that links open telemetry spans across threads - * - *

Any tasks submitted to the executor will point to the parent context that was present at the - * time of task submission. Or, if no parent context was present tasks will create spans as they - * normally would (or would not) without this executor. - */ -class OtelContextPassingExecutor implements Executor { - private final Executor underlying; - - public OtelContextPassingExecutor(Executor executor) { - this.underlying = executor; - } - - @Override - public void execute(@NonNull Runnable command) { - var context = Context.current(); - underlying.execute( - () -> { - try (var ignored = context.makeCurrent()) { - command.run(); - } - }); - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/TracingProxy.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/TracingProxy.java deleted file mode 100644 index 5bbb4cd8..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/TracingProxy.java +++ /dev/null @@ -1,50 +0,0 @@ -package dev.braintrust.instrumentation.langchain; - -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Scope; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.Proxy; -import org.jspecify.annotations.NonNull; - -class TracingProxy { - /** - * Use a java {@link Proxy} to wrap a service interface methods with spans - * - *

Each interface method will create a span with the same name as the method - */ - @SuppressWarnings("unchecked") - public static @NonNull T create(Class serviceInterface, T service, Tracer tracer) { - return (T) - Proxy.newProxyInstance( - serviceInterface.getClassLoader(), - new Class[] {serviceInterface}, - (proxy, method, args) -> { - // Skip Object methods (equals, hashCode, toString) - if (method.getDeclaringClass() == Object.class) { - return method.invoke(service, args); - } - - Span span = tracer.spanBuilder(method.getName()).startSpan(); - try (Scope ignored = span.makeCurrent()) { - // Use setAccessible to handle non-public interfaces - method.setAccessible(true); - return method.invoke(service, args); - } catch (InvocationTargetException e) { - Throwable cause = e.getCause(); - span.setStatus(StatusCode.ERROR, cause.getMessage()); - span.recordException(cause); - throw cause; - } catch (Exception e) { - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.recordException(e); - throw e; - } finally { - span.end(); - } - }); - } - - private TracingProxy() {} -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/TracingToolExecutor.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/TracingToolExecutor.java deleted file mode 100644 index dd6dd5ee..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/TracingToolExecutor.java +++ /dev/null @@ -1,78 +0,0 @@ -package dev.braintrust.instrumentation.langchain; - -import dev.langchain4j.agent.tool.ToolExecutionRequest; -import dev.langchain4j.invocation.InvocationContext; -import dev.langchain4j.service.tool.ToolExecutionResult; -import dev.langchain4j.service.tool.ToolExecutor; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Scope; -import javax.annotation.Nullable; -import lombok.extern.slf4j.Slf4j; - -/** A ToolExecutor wrapper that creates a span around tool execution */ -@Slf4j -class TracingToolExecutor implements ToolExecutor { - static final String TYPE_TOOL_JSON = "{\"type\":\"tool\"}"; - - private final ToolExecutor delegate; - private final String toolName; - private final Tracer tracer; - - TracingToolExecutor(ToolExecutor delegate, String toolName, Tracer tracer) { - this.delegate = delegate; - this.toolName = toolName; - this.tracer = tracer; - } - - @Override - public String execute(ToolExecutionRequest request, Object memoryId) { - Span span = tracer.spanBuilder(toolName).startSpan(); - try (Scope ignored = span.makeCurrent()) { - String result = delegate.execute(request, memoryId); - setSpanAttributes(span, request, result); - return result; - } catch (Exception e) { - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.recordException(e); - throw e; - } finally { - span.end(); - } - } - - @Override - public ToolExecutionResult executeWithContext( - ToolExecutionRequest request, InvocationContext context) { - Span span = tracer.spanBuilder(toolName).startSpan(); - try (Scope ignored = span.makeCurrent()) { - ToolExecutionResult result = delegate.executeWithContext(request, context); - setSpanAttributes(span, request, result.resultText()); - return result; - } catch (Exception e) { - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.recordException(e); - throw e; - } finally { - span.end(); - } - } - - private void setSpanAttributes( - Span span, ToolExecutionRequest request, @Nullable String toolCallResult) { - try { - span.setAttribute("braintrust.span_attributes", TYPE_TOOL_JSON); - - String args = request.arguments(); - if (args != null && !args.isEmpty()) { - span.setAttribute("braintrust.input_json", args); - } - if (toolCallResult != null) { - span.setAttribute("braintrust.output", toolCallResult); - } - } catch (Exception e) { - log.debug("Failed to set tool span attributes", e); - } - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/WrappedHttpClient.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/WrappedHttpClient.java deleted file mode 100644 index 4e73f422..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/WrappedHttpClient.java +++ /dev/null @@ -1,249 +0,0 @@ -package dev.braintrust.instrumentation.langchain; - -import static dev.braintrust.json.BraintrustJsonMapper.toJson; - -import com.fasterxml.jackson.databind.JsonNode; -import dev.braintrust.instrumentation.InstrumentationSemConv; -import dev.braintrust.json.BraintrustJsonMapper; -import dev.braintrust.trace.BraintrustTracing; -import dev.langchain4j.exception.HttpException; -import dev.langchain4j.http.client.HttpClient; -import dev.langchain4j.http.client.HttpRequest; -import dev.langchain4j.http.client.SuccessfulHttpResponse; -import dev.langchain4j.http.client.sse.ServerSentEvent; -import dev.langchain4j.http.client.sse.ServerSentEventContext; -import dev.langchain4j.http.client.sse.ServerSentEventListener; -import dev.langchain4j.http.client.sse.ServerSentEventParser; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Scope; -import java.net.URI; -import java.util.Arrays; -import java.util.List; -import java.util.concurrent.atomic.AtomicLong; -import lombok.extern.slf4j.Slf4j; - -@Slf4j -class WrappedHttpClient implements HttpClient { - private final Tracer tracer; - private final HttpClient underlying; - private final BraintrustLangchain.Options options; - - public WrappedHttpClient( - OpenTelemetry openTelemetry, - HttpClient underlying, - BraintrustLangchain.Options options) { - this.tracer = BraintrustTracing.getTracer(openTelemetry); - this.underlying = underlying; - this.options = options; - } - - @Override - public SuccessfulHttpResponse execute(HttpRequest request) - throws HttpException, RuntimeException { - Span span = - tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME) - .setSpanKind(SpanKind.CLIENT) - .startSpan(); - try (Scope scope = span.makeCurrent()) { - tagRequest(span, request); - var response = underlying.execute(request); - // Non-streaming: time_to_first_token is not meaningful - InstrumentationSemConv.tagLLMSpanResponse( - span, options.providerName(), response.body()); - return response; - } catch (Throwable t) { - InstrumentationSemConv.tagLLMSpanResponse(span, t); - throw t; - } finally { - span.end(); - } - } - - @Override - public void execute(HttpRequest request, ServerSentEventListener listener) { - if (listener instanceof WrappedServerSentEventListener) { - // already instrumented - underlying.execute(request, listener); - return; - } - Span span = - tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME) - .setSpanKind(SpanKind.CLIENT) - .startSpan(); - try (Scope ignored = span.makeCurrent()) { - tagRequest(span, request); - underlying.execute( - request, - new WrappedServerSentEventListener(listener, span, options.providerName())); - } catch (Throwable t) { - InstrumentationSemConv.tagLLMSpanResponse(span, t); - span.end(); - throw t; - } - } - - @Override - public void execute( - HttpRequest request, ServerSentEventParser parser, ServerSentEventListener listener) { - if (listener instanceof WrappedServerSentEventListener) { - // already instrumented - underlying.execute(request, parser, listener); - return; - } - Span span = - tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME) - .setSpanKind(SpanKind.CLIENT) - .startSpan(); - try (Scope ignored = span.makeCurrent()) { - tagRequest(span, request); - underlying.execute( - request, - parser, - new WrappedServerSentEventListener(listener, span, options.providerName())); - } catch (Throwable t) { - InstrumentationSemConv.tagLLMSpanResponse(span, t); - span.end(); - throw t; - } - } - - private void tagRequest(Span span, HttpRequest request) { - try { - URI uri = new URI(request.url()); - String baseUrl = uri.getScheme() + "://" + uri.getAuthority(); - List pathSegments = - Arrays.stream(uri.getPath().split("/")).filter(s -> !s.isEmpty()).toList(); - InstrumentationSemConv.tagLLMSpanRequest( - span, options.providerName(), baseUrl, pathSegments, "POST", request.body()); - } catch (Exception e) { - log.debug("Failed to tag request span", e); - } - } - - /** - * Wraps a {@link ServerSentEventListener} to keep the span open for the duration of the stream, - * accumulate SSE chunks, and finalize span tagging via {@link InstrumentationSemConv} when the - * stream closes or errors. - */ - private static class WrappedServerSentEventListener implements ServerSentEventListener { - private final ServerSentEventListener delegate; - private final Span span; - private final String providerName; - private final long startNanos = System.nanoTime(); - private final AtomicLong timeToFirstTokenNanos = new AtomicLong(); - private final StringBuilder contentBuffer = new StringBuilder(); - private String finishReason = null; - private JsonNode usageData = null; - - WrappedServerSentEventListener( - ServerSentEventListener delegate, Span span, String providerName) { - this.delegate = delegate; - this.span = span; - this.providerName = providerName; - } - - @Override - public void onOpen(SuccessfulHttpResponse response) { - try (Scope ignored = span.makeCurrent()) { - delegate.onOpen(response); - } - } - - @Override - public void onEvent(ServerSentEvent event, ServerSentEventContext context) { - try (Scope ignored = span.makeCurrent()) { - accumulateChunk(event.data()); - delegate.onEvent(event, context); - } - } - - @Override - public void onEvent(ServerSentEvent event) { - try (Scope ignored = span.makeCurrent()) { - accumulateChunk(event.data()); - delegate.onEvent(event); - } - } - - @Override - public void onError(Throwable error) { - try (Scope ignored = span.makeCurrent()) { - delegate.onError(error); - } finally { - InstrumentationSemConv.tagLLMSpanResponse(span, error); - span.end(); - } - } - - @Override - public void onClose() { - try (Scope ignored = span.makeCurrent()) { - delegate.onClose(); - } finally { - finalizeSpan(); - span.end(); - } - } - - private void accumulateChunk(String data) { - if (data == null || data.isEmpty() || "[DONE]".equals(data)) return; - try { - if (timeToFirstTokenNanos.get() == 0L) { - // conditional so we don't make unnecessary calls to nano time - timeToFirstTokenNanos.compareAndExchange(0L, System.nanoTime() - startNanos); - } - JsonNode chunk = BraintrustJsonMapper.get().readTree(data); - if (chunk.has("choices") && chunk.get("choices").size() > 0) { - JsonNode choice = chunk.get("choices").get(0); - if (choice.has("delta")) { - JsonNode delta = choice.get("delta"); - if (delta.has("content")) { - contentBuffer.append(delta.get("content").asText()); - } - } - if (choice.has("finish_reason") && !choice.get("finish_reason").isNull()) { - finishReason = choice.get("finish_reason").asText(); - } - } - if (chunk.has("usage") && !chunk.get("usage").isNull()) { - usageData = chunk.get("usage"); - } - } catch (Exception e) { - log.debug("Failed to parse SSE chunk: {}", data, e); - } - } - - private void finalizeSpan() { - try { - // Reconstruct a minimal response JSON that tagLLMSpanResponse already knows - // how to parse: {"choices":[{"index":0,"finish_reason":"...","message":{...}}], - // "usage":{...}} - var root = BraintrustJsonMapper.get().createObjectNode(); - - var choicesArray = BraintrustJsonMapper.get().createArrayNode(); - var choice = BraintrustJsonMapper.get().createObjectNode(); - choice.put("index", 0); - if (finishReason != null) choice.put("finish_reason", finishReason); - var message = BraintrustJsonMapper.get().createObjectNode(); - message.put("role", "assistant"); - message.put("content", contentBuffer.toString()); - choice.set("message", message); - choicesArray.add(choice); - root.set("choices", choicesArray); - - if (usageData != null) { - root.set("usage", usageData); - } - - long ttft = timeToFirstTokenNanos.get(); - InstrumentationSemConv.tagLLMSpanResponse( - span, providerName, toJson(root), ttft == 0L ? null : ttft); - } catch (Exception e) { - log.debug("Failed to finalize streaming span", e); - } - } - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/WrappedHttpClientBuilder.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/WrappedHttpClientBuilder.java deleted file mode 100644 index a78aa39e..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/langchain/WrappedHttpClientBuilder.java +++ /dev/null @@ -1,48 +0,0 @@ -package dev.braintrust.instrumentation.langchain; - -import dev.langchain4j.http.client.HttpClient; -import dev.langchain4j.http.client.HttpClientBuilder; -import io.opentelemetry.api.OpenTelemetry; -import java.time.Duration; - -class WrappedHttpClientBuilder implements HttpClientBuilder { - private final OpenTelemetry openTelemetry; - private final HttpClientBuilder underlying; - private final BraintrustLangchain.Options options; - - public WrappedHttpClientBuilder( - OpenTelemetry openTelemetry, - HttpClientBuilder underlying, - BraintrustLangchain.Options options) { - this.openTelemetry = openTelemetry; - this.underlying = underlying; - this.options = options; - } - - @Override - public Duration connectTimeout() { - return underlying.connectTimeout(); - } - - @Override - public HttpClientBuilder connectTimeout(Duration timeout) { - underlying.connectTimeout(timeout); - return this; - } - - @Override - public Duration readTimeout() { - return underlying.readTimeout(); - } - - @Override - public HttpClientBuilder readTimeout(Duration timeout) { - underlying.readTimeout(timeout); - return this; - } - - @Override - public HttpClient build() { - return new WrappedHttpClient(openTelemetry, underlying.build(), options); - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java deleted file mode 100644 index abee01e3..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java +++ /dev/null @@ -1,165 +0,0 @@ -package dev.braintrust.instrumentation.openai; - -import com.openai.client.OpenAIClient; -import com.openai.core.ClientOptions; -import com.openai.core.ObjectMappers; -import com.openai.core.http.HttpClient; -import com.openai.models.chat.completions.ChatCompletionCreateParams; -import dev.braintrust.prompt.BraintrustPrompt; -import io.opentelemetry.api.OpenTelemetry; -import java.lang.reflect.Field; -import java.lang.reflect.Modifier; -import java.util.HashMap; -import java.util.Map; -import java.util.function.BiConsumer; -import java.util.function.Consumer; -import kotlin.Lazy; -import lombok.SneakyThrows; -import lombok.extern.slf4j.Slf4j; - -/** Braintrust OpenAI client instrumentation. */ -@Slf4j -public class BraintrustOpenAI { - - /** Instrument openai client with braintrust traces */ - public static OpenAIClient wrapOpenAI(OpenTelemetry openTelemetry, OpenAIClient openAIClient) { - try { - instrumentHttpClient(openTelemetry, openAIClient); - return openAIClient; - } catch (Exception e) { - log.error("failed to apply openai instrumentation", e); - return openAIClient; - } - } - - @SneakyThrows - public static ChatCompletionCreateParams buildChatCompletionsPrompt( - BraintrustPrompt prompt, Map parameters) { - var promptMap = new HashMap<>(prompt.getOptions()); - promptMap.put("messages", prompt.renderMessages(parameters)); - var promptJson = ObjectMappers.jsonMapper().writeValueAsString(promptMap); - - var body = - ObjectMappers.jsonMapper() - .readValue(promptJson, ChatCompletionCreateParams.Body.class); - - return ChatCompletionCreateParams.builder() - .body(body) - .additionalBodyProperties(Map.of()) - .build(); - } - - private static void instrumentHttpClient( - OpenTelemetry openTelemetry, OpenAIClient openAIClient) { - forAllFields( - openAIClient, - fieldName -> { - try { - var field = getField(openAIClient, fieldName); - if (field instanceof ClientOptions clientOptions) { - instrumentClientOptions( - openTelemetry, clientOptions, "originalHttpClient"); - instrumentClientOptions(openTelemetry, clientOptions, "httpClient"); - } else { - if (field instanceof Lazy lazyField) { - var resolved = lazyField.getValue(); - forAllFieldsOfType( - resolved, - ClientOptions.class, - (clientOptions, subfieldName) -> - instrumentClientOptions( - openTelemetry, - clientOptions, - subfieldName)); - } else { - forAllFieldsOfType( - field, - ClientOptions.class, - (clientOptions, subfieldName) -> - instrumentClientOptions( - openTelemetry, - clientOptions, - subfieldName)); - } - } - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } - }); - } - - private static void forAllFields(Object object, Consumer consumer) { - if (object == null || consumer == null) return; - - Class clazz = object.getClass(); - while (clazz != null && clazz != Object.class) { - for (Field field : clazz.getDeclaredFields()) { - if (field.isSynthetic()) continue; - if (Modifier.isStatic(field.getModifiers())) continue; - - consumer.accept(field.getName()); - } - clazz = clazz.getSuperclass(); - } - } - - private static void forAllFieldsOfType( - Object object, Class targetClazz, BiConsumer consumer) { - forAllFields( - object, - fieldName -> { - try { - if (targetClazz.isAssignableFrom(object.getClass())) { - consumer.accept(getField(object, fieldName), fieldName); - } - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } - }); - } - - private static void instrumentClientOptions( - OpenTelemetry openTelemetry, ClientOptions clientOptions, String fieldName) { - try { - HttpClient httpClient = getField(clientOptions, fieldName); - if (!(httpClient instanceof TracingHttpClient)) { - setPrivateField( - clientOptions, fieldName, new TracingHttpClient(openTelemetry, httpClient)); - } - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } - } - - @SuppressWarnings("unchecked") - private static T getField(Object obj, String fieldName) - throws ReflectiveOperationException { - Class clazz = obj.getClass(); - while (clazz != null) { - try { - java.lang.reflect.Field field = clazz.getDeclaredField(fieldName); - field.setAccessible(true); - return (T) field.get(obj); - } catch (NoSuchFieldException e) { - clazz = clazz.getSuperclass(); - } - } - throw new NoSuchFieldException(fieldName); - } - - private static void setPrivateField(Object obj, String fieldName, Object value) - throws ReflectiveOperationException { - Class clazz = obj.getClass(); - while (clazz != null) { - try { - java.lang.reflect.Field field = clazz.getDeclaredField(fieldName); - field.setAccessible(true); - field.set(obj, value); - return; - } catch (NoSuchFieldException e) { - clazz = clazz.getSuperclass(); - } - } - throw new NoSuchFieldException(fieldName); - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/openai/TracingHttpClient.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/openai/TracingHttpClient.java deleted file mode 100644 index 37dc40f7..00000000 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/openai/TracingHttpClient.java +++ /dev/null @@ -1,371 +0,0 @@ -package dev.braintrust.instrumentation.openai; - -import com.openai.core.RequestOptions; -import com.openai.core.http.HttpClient; -import com.openai.core.http.HttpRequest; -import com.openai.core.http.HttpRequestBody; -import com.openai.core.http.HttpResponse; -import com.openai.helpers.ChatCompletionAccumulator; -import com.openai.models.chat.completions.ChatCompletionChunk; -import dev.braintrust.instrumentation.InstrumentationSemConv; -import dev.braintrust.json.BraintrustJsonMapper; -import dev.braintrust.trace.BraintrustTracing; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.Tracer; -import java.io.BufferedReader; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.InputStream; -import java.io.InputStreamReader; -import java.io.OutputStream; -import java.nio.charset.StandardCharsets; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; -import lombok.extern.slf4j.Slf4j; -import org.jspecify.annotations.NonNull; - -@Slf4j -public class TracingHttpClient implements HttpClient { - private final Tracer tracer; - private final HttpClient underlying; - - public TracingHttpClient(OpenTelemetry openTelemetry, HttpClient underlying) { - this.tracer = BraintrustTracing.getTracer(openTelemetry); - this.underlying = underlying; - } - - @Override - public void close() { - underlying.close(); - } - - @Override - public @NonNull HttpResponse execute( - @NonNull HttpRequest httpRequest, @NonNull RequestOptions requestOptions) { - var span = tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME).startSpan(); - try (var ignored = span.makeCurrent()) { - // Buffer the request body so we can (a) read its bytes for the span attribute and - // (b) supply a fresh, repeatable body to the underlying client — avoiding any - // one-shot stream consumption issue. - var bufferedRequest = bufferRequestBody(httpRequest); - - String inputJson = - bufferedRequest.body() != null - ? readBodyAsString(bufferedRequest.body()) - : null; - - InstrumentationSemConv.tagLLMSpanRequest( - span, - InstrumentationSemConv.PROVIDER_NAME_OPENAI, - bufferedRequest.baseUrl(), - bufferedRequest.pathSegments(), - bufferedRequest.method().name(), - inputJson); - var response = underlying.execute(bufferedRequest, requestOptions); - // Always tee the response body. onStreamClosed() detects whether the collected - // bytes are SSE or plain JSON and tags the span accordingly. - return new TeeingStreamHttpResponse(response, span); - } catch (Exception e) { - InstrumentationSemConv.tagLLMSpanResponse(span, e); - span.end(); - throw e; - } - } - - @Override - public @NonNull CompletableFuture executeAsync( - @NonNull HttpRequest httpRequest, @NonNull RequestOptions requestOptions) { - var span = tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME).startSpan(); - try { - var bufferedRequest = bufferRequestBody(httpRequest); - String inputJson = - bufferedRequest.body() != null - ? readBodyAsString(bufferedRequest.body()) - : null; - InstrumentationSemConv.tagLLMSpanRequest( - span, - InstrumentationSemConv.PROVIDER_NAME_OPENAI, - bufferedRequest.baseUrl(), - bufferedRequest.pathSegments(), - bufferedRequest.method().name(), - inputJson); - return underlying - .executeAsync(bufferedRequest, requestOptions) - .thenApply( - response -> (HttpResponse) new TeeingStreamHttpResponse(response, span)) - .whenComplete( - (response, t) -> { - if (t != null) { - // this means the future itself failed - InstrumentationSemConv.tagLLMSpanResponse(span, t); - span.end(); - } - }); - } catch (Exception e) { - InstrumentationSemConv.tagLLMSpanResponse(span, e); - span.end(); - throw e; - } - } - - /** - * Captures the request body into an in-memory byte array and returns a new {@link HttpRequest} - * backed by those bytes. The original body stream is consumed exactly once here; the returned - * request uses a {@link HttpRequestBody} that is always {@link HttpRequestBody#repeatable() - * repeatable}, so the underlying client can read it safely (including on retry). - * - *

If the original body is {@code null} or already in-memory (repeatable), the cost is just - * one extra copy of the bytes — acceptable for observability. - */ - private static HttpRequest bufferRequestBody(HttpRequest request) { - HttpRequestBody originalBody = request.body(); - if (originalBody == null) { - return request; - } - var baos = new ByteArrayOutputStream(); - originalBody.writeTo(baos); - byte[] bytes = baos.toByteArray(); - String contentType = originalBody.contentType(); - - HttpRequestBody bufferedBody = - new HttpRequestBody() { - @Override - public void writeTo(OutputStream outputStream) { - try { - outputStream.write(bytes); - } catch (java.io.IOException e) { - throw new RuntimeException(e); - } - } - - @Override - public String contentType() { - return contentType; - } - - @Override - public long contentLength() { - return bytes.length; - } - - @Override - public boolean repeatable() { - return true; - } - - @Override - public void close() {} - }; - - return request.toBuilder().body(bufferedBody).build(); - } - - private static String readBodyAsString(HttpRequestBody body) { - // Body was already buffered by bufferRequestBody, so writeTo is safe to call again. - var baos = new ByteArrayOutputStream((int) Math.max(body.contentLength(), 0)); - body.writeTo(baos); - return baos.toString(StandardCharsets.UTF_8); - } - - /** - * Tags the span from bytes collected by {@link TeeingStreamHttpResponse}. Auto-detects whether - * the bytes are an SSE stream (first non-empty line starts with {@code "data: "}) or a plain - * JSON response, and parses accordingly. - */ - private static void tagSpanFromBuffer(Span span, byte[] bytes, Long timeToFirstTokenNanos) { - if (bytes.length == 0) return; - try { - String firstLine = firstNonEmptyLine(bytes); - if (firstLine != null && firstLine.startsWith("data:")) { - tagSpanFromSseBytes(span, bytes, timeToFirstTokenNanos); - } else { - InstrumentationSemConv.tagLLMSpanResponse( - span, - InstrumentationSemConv.PROVIDER_NAME_OPENAI, - new String(bytes, StandardCharsets.UTF_8)); - } - } catch (Exception e) { - log.error("Could not tag span from response buffer", e); - } - } - - private static String firstNonEmptyLine(byte[] bytes) { - int start = 0; - for (int i = 0; i <= bytes.length; i++) { - if (i == bytes.length || bytes[i] == '\n') { - String line = new String(bytes, start, i - start, StandardCharsets.UTF_8).strip(); - if (!line.isEmpty()) return line; - start = i + 1; - } - } - return null; - } - - /** - * Parses SSE wire bytes, feeds each {@code data:} chunk through {@link - * ChatCompletionAccumulator}, then tags the span with the reassembled output JSON. - */ - private static void tagSpanFromSseBytes( - Span span, byte[] sseBytes, Long timeToFirstTokenNanos) { - try { - var accumulator = ChatCompletionAccumulator.create(); - var reader = - new BufferedReader( - new InputStreamReader( - new ByteArrayInputStream(sseBytes), StandardCharsets.UTF_8)); - String line; - while ((line = reader.readLine()) != null) { - if (!line.startsWith("data:")) continue; - String data = line.substring("data:".length()).strip(); - if (data.isEmpty() || data.equals("[DONE]")) continue; - ChatCompletionChunk chunk = - BraintrustJsonMapper.get().readValue(data, ChatCompletionChunk.class); - accumulator.accumulate(chunk); - } - var chatCompletion = accumulator.chatCompletion(); - InstrumentationSemConv.tagLLMSpanResponse( - span, - InstrumentationSemConv.PROVIDER_NAME_OPENAI, - BraintrustJsonMapper.toJson(chatCompletion), - timeToFirstTokenNanos); - } catch (Exception e) { - log.error("Could not parse SSE buffer to tag streaming span output", e); - } - } - - /** - * {@link HttpResponse} wrapper for streaming (SSE) responses. Its {@link #body()} returns a tee - * {@link InputStream} that copies every byte the caller reads into an in-memory buffer. When - * the stream is fully consumed and {@link #close()} is called, the accumulated bytes are - * available via {@link #collectedBytes()} for span tagging. - */ - private static final class TeeingStreamHttpResponse implements HttpResponse { - private final HttpResponse delegate; - private final Span span; - private final long spanStartNanos = System.nanoTime(); - private final AtomicLong timeToFirstTokenNanos = new AtomicLong(); - private final ByteArrayOutputStream teeBuffer = new ByteArrayOutputStream(); - private final InputStream teeStream; - - TeeingStreamHttpResponse(HttpResponse delegate, Span span) { - this.delegate = delegate; - this.span = span; - this.teeStream = - new TeeInputStream( - delegate.body(), teeBuffer, this::onFirstByte, this::onStreamClosed); - } - - private void onFirstByte() { - timeToFirstTokenNanos.set(System.nanoTime() - spanStartNanos); - } - - /** Called back by {@link TeeInputStream} when the stream is fully drained or closed. */ - private void onStreamClosed() { - try { - // Synchronize on teeBuffer to ensure any write() that was in-flight on a - // concurrent read thread has fully completed before we snapshot the bytes. - byte[] bytes; - synchronized (teeBuffer) { - bytes = teeBuffer.toByteArray(); - } - tagSpanFromBuffer(span, bytes, timeToFirstTokenNanos.get()); - } finally { - span.end(); - } - } - - byte[] collectedBytes() { - return teeBuffer.toByteArray(); - } - - @Override - public int statusCode() { - return delegate.statusCode(); - } - - @Override - public com.openai.core.http.Headers headers() { - return delegate.headers(); - } - - @Override - public InputStream body() { - return teeStream; - } - - @Override - public void close() { - try { - teeStream.close(); // triggers onStreamClosed if not already fired (e.g. abandoned - // stream) - } catch (java.io.IOException ignored) { - } - delegate.close(); - } - } - - /** - * An {@link InputStream} that copies every byte read from {@code source} into {@code sink}, and - * fires {@code onClose} exactly once when the stream reaches EOF or is explicitly closed. - */ - private static final class TeeInputStream extends InputStream { - private final InputStream source; - private final OutputStream sink; - private final Runnable onFirstByte; - private final Runnable onClose; - private final AtomicBoolean firstByteSeen = new AtomicBoolean(false); - private final AtomicBoolean closed = new AtomicBoolean(false); - - TeeInputStream( - InputStream source, OutputStream sink, Runnable onFirstByte, Runnable onClose) { - this.source = source; - this.sink = sink; - this.onFirstByte = onFirstByte; - this.onClose = onClose; - } - - @Override - public int read() throws java.io.IOException { - int b = source.read(); - if (b == -1) { - notifyClosed(); - } else { - notifyFirstByte(); - sink.write(b); - } - return b; - } - - @Override - public int read(byte[] buf, int off, int len) throws java.io.IOException { - int n = source.read(buf, off, len); - if (n == -1) { - notifyClosed(); - } else { - notifyFirstByte(); - sink.write(buf, off, n); - } - return n; - } - - @Override - public void close() throws java.io.IOException { - notifyClosed(); - source.close(); - } - - private void notifyFirstByte() { - if (!firstByteSeen.compareAndExchange(false, true)) { - onFirstByte.run(); - } - } - - private void notifyClosed() { - if (!closed.compareAndExchange(false, true)) { - onClose.run(); - } - } - } -} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java index 1676d4d2..0ff98c89 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java @@ -10,9 +10,9 @@ import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; -import org.jetbrains.annotations.NotNull; /** * Custom span processor that enriches spans with Braintrust-specific attributes. Supports parent @@ -34,7 +34,7 @@ public class BraintrustSpanProcessor implements SpanProcessor { } @Override - public void onStart(@NotNull Context parentContext, ReadWriteSpan span) { + public void onStart(@Nonnull Context parentContext, ReadWriteSpan span) { log.debug("OnStart: span={}, parent={}", span.getName(), parentContext); // Check if span already has a parent attribute diff --git a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropicTest.java b/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropicTest.java deleted file mode 100644 index 78a9f507..00000000 --- a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropicTest.java +++ /dev/null @@ -1,469 +0,0 @@ -package dev.braintrust.instrumentation.anthropic; - -import static org.junit.jupiter.api.Assertions.*; - -import com.anthropic.client.AnthropicClient; -import com.anthropic.client.okhttp.AnthropicOkHttpClient; -import com.anthropic.models.messages.MessageCreateParams; -import com.anthropic.models.messages.Model; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import dev.braintrust.TestHarness; -import io.opentelemetry.api.common.AttributeKey; -import java.util.concurrent.TimeUnit; -import lombok.SneakyThrows; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -public class BraintrustAnthropicTest { - private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); - - private TestHarness testHarness; - - @BeforeEach - void beforeEach() { - testHarness = TestHarness.setup(); - } - - @Test - @SneakyThrows - void testWrapAnthropic() { - AnthropicClient anthropicClient = - AnthropicOkHttpClient.builder() - .baseUrl(testHarness.anthropicBaseUrl()) - .apiKey(testHarness.anthropicApiKey()) - .build(); - - anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); - - var request = - MessageCreateParams.builder() - .model(Model.CLAUDE_3_HAIKU_20240307) - .system("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .maxTokens(50) - .temperature(0.0) - .build(); - - var response = anthropicClient.messages().create(request); - - // Verify the response - assertNotNull(response); - assertNotNull(response.id()); - var contentBlock = response.content().get(0); - assertTrue(contentBlock.isText()); - assertNotNull(contentBlock.asText().text()); - - // Verify spans were exported - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("anthropic.messages.create", span.getName()); - - // Verify span_attributes - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - // Verify metadata - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("anthropic", metadata.get("provider").asText()); - assertTrue( - metadata.get("model").asText().startsWith("claude-3-haiku"), - "model should start with claude-3-haiku"); - - // Verify input - String inputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")); - assertNotNull(inputJson); - JsonNode input = JSON_MAPPER.readTree(inputJson); - assertTrue(input.isArray()); - assertTrue(input.size() > 0); - - // Verify output — full Message object - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - JsonNode outputMessage = JSON_MAPPER.readTree(outputJson); - assertNotNull(outputMessage.get("id")); - assertEquals("message", outputMessage.get("type").asText()); - assertEquals("assistant", outputMessage.get("role").asText()); - assertNotNull(outputMessage.get("content").get(0).get("text")); - assertTrue(outputMessage.get("usage").get("output_tokens").asInt() > 0); - assertTrue(outputMessage.get("usage").get("input_tokens").asInt() > 0); - - // Verify metrics — tokens; non-streaming should NOT have time_to_first_token - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens"), "prompt_tokens should be present"); - assertTrue(metrics.has("completion_tokens"), "completion_tokens should be present"); - assertTrue(metrics.has("tokens"), "tokens should be present"); - assertFalse( - metrics.has("time_to_first_token"), - "time_to_first_token should not be present for non-streaming"); - } - - @Test - @SneakyThrows - void testWrapAnthropicStreaming() { - AnthropicClient anthropicClient = - AnthropicOkHttpClient.builder() - .baseUrl(testHarness.anthropicBaseUrl()) - .apiKey(testHarness.anthropicApiKey()) - .build(); - - anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); - - var request = - MessageCreateParams.builder() - .model(Model.CLAUDE_3_HAIKU_20240307) - .system("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .maxTokens(50) - .temperature(0.0) - .build(); - - StringBuilder fullResponse = new StringBuilder(); - try (var stream = anthropicClient.messages().createStreaming(request)) { - stream.stream() - .forEach( - event -> { - if (event.contentBlockDelta().isPresent()) { - var delta = event.contentBlockDelta().get().delta(); - if (delta.text().isPresent()) { - fullResponse.append(delta.text().get().text()); - } - } - }); - } - - assertFalse(fullResponse.toString().isEmpty()); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("anthropic.messages.create", span.getName()); - - // Verify metadata - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("anthropic", metadata.get("provider").asText()); - - // Verify input - assertNotNull(span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json"))); - - // Verify output — full Message object assembled from SSE stream - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - JsonNode outputMessage = JSON_MAPPER.readTree(outputJson); - assertEquals("assistant", outputMessage.get("role").asText()); - assertFalse( - outputMessage.get("content").get(0).get("text").asText().isEmpty(), - "content should not be empty"); - - // Verify metrics — tokens and time_to_first_token - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens"), "prompt_tokens should be present"); - assertTrue(metrics.has("completion_tokens"), "completion_tokens should be present"); - assertTrue(metrics.has("time_to_first_token"), "time_to_first_token should be present"); - assertTrue( - metrics.get("time_to_first_token").asDouble() >= 0.0, - "time_to_first_token should be non-negative"); - } - - @Test - @SneakyThrows - void testWrapAnthropicAsync() { - AnthropicClient anthropicClient = - AnthropicOkHttpClient.builder() - .baseUrl(testHarness.anthropicBaseUrl()) - .apiKey(testHarness.anthropicApiKey()) - .build(); - - anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); - - var request = - MessageCreateParams.builder() - .model(Model.CLAUDE_3_HAIKU_20240307) - .system("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .maxTokens(50) - .temperature(0.0) - .build(); - - var response = anthropicClient.async().messages().create(request).get(); - - assertNotNull(response); - assertNotNull(response.id()); - var contentBlock = response.content().get(0); - assertTrue(contentBlock.isText()); - assertNotNull(contentBlock.asText().text()); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("anthropic.messages.create", span.getName()); - - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("anthropic", metadata.get("provider").asText()); - - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - JsonNode outputMessage = JSON_MAPPER.readTree(outputJson); - assertEquals("message", outputMessage.get("type").asText()); - assertEquals("assistant", outputMessage.get("role").asText()); - assertNotNull(outputMessage.get("content").get(0).get("text")); - - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens")); - assertTrue(metrics.has("completion_tokens")); - assertTrue(metrics.has("tokens")); - assertFalse( - metrics.has("time_to_first_token"), - "time_to_first_token should not be present for non-streaming"); - } - - @Test - @SneakyThrows - void testWrapAnthropicAsyncStreaming() { - AnthropicClient anthropicClient = - AnthropicOkHttpClient.builder() - .baseUrl(testHarness.anthropicBaseUrl()) - .apiKey(testHarness.anthropicApiKey()) - .build(); - - anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); - - var request = - MessageCreateParams.builder() - .model(Model.CLAUDE_3_HAIKU_20240307) - .system("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .maxTokens(50) - .temperature(0.0) - .build(); - - var fullResponse = new StringBuilder(); - var stream = anthropicClient.async().messages().createStreaming(request); - stream.subscribe( - event -> { - if (event.contentBlockDelta().isPresent()) { - var delta = event.contentBlockDelta().get().delta(); - if (delta.text().isPresent()) { - fullResponse.append(delta.text().get().text()); - } - } - }); - stream.onCompleteFuture().get(30, TimeUnit.SECONDS); - - assertFalse(fullResponse.toString().isEmpty()); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("anthropic.messages.create", span.getName()); - - assertNotNull(span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json"))); - - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - JsonNode outputMessage = JSON_MAPPER.readTree(outputJson); - assertEquals("assistant", outputMessage.get("role").asText()); - assertFalse(outputMessage.get("content").get(0).get("text").asText().isEmpty()); - - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens")); - assertTrue(metrics.has("completion_tokens")); - assertTrue( - metrics.has("time_to_first_token"), - "time_to_first_token should be present for streaming"); - assertTrue(metrics.get("time_to_first_token").asDouble() >= 0.0); - } - - @Test - @SneakyThrows - void testWrapAnthropicBeta() { - AnthropicClient anthropicClient = - AnthropicOkHttpClient.builder() - .baseUrl(testHarness.anthropicBaseUrl()) - .apiKey(testHarness.anthropicApiKey()) - .build(); - - anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); - - var request = - com.anthropic.models.beta.messages.MessageCreateParams.builder() - .model(Model.CLAUDE_3_HAIKU_20240307) - .system("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .maxTokens(50) - .temperature(0.0) - .build(); - - var response = anthropicClient.beta().messages().create(request); - - assertNotNull(response); - assertNotNull(response.id()); - var contentBlock = response.content().get(0); - assertTrue(contentBlock.isText()); - assertNotNull(contentBlock.asText().text()); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("anthropic.messages.create", span.getName()); - - // Verify span_attributes - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - // Verify metadata - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("anthropic", metadata.get("provider").asText()); - assertTrue( - metadata.get("model").asText().startsWith("claude-3-haiku"), - "model should start with claude-3-haiku"); - - // Verify input - String inputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")); - assertNotNull(inputJson); - JsonNode input = JSON_MAPPER.readTree(inputJson); - assertTrue(input.isArray()); - assertTrue(input.size() > 0); - - // Verify output — full BetaMessage object - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - JsonNode outputMessage = JSON_MAPPER.readTree(outputJson); - assertNotNull(outputMessage.get("id")); - assertEquals("message", outputMessage.get("type").asText()); - assertEquals("assistant", outputMessage.get("role").asText()); - assertNotNull(outputMessage.get("content").get(0).get("text")); - assertTrue(outputMessage.get("usage").get("output_tokens").asInt() > 0); - assertTrue(outputMessage.get("usage").get("input_tokens").asInt() > 0); - - // Verify metrics — tokens; non-streaming should NOT have time_to_first_token - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens"), "prompt_tokens should be present"); - assertTrue(metrics.has("completion_tokens"), "completion_tokens should be present"); - assertTrue(metrics.has("tokens"), "tokens should be present"); - assertFalse( - metrics.has("time_to_first_token"), - "time_to_first_token should not be present for non-streaming"); - } - - @Test - @SneakyThrows - void testWrapAnthropicBetaStreaming() { - AnthropicClient anthropicClient = - AnthropicOkHttpClient.builder() - .baseUrl(testHarness.anthropicBaseUrl()) - .apiKey(testHarness.anthropicApiKey()) - .build(); - - anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); - - var request = - com.anthropic.models.beta.messages.MessageCreateParams.builder() - .model(Model.CLAUDE_3_HAIKU_20240307) - .system("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .maxTokens(50) - .temperature(0.0) - .build(); - - StringBuilder fullResponse = new StringBuilder(); - try (var stream = anthropicClient.beta().messages().createStreaming(request)) { - stream.stream() - .forEach( - event -> { - if (event.contentBlockDelta().isPresent()) { - var delta = event.contentBlockDelta().get(); - if (delta.delta().text().isPresent()) { - fullResponse.append(delta.delta().text().get().text()); - } - } - }); - } - - assertFalse(fullResponse.toString().isEmpty()); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("anthropic.messages.create", span.getName()); - - // Verify metadata - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("anthropic", metadata.get("provider").asText()); - - // Verify input - assertNotNull(span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json"))); - - // Verify output — full BetaMessage object assembled from SSE stream - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - JsonNode outputMessage = JSON_MAPPER.readTree(outputJson); - assertEquals("assistant", outputMessage.get("role").asText()); - assertFalse( - outputMessage.get("content").get(0).get("text").asText().isEmpty(), - "content should not be empty"); - - // Verify metrics — tokens and time_to_first_token - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens"), "prompt_tokens should be present"); - assertTrue(metrics.has("completion_tokens"), "completion_tokens should be present"); - assertTrue(metrics.has("time_to_first_token"), "time_to_first_token should be present"); - assertTrue( - metrics.get("time_to_first_token").asDouble() >= 0.0, - "time_to_first_token should be non-negative"); - } -} diff --git a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/genai/BraintrustGenAITest.java b/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/genai/BraintrustGenAITest.java deleted file mode 100644 index e25a74c5..00000000 --- a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/genai/BraintrustGenAITest.java +++ /dev/null @@ -1,183 +0,0 @@ -package dev.braintrust.instrumentation.genai; - -import static org.junit.jupiter.api.Assertions.*; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.genai.Client; -import com.google.genai.types.GenerateContentConfig; -import com.google.genai.types.HttpOptions; -import dev.braintrust.TestHarness; -import io.opentelemetry.api.common.AttributeKey; -import lombok.SneakyThrows; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -public class BraintrustGenAITest { - private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); - - private TestHarness testHarness; - - @BeforeEach - void beforeEach() { - testHarness = TestHarness.setup(); - } - - @Test - @SneakyThrows - void testWrapGemini() { - // Create Gemini client using VCR - HttpOptions httpOptions = - HttpOptions.builder().baseUrl(testHarness.googleBaseUrl()).build(); - - // Wrap with Braintrust instrumentation - var geminiClient = - BraintrustGenAI.wrap( - testHarness.openTelemetry(), - new Client.Builder() - .apiKey(testHarness.googleApiKey()) - .httpOptions(httpOptions)); - - var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build(); - - var response = - geminiClient.models.generateContent( - "gemini-2.0-flash-lite", "What is the capital of France?", config); - - // Verify the response - assertNotNull(response); - assertNotNull(response.text()); - - // Verify spans were exported - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size(), "Expected exactly 1 span to be created"); - var span = spans.get(0); - - // Verify span name matches the operation - assertEquals("generate_content", span.getName()); - - // Verify braintrust.metadata contains provider and model - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson, "braintrust.metadata should be set"); - var metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("gemini", metadata.get("provider").asText()); - assertEquals("gemini-2.0-flash-lite", metadata.get("model").asText()); - assertEquals(0.0, metadata.get("temperature").asDouble()); - assertEquals(50, metadata.get("maxOutputTokens").asInt()); - - // Verify braintrust.metrics contains token counts - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson, "braintrust.metrics should be set"); - var metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.get("prompt_tokens").asInt() > 0, "prompt_tokens should be > 0"); - assertTrue(metrics.get("completion_tokens").asInt() > 0, "completion_tokens should be > 0"); - assertTrue(metrics.get("tokens").asInt() > 0, "tokens should be > 0"); - - // Verify braintrust.span_attributes marks this as an LLM span - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson, "braintrust.span_attributes should be set"); - var spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - // Verify braintrust.input_json contains the request - String inputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")); - assertNotNull(inputJson, "braintrust.input_json should be set"); - var input = JSON_MAPPER.readTree(inputJson); - assertEquals("gemini-2.0-flash-lite", input.get("model").asText()); - assertTrue(input.has("contents"), "input should have contents"); - assertTrue(input.has("config"), "input should have config"); - - // Verify braintrust.output_json contains the response - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson, "braintrust.output_json should be set"); - var output = JSON_MAPPER.readTree(outputJson); - assertTrue(output.has("candidates"), "output should have candidates"); - assertNotNull(output.get("candidates").get(0).get("finishReason")); - assertNotNull( - output.get("candidates").get(0).get("content").get("parts").get(0).get("text")); - } - - @Test - @SneakyThrows - void testWrapGeminiAsync() { - // Create Gemini client using VCR - HttpOptions httpOptions = - HttpOptions.builder().baseUrl(testHarness.googleBaseUrl()).build(); - - // Wrap with Braintrust instrumentation - var geminiClient = - BraintrustGenAI.wrap( - testHarness.openTelemetry(), - new Client.Builder() - .apiKey(testHarness.googleApiKey()) - .httpOptions(httpOptions)); - - var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build(); - - // Call async method and wait for completion - var responseFuture = - geminiClient.async.models.generateContent( - "gemini-2.0-flash-lite", "What is the capital of Germany?", config); - - var response = responseFuture.get(); // Wait for completion - - // Verify the response - assertNotNull(response); - assertNotNull(response.text()); - - // Verify spans were exported - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size(), "Expected exactly 1 span to be created"); - var span = spans.get(0); - - // Verify span name matches the operation - assertEquals("generate_content", span.getName()); - - // Verify braintrust.metadata contains provider and model - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson, "braintrust.metadata should be set"); - var metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("gemini", metadata.get("provider").asText()); - assertEquals("gemini-2.0-flash-lite", metadata.get("model").asText()); - assertEquals(0.0, metadata.get("temperature").asDouble()); - assertEquals(50, metadata.get("maxOutputTokens").asInt()); - - // Verify braintrust.metrics contains token counts - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson, "braintrust.metrics should be set"); - var metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.get("prompt_tokens").asInt() > 0, "prompt_tokens should be > 0"); - assertTrue(metrics.get("completion_tokens").asInt() > 0, "completion_tokens should be > 0"); - assertTrue(metrics.get("tokens").asInt() > 0, "tokens should be > 0"); - - // Verify braintrust.span_attributes marks this as an LLM span - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson, "braintrust.span_attributes should be set"); - var spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - // Verify braintrust.input_json contains the request - String inputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")); - assertNotNull(inputJson, "braintrust.input_json should be set"); - var input = JSON_MAPPER.readTree(inputJson); - assertEquals("gemini-2.0-flash-lite", input.get("model").asText()); - assertTrue(input.has("contents"), "input should have contents"); - assertTrue(input.has("config"), "input should have config"); - - // Verify braintrust.output_json contains the response - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson, "braintrust.output_json should be set"); - var output = JSON_MAPPER.readTree(outputJson); - assertTrue(output.has("candidates"), "output should have candidates"); - assertNotNull(output.get("candidates").get(0).get("finishReason")); - assertNotNull( - output.get("candidates").get(0).get("content").get("parts").get(0).get("text")); - } -} diff --git a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/langchain/BraintrustLangchainTest.java b/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/langchain/BraintrustLangchainTest.java deleted file mode 100644 index 99e82fef..00000000 --- a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/langchain/BraintrustLangchainTest.java +++ /dev/null @@ -1,353 +0,0 @@ -package dev.braintrust.instrumentation.langchain; - -import static org.junit.jupiter.api.Assertions.*; - -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import dev.braintrust.TestHarness; -import dev.langchain4j.agent.tool.Tool; -import dev.langchain4j.data.message.UserMessage; -import dev.langchain4j.model.chat.ChatModel; -import dev.langchain4j.model.chat.StreamingChatModel; -import dev.langchain4j.model.chat.response.ChatResponse; -import dev.langchain4j.model.chat.response.StreamingChatResponseHandler; -import dev.langchain4j.model.openai.OpenAiChatModel; -import dev.langchain4j.model.openai.OpenAiStreamingChatModel; -import dev.langchain4j.service.AiServices; -import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.sdk.trace.data.SpanData; -import java.util.List; -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicInteger; -import lombok.SneakyThrows; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -public class BraintrustLangchainTest { - - private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); - - private TestHarness testHarness; - - @BeforeEach - void beforeEach() { - testHarness = TestHarness.setup(); - } - - @Test - @SneakyThrows - void typeToolJsonCorrect() { - assertEquals( - JSON_MAPPER.writeValueAsString(Map.of("type", "tool")), - TracingToolExecutor.TYPE_TOOL_JSON); - } - - @Test - @SneakyThrows - void testSyncChatCompletion() { - ChatModel model = - BraintrustLangchain.wrap( - testHarness.openTelemetry(), - OpenAiChatModel.builder() - .apiKey(testHarness.openAiApiKey()) - .baseUrl(testHarness.openAiBaseUrl()) - .modelName("gpt-4o-mini") - .temperature(0.0)); - - // Execute chat request - var message = UserMessage.from("What is the capital of France?"); - var response = model.chat(message); - - // Verify the response - assertNotNull(response); - assertNotNull(response.aiMessage().text()); - - // Verify spans were exported - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size(), "Expected one span for sync chat completion"); - var span = spans.get(0); - - // Verify span name - assertEquals("Chat Completion", span.getName(), "Span name should be 'Chat Completion'"); - - // Verify span attributes - var attributes = span.getAttributes(); - var braintrustSpanAttributesJson = - attributes.get(AttributeKey.stringKey("braintrust.span_attributes")); - - // Verify span type - JsonNode spanAttributes = JSON_MAPPER.readTree(braintrustSpanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText(), "Span type should be 'llm'"); - - // Verify metadata - String metadataJson = attributes.get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson, "Metadata should be present"); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("openai", metadata.get("provider").asText(), "Provider should be 'openai'"); - assertEquals( - "gpt-4o-mini", metadata.get("model").asText(), "Model should be 'gpt-4o-mini'"); - - // Verify metrics - String metricsJson = attributes.get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson, "Metrics should be present"); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.get("tokens").asLong() > 0, "Total tokens should be > 0"); - assertTrue(metrics.get("prompt_tokens").asLong() > 0, "Prompt tokens should be > 0"); - assertTrue( - metrics.get("completion_tokens").asLong() > 0, "Completion tokens should be > 0"); - assertFalse( - metrics.has("time_to_first_token"), - "time_to_first_token should not be present for non-streaming"); - - // Verify input - String inputJson = attributes.get(AttributeKey.stringKey("braintrust.input_json")); - assertNotNull(inputJson, "Input should be present"); - JsonNode input = JSON_MAPPER.readTree(inputJson); - assertTrue(input.isArray(), "Input should be an array"); - assertTrue(input.size() > 0, "Input array should not be empty"); - assertTrue( - input.get(0).get("content").asText().contains("What is the capital of France"), - "Input should contain the user message"); - - // Verify output - String outputJson = attributes.get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson, "Output should be present"); - JsonNode output = JSON_MAPPER.readTree(outputJson); - assertTrue(output.isArray(), "Output should be an array"); - assertTrue(output.size() > 0, "Output array should not be empty"); - assertNotNull( - output.get(0).get("message").get("content"), - "Output should contain assistant response content"); - } - - @Test - @SneakyThrows - void testStreamingChatCompletion() { - var tracer = testHarness.openTelemetry().getTracer("test-tracer"); - - StreamingChatModel model = - BraintrustLangchain.wrap( - testHarness.openTelemetry(), - OpenAiStreamingChatModel.builder() - .apiKey(testHarness.openAiApiKey()) - .baseUrl(testHarness.openAiBaseUrl()) - .modelName("gpt-4o-mini") - .temperature(0.0)); - - // Execute streaming chat request - var future = new CompletableFuture(); - var responseBuilder = new StringBuilder(); - var callbackCount = new AtomicInteger(0); - - model.chat( - "What is the capital of France?", - new StreamingChatResponseHandler() { - @Override - public void onPartialResponse(String token) { - // Create a child span during the callback to verify parenting - Span childSpan = - tracer.spanBuilder( - "callback-span-" + callbackCount.incrementAndGet()) - .startSpan(); - childSpan.end(); - responseBuilder.append(token); - } - - @Override - public void onCompleteResponse(ChatResponse response) { - future.complete(response); - } - - @Override - public void onError(Throwable error) { - future.completeExceptionally(error); - } - }); - - // Wait for completion - var response = future.get(); - - // Verify the response - assertNotNull(response); - assertFalse(responseBuilder.toString().isEmpty(), "Response should not be empty"); - - // We expect at least 2 spans: 1 LLM span + at least 1 callback span - int expectedMinSpans = 1 + callbackCount.get(); - var spans = testHarness.awaitExportedSpans(expectedMinSpans); - assertTrue( - spans.size() >= expectedMinSpans, - "Expected at least " + expectedMinSpans + " spans, got " + spans.size()); - - // Find the LLM span and callback spans - SpanData llmSpan = null; - List callbackSpans = new java.util.ArrayList<>(); - - for (var span : spans) { - if (span.getName().equals("Chat Completion")) { - llmSpan = span; - } else if (span.getName().startsWith("callback-span-")) { - callbackSpans.add(span); - } - } - - assertNotNull(llmSpan, "Should have an LLM span named 'Chat Completion'"); - assertEquals( - callbackCount.get(), - callbackSpans.size(), - "Should have one callback span per onPartialResponse invocation"); - - // Verify all callback spans are parented under the LLM span - String llmSpanId = llmSpan.getSpanId(); - for (var callbackSpan : callbackSpans) { - assertEquals( - llmSpanId, - callbackSpan.getParentSpanId(), - "Callback span '" - + callbackSpan.getName() - + "' should be parented under LLM span"); - } - - // Verify LLM span attributes - var attributes = llmSpan.getAttributes(); - - var braintrustSpanAttributesJson = - attributes.get(AttributeKey.stringKey("braintrust.span_attributes")); - - // Verify span type - JsonNode spanAttributes = JSON_MAPPER.readTree(braintrustSpanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText(), "Span type should be 'llm'"); - - // Verify metadata - String metadataJson = attributes.get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson, "Metadata should be present"); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("openai", metadata.get("provider").asText(), "Provider should be 'openai'"); - assertEquals( - "gpt-4o-mini", metadata.get("model").asText(), "Model should be 'gpt-4o-mini'"); - - // Verify metrics for streaming - String metricsJson = attributes.get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson, "Metrics should be present"); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.get("tokens").asLong() > 0, "Total tokens should be > 0"); - assertTrue(metrics.get("prompt_tokens").asLong() > 0, "Prompt tokens should be > 0"); - assertTrue( - metrics.get("completion_tokens").asLong() > 0, "Completion tokens should be > 0"); - assertTrue( - metrics.has("time_to_first_token"), - "Metrics should contain time_to_first_token for streaming"); - assertTrue( - metrics.get("time_to_first_token").isNumber(), - "time_to_first_token should be a number"); - - // Verify input - String inputJson = attributes.get(AttributeKey.stringKey("braintrust.input_json")); - assertNotNull(inputJson, "Input should be present"); - JsonNode input = JSON_MAPPER.readTree(inputJson); - assertTrue(input.isArray(), "Input should be an array"); - assertTrue(input.size() > 0, "Input array should not be empty"); - assertTrue( - input.get(0).get("content").asText().contains("What is the capital of France"), - "Input should contain the user message"); - - // Verify output (streaming reconstructs the output) - String outputJson = attributes.get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson, "Output should be present"); - JsonNode output = JSON_MAPPER.readTree(outputJson); - assertTrue(output.isArray(), "Output should be an array"); - assertTrue(output.size() > 0, "Output array should not be empty"); - JsonNode choice = output.get(0); - assertNotNull( - choice.get("message").get("content"), - "Output should contain the complete streamed response"); - assertNotNull(choice.get("finish_reason"), "Output should have finish_reason"); - } - - @Test - @SneakyThrows - void testAiServicesWithTools() { - Assistant assistant = - BraintrustLangchain.wrap( - testHarness.openTelemetry(), - AiServices.builder(Assistant.class) - .chatModel( - OpenAiChatModel.builder() - .apiKey(testHarness.openAiApiKey()) - .baseUrl(testHarness.openAiBaseUrl()) - .modelName("gpt-4o-mini") - .temperature(0.0) - .build()) - .tools(new WeatherTools()) - .executeToolsConcurrently()); - - // This should trigger two (concurrent) tool calls - var response = assistant.chat("is it hotter in Paris or New York right now?"); - - // Verify the response - assertNotNull(response); - - // Verify spans were exported - should have at least: - // - one AI service method span ("chat") - // - at least two LLM calls (one to request the tool calls, and another to analyze) - // - at least two tool call spans - var spans = testHarness.awaitExportedSpans(3); - assertTrue(spans.size() >= 3, "Expected at least 3 spans for AI Services with tools"); - - // Verify we have the expected span types - int numServiceMethodSpans = 0; - int numLLMSpans = 0; - int numToolCallSpans = 0; - - for (var span : spans) { - String spanName = span.getName(); - var attributes = span.getAttributes(); - - if (spanName.equals("chat")) { - numServiceMethodSpans++; - } else if (spanName.equals("Chat Completion")) { - numLLMSpans++; - // Verify LLM span has proper attributes - var spanAttributesJson = - attributes.get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson, "LLM span should have span_attributes"); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals( - "llm", spanAttributes.get("type").asText(), "Span type should be 'llm'"); - } else if (spanName.equals("getWeather")) { - numToolCallSpans++; - // Verify tool span has proper attributes - var spanAttributesJson = - attributes.get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson, "Tool span should have span_attributes"); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals( - "tool", spanAttributes.get("type").asText(), "Span type should be 'tool'"); - } - } - assertEquals(1, numServiceMethodSpans, "should be exactly one service call"); - assertTrue(numLLMSpans >= 2, "should be at least two llm spans"); - assertTrue(numToolCallSpans >= 2, "should be at least two tool call spans"); - } - - /** AI Service interface for the assistant */ - interface Assistant { - String chat(String userMessage); - } - - /** Example tool class with weather-related methods */ - public static class WeatherTools { - @Tool("Get current weather for a location") - public String getWeather(String location) { - return String.format("The weather in %s is sunny with 72°F temperature.", location); - } - - @Tool("Get weather forecast for next N days") - public String getForecast(String location, int days) { - return String.format( - "The %d-day forecast for %s: Mostly sunny with temperatures between 65-75°F.", - days, location); - } - } -} diff --git a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/langchain/TestTools.java b/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/langchain/TestTools.java deleted file mode 100644 index b289611e..00000000 --- a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/langchain/TestTools.java +++ /dev/null @@ -1,22 +0,0 @@ -package dev.braintrust.instrumentation.langchain; - -import dev.langchain4j.agent.tool.Tool; - -public class TestTools { - - @Tool("Get weather for a location") - public String getWeather(String location) { - return String.format( - "{\"location\":\"%s\",\"temperature\":72,\"condition\":\"sunny\"}", location); - } - - @Tool("Calculate sum") - public int calculateSum(int a, int b) { - return a + b; - } - - @Tool("Tool that throws exception") - public String throwError() { - throw new RuntimeException("Intentional error for testing"); - } -} diff --git a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java b/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java deleted file mode 100644 index 9401241c..00000000 --- a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java +++ /dev/null @@ -1,401 +0,0 @@ -package dev.braintrust.instrumentation.openai; - -import static org.junit.jupiter.api.Assertions.*; - -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.openai.client.OpenAIClient; -import com.openai.client.okhttp.OpenAIOkHttpClient; -import com.openai.models.ChatModel; -import com.openai.models.Reasoning; -import com.openai.models.ReasoningEffort; -import com.openai.models.chat.completions.*; -import com.openai.models.responses.EasyInputMessage; -import com.openai.models.responses.ResponseCreateParams; -import com.openai.models.responses.ResponseInputItem; -import dev.braintrust.TestHarness; -import io.opentelemetry.api.common.AttributeKey; -import java.util.List; -import java.util.concurrent.TimeUnit; -import lombok.SneakyThrows; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -public class BraintrustOpenAITest { - private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); - - private TestHarness testHarness; - - @BeforeEach - void beforeEach() { - testHarness = TestHarness.setup(); - } - - @Test - @SneakyThrows - void testWrapOpenAi() { - // Create OpenAI client using TestHarness configuration - // TestHarness automatically provides the correct base URL and API key - OpenAIClient openAIClient = - OpenAIOkHttpClient.builder() - .baseUrl(testHarness.openAiBaseUrl()) - .apiKey(testHarness.openAiApiKey()) - .build(); - - // Wrap with Braintrust instrumentation - openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); - - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O_MINI) - .addSystemMessage("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .temperature(0.0) - .build(); - - var response = openAIClient.chat().completions().create(request); - - // Verify the response (same assertions work for both modes) - assertNotNull(response); - assertNotNull(response.id()); - assertTrue(response.choices().get(0).message().content().isPresent()); - String content = response.choices().get(0).message().content().get(); - assertTrue(content.toLowerCase().contains("paris"), "Response should mention Paris"); - - // Verify spans were exported - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - // Verify span name matches other SDKs - assertEquals("Chat Completion", span.getName()); - - // Verify span_attributes JSON - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - // Verify braintrust.metadata JSON - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("openai", metadata.get("provider").asText()); - assertTrue( - metadata.get("model").asText().startsWith("gpt-4o-mini"), - "model should start with gpt-4o-mini"); - - // Verify braintrust.metrics JSON (tokens) - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens"), "prompt_tokens should be present"); - assertTrue( - metrics.get("prompt_tokens").asInt() >= 0, "prompt_tokens should be non-negative"); - assertTrue(metrics.has("completion_tokens"), "completion_tokens should be present"); - assertTrue( - metrics.get("completion_tokens").asInt() >= 0, - "completion_tokens should be non-negative"); - assertTrue(metrics.has("tokens"), "tokens should be present"); - assertTrue(metrics.get("tokens").asInt() >= 0, "tokens should be non-negative"); - assertFalse( - metrics.has("time_to_first_token"), - "time_to_first_token should not be present for non-streaming"); - - // Verify output (choices array) - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - var outputChoices = JSON_MAPPER.readTree(outputJson); - assertEquals(1, outputChoices.size()); - var choice = outputChoices.get(0); - assertEquals("assistant", choice.get("message").get("role").asText()); - assertNotNull(choice.get("finish_reason")); - } - - @Test - @SneakyThrows - void testWrapOpenAiStreaming() { - OpenAIClient openAIClient = - OpenAIOkHttpClient.builder() - .baseUrl(testHarness.openAiBaseUrl()) - .apiKey(testHarness.openAiApiKey()) - .build(); - - openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); - - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O_MINI) - .addSystemMessage("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .temperature(0.0) - .streamOptions( - ChatCompletionStreamOptions.builder().includeUsage(true).build()) - .build(); - - // Consume the stream - StringBuilder fullResponse = new StringBuilder(); - try (var stream = openAIClient.chat().completions().createStreaming(request)) { - stream.stream() - .forEach( - chunk -> { - if (!chunk.choices().isEmpty()) { - chunk.choices() - .get(0) - .delta() - .content() - .ifPresent(fullResponse::append); - } - }); - } - - // Verify the response - assertFalse(fullResponse.isEmpty(), "Should have received streaming response"); - assertTrue( - fullResponse.toString().toLowerCase().contains("paris"), - "Response should mention Paris"); - - // Verify spans - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("Chat Completion", span.getName()); - - // Verify braintrust.metadata has provider=openai - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("openai", metadata.get("provider").asText()); - - // Verify time_to_first_token is in braintrust.metrics JSON - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("time_to_first_token"), "time_to_first_token should be present"); - assertTrue( - metrics.get("time_to_first_token").asDouble() >= 0.0, - "time_to_first_token should be non-negative"); - } - - @Test - @SneakyThrows - void testWrapOpenAiResponses() { - OpenAIClient openAIClient = - OpenAIOkHttpClient.builder() - .baseUrl(testHarness.openAiBaseUrl()) - .apiKey(testHarness.openAiApiKey()) - .build(); - - openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); - - var inputMsg = - EasyInputMessage.builder() - .role(EasyInputMessage.Role.USER) - .content("What is the capital of France? Reply in one word.") - .build(); - - var request = - ResponseCreateParams.builder() - .model("o4-mini") - .reasoning( - Reasoning.builder() - .effort(ReasoningEffort.LOW) - .summary(Reasoning.Summary.AUTO) - .build()) - .inputOfResponse(List.of(ResponseInputItem.ofEasyInputMessage(inputMsg))) - .build(); - - var response = openAIClient.responses().create(request); - - assertNotNull(response); - assertNotNull(response.id()); - assertFalse(response.output().isEmpty(), "Response should have output items"); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - // Span name for /v1/responses endpoint - assertEquals("responses", span.getName()); - - // span_attributes type=llm - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - // metadata: provider and model - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("openai", metadata.get("provider").asText()); - assertTrue(metadata.get("model").asText().startsWith("o4-mini")); - - // input_json: captured from "input" array - String inputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")); - assertNotNull(inputJson, "braintrust.input_json should be set"); - JsonNode inputItems = JSON_MAPPER.readTree(inputJson); - assertTrue(inputItems.isArray() && inputItems.size() > 0); - assertEquals("user", inputItems.get(0).get("role").asText()); - - // output_json: captured from "output" array - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson, "braintrust.output_json should be set"); - JsonNode outputItems = JSON_MAPPER.readTree(outputJson); - assertTrue(outputItems.isArray() && outputItems.size() > 0); - - // metrics: tokens from Responses API usage fields - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens"), "prompt_tokens should be present"); - assertTrue(metrics.get("prompt_tokens").asInt() >= 0); - assertTrue(metrics.has("completion_tokens"), "completion_tokens should be present"); - assertTrue(metrics.get("completion_tokens").asInt() >= 0); - assertTrue(metrics.has("tokens"), "tokens should be present"); - assertTrue(metrics.get("tokens").asInt() >= 0); - assertTrue( - metrics.has("completion_reasoning_tokens"), - "completion_reasoning_tokens should be present"); - assertTrue(metrics.get("completion_reasoning_tokens").asInt() >= 0); - } - - @Test - @SneakyThrows - void testWrapOpenAiAsync() { - OpenAIClient openAIClient = - OpenAIOkHttpClient.builder() - .baseUrl(testHarness.openAiBaseUrl()) - .apiKey(testHarness.openAiApiKey()) - .build(); - - openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); - - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O_MINI) - .addSystemMessage("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .temperature(0.0) - .build(); - - var response = openAIClient.async().chat().completions().create(request).get(); - - assertNotNull(response); - assertNotNull(response.id()); - assertTrue(response.choices().get(0).message().content().isPresent()); - String content = response.choices().get(0).message().content().get(); - assertTrue(content.toLowerCase().contains("paris"), "Response should mention Paris"); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("Chat Completion", span.getName()); - - String spanAttributesJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); - assertNotNull(spanAttributesJson); - JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); - assertEquals("llm", spanAttributes.get("type").asText()); - - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("openai", metadata.get("provider").asText()); - - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - var outputChoices = JSON_MAPPER.readTree(outputJson); - assertEquals(1, outputChoices.size()); - assertEquals("assistant", outputChoices.get(0).get("message").get("role").asText()); - - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens")); - assertTrue(metrics.has("completion_tokens")); - assertTrue(metrics.has("tokens")); - assertFalse( - metrics.has("time_to_first_token"), - "time_to_first_token should not be present for non-streaming"); - } - - @Test - @SneakyThrows - void testWrapOpenAiAsyncStreaming() { - OpenAIClient openAIClient = - OpenAIOkHttpClient.builder() - .baseUrl(testHarness.openAiBaseUrl()) - .apiKey(testHarness.openAiApiKey()) - .build(); - - openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); - - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O_MINI) - .addSystemMessage("You are a helpful assistant") - .addUserMessage("What is the capital of France?") - .temperature(0.0) - .streamOptions( - ChatCompletionStreamOptions.builder().includeUsage(true).build()) - .build(); - - var fullResponse = new StringBuilder(); - var stream = openAIClient.async().chat().completions().createStreaming(request); - stream.subscribe( - chunk -> { - if (!chunk.choices().isEmpty()) { - chunk.choices().get(0).delta().content().ifPresent(fullResponse::append); - } - }); - stream.onCompleteFuture().get(30, TimeUnit.SECONDS); - - assertFalse(fullResponse.toString().isEmpty()); - assertTrue(fullResponse.toString().toLowerCase().contains("paris")); - - var spans = testHarness.awaitExportedSpans(); - assertEquals(1, spans.size()); - var span = spans.get(0); - - assertEquals("Chat Completion", span.getName()); - - String metadataJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata")); - assertNotNull(metadataJson); - JsonNode metadata = JSON_MAPPER.readTree(metadataJson); - assertEquals("openai", metadata.get("provider").asText()); - - assertNotNull(span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json"))); - - String outputJson = - span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); - assertNotNull(outputJson); - var outputChoices = JSON_MAPPER.readTree(outputJson); - assertEquals(1, outputChoices.size()); - assertEquals("assistant", outputChoices.get(0).get("message").get("role").asText()); - - String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); - assertNotNull(metricsJson); - JsonNode metrics = JSON_MAPPER.readTree(metricsJson); - assertTrue(metrics.has("prompt_tokens")); - assertTrue(metrics.has("completion_tokens")); - assertTrue(metrics.has("tokens")); - assertTrue( - metrics.has("time_to_first_token"), - "time_to_first_token should be present for streaming"); - assertTrue(metrics.get("time_to_first_token").asDouble() >= 0.0); - } -} diff --git a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg b/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg deleted file mode 100644 index f0164667..00000000 Binary files a/braintrust-sdk/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg and /dev/null differ diff --git a/examples/build.gradle b/examples/build.gradle index fe166f90..8bbe1a4d 100644 --- a/examples/build.gradle +++ b/examples/build.gradle @@ -20,6 +20,11 @@ def braintrustLogLevel = System.getenv('BRAINTRUST_LOG_LEVEL') ?: 'info' dependencies { implementation project(':braintrust-sdk') + implementation project(':braintrust-sdk:instrumentation:openai_2_8_0') + implementation project(':braintrust-sdk:instrumentation:anthropic_2_2_0') + implementation project(':braintrust-sdk:instrumentation:genai_1_18_0') + implementation project(':braintrust-sdk:instrumentation:langchain_1_8_0') + implementation project(':braintrust-sdk:instrumentation:springai_1_0_0') runtimeOnly "org.slf4j:slf4j-simple:${slf4jVersion}" // To run otel examples implementation "io.opentelemetry:opentelemetry-exporter-otlp:${otelVersion}" diff --git a/settings.gradle b/settings.gradle index fae6cad9..97cccff0 100644 --- a/settings.gradle +++ b/settings.gradle @@ -11,11 +11,11 @@ include 'braintrust-java-agent' include 'braintrust-java-agent:bootstrap' include 'braintrust-java-agent:internal' include 'braintrust-java-agent:instrumenter' -include 'braintrust-java-agent:instrumentation:openai_2_8_0' -include 'braintrust-java-agent:instrumentation:anthropic_2_2_0' -include 'braintrust-java-agent:instrumentation:genai_1_18_0' -include 'braintrust-java-agent:instrumentation:langchain_1_8_0' -include 'braintrust-java-agent:instrumentation:springai_1_0_0' +include 'braintrust-sdk:instrumentation:openai_2_8_0' +include 'braintrust-sdk:instrumentation:anthropic_2_2_0' +include 'braintrust-sdk:instrumentation:genai_1_18_0' +include 'braintrust-sdk:instrumentation:langchain_1_8_0' +include 'braintrust-sdk:instrumentation:springai_1_0_0' include 'braintrust-java-agent:smoke-test:test-instrumentation' include 'braintrust-java-agent:smoke-test:dd-agent' include 'braintrust-java-agent:smoke-test:otel-agent'