Skip to content

Commit e70508a

Browse files
realarkrealark
authored
Merge pull request #30 from braintrustdata/ark/langchain-ai-services
Instrument Langchain4j AI Services
2 parents 699b42d + c4f4249 commit e70508a

33 files changed

Lines changed: 971 additions & 188 deletions

File tree

‎README.md‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,65 @@ var request =
8282
var response = openAIClient.chat().completions().create(request);
8383
```
8484

85+
### LangChain4j Instrumentation
86+
87+
```java
88+
var braintrust = Braintrust.get();
89+
var openTelemetry = braintrust.openTelemetryCreate();
90+
91+
// Wrap the chat model to trace LLM calls
92+
ChatModel model = BraintrustLangchain.wrap(
93+
openTelemetry,
94+
OpenAiChatModel.builder()
95+
.apiKey(System.getenv("OPENAI_API_KEY"))
96+
.modelName("gpt-4o-mini")
97+
.temperature(0.0));
98+
99+
var response = model.chat("What is the capital of France?");
100+
```
101+
102+
#### Tool Wrapping
103+
104+
Use `BraintrustLangchain.wrapTools()` to automatically trace tool executions in your LangChain4j agents:
105+
106+
```java
107+
// Create your tool class
108+
public class WeatherTools {
109+
@Tool("Get current weather for a location")
110+
public String getWeather(String location) {
111+
return "The weather in " + location + " is sunny.";
112+
}
113+
}
114+
115+
// Wrap tools to create spans for each tool execution
116+
WeatherTools tools = new WeatherTools();
117+
WeatherTools instrumentedTools = BraintrustLangchain.wrapTools(openTelemetry, tools);
118+
119+
// Use instrumented tools in your AI service
120+
Assistant assistant = AiServices.builder(Assistant.class)
121+
.chatModel(model)
122+
.tools(instrumentedTools)
123+
.build();
124+
```
125+
126+
Each tool call will automatically create an OpenTelemetry span in Braintrust with:
127+
- Tool name and parameters
128+
- Execution duration
129+
- Return values
130+
- Any exceptions thrown
131+
132+
**Note:** For proper display in the Braintrust UI, ensure parent spans (conversation, turn, etc.) also set the required Braintrust attributes:
133+
```java
134+
var span = tracer.spanBuilder("my-span").startSpan();
135+
span.setAttribute("braintrust.span_attributes", "{\"type\":\"task\",\"name\":\"my-span\"}");
136+
span.setAttribute("braintrust.input_json", "{\"user_message\":\"...\"}");
137+
// ... do work ...
138+
span.setAttribute("braintrust.output_json", "{\"result\":\"...\"}");
139+
span.end();
140+
```
141+
142+
See [LangchainToolWrappingExample.java](./examples/src/main/java/dev/braintrust/examples/LangchainToolWrappingExample.java) for a complete example with proper span hierarchy.
143+
85144
## Running Examples
86145

87146
Example source code can be found [here](./examples/src/main/java/dev/braintrust/examples)

‎examples/build.gradle‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,11 +159,25 @@ task runRemoteEval(type: JavaExec) {
159159
}
160160
}
161161

162-
task runLangchain(type: JavaExec) {
162+
task runLangchainSimple(type: JavaExec) {
163163
group = 'Braintrust SDK Examples'
164164
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'
165165
classpath = sourceSets.main.runtimeClasspath
166-
mainClass = 'dev.braintrust.examples.LangchainExample'
166+
mainClass = 'dev.braintrust.examples.LangchainSimpleExample'
167+
systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel
168+
debugOptions {
169+
enabled = true
170+
port = 5566
171+
server = true
172+
suspend = false
173+
}
174+
}
175+
176+
task runLangchainAIServices(type: JavaExec) {
177+
group = 'Braintrust SDK Examples'
178+
description = 'Run the LangChain4j AI Services example. NOTE: this requires OPENAI_API_KEY to be exported and will make a small call to openai, using your tokens'
179+
classpath = sourceSets.main.runtimeClasspath
180+
mainClass = 'dev.braintrust.examples.LangchainAIServicesExample'
167181
systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel
168182
debugOptions {
169183
enabled = true

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package dev.braintrust.examples;
2+
3+
import dev.braintrust.Braintrust;
4+
import dev.braintrust.instrumentation.langchain.BraintrustLangchain;
5+
import dev.langchain4j.agent.tool.Tool;
6+
import dev.langchain4j.model.openai.OpenAiChatModel;
7+
import dev.langchain4j.service.AiServices;
8+
9+
public class LangchainAIServicesExample {
10+
11+
public static void main(String[] args) throws Exception {
12+
var braintrust = Braintrust.get();
13+
var openTelemetry = braintrust.openTelemetryCreate();
14+
15+
Assistant assistant =
16+
BraintrustLangchain.wrap(
17+
openTelemetry,
18+
AiServices.builder(Assistant.class)
19+
.chatModel(
20+
OpenAiChatModel.builder()
21+
.apiKey(System.getenv("OPENAI_API_KEY"))
22+
.modelName("gpt-4o-mini")
23+
.temperature(0.0)
24+
.build())
25+
.tools(new WeatherTools())
26+
.executeToolsConcurrently());
27+
28+
var rootSpan =
29+
openTelemetry
30+
.getTracer("my-instrumentation")
31+
.spanBuilder("langchain4j-ai-services-example")
32+
.startSpan();
33+
try (var ignored = rootSpan.makeCurrent()) {
34+
// response 1 should do a concurrent tool call
35+
var response1 = assistant.chat("is it hotter in Paris or New York right now?");
36+
System.out.println("response1: " + response1);
37+
var response2 = assistant.chat("what's the five day forecast for San Francisco?");
38+
System.out.println("response2: " + response2);
39+
} finally {
40+
rootSpan.end();
41+
}
42+
var url =
43+
braintrust.projectUri()
44+
+ "/logs?r=%s&s=%s"
45+
.formatted(
46+
rootSpan.getSpanContext().getTraceId(),
47+
rootSpan.getSpanContext().getSpanId());
48+
System.out.println(
49+
"\n\n Example complete! View your data in Braintrust: %s\n".formatted(url));
50+
}
51+
52+
/** AI Service interface for the assistant */
53+
interface Assistant {
54+
String chat(String userMessage);
55+
}
56+
57+
/** Example tool class with weather-related methods */
58+
public static class WeatherTools {
59+
@Tool("Get current weather for a location")
60+
public String getWeather(String location) {
61+
return String.format("The weather in %s is sunny with 72°F temperature.", location);
62+
}
63+
64+
@Tool("Get weather forecast for next N days")
65+
public String getForecast(String location, int days) {
66+
return String.format(
67+
"The %d-day forecast for %s: Mostly sunny with temperatures between 65-75°F.",
68+
days, location);
69+
}
70+
}
71+
}

examples/src/main/java/dev/braintrust/examples/LangchainExample.java renamed to examples/src/main/java/dev/braintrust/examples/LangchainSimpleExample.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import dev.langchain4j.model.openai.OpenAiChatModel;
88

99
/** Basic OTel + LangChain4j instrumentation example */
10-
public class LangchainExample {
10+
public class LangchainSimpleExample {
1111

1212
public static void main(String[] args) throws Exception {
1313
if (null == System.getenv("OPENAI_API_KEY")) {
@@ -46,8 +46,7 @@ public static void main(String[] args) throws Exception {
4646
}
4747

4848
private static void chatExample(ChatModel model) {
49-
var message = UserMessage.from("What is the capital of France?");
50-
var response = model.chat(message);
49+
var response = model.chat(UserMessage.from("What is the capital of France?"));
5150
System.out.println(
5251
"\n~~~ LANGCHAIN4J CHAT RESPONSE: %s\n".formatted(response.aiMessage().text()));
5352
}

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

Lines changed: 143 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,66 +1,178 @@
11
package dev.braintrust.instrumentation.langchain;
22

3-
import dev.langchain4j.http.client.HttpClientBuilder;
4-
import dev.langchain4j.http.client.HttpClientBuilderLoader;
53
import dev.langchain4j.model.openai.OpenAiChatModel;
64
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
5+
import dev.langchain4j.service.AiServiceContext;
6+
import dev.langchain4j.service.AiServices;
7+
import dev.langchain4j.service.tool.ToolExecutor;
78
import io.opentelemetry.api.OpenTelemetry;
9+
import io.opentelemetry.api.trace.Tracer;
10+
import java.util.Map;
811
import lombok.extern.slf4j.Slf4j;
912

1013
/** Braintrust LangChain4j client instrumentation. */
1114
@Slf4j
1215
public final class BraintrustLangchain {
16+
17+
private static final String INSTRUMENTATION_NAME = "braintrust-langchain4j";
18+
19+
@SuppressWarnings("unchecked")
20+
public static <T> T wrap(OpenTelemetry openTelemetry, AiServices<T> aiServices) {
21+
try {
22+
AiServiceContext context = getPrivateField(aiServices, "context");
23+
Tracer tracer = openTelemetry.getTracer(INSTRUMENTATION_NAME);
24+
25+
// ////// CREATE A LLM SPAN FOR EACH CALL TO AI PROVIDER
26+
var chatModel = context.chatModel;
27+
var streamingChatModel = context.streamingChatModel;
28+
if (chatModel != null) {
29+
if (chatModel instanceof OpenAiChatModel oaiModel) {
30+
aiServices.chatModel(wrap(openTelemetry, oaiModel));
31+
} else {
32+
log.warn(
33+
"unsupported model: {}. LLM calls will not be instrumented",
34+
chatModel.getClass().getName());
35+
}
36+
// intentional fall-through
37+
} else if (streamingChatModel != null) {
38+
if (streamingChatModel instanceof OpenAiStreamingChatModel oaiModel) {
39+
aiServices.streamingChatModel(wrap(openTelemetry, oaiModel));
40+
} else {
41+
log.warn(
42+
"unsupported model: {}. LLM calls will not be instrumented",
43+
streamingChatModel.getClass().getName());
44+
}
45+
// intentional fall-through
46+
} else {
47+
// langchain is going to fail to build. don't apply instrumentation.
48+
throw new RuntimeException("model or chat model must be set");
49+
}
50+
51+
if (context.toolService != null) {
52+
// ////// CREATE A SPAN FOR EACH TOOL CALL
53+
for (Map.Entry<String, ToolExecutor> entry :
54+
context.toolService.toolExecutors().entrySet()) {
55+
String toolName = entry.getKey();
56+
ToolExecutor original = entry.getValue();
57+
entry.setValue(new TracingToolExecutor(original, toolName, tracer));
58+
}
59+
60+
// ////// LINK SPANS ACROSS CONCURRENT TOOL CALLS
61+
var underlyingExecutor = context.toolService.executor();
62+
if (underlyingExecutor != null) {
63+
aiServices.executeToolsConcurrently(
64+
new OtelContextPassingExecutor(underlyingExecutor));
65+
}
66+
}
67+
68+
// ////// CREATE A SPAN ON SERVICE METHOD INVOKE
69+
T service = aiServices.build();
70+
Class<T> serviceInterface = (Class<T>) context.aiServiceClass;
71+
return TracingProxy.create(serviceInterface, service, tracer);
72+
} catch (Exception e) {
73+
log.warn("failed to apply langchain AI services instrumentation", e);
74+
return aiServices.build();
75+
}
76+
}
77+
1378
/** Instrument langchain openai chat model with braintrust traces */
1479
public static OpenAiChatModel wrap(
1580
OpenTelemetry otel, OpenAiChatModel.OpenAiChatModelBuilder builder) {
81+
return wrap(otel, builder.build());
82+
}
83+
84+
private static OpenAiChatModel wrap(OpenTelemetry otel, OpenAiChatModel model) {
1685
try {
17-
HttpClientBuilder underlyingHttpClient = getPrivateField(builder, "httpClientBuilder");
18-
if (underlyingHttpClient == null) {
19-
underlyingHttpClient = HttpClientBuilderLoader.loadHttpClientBuilder();
86+
// Get the internal OpenAiClient from the chat model
87+
Object internalClient = getPrivateField(model, "client");
88+
89+
// Get the HttpClient from the internal client
90+
dev.langchain4j.http.client.HttpClient httpClient =
91+
getPrivateField(internalClient, "httpClient");
92+
93+
if (httpClient instanceof WrappedHttpClient) {
94+
log.debug("model already instrumented. skipping: {}", httpClient.getClass());
95+
return model;
2096
}
21-
HttpClientBuilder wrappedHttpClient =
22-
wrap(otel, underlyingHttpClient, new Options("openai"));
23-
return builder.httpClientBuilder(wrappedHttpClient).build();
97+
98+
// Wrap the HttpClient with our instrumented version
99+
dev.langchain4j.http.client.HttpClient wrappedHttpClient =
100+
new WrappedHttpClient(otel, httpClient, new Options("openai"));
101+
102+
setPrivateField(internalClient, "httpClient", wrappedHttpClient);
103+
104+
return model;
24105
} catch (Exception e) {
25-
log.warn(
26-
"Braintrust instrumentation could not be applied to OpenAiChatModel builder",
27-
e);
28-
return builder.build();
106+
log.warn("failed to instrument OpenAiChatModel", e);
107+
return model;
29108
}
30109
}
31110

32111
/** Instrument langchain openai chat model with braintrust traces */
33112
public static OpenAiStreamingChatModel wrap(
34113
OpenTelemetry otel, OpenAiStreamingChatModel.OpenAiStreamingChatModelBuilder builder) {
114+
return wrap(otel, builder.build());
115+
}
116+
117+
public static OpenAiStreamingChatModel wrap(
118+
OpenTelemetry otel, OpenAiStreamingChatModel model) {
35119
try {
36-
HttpClientBuilder underlyingHttpClient = getPrivateField(builder, "httpClientBuilder");
37-
if (underlyingHttpClient == null) {
38-
underlyingHttpClient = HttpClientBuilderLoader.loadHttpClientBuilder();
120+
// Get the internal OpenAiClient from the streaming chat model
121+
Object internalClient = getPrivateField(model, "client");
122+
123+
// Get the HttpClient from the internal client
124+
dev.langchain4j.http.client.HttpClient httpClient =
125+
getPrivateField(internalClient, "httpClient");
126+
127+
if (httpClient instanceof WrappedHttpClient) {
128+
log.debug("model already instrumented. skipping: {}", httpClient.getClass());
129+
return model;
39130
}
40-
HttpClientBuilder wrappedHttpClient =
41-
wrap(otel, underlyingHttpClient, new Options("openai"));
42-
return builder.httpClientBuilder(wrappedHttpClient).build();
131+
132+
// Wrap the HttpClient with our instrumented version
133+
dev.langchain4j.http.client.HttpClient wrappedHttpClient =
134+
new WrappedHttpClient(otel, httpClient, new Options("openai"));
135+
136+
setPrivateField(internalClient, "httpClient", wrappedHttpClient);
137+
138+
return model;
43139
} catch (Exception e) {
44-
log.warn(
45-
"Braintrust instrumentation could not be applied to OpenAiStreamingChatModel"
46-
+ " builder",
47-
e);
48-
return builder.build();
140+
log.warn("failed to instrument OpenAiStreamingChatModel", e);
141+
return model;
49142
}
50143
}
51144

52-
private static HttpClientBuilder wrap(
53-
OpenTelemetry otel, HttpClientBuilder builder, Options options) {
54-
return new WrappedHttpClientBuilder(otel, builder, options);
55-
}
56-
57145
public record Options(String providerName) {}
58146

59147
@SuppressWarnings("unchecked")
60148
private static <T> T getPrivateField(Object obj, String fieldName)
61149
throws ReflectiveOperationException {
62-
java.lang.reflect.Field field = obj.getClass().getDeclaredField(fieldName);
63-
field.setAccessible(true);
64-
return (T) field.get(obj);
150+
Class<?> clazz = obj.getClass();
151+
while (clazz != null) {
152+
try {
153+
java.lang.reflect.Field field = clazz.getDeclaredField(fieldName);
154+
field.setAccessible(true);
155+
return (T) field.get(obj);
156+
} catch (NoSuchFieldException e) {
157+
clazz = clazz.getSuperclass();
158+
}
159+
}
160+
throw new NoSuchFieldException(fieldName);
161+
}
162+
163+
private static void setPrivateField(Object obj, String fieldName, Object value)
164+
throws ReflectiveOperationException {
165+
Class<?> clazz = obj.getClass();
166+
while (clazz != null) {
167+
try {
168+
java.lang.reflect.Field field = clazz.getDeclaredField(fieldName);
169+
field.setAccessible(true);
170+
field.set(obj, value);
171+
return;
172+
} catch (NoSuchFieldException e) {
173+
clazz = clazz.getSuperclass();
174+
}
175+
}
176+
throw new NoSuchFieldException(fieldName);
65177
}
66178
}

0 commit comments

Comments
 (0)