Skip to content

Commit 1850e79

Browse files
ygreeygree
committed
Reorg/rename decorator functions
Embeddings WIP
1 parent c26909b commit 1850e79

11 files changed

+131
-125
lines changed

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,8 @@ public static AgentScope enter(
5555
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
5656
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5757
DECORATE.afterStart(span);
58-
DECORATE.decorateWithClientOptions(span, clientOptions);
59-
DECORATE.decorateChatCompletion(span, params, false);
58+
DECORATE.withClientOptions(span, clientOptions);
59+
DECORATE.withChatCompletionCreateParams(span, params, false);
6060
return activateSpan(span);
6161
}
6262

@@ -71,9 +71,7 @@ public static void exit(
7171
DECORATE.onError(span, err);
7272
}
7373
if (future != null) {
74-
future =
75-
ResponseWrappers.wrapFutureResponse(
76-
future, span, DECORATE::decorateWithChatCompletion);
74+
future = ResponseWrappers.wrapFutureResponse(future, span, DECORATE::withChatCompletion);
7775
} else {
7876
span.finish();
7977
}
@@ -90,8 +88,8 @@ public static AgentScope enter(
9088
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
9189
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
9290
DECORATE.afterStart(span);
93-
DECORATE.decorateWithClientOptions(span, clientOptions);
94-
DECORATE.decorateChatCompletion(span, params, true);
91+
DECORATE.withClientOptions(span, clientOptions);
92+
DECORATE.withChatCompletionCreateParams(span, params, true);
9593
return activateSpan(span);
9694
}
9795

@@ -109,7 +107,7 @@ public static void exit(
109107
if (future != null) {
110108
future =
111109
ResponseWrappers.wrapFutureStreamResponse(
112-
future, span, DECORATE::decorateWithChatCompletionChunks);
110+
future, span, DECORATE::withChatCompletionChunks);
113111
} else {
114112
span.finish();
115113
}

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ public static AgentScope enter(
5454
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
5555
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5656
DECORATE.afterStart(span);
57-
DECORATE.decorateWithClientOptions(span, clientOptions);
58-
DECORATE.decorateChatCompletion(span, params, false);
57+
DECORATE.withClientOptions(span, clientOptions);
58+
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::decorateWithChatCompletion);
75+
response, span, OpenAiDecorator.DECORATE::withChatCompletion);
7676
}
7777
DECORATE.beforeFinish(span);
7878
} finally {
@@ -90,8 +90,8 @@ public static AgentScope enter(
9090
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
9191
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
9292
DECORATE.afterStart(span);
93-
DECORATE.decorateWithClientOptions(span, clientOptions);
94-
DECORATE.decorateChatCompletion(span, params, true);
93+
DECORATE.withClientOptions(span, clientOptions);
94+
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::decorateWithChatCompletionChunks);
112+
response, span, DECORATE::withChatCompletionChunks);
113113
} else {
114114
span.finish();
115115
}

‎dd-java-agent/instrumentation/openai-java/openai-java-1.0/src/main/java/datadog/trace/instrumentation/openai_java/CompletionServiceAsyncInstrumentation.java‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@ public static AgentScope enter(
5050
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
5151
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5252
DECORATE.afterStart(span);
53-
DECORATE.decorateWithClientOptions(span, clientOptions);
54-
DECORATE.decorateCompletion(span, params);
53+
DECORATE.withClientOptions(span, clientOptions);
54+
DECORATE.withCompletionCreateParams(span, params);
5555
return activateSpan(span);
5656
}
5757

@@ -66,8 +66,7 @@ public static void exit(
6666
DECORATE.onError(span, err);
6767
}
6868
if (future != null) {
69-
future =
70-
ResponseWrappers.wrapFutureResponse(future, span, DECORATE::decorateWithCompletion);
69+
future = ResponseWrappers.wrapFutureResponse(future, span, DECORATE::withCompletion);
7170
} else {
7271
span.finish();
7372
}
@@ -85,8 +84,8 @@ public static AgentScope enter(
8584
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
8685
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
8786
DECORATE.afterStart(span);
88-
DECORATE.decorateWithClientOptions(span, clientOptions);
89-
DECORATE.decorateCompletion(span, params);
87+
DECORATE.withClientOptions(span, clientOptions);
88+
DECORATE.withCompletionCreateParams(span, params);
9089
return activateSpan(span);
9190
}
9291

@@ -103,8 +102,7 @@ public static void exit(
103102
}
104103
if (future != null) {
105104
future =
106-
ResponseWrappers.wrapFutureStreamResponse(
107-
future, span, DECORATE::decorateWithCompletions);
105+
ResponseWrappers.wrapFutureStreamResponse(future, span, DECORATE::withCompletions);
108106
} else {
109107
span.finish();
110108
}

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,8 @@ public static AgentScope enter(
5353
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5454
// TODO get span from context?
5555
DECORATE.afterStart(span);
56-
DECORATE.decorateWithClientOptions(span, clientOptions);
57-
DECORATE.decorateCompletion(span, params);
56+
DECORATE.withClientOptions(span, clientOptions);
57+
DECORATE.withCompletionCreateParams(span, params);
5858

5959
llmScope = LLMObsContext.attach(span.context());
6060
// TODO should the agent span be activated via the context api or keep separate?
@@ -75,7 +75,7 @@ public static void exit(
7575
if (response != null) {
7676
response =
7777
ResponseWrappers.wrapResponse(
78-
response, span, OpenAiDecorator.DECORATE::decorateWithCompletion);
78+
response, span, OpenAiDecorator.DECORATE::withCompletion);
7979
}
8080
DECORATE.beforeFinish(span);
8181
} finally {
@@ -94,8 +94,8 @@ public static AgentScope enter(
9494
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
9595
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
9696
DECORATE.afterStart(span);
97-
DECORATE.decorateWithClientOptions(span, clientOptions);
98-
DECORATE.decorateCompletion(span, params);
97+
DECORATE.withClientOptions(span, clientOptions);
98+
DECORATE.withCompletionCreateParams(span, params);
9999
return activateSpan(span);
100100
}
101101

@@ -110,9 +110,7 @@ public static void exit(
110110
DECORATE.onError(span, err);
111111
}
112112
if (response != null) {
113-
response =
114-
ResponseWrappers.wrapStreamResponse(
115-
response, span, DECORATE::decorateWithCompletions);
113+
response = ResponseWrappers.wrapStreamResponse(response, span, DECORATE::withCompletions);
116114
} else {
117115
span.finish();
118116
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,8 @@ public static AgentScope enter(
4141
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
4242
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
4343
DECORATE.afterStart(span);
44-
DECORATE.decorateWithClientOptions(span, clientOptions);
45-
DECORATE.decorateEmbedding(span, params);
44+
DECORATE.withClientOptions(span, clientOptions);
45+
DECORATE.withEmbeddingCreateParams(span, params);
4646
return activateSpan(span);
4747
}
4848

@@ -59,7 +59,7 @@ public static void exit(
5959
if (response != null) {
6060
response =
6161
ResponseWrappers.wrapResponse(
62-
response, span, OpenAiDecorator.DECORATE::decorateWithEmbedding);
62+
response, span, OpenAiDecorator.DECORATE::withCreateEmbeddingResponse);
6363
}
6464
DECORATE.beforeFinish(span);
6565
} finally {

0 commit comments

Comments
 (0)