Skip to content

Commit 42784e2

Browse files
nayeem-kamalgary-huang
nayeem-kamal
authored and
gary-huang
committed
added chatcompletion llmobs spans
1 parent 378d442 commit 42784e2

File tree

10 files changed

+188
-48
lines changed

10 files changed

+188
-48
lines changed

‎dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import java.lang.instrument.Instrumentation;
1313
import java.util.Map;
1414
import java.util.concurrent.TimeUnit;
15-
import org.jetbrains.annotations.Nullable;
15+
import javax.annotation.Nullable;
1616
import org.slf4j.Logger;
1717
import org.slf4j.LoggerFactory;
1818

@@ -22,7 +22,7 @@ public class LLMObsSystem {
2222

2323
private static final String CUSTOM_MODEL_VAL = "custom";
2424

25-
public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
25+
public static void start(@Nullable Instrumentation inst, SharedCommunicationObjects sco) {
2626
Config config = Config.get();
2727
if (!config.isLlmObsEnabled()) {
2828
LOGGER.debug("LLM Observability is disabled");

‎dd-java-agent/instrumentation/openai/build.gradle

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@ muzzle {
88
}
99

1010
// Instrumentation is currently in preview. Additional testing will be implemented before enabling this by default.
11-
//minimumBranchCoverage = 0.0
12-
//minimumInstructionCoverage = 0.0
1311

1412
apply from: "$rootDir/gradle/java.gradle"
1513

‎dd-java-agent/instrumentation/openai/src/main/java/datadog/trace/instrumentation/openai/ChatCompletionServiceInstrumentation.java

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,16 @@
55
import static net.bytebuddy.matcher.ElementMatchers.*;
66
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
77

8+
import com.openai.models.chat.completions.ChatCompletion;
89
import com.openai.models.chat.completions.ChatCompletionCreateParams;
910
import com.openai.services.blocking.CompletionService;
1011
import datadog.trace.agent.tooling.Instrumenter;
1112
import datadog.trace.agent.tooling.InstrumenterModule;
13+
import datadog.trace.api.InstrumenterConfig;
14+
import datadog.trace.api.llmobs.LLMObsSpan;
1215
import datadog.trace.bootstrap.InstrumentationContext;
1316
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17+
import java.util.Collections;
1418
import java.util.HashMap;
1519
import java.util.Map;
1620
import net.bytebuddy.asm.Advice;
@@ -23,6 +27,11 @@ public ChatCompletionServiceInstrumentation() {
2327
super("openai-client");
2428
}
2529

30+
@Override
31+
protected boolean defaultEnabled() {
32+
return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
33+
}
34+
2635
@Override
2736
public String[] helperClassNames() {
2837
return new String[] {
@@ -34,7 +43,7 @@ public String[] helperClassNames() {
3443
public Map<String, String> contextStore() {
3544
Map<String, String> contextStores = new HashMap<>(1);
3645
contextStores.put(
37-
"com.openai.services.blocking.CompletionService", OpenAIClientInfo.class.getName());
46+
"com.openai.services.blocking.ChatCompletionService", OpenAIClientInfo.class.getName());
3847
return contextStores;
3948
}
4049

@@ -47,7 +56,7 @@ public void methodAdvice(MethodTransformer transformer) {
4756
.and(
4857
takesArgument(
4958
0, named("com.openai.models.completions.ChatCompletionCreateParams"))),
50-
getClass().getName() + "$CompletionServiceAdvice");
59+
getClass().getName() + "$ChatCompletionServiceAdvice");
5160
}
5261

5362
@Override
@@ -60,24 +69,37 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
6069
return implementsInterface(named(hierarchyMarkerType()));
6170
}
6271

63-
public static class CompletionServiceAdvice {
72+
public static class ChatCompletionServiceAdvice {
6473
@Advice.OnMethodEnter(suppress = Throwable.class)
65-
public static AgentScope methodEnter(
74+
public static Map<String, Object> methodEnter(
6675
@Advice.Argument(0) final ChatCompletionCreateParams params) {
67-
68-
return OpenAIClientDecorator.DECORATE.startChatCompletionSpan(params);
76+
Map<String, Object> spans = new HashMap<>();
77+
spans.put(
78+
"datadog.trace.api.llmobs.LLMObsSpan",
79+
OpenAIClientDecorator.DECORATE.startLLMObsChatCompletionSpan(params));
80+
spans.put(
81+
"datadog.trace.bootstrap.instrumentation.api.AgentScope",
82+
OpenAIClientDecorator.DECORATE.startChatCompletionSpan(params));
83+
return spans;
6984
}
7085

7186
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
7287
public static void methodExit(
73-
@Advice.Enter final AgentScope span,
88+
@Advice.Enter final Map<String, Object> spans,
7489
@Advice.This final CompletionService completionService,
7590
@Advice.Return final Object result,
7691
@Advice.Thrown final Throwable throwable) {
7792
OpenAIClientInfo info =
7893
InstrumentationContext.get(CompletionService.class, OpenAIClientInfo.class)
7994
.get(completionService);
80-
OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable);
95+
OpenAIClientDecorator.DECORATE.finishLLMObsChatCompletionSpan(
96+
(LLMObsSpan) spans.get("datadog.trace.api.llmobs.LLMObsSpan"),
97+
(ChatCompletion) result,
98+
throwable);
99+
OpenAIClientDecorator.DECORATE.finishSpan(
100+
(AgentScope) spans.get("datadog.trace.bootstrap.instrumentation.api.AgentScope"),
101+
result,
102+
throwable);
81103
}
82104
}
83105
}

‎dd-java-agent/instrumentation/openai/src/main/java/datadog/trace/instrumentation/openai/CompletionServiceInstrumentation.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,10 @@
99
import com.openai.services.blocking.CompletionService;
1010
import datadog.trace.agent.tooling.Instrumenter;
1111
import datadog.trace.agent.tooling.InstrumenterModule;
12+
import datadog.trace.api.InstrumenterConfig;
1213
import datadog.trace.bootstrap.InstrumentationContext;
1314
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
15+
import java.util.Collections;
1416
import java.util.HashMap;
1517
import java.util.Map;
1618
import net.bytebuddy.asm.Advice;
@@ -23,6 +25,11 @@ public CompletionServiceInstrumentation() {
2325
super("openai-client");
2426
}
2527

28+
@Override
29+
protected boolean defaultEnabled() {
30+
return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
31+
}
32+
2633
@Override
2734
public String[] helperClassNames() {
2835
return new String[] {

‎dd-java-agent/instrumentation/openai/src/main/java/datadog/trace/instrumentation/openai/EmbeddingServiceInstrumentation.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@
1010
import com.openai.models.embeddings.EmbeddingCreateParams;
1111
import datadog.trace.agent.tooling.Instrumenter;
1212
import datadog.trace.agent.tooling.InstrumenterModule;
13+
import datadog.trace.api.InstrumenterConfig;
1314
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
15+
import java.util.Collections;
1416
import net.bytebuddy.asm.Advice;
1517
import net.bytebuddy.description.type.TypeDescription;
1618
import net.bytebuddy.matcher.ElementMatcher;
@@ -23,6 +25,11 @@ public EmbeddingServiceInstrumentation() {
2325
super("openai", "openai-java");
2426
}
2527

28+
@Override
29+
protected boolean defaultEnabled() {
30+
return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
31+
}
32+
2633
@Override
2734
public String[] helperClassNames() {
2835
return new String[] {

‎dd-java-agent/instrumentation/openai/src/main/java/datadog/trace/instrumentation/openai/OpenAIClientDecorator.java

Lines changed: 126 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@
1010
import com.openai.models.embeddings.CreateEmbeddingResponse;
1111
import com.openai.models.embeddings.Embedding;
1212
import com.openai.models.embeddings.EmbeddingCreateParams;
13+
import datadog.trace.api.Config;
14+
import datadog.trace.api.llmobs.LLMObs;
15+
import datadog.trace.api.llmobs.LLMObsSpan;
1316
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1417
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1518
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
@@ -20,6 +23,8 @@
2023
import java.util.Optional;
2124

2225
public class OpenAIClientDecorator extends ClientDecorator {
26+
private static final String mlApp = Config.get().getLlmObsMlApp();
27+
private static final String mlProvider = "openai";
2328
private static final String COMPONENT_NAME = "openai";
2429
private static final UTF8BytesString OPENAI_REQUEST = UTF8BytesString.create("openai.request");
2530

@@ -42,14 +47,14 @@ protected CharSequence spanType() {
4247

4348
@Override
4449
protected String service() {
45-
return null; // Use default service name
50+
return null;
4651
}
4752

4853
public AgentScope startChatCompletionSpan(ChatCompletionCreateParams params) {
4954
AgentSpan span = startSpan(OPENAI_REQUEST);
5055
span.setTag("openai.request.endpoint", "/chat/completions");
5156
span.setResourceName("chat.completions.create");
52-
span.setTag("openai.provider", "openai");
57+
span.setTag("openai.request.provider", "openai");
5358
extractChatCompletionRequestData(span, params);
5459
afterStart(span);
5560
return activateSpan(span);
@@ -59,7 +64,7 @@ public AgentScope startLLMChatCompletionSpan(ChatCompletionCreateParams params)
5964
AgentSpan span = startSpan(OPENAI_REQUEST);
6065
span.setTag("openai.request.endpoint", "/chat/completions");
6166
span.setResourceName("chat.completions.create");
62-
span.setTag("openai.provider", "openai");
67+
span.setTag("openai.request.provider", "openai");
6368
extractChatCompletionRequestData(span, params);
6469
afterStart(span);
6570
return activateSpan(span);
@@ -88,6 +93,115 @@ public AgentScope startEmbeddingSpan(EmbeddingCreateParams params) {
8893
return activateSpan(span);
8994
}
9095

96+
public LLMObsSpan startLLMObsChatCompletionSpan(ChatCompletionCreateParams params) {
97+
String modelName = params.model().toString();
98+
99+
LLMObsSpan llmObsSpan =
100+
LLMObs.startLLMSpan("chat_completion", modelName, mlProvider, mlApp, null);
101+
102+
// Extract and set input data from chat completion params
103+
// May need to be reformatted
104+
StringBuilder inputData = new StringBuilder();
105+
List<ChatCompletionMessageParam> messages = params.messages();
106+
for (int i = 0; i < messages.size(); i++) {
107+
ChatCompletionMessageParam messageParam = messages.get(i);
108+
if (i > 0) {
109+
inputData.append(" | ");
110+
}
111+
112+
if (messageParam.isUser()) {
113+
inputData.append("User: ").append(messageParam.asUser().content());
114+
} else if (messageParam.isAssistant()) {
115+
inputData.append("Assistant: ").append(messageParam.asAssistant().content());
116+
} else if (messageParam.isDeveloper()) {
117+
inputData.append("Developer: ").append(messageParam.asDeveloper().content());
118+
} else if (messageParam.isSystem()) {
119+
inputData.append("System: ").append(messageParam.asSystem().content());
120+
} else if (messageParam.isTool()) {
121+
inputData.append("Tool: ").append(messageParam.asTool().content());
122+
}
123+
}
124+
125+
if (inputData.length() > 0) {
126+
llmObsSpan.annotateIO(inputData.toString(), null); // No output yet, will be set in response
127+
}
128+
129+
java.util.Map<String, Object> metadata = new java.util.HashMap<>();
130+
metadata.put("endpoint", "/chat/completions");
131+
metadata.put("provider", "openai");
132+
metadata.put("model", modelName);
133+
134+
params.maxTokens().ifPresent(tokens -> metadata.put("max_tokens", tokens));
135+
params.temperature().ifPresent(temp -> metadata.put("temperature", temp));
136+
137+
llmObsSpan.setMetadata(metadata);
138+
139+
return llmObsSpan;
140+
}
141+
142+
public void finishLLMObsChatCompletionSpan(
143+
LLMObsSpan llmObsSpan, ChatCompletion response, Throwable throwable) {
144+
try {
145+
if (throwable != null) {
146+
// Set error information
147+
java.util.Map<String, Object> errorMetadata = new java.util.HashMap<>();
148+
errorMetadata.put("error.type", throwable.getClass().getSimpleName());
149+
errorMetadata.put("error.message", throwable.getMessage());
150+
llmObsSpan.setMetadata(errorMetadata);
151+
} else if (response != null) {
152+
StringBuilder outputData = new StringBuilder();
153+
List<ChatCompletion.Choice> choices = response.choices();
154+
155+
for (int i = 0; i < choices.size(); i++) {
156+
ChatCompletion.Choice choice = choices.get(i);
157+
ChatCompletionMessage message = choice.message();
158+
159+
if (i > 0) {
160+
outputData.append(" | ");
161+
}
162+
163+
// Extract content
164+
Optional<String> content = message.content();
165+
content.ifPresent(s -> outputData.append("Assistant: ").append(s));
166+
167+
// Extract tool calls if present
168+
Optional<List<ChatCompletionMessageToolCall>> toolCalls = message.toolCalls();
169+
if (toolCalls.isPresent() && !toolCalls.get().isEmpty()) {
170+
content.ifPresent(s -> outputData.append(" | "));
171+
outputData.append("Tool calls: ");
172+
for (int j = 0; j < toolCalls.get().size(); j++) {
173+
ChatCompletionMessageToolCall call = toolCalls.get().get(j);
174+
if (j > 0) {
175+
outputData.append(", ");
176+
}
177+
outputData
178+
.append(call.function().name())
179+
.append("(")
180+
.append(call.function().arguments())
181+
.append(")");
182+
}
183+
}
184+
}
185+
186+
if (outputData.length() > 0) {
187+
llmObsSpan.annotateIO(null, outputData.toString());
188+
}
189+
java.util.Map<String, Object> responseMetadata = new java.util.HashMap<>();
190+
responseMetadata.put("response.choices_count", choices.size());
191+
192+
llmObsSpan.setMetadata(responseMetadata);
193+
}
194+
} catch (Exception e) {
195+
java.util.Map<String, Object> errorMetadata = new java.util.HashMap<>();
196+
errorMetadata.put("error.type", "ResponseProcessingError");
197+
errorMetadata.put("error.message", "Failed to process response: " + e.getMessage());
198+
llmObsSpan.setMetadata(errorMetadata);
199+
} finally {
200+
// Always finish the span
201+
llmObsSpan.finish();
202+
}
203+
}
204+
91205
public void finishSpan(AgentScope scope, Object result, Throwable throwable) {
92206

93207
AgentSpan span = scope.span();
@@ -107,8 +221,6 @@ public void finishSpan(AgentScope scope, Object result, Throwable throwable) {
107221

108222
private void extractChatCompletionRequestData(AgentSpan span, ChatCompletionCreateParams params) {
109223

110-
// Extract model
111-
112224
span.setTag("openai.model.name", params.model().toString());
113225

114226
// Extract messages
@@ -124,11 +236,7 @@ private void extractChatCompletionRequestData(AgentSpan span, ChatCompletionCrea
124236

125237
private void extractCompletionRequestData(
126238
AgentSpan span, CompletionCreateParams completionParams) {
127-
// Extract model
128-
CompletionCreateParams.Model model = completionParams.model();
129-
if (model != null) {
130-
span.setTag("openai.model.name", model.toString());
131-
}
239+
span.setTag("openai.model.name", completionParams.model().toString());
132240

133241
// Extract prompt
134242
Optional<CompletionCreateParams.Prompt> prompt = completionParams.prompt();
@@ -156,18 +264,13 @@ private void extractCompletionRequestData(
156264

157265
private void extractEmbeddingRequestData(AgentSpan span, EmbeddingCreateParams embeddingParams) {
158266

159-
// Extract model
160-
Object model = embeddingParams.model();
161-
if (model != null) {
162-
span.setTag("openai.model.name", model.toString());
163-
}
164-
267+
span.setTag("openai.model.name", embeddingParams.model().toString());
165268
// Extract input
166269
EmbeddingCreateParams.Input input = embeddingParams.input();
167270
int inputIndex = 0;
168271
List<String> inputStrings = input.asArrayOfStrings();
169272
for (String inputItem : inputStrings) {
170-
span.setTag("openai.request.input." + inputIndex, input.asString());
273+
span.setTag("openai.request.input." + inputIndex, inputItem);
171274
inputIndex++;
172275
}
173276
}
@@ -205,25 +308,19 @@ private void extractMessageData(
205308

206309
private void extractChatCompletionParameters(AgentSpan span, ChatCompletionCreateParams params) {
207310
// Extract max_tokens
208-
Optional<Long> maxTokens = params.maxTokens();
209-
maxTokens.ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens));
311+
params
312+
.maxCompletionTokens()
313+
.ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens));
210314

211315
// Extract temperature
212-
Optional<Double> temperature = params.temperature();
213-
temperature.ifPresent(temp -> span.setTag("openai.request.temperature", temp));
316+
params.temperature().ifPresent(temp -> span.setTag("openai.request.temperature", temp));
214317
}
215318

216319
private void extractCompletionParameters(AgentSpan span, CompletionCreateParams params) {
217320
// Extract max_tokens
218-
if (params.maxTokens().isPresent()) {
219-
span.setTag("openai.request.max_tokens", params.maxTokens().get());
220-
}
221-
321+
params.maxTokens().ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens));
222322
// Extract temperature
223-
Optional<Double> temperature = params.temperature();
224-
if (temperature.isPresent()) {
225-
span.setTag("openai.request.temperature", temperature.get());
226-
}
323+
params.temperature().ifPresent(temp -> span.setTag("openai.request.temperature", temp));
227324
}
228325

229326
private void extractResponseData(AgentSpan span, Object result) {

0 commit comments

Comments
 (0)