Skip to content

Commit 0f0b0a1

Browse files
ygreeygree
committed
Split to modules b/o exceeded Muzzle.create limit
Failed to transform at least one type: [class datadog.trace.instrumentation.openai_java.OpenAiModule:[net.bytebuddy.jar.asm.MethodTooLargeException: Method too large: datadog/trace/instrumentation/openai_java/OpenAiModule$Muzzle.create ()Ldatadog/trace/agent/tooling/muzzle/ReferenceMatcher;]]
1 parent c29ee56 commit 0f0b0a1

21 files changed

+814
-533
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ muzzle {
99
group = "com.openai"
1010
module = "openai-java"
1111
versions = "[$minVer,)"
12-
assertInverse = true
12+
// assertInverse = true //TODO fix after module split
1313
}
1414
}
1515

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionDecorator.java‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
package datadog.trace.instrumentation.openai_java;
2+
3+
import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.REQUEST_MODEL;
4+
import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.RESPONSE_MODEL;
5+
6+
import com.openai.helpers.ChatCompletionAccumulator;
7+
import com.openai.models.chat.completions.ChatCompletion;
8+
import com.openai.models.chat.completions.ChatCompletionChunk;
9+
import com.openai.models.chat.completions.ChatCompletionCreateParams;
10+
import com.openai.models.chat.completions.ChatCompletionMessage;
11+
import com.openai.models.chat.completions.ChatCompletionMessageParam;
12+
import com.openai.models.chat.completions.ChatCompletionMessageToolCall;
13+
import com.openai.models.completions.CompletionUsage;
14+
import datadog.trace.api.llmobs.LLMObs;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
16+
import datadog.trace.bootstrap.instrumentation.api.Tags;
17+
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
18+
import java.util.ArrayList;
19+
import java.util.Collections;
20+
import java.util.HashMap;
21+
import java.util.List;
22+
import java.util.Map;
23+
import java.util.Optional;
24+
import java.util.stream.Collectors;
25+
26+
public class ChatCompletionDecorator {
27+
public static final ChatCompletionDecorator DECORATE = new ChatCompletionDecorator();
28+
private static final CharSequence CHAT_COMPLETIONS_CREATE =
29+
UTF8BytesString.create("createChatCompletion");
30+
31+
public void withChatCompletionCreateParams(
32+
AgentSpan span, ChatCompletionCreateParams params, boolean stream) {
33+
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
34+
span.setResourceName(CHAT_COMPLETIONS_CREATE);
35+
span.setTag("openai.request.endpoint", "v1/chat/completions");
36+
span.setTag("openai.request.method", "POST");
37+
if (params == null) {
38+
return;
39+
}
40+
params.model()._value().asString().ifPresent(str -> span.setTag(REQUEST_MODEL, str));
41+
42+
span.setTag(
43+
"_ml_obs_tag.input",
44+
params.messages().stream()
45+
.map(ChatCompletionDecorator::llmMessage)
46+
.collect(Collectors.toList()));
47+
48+
Map<String, Object> metadata = new HashMap<>();
49+
// maxTokens is deprecated but integration tests missing to provide maxCompletionTokens
50+
params.maxTokens().ifPresent(v -> metadata.put("max_tokens", v));
51+
params.temperature().ifPresent(v -> metadata.put("temperature", v));
52+
if (stream) {
53+
metadata.put("stream", true);
54+
}
55+
params
56+
.streamOptions()
57+
.ifPresent(
58+
v -> {
59+
if (v.includeUsage().orElse(false)) {
60+
metadata.put("stream_options", Collections.singletonMap("include_usage", true));
61+
}
62+
});
63+
span.setTag("_ml_obs_tag.metadata", metadata);
64+
}
65+
66+
private static LLMObs.LLMMessage llmMessage(ChatCompletionMessageParam m) {
67+
String role = "unknown";
68+
String content = null;
69+
if (m.isAssistant()) {
70+
role = "assistant";
71+
content = m.asAssistant().content().map(v -> v.text().orElse(null)).orElse(null);
72+
} else if (m.isDeveloper()) {
73+
role = "developer";
74+
content = m.asDeveloper().content().text().orElse(null);
75+
} else if (m.isSystem()) {
76+
role = "system";
77+
content = m.asSystem().content().text().orElse(null);
78+
} else if (m.isTool()) {
79+
role = "tool";
80+
content = m.asTool().content().text().orElse(null);
81+
} else if (m.isUser()) {
82+
role = "user";
83+
content = m.asUser().content().text().orElse(null);
84+
}
85+
return LLMObs.LLMMessage.from(role, content);
86+
}
87+
88+
public void withChatCompletion(AgentSpan span, ChatCompletion completion) {
89+
String modelName = completion.model();
90+
span.setTag(RESPONSE_MODEL, modelName);
91+
span.setTag("_ml_obs_tag.model_name", modelName);
92+
span.setTag("_ml_obs_tag.model_provider", "openai");
93+
94+
List<LLMObs.LLMMessage> output =
95+
completion.choices().stream()
96+
.map(ChatCompletionDecorator::llmMessage)
97+
.collect(Collectors.toList());
98+
span.setTag("_ml_obs_tag.output", output);
99+
100+
completion.usage().ifPresent(usage -> withCompletionUsage(span, usage));
101+
}
102+
103+
private static void withCompletionUsage(AgentSpan span, CompletionUsage usage) {
104+
span.setTag("_ml_obs_metric.input_tokens", usage.promptTokens());
105+
span.setTag("_ml_obs_metric.output_tokens", usage.completionTokens());
106+
span.setTag("_ml_obs_metric.total_tokens", usage.totalTokens());
107+
}
108+
109+
private static LLMObs.LLMMessage llmMessage(ChatCompletion.Choice choice) {
110+
ChatCompletionMessage msg = choice.message();
111+
Optional<?> roleOpt = msg._role().asString();
112+
String role = "unknown";
113+
if (roleOpt.isPresent()) {
114+
role = String.valueOf(roleOpt.get());
115+
}
116+
String content = msg.content().orElse(null);
117+
118+
Optional<List<ChatCompletionMessageToolCall>> toolCallsOpt = msg.toolCalls();
119+
if (toolCallsOpt.isPresent() && !toolCallsOpt.get().isEmpty()) {
120+
List<LLMObs.ToolCall> toolCalls = new ArrayList<>();
121+
for (ChatCompletionMessageToolCall toolCall : toolCallsOpt.get()) {
122+
LLMObs.ToolCall llmObsToolCall = ToolCallExtractor.getToolCall(toolCall);
123+
if (llmObsToolCall != null) {
124+
toolCalls.add(llmObsToolCall);
125+
}
126+
}
127+
128+
if (!toolCalls.isEmpty()) {
129+
return LLMObs.LLMMessage.from(role, content, toolCalls);
130+
}
131+
}
132+
133+
return LLMObs.LLMMessage.from(role, content);
134+
}
135+
136+
public void withChatCompletionChunks(AgentSpan span, List<ChatCompletionChunk> chunks) {
137+
ChatCompletionAccumulator accumulator = ChatCompletionAccumulator.create();
138+
for (ChatCompletionChunk chunk : chunks) {
139+
accumulator.accumulate(chunk);
140+
}
141+
ChatCompletion chatCompletion = accumulator.chatCompletion();
142+
withChatCompletion(span, chatCompletion);
143+
}
144+
}

dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiModule.java renamed to dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionModule.java

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,14 +7,15 @@
77
import java.util.List;
88

99
@AutoService(InstrumenterModule.class)
10-
public class OpenAiModule extends InstrumenterModule.Tracing {
11-
public OpenAiModule() {
10+
public class ChatCompletionModule extends InstrumenterModule.Tracing {
11+
public ChatCompletionModule() {
1212
super("openai-java");
1313
}
1414

1515
@Override
1616
public String[] helperClassNames() {
1717
return new String[] {
18+
packageName + ".ChatCompletionDecorator",
1819
packageName + ".OpenAiDecorator",
1920
packageName + ".ResponseWrappers",
2021
packageName + ".ResponseWrappers$DDHttpResponseFor",
@@ -30,11 +31,6 @@ public String[] helperClassNames() {
3031
public List<Instrumenter> typeInstrumentations() {
3132
return Arrays.asList(
3233
new ChatCompletionServiceAsyncInstrumentation(),
33-
new ChatCompletionServiceInstrumentation(),
34-
new CompletionServiceAsyncInstrumentation(),
35-
new CompletionServiceInstrumentation(),
36-
new EmbeddingServiceInstrumentation(),
37-
new ResponseServiceAsyncInstrumentation(),
38-
new ResponseServiceInstrumentation());
34+
new ChatCompletionServiceInstrumentation());
3935
}
4036
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionServiceAsyncInstrumentation.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ public static AgentScope enter(
5656
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5757
DECORATE.afterStart(span);
5858
DECORATE.withClientOptions(span, clientOptions);
59-
DECORATE.withChatCompletionCreateParams(span, params, false);
59+
ChatCompletionDecorator.DECORATE.withChatCompletionCreateParams(span, params, false);
6060
return activateSpan(span);
6161
}
6262

@@ -71,7 +71,9 @@ public static void exit(
7171
DECORATE.onError(span, err);
7272
}
7373
if (future != null) {
74-
future = ResponseWrappers.wrapFutureResponse(future, span, DECORATE::withChatCompletion);
74+
future =
75+
ResponseWrappers.wrapFutureResponse(
76+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
7577
} else {
7678
span.finish();
7779
}
@@ -89,7 +91,7 @@ public static AgentScope enter(
8991
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
9092
DECORATE.afterStart(span);
9193
DECORATE.withClientOptions(span, clientOptions);
92-
DECORATE.withChatCompletionCreateParams(span, params, true);
94+
ChatCompletionDecorator.DECORATE.withChatCompletionCreateParams(span, params, true);
9395
return activateSpan(span);
9496
}
9597

@@ -107,7 +109,7 @@ public static void exit(
107109
if (future != null) {
108110
future =
109111
ResponseWrappers.wrapFutureStreamResponse(
110-
future, span, DECORATE::withChatCompletionChunks);
112+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
111113
} else {
112114
span.finish();
113115
}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ public static AgentScope enter(
5555
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5656
DECORATE.afterStart(span);
5757
DECORATE.withClientOptions(span, clientOptions);
58-
DECORATE.withChatCompletionCreateParams(span, params, false);
58+
ChatCompletionDecorator.DECORATE.withChatCompletionCreateParams(span, params, false);
5959
return activateSpan(span);
6060
}
6161

@@ -72,7 +72,7 @@ public static void exit(
7272
if (response != null) {
7373
response =
7474
ResponseWrappers.wrapResponse(
75-
response, span, OpenAiDecorator.DECORATE::withChatCompletion);
75+
response, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
7676
}
7777
DECORATE.beforeFinish(span);
7878
} finally {
@@ -91,7 +91,7 @@ public static AgentScope enter(
9191
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
9292
DECORATE.afterStart(span);
9393
DECORATE.withClientOptions(span, clientOptions);
94-
DECORATE.withChatCompletionCreateParams(span, params, true);
94+
ChatCompletionDecorator.DECORATE.withChatCompletionCreateParams(span, params, true);
9595
return activateSpan(span);
9696
}
9797

@@ -109,7 +109,7 @@ public static void exit(
109109
if (response != null) {
110110
response =
111111
ResponseWrappers.wrapStreamResponse(
112-
response, span, DECORATE::withChatCompletionChunks);
112+
response, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
113113
} else {
114114
span.finish();
115115
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CompletionDecorator.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package datadog.trace.instrumentation.openai_java;
2+
3+
import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.REQUEST_MODEL;
4+
import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.RESPONSE_MODEL;
5+
6+
import com.openai.models.completions.Completion;
7+
import com.openai.models.completions.CompletionCreateParams;
8+
import com.openai.models.completions.CompletionUsage;
9+
import datadog.trace.api.llmobs.LLMObs;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.bootstrap.instrumentation.api.Tags;
12+
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
13+
import java.util.Collections;
14+
import java.util.HashMap;
15+
import java.util.List;
16+
import java.util.Map;
17+
import java.util.stream.Collectors;
18+
19+
public class CompletionDecorator {
20+
public static final CompletionDecorator DECORATE = new CompletionDecorator();
21+
22+
private static final CharSequence COMPLETIONS_CREATE = UTF8BytesString.create("createCompletion");
23+
24+
public void withCompletionCreateParams(AgentSpan span, CompletionCreateParams params) {
25+
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
26+
27+
span.setResourceName(COMPLETIONS_CREATE);
28+
span.setTag("openai.request.endpoint", "v1/completions");
29+
span.setTag("openai.request.method", "POST");
30+
if (params == null) {
31+
return;
32+
}
33+
34+
params.model()._value().asString().ifPresent(str -> span.setTag(REQUEST_MODEL, str));
35+
params
36+
.prompt()
37+
.flatMap(p -> p.string())
38+
.ifPresent(
39+
input ->
40+
span.setTag(
41+
"_ml_obs_tag.input",
42+
Collections.singletonList(LLMObs.LLMMessage.from(null, input))));
43+
44+
Map<String, Object> metadata = new HashMap<>();
45+
params.maxTokens().ifPresent(v -> metadata.put("max_tokens", v));
46+
params.temperature().ifPresent(v -> metadata.put("temperature", v));
47+
span.setTag("_ml_obs_tag.metadata", metadata);
48+
}
49+
50+
public void withCompletion(AgentSpan span, Completion completion) {
51+
String modelName = completion.model();
52+
span.setTag(RESPONSE_MODEL, modelName);
53+
span.setTag("_ml_obs_tag.model_name", modelName);
54+
span.setTag("_ml_obs_tag.model_provider", "openai");
55+
56+
List<LLMObs.LLMMessage> output =
57+
completion.choices().stream()
58+
.map(v -> LLMObs.LLMMessage.from(null, v.text()))
59+
.collect(Collectors.toList());
60+
span.setTag("_ml_obs_tag.output", output);
61+
62+
completion.usage().ifPresent(usage -> withCompletionUsage(span, usage));
63+
}
64+
65+
public void withCompletions(AgentSpan span, List<Completion> completions) {
66+
if (!completions.isEmpty()) {
67+
withCompletion(span, completions.get(0));
68+
}
69+
}
70+
71+
private static void withCompletionUsage(AgentSpan span, CompletionUsage usage) {
72+
span.setTag("_ml_obs_metric.input_tokens", usage.promptTokens());
73+
span.setTag("_ml_obs_metric.output_tokens", usage.completionTokens());
74+
span.setTag("_ml_obs_metric.total_tokens", usage.totalTokens());
75+
}
76+
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CompletionModule.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package datadog.trace.instrumentation.openai_java;
2+
3+
import com.google.auto.service.AutoService;
4+
import datadog.trace.agent.tooling.Instrumenter;
5+
import datadog.trace.agent.tooling.InstrumenterModule;
6+
import java.util.Arrays;
7+
import java.util.List;
8+
9+
@AutoService(InstrumenterModule.class)
10+
public class CompletionModule extends InstrumenterModule.Tracing {
11+
public CompletionModule() {
12+
super("openai-java");
13+
}
14+
15+
@Override
16+
public String[] helperClassNames() {
17+
return new String[] {
18+
packageName + ".CompletionDecorator",
19+
packageName + ".OpenAiDecorator",
20+
packageName + ".ResponseWrappers",
21+
packageName + ".ResponseWrappers$DDHttpResponseFor",
22+
packageName + ".ResponseWrappers$1",
23+
packageName + ".ResponseWrappers$2",
24+
packageName + ".ResponseWrappers$2$1"
25+
};
26+
}
27+
28+
@Override
29+
public List<Instrumenter> typeInstrumentations() {
30+
return Arrays.asList(
31+
new CompletionServiceAsyncInstrumentation(), new CompletionServiceInstrumentation());
32+
}
33+
}

0 commit comments

Comments
 (0)