Skip to content

Commit f8d8ce5

Browse files
ygreeygree
committed
Do not decorate LLMObs specific tags when it's disabled
1 parent 1ca2ae0 commit f8d8ce5

File tree

5 files changed

+83
-24
lines changed

5 files changed

+83
-24
lines changed

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

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
import com.openai.models.chat.completions.ChatCompletionMessage;
1111
import com.openai.models.chat.completions.ChatCompletionMessageParam;
1212
import com.openai.models.chat.completions.ChatCompletionMessageToolCall;
13-
import com.openai.models.completions.CompletionUsage;
13+
import datadog.trace.api.Config;
1414
import datadog.trace.api.llmobs.LLMObs;
1515
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1616
import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -28,12 +28,18 @@ public class ChatCompletionDecorator {
2828
private static final CharSequence CHAT_COMPLETIONS_CREATE =
2929
UTF8BytesString.create("createChatCompletion");
3030

31+
private final boolean llmObsEnabled = Config.get().isLlmObsEnabled();
32+
3133
public void withChatCompletionCreateParams(
3234
AgentSpan span, ChatCompletionCreateParams params, boolean stream) {
33-
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
3435
span.setResourceName(CHAT_COMPLETIONS_CREATE);
3536
span.setTag("openai.request.endpoint", "v1/chat/completions");
3637
span.setTag("openai.request.method", "POST");
38+
if (!llmObsEnabled) {
39+
return;
40+
}
41+
42+
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
3743
if (params == null) {
3844
return;
3945
}
@@ -86,6 +92,9 @@ private static LLMObs.LLMMessage llmMessage(ChatCompletionMessageParam m) {
8692
}
8793

8894
public void withChatCompletion(AgentSpan span, ChatCompletion completion) {
95+
if (!llmObsEnabled) {
96+
return;
97+
}
8998
String modelName = completion.model();
9099
span.setTag(RESPONSE_MODEL, modelName);
91100
span.setTag("_ml_obs_tag.model_name", modelName);
@@ -97,13 +106,14 @@ public void withChatCompletion(AgentSpan span, ChatCompletion completion) {
97106
.collect(Collectors.toList());
98107
span.setTag("_ml_obs_tag.output", output);
99108

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());
109+
completion
110+
.usage()
111+
.ifPresent(
112+
usage -> {
113+
span.setTag("_ml_obs_metric.input_tokens", usage.promptTokens());
114+
span.setTag("_ml_obs_metric.output_tokens", usage.completionTokens());
115+
span.setTag("_ml_obs_metric.total_tokens", usage.totalTokens());
116+
});
107117
}
108118

109119
private static LLMObs.LLMMessage llmMessage(ChatCompletion.Choice choice) {
@@ -134,6 +144,9 @@ private static LLMObs.LLMMessage llmMessage(ChatCompletion.Choice choice) {
134144
}
135145

136146
public void withChatCompletionChunks(AgentSpan span, List<ChatCompletionChunk> chunks) {
147+
if (!llmObsEnabled) {
148+
return;
149+
}
137150
ChatCompletionAccumulator accumulator = ChatCompletionAccumulator.create();
138151
for (ChatCompletionChunk chunk : chunks) {
139152
accumulator.accumulate(chunk);

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

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
import com.openai.models.completions.Completion;
77
import com.openai.models.completions.CompletionCreateParams;
8-
import com.openai.models.completions.CompletionUsage;
8+
import datadog.trace.api.Config;
99
import datadog.trace.api.llmobs.LLMObs;
1010
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1111
import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -21,12 +21,17 @@ public class CompletionDecorator {
2121

2222
private static final CharSequence COMPLETIONS_CREATE = UTF8BytesString.create("createCompletion");
2323

24-
public void withCompletionCreateParams(AgentSpan span, CompletionCreateParams params) {
25-
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
24+
private final boolean llmObsEnabled = Config.get().isLlmObsEnabled();
2625

26+
public void withCompletionCreateParams(AgentSpan span, CompletionCreateParams params) {
2727
span.setResourceName(COMPLETIONS_CREATE);
2828
span.setTag("openai.request.endpoint", "v1/completions");
2929
span.setTag("openai.request.method", "POST");
30+
if (!llmObsEnabled) {
31+
return;
32+
}
33+
34+
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
3035
if (params == null) {
3136
return;
3237
}
@@ -48,6 +53,10 @@ public void withCompletionCreateParams(AgentSpan span, CompletionCreateParams pa
4853
}
4954

5055
public void withCompletion(AgentSpan span, Completion completion) {
56+
if (!llmObsEnabled) {
57+
return;
58+
}
59+
5160
String modelName = completion.model();
5261
span.setTag(RESPONSE_MODEL, modelName);
5362
span.setTag("_ml_obs_tag.model_name", modelName);
@@ -59,18 +68,23 @@ public void withCompletion(AgentSpan span, Completion completion) {
5968
.collect(Collectors.toList());
6069
span.setTag("_ml_obs_tag.output", output);
6170

62-
completion.usage().ifPresent(usage -> withCompletionUsage(span, usage));
71+
completion
72+
.usage()
73+
.ifPresent(
74+
usage -> {
75+
span.setTag("_ml_obs_metric.input_tokens", usage.promptTokens());
76+
span.setTag("_ml_obs_metric.output_tokens", usage.completionTokens());
77+
span.setTag("_ml_obs_metric.total_tokens", usage.totalTokens());
78+
});
6379
}
6480

6581
public void withCompletions(AgentSpan span, List<Completion> completions) {
82+
if (!llmObsEnabled) {
83+
return;
84+
}
85+
6686
if (!completions.isEmpty()) {
6787
withCompletion(span, completions.get(0));
6888
}
6989
}
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-
}
7690
}

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import com.openai.models.embeddings.CreateEmbeddingResponse;
77
import com.openai.models.embeddings.Embedding;
88
import com.openai.models.embeddings.EmbeddingCreateParams;
9+
import datadog.trace.api.Config;
910
import datadog.trace.api.llmobs.LLMObs;
1011
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1112
import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -22,11 +23,17 @@ public class EmbeddingDecorator {
2223

2324
private static final CharSequence EMBEDDINGS_CREATE = UTF8BytesString.create("createEmbedding");
2425

26+
private final boolean llmObsEnabled = Config.get().isLlmObsEnabled();
27+
2528
public void withEmbeddingCreateParams(AgentSpan span, EmbeddingCreateParams params) {
26-
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_EMBEDDING_SPAN_KIND);
2729
span.setResourceName(EMBEDDINGS_CREATE);
2830
span.setTag("openai.request.endpoint", "v1/embeddings");
2931
span.setTag("openai.request.method", "POST");
32+
if (!llmObsEnabled) {
33+
return;
34+
}
35+
36+
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_EMBEDDING_SPAN_KIND);
3037
if (params == null) {
3138
return;
3239
}
@@ -52,6 +59,10 @@ private List<LLMObs.Document> embeddingDocuments(EmbeddingCreateParams.Input inp
5259
}
5360

5461
public void withCreateEmbeddingResponse(AgentSpan span, CreateEmbeddingResponse response) {
62+
if (!llmObsEnabled) {
63+
return;
64+
}
65+
5566
String modelName = response.model();
5667
span.setTag(RESPONSE_MODEL, modelName);
5768
span.setTag("_ml_obs_tag.model_name", modelName);

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.openai.core.ClientOptions;
44
import com.openai.core.http.Headers;
5+
import datadog.trace.api.Config;
56
import datadog.trace.api.llmobs.LLMObsContext;
67
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
78
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -22,6 +23,8 @@ public class OpenAiDecorator extends ClientDecorator {
2223

2324
private static final CharSequence COMPONENT_NAME = UTF8BytesString.create("openai");
2425

26+
private final boolean llmObsEnabled = Config.get().isLlmObsEnabled();
27+
2528
public AgentSpan startSpan(ClientOptions clientOptions) {
2629
AgentSpan span = AgentTracer.startSpan(INSTRUMENTATION_NAME, SPAN_NAME);
2730
afterStart(span);
@@ -61,17 +64,20 @@ protected CharSequence component() {
6164

6265
@Override
6366
public AgentSpan afterStart(AgentSpan span) {
64-
// TODO only if llmobs enabled
65-
span.setTag("_ml_obs_tag.parent_id", LLMObsContext.parentSpanId());
67+
if (llmObsEnabled) {
68+
span.setTag("_ml_obs_tag.parent_id", LLMObsContext.parentSpanId());
69+
}
6670
return super.afterStart(span);
6771
}
6872

6973
public void withHttpResponse(AgentSpan span, Headers headers) {
74+
if (!llmObsEnabled) {
75+
return;
76+
}
7077
List<String> values = headers.values("openai-organization");
7178
if (!values.isEmpty()) {
7279
span.setTag(OPENAI_ORGANIZATION_NAME, values.get(0));
7380
}
74-
7581
setMetricFromHeader(
7682
span,
7783
"openai.organization.ratelimit.requests.limit",

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.openai.models.responses.ResponseReasoningItem;
1919
import com.openai.models.responses.ResponseStreamEvent;
2020
import datadog.json.JsonWriter;
21+
import datadog.trace.api.Config;
2122
import datadog.trace.api.llmobs.LLMObs;
2223
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2324
import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -34,11 +35,17 @@ public class ResponseDecorator {
3435

3536
private static final CharSequence RESPONSES_CREATE = UTF8BytesString.create("createResponse");
3637

38+
private final boolean llmObsEnabled = Config.get().isLlmObsEnabled();
39+
3740
public void withResponseCreateParams(AgentSpan span, ResponseCreateParams params) {
38-
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
3941
span.setResourceName(RESPONSES_CREATE);
4042
span.setTag("openai.request.endpoint", "v1/responses");
4143
span.setTag("openai.request.method", "POST");
44+
if (!llmObsEnabled) {
45+
return;
46+
}
47+
48+
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
4249
if (params == null) {
4350
return;
4451
}
@@ -366,6 +373,10 @@ public void withResponse(AgentSpan span, Response response) {
366373
}
367374

368375
public void withResponseStreamEvents(AgentSpan span, List<ResponseStreamEvent> events) {
376+
if (!llmObsEnabled) {
377+
return;
378+
}
379+
369380
for (ResponseStreamEvent event : events) {
370381
if (event.isCompleted()) {
371382
Response response = event.asCompleted().response();
@@ -381,6 +392,10 @@ public void withResponseStreamEvents(AgentSpan span, List<ResponseStreamEvent> e
381392
}
382393

383394
private void withResponse(AgentSpan span, Response response, boolean stream) {
395+
if (!llmObsEnabled) {
396+
return;
397+
}
398+
384399
String modelName = extractResponseModel(response._model());
385400
span.setTag(RESPONSE_MODEL, modelName);
386401
span.setTag("_ml_obs_tag.model_name", modelName);

0 commit comments

Comments
 (0)