Skip to content

Commit 0ee6b46

Browse files
nayeem-kamalnayeem-kamal
committed
added apm openai instrumentation
1 parent aa9dd27 commit 0ee6b46

File tree

9 files changed

+733
-0
lines changed

9 files changed

+733
-0
lines changed

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
muzzle {
2+
pass {
3+
group = "com.openai"
4+
module = "openai-java"
5+
versions = "[2.8.0,)"
6+
assertInverse = true
7+
}
8+
}
9+
10+
// Instrumentation is currently in preview. Additional testing will be implemented before enabling this by default.
11+
//minimumBranchCoverage = 0.0
12+
//minimumInstructionCoverage = 0.0
13+
14+
apply from: "$rootDir/gradle/java.gradle"
15+
16+
addTestSuiteForDir('latestDepTest', 'test')
17+
18+
dependencies {
19+
compileOnly(group: 'com.openai', name: 'openai-java', version: '2.8.1')
20+
21+
testImplementation(group: 'com.openai', name: 'openai-java') {
22+
version {
23+
strictly '[2.8.0,)'
24+
}
25+
}
26+
27+
latestDepTestImplementation group: 'com.openai', name: 'openai-java', version: '+'
28+
29+
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
# This is a Gradle generated file for dependency locking.
2+
# This file is expected to be part of source control.
3+
# Manual entries are allowed.
4+
empty=

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

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package datadog.trace.instrumentation.openai;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.*;
6+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
7+
8+
import com.openai.models.chat.completions.ChatCompletionCreateParams;
9+
import com.openai.services.blocking.CompletionService;
10+
import datadog.trace.agent.tooling.Instrumenter;
11+
import datadog.trace.agent.tooling.InstrumenterModule;
12+
import datadog.trace.bootstrap.InstrumentationContext;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
14+
import java.util.HashMap;
15+
import java.util.Map;
16+
import net.bytebuddy.asm.Advice;
17+
import net.bytebuddy.description.type.TypeDescription;
18+
import net.bytebuddy.matcher.ElementMatcher;
19+
20+
public class ChatCompletionServiceInstrumentation extends InstrumenterModule.Tracing
21+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
22+
public ChatCompletionServiceInstrumentation() {
23+
super("openai-client");
24+
}
25+
26+
@Override
27+
public String[] helperClassNames() {
28+
return new String[] {
29+
packageName + ".OpenAIClientInfo",
30+
};
31+
}
32+
33+
@Override
34+
public Map<String, String> contextStore() {
35+
Map<String, String> contextStores = new HashMap<>(1);
36+
contextStores.put(
37+
"com.openai.services.blocking.CompletionService", OpenAIClientInfo.class.getName());
38+
return contextStores;
39+
}
40+
41+
@Override
42+
public void methodAdvice(MethodTransformer transformer) {
43+
transformer.applyAdvice(
44+
isMethod()
45+
.and(isPublic())
46+
.and(named("create"))
47+
.and(
48+
takesArgument(
49+
0, named("com.openai.models.completions.ChatCompletionCreateParams"))),
50+
getClass().getName() + "$CompletionServiceAdvice");
51+
}
52+
53+
@Override
54+
public String hierarchyMarkerType() {
55+
return "com.openai.services.blocking.chat.ChatCompletionService";
56+
}
57+
58+
@Override
59+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
60+
return implementsInterface(named(hierarchyMarkerType()));
61+
}
62+
63+
public static class CompletionServiceAdvice {
64+
@Advice.OnMethodEnter(suppress = Throwable.class)
65+
public static AgentScope methodEnter(
66+
@Advice.Argument(0) final ChatCompletionCreateParams params) {
67+
68+
return OpenAIClientDecorator.DECORATE.startChatCompletionSpan(params);
69+
}
70+
71+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
72+
public static void methodExit(
73+
@Advice.Enter final AgentScope span,
74+
@Advice.This final CompletionService completionService,
75+
@Advice.Return final Object result,
76+
@Advice.Thrown final Throwable throwable) {
77+
OpenAIClientInfo info =
78+
InstrumentationContext.get(CompletionService.class, OpenAIClientInfo.class)
79+
.get(completionService);
80+
OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable);
81+
}
82+
}
83+
}

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package datadog.trace.instrumentation.openai;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.*;
6+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
7+
8+
import com.openai.models.completions.CompletionCreateParams;
9+
import com.openai.services.blocking.CompletionService;
10+
import datadog.trace.agent.tooling.Instrumenter;
11+
import datadog.trace.agent.tooling.InstrumenterModule;
12+
import datadog.trace.bootstrap.InstrumentationContext;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
14+
import java.util.HashMap;
15+
import java.util.Map;
16+
import net.bytebuddy.asm.Advice;
17+
import net.bytebuddy.description.type.TypeDescription;
18+
import net.bytebuddy.matcher.ElementMatcher;
19+
20+
public class CompletionServiceInstrumentation extends InstrumenterModule.Tracing
21+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
22+
public CompletionServiceInstrumentation() {
23+
super("openai-client");
24+
}
25+
26+
@Override
27+
public String[] helperClassNames() {
28+
return new String[] {
29+
packageName + ".OpenAIClientInfo",
30+
};
31+
}
32+
33+
@Override
34+
public Map<String, String> contextStore() {
35+
Map<String, String> contextStores = new HashMap<>(1);
36+
contextStores.put(
37+
"com.openai.services.blocking.CompletionService", OpenAIClientInfo.class.getName());
38+
return contextStores;
39+
}
40+
41+
@Override
42+
public void methodAdvice(MethodTransformer transformer) {
43+
transformer.applyAdvice(
44+
isMethod()
45+
.and(isPublic())
46+
.and(named("create"))
47+
.and(takesArgument(0, named("com.openai.models.completions.CompletionCreateParams"))),
48+
getClass().getName() + "$CompletionServiceAdvice");
49+
}
50+
51+
@Override
52+
public String hierarchyMarkerType() {
53+
return "com.openai.services.blocking.CompletionService";
54+
}
55+
56+
@Override
57+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
58+
return implementsInterface(named(hierarchyMarkerType()));
59+
}
60+
61+
public static class CompletionServiceAdvice {
62+
@Advice.OnMethodEnter(suppress = Throwable.class)
63+
public static AgentScope methodEnter(@Advice.Argument(0) final CompletionCreateParams params) {
64+
65+
return OpenAIClientDecorator.DECORATE.startCompletionSpan(params);
66+
}
67+
68+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
69+
public static void methodExit(
70+
@Advice.Enter final AgentScope span,
71+
@Advice.This final CompletionService completionService,
72+
@Advice.Return final Object result,
73+
@Advice.Thrown final Throwable throwable) {
74+
OpenAIClientInfo info =
75+
InstrumentationContext.get(CompletionService.class, OpenAIClientInfo.class)
76+
.get(completionService);
77+
OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable);
78+
}
79+
}
80+
}

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package datadog.trace.instrumentation.openai;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
7+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
8+
9+
import com.google.auto.service.AutoService;
10+
import com.openai.models.embeddings.EmbeddingCreateParams;
11+
import datadog.trace.agent.tooling.Instrumenter;
12+
import datadog.trace.agent.tooling.InstrumenterModule;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
14+
import net.bytebuddy.asm.Advice;
15+
import net.bytebuddy.description.type.TypeDescription;
16+
import net.bytebuddy.matcher.ElementMatcher;
17+
18+
@AutoService(InstrumenterModule.class)
19+
public class EmbeddingServiceInstrumentation extends InstrumenterModule.Tracing
20+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
21+
22+
public EmbeddingServiceInstrumentation() {
23+
super("openai", "openai-java");
24+
}
25+
26+
@Override
27+
public String[] helperClassNames() {
28+
return new String[] {
29+
packageName + ".OpenAIDecorator",
30+
};
31+
}
32+
33+
@Override
34+
public void methodAdvice(MethodTransformer transformer) {
35+
// Instrument embedding creation methods
36+
transformer.applyAdvice(
37+
isMethod()
38+
.and(named("create"))
39+
.and(isPublic())
40+
.and(takesArgument(0, named("com.openai.models.embeddings.EmbeddingCreateParams"))),
41+
EmbeddingServiceInstrumentation.class.getName() + "$EmbeddingAdvice");
42+
}
43+
44+
@Override
45+
public String hierarchyMarkerType() {
46+
return "com.openai.services.embedding.EmbeddingService";
47+
}
48+
49+
@Override
50+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
51+
return implementsInterface(named(hierarchyMarkerType()));
52+
}
53+
54+
public static class EmbeddingAdvice {
55+
@Advice.OnMethodEnter(suppress = Throwable.class)
56+
public static AgentScope onEnter(
57+
@Advice.Argument(0) final EmbeddingCreateParams embeddingParams) {
58+
59+
return OpenAIClientDecorator.DECORATE.startEmbeddingSpan(embeddingParams);
60+
}
61+
62+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
63+
public static void onExit(
64+
@Advice.Enter final AgentScope span,
65+
@Advice.Return final Object result,
66+
@Advice.Thrown final Throwable throwable) {
67+
68+
OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable);
69+
}
70+
}
71+
}

0 commit comments

Comments
 (0)