Merge pull request #24 from braintrustdata/ark/langchain4j-instrumentation

Instrumentation for Lanchain4j's OpenAI Chat Model

Author: realark

build.gradle

@@ -46,6 +46,7 @@ ext {
   jacksonVersion = '2.16.1'
   junitVersion = '5.11.4'
   slf4jVersion = '2.0.17'
+  langchainVersion = '1.9.1'
 }
 
 dependencies {
@@ -86,6 +87,14 @@ dependencies {
   // 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}"
 }
 
 /**

examples/build.gradle

@@ -35,6 +35,9 @@ dependencies {
   implementation('org.springframework.boot:spring-boot-starter:3.4.1') {
     exclude group: 'org.springframework.boot', module: 'spring-boot-starter-logging'
   }
+  // to run langchain4j examples
+  implementation 'dev.langchain4j:langchain4j:1.9.1'
+  implementation 'dev.langchain4j:langchain4j-open-ai:1.9.1'
 }
 
 application {
@@ -142,7 +145,6 @@ task runSpringAI(type: JavaExec) {
   }
 }
 
-
 task runRemoteEval(type: JavaExec) {
   group = 'Braintrust SDK Examples'
   description = 'Run the remote eval example'
@@ -156,3 +158,17 @@ task runRemoteEval(type: JavaExec) {
     suspend = false
   }
 }
+
+task runLangchain(type: JavaExec) {
+  group = 'Braintrust SDK Examples'
+  description = 'Run the LangChain4j instrumentation example. NOTE: this requires OPENAI_API_KEY to be exported and will make a small call to openai, using your tokens'
+  classpath = sourceSets.main.runtimeClasspath
+  mainClass = 'dev.braintrust.examples.LangchainExample'
+  systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel
+  debugOptions {
+    enabled = true
+    port = 5566
+    server = true
+    suspend = false
+  }
+}

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

@@ -0,0 +1,54 @@
+package dev.braintrust.examples;
+
+import dev.braintrust.Braintrust;
+import dev.braintrust.instrumentation.langchain.BraintrustLangchain;
+import dev.langchain4j.data.message.UserMessage;
+import dev.langchain4j.model.chat.ChatModel;
+import dev.langchain4j.model.openai.OpenAiChatModel;
+
+/** Basic OTel + LangChain4j instrumentation example */
+public class LangchainExample {
+
+    public static void main(String[] args) throws Exception {
+        if (null == System.getenv("OPENAI_API_KEY")) {
+            System.err.println(
+                    "\nWARNING envar OPENAI_API_KEY not found. This example will likely fail.\n");
+        }
+        var braintrust = Braintrust.get();
+        var openTelemetry = braintrust.openTelemetryCreate();
+
+        ChatModel model =
+                BraintrustLangchain.wrap(
+                        openTelemetry,
+                        OpenAiChatModel.builder()
+                                .apiKey(System.getenv("OPENAI_API_KEY"))
+                                .modelName("gpt-4o-mini")
+                                .temperature(0.0));
+
+        var rootSpan =
+                openTelemetry
+                        .getTracer("my-instrumentation")
+                        .spanBuilder("langchain4j-instrumentation-example")
+                        .startSpan();
+        try (var ignored = rootSpan.makeCurrent()) {
+            chatExample(model);
+        } finally {
+            rootSpan.end();
+        }
+        var url =
+                braintrust.projectUri()
+                        + "/logs?r=%s&s=%s"
+                                .formatted(
+                                        rootSpan.getSpanContext().getTraceId(),
+                                        rootSpan.getSpanContext().getSpanId());
+        System.out.println(
+                "\n\n  Example complete! View your data in Braintrust: %s\n".formatted(url));
+    }
+
+    private static void chatExample(ChatModel model) {
+        var message = UserMessage.from("What is the capital of France?");
+        var response = model.chat(message);
+        System.out.println(
+                "\n~~~ LANGCHAIN4J CHAT RESPONSE: %s\n".formatted(response.aiMessage().text()));
+    }
+}

src/main/java/dev/braintrust/instrumentation/langchain/BraintrustLangchain.java

@@ -0,0 +1,66 @@
+package dev.braintrust.instrumentation.langchain;
+
+import dev.langchain4j.http.client.HttpClientBuilder;
+import dev.langchain4j.http.client.HttpClientBuilderLoader;
+import dev.langchain4j.model.openai.OpenAiChatModel;
+import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
+import io.opentelemetry.api.OpenTelemetry;
+import lombok.extern.slf4j.Slf4j;
+
+/** Braintrust LangChain4j client instrumentation. */
+@Slf4j
+public final class BraintrustLangchain {
+    /** Instrument langchain openai chat model with braintrust traces */
+    public static OpenAiChatModel wrap(
+            OpenTelemetry otel, OpenAiChatModel.OpenAiChatModelBuilder builder) {
+        try {
+            HttpClientBuilder underlyingHttpClient = getPrivateField(builder, "httpClientBuilder");
+            if (underlyingHttpClient == null) {
+                underlyingHttpClient = HttpClientBuilderLoader.loadHttpClientBuilder();
+            }
+            HttpClientBuilder wrappedHttpClient =
+                    wrap(otel, underlyingHttpClient, new Options("openai"));
+            return builder.httpClientBuilder(wrappedHttpClient).build();
+        } catch (Exception e) {
+            log.warn(
+                    "Braintrust instrumentation could not be applied to OpenAiChatModel builder",
+                    e);
+            return builder.build();
+        }
+    }
+
+    /** Instrument langchain openai chat model with braintrust traces */
+    public static OpenAiStreamingChatModel wrap(
+            OpenTelemetry otel, OpenAiStreamingChatModel.OpenAiStreamingChatModelBuilder builder) {
+        try {
+            HttpClientBuilder underlyingHttpClient = getPrivateField(builder, "httpClientBuilder");
+            if (underlyingHttpClient == null) {
+                underlyingHttpClient = HttpClientBuilderLoader.loadHttpClientBuilder();
+            }
+            HttpClientBuilder wrappedHttpClient =
+                    wrap(otel, underlyingHttpClient, new Options("openai"));
+            return builder.httpClientBuilder(wrappedHttpClient).build();
+        } catch (Exception e) {
+            log.warn(
+                    "Braintrust instrumentation could not be applied to OpenAiStreamingChatModel"
+                            + " builder",
+                    e);
+            return builder.build();
+        }
+    }
+
+    private static HttpClientBuilder wrap(
+            OpenTelemetry otel, HttpClientBuilder builder, Options options) {
+        return new WrappedHttpClientBuilder(otel, builder, options);
+    }
+
+    public record Options(String providerName) {}
+
+    @SuppressWarnings("unchecked")
+    private static <T> T getPrivateField(Object obj, String fieldName)
+            throws ReflectiveOperationException {
+        java.lang.reflect.Field field = obj.getClass().getDeclaredField(fieldName);
+        field.setAccessible(true);
+        return (T) field.get(obj);
+    }
+}