Skip to content

Commit 8718f94

Browse files
ygreeygree
committed
Clean up HttpResponse wrappers
1 parent 2d127f2 commit 8718f94

15 files changed

+249
-186
lines changed

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,9 @@ public String[] helperClassNames() {
1717
return new String[] {
1818
packageName + ".ChatCompletionDecorator",
1919
packageName + ".OpenAiDecorator",
20-
packageName + ".HttpResponseWrappers",
21-
packageName + ".HttpResponseWrappers$DDHttpResponseFor",
22-
packageName + ".HttpResponseWrappers$1",
23-
packageName + ".HttpResponseWrappers$2",
24-
packageName + ".HttpResponseWrappers$2$1",
20+
packageName + ".HttpResponseWrapper",
21+
packageName + ".HttpStreamResponseWrapper",
22+
packageName + ".HttpStreamResponseStreamWrapper",
2523
packageName + ".ToolCallExtractor",
2624
packageName + ".ToolCallExtractor$1"
2725
};

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,8 @@ public static void exit(
6969
}
7070
if (future != null) {
7171
future =
72-
HttpResponseWrappers.wrapFutureHttpResponse(
73-
future, span, ChatCompletionDecorator.DECORATE::withChatCompletion, DECORATE::onError);
72+
HttpResponseWrapper.wrapFuture(
73+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
7474
} else {
7575
span.finish();
7676
}
@@ -103,8 +103,8 @@ public static void exit(
103103
}
104104
if (future != null) {
105105
future =
106-
HttpResponseWrappers.wrapFutureHttpResponseStream(
107-
future, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks, DECORATE::onError);
106+
HttpStreamResponseWrapper.wrapFuture(
107+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
108108
} else {
109109
span.finish();
110110
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.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
@@ -68,13 +68,13 @@ public static void exit(
6868
}
6969
if (response != null) {
7070
response =
71-
HttpResponseWrappers.wrapHttpResponse(
72-
response, span, ChatCompletionDecorator.DECORATE::withChatCompletion, DECORATE::onError); // TODO withChatCompletion may not be called if parsing failed, or never called b/o it's lazy
71+
HttpResponseWrapper.wrap(
72+
response, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
7373
}
7474
DECORATE.beforeFinish(span);
7575
} finally {
7676
scope.close();
77-
span.finish(); // TODO so we finish here but decorate withChatCompletion later? should move
77+
span.finish();
7878
}
7979
}
8080
}
@@ -103,12 +103,12 @@ public static void exit(
103103
}
104104
if (response != null) {
105105
response =
106-
HttpResponseWrappers.wrapHttpResponseStream(
107-
response, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks, DECORATE::onError);
106+
HttpStreamResponseWrapper.wrap(
107+
response, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
108108
} else {
109+
DECORATE.beforeFinish(span);
109110
span.finish();
110111
}
111-
DECORATE.beforeFinish(span);
112112
} finally {
113113
scope.close();
114114
}

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,9 @@ public String[] helperClassNames() {
1717
return new String[] {
1818
packageName + ".CompletionDecorator",
1919
packageName + ".OpenAiDecorator",
20-
packageName + ".HttpResponseWrappers",
21-
packageName + ".HttpResponseWrappers$DDHttpResponseFor",
22-
packageName + ".HttpResponseWrappers$1",
23-
packageName + ".HttpResponseWrappers$2",
24-
packageName + ".HttpResponseWrappers$2$1"
20+
packageName + ".HttpResponseWrapper",
21+
packageName + ".HttpStreamResponseWrapper",
22+
packageName + ".HttpStreamResponseStreamWrapper",
2523
};
2624
}
2725

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,8 @@ public static void exit(
6464
}
6565
if (future != null) {
6666
future =
67-
HttpResponseWrappers.wrapFutureHttpResponse(
68-
future, span, CompletionDecorator.DECORATE::withCompletion, DECORATE::onError);
67+
HttpResponseWrapper.wrapFuture(
68+
future, span, CompletionDecorator.DECORATE::withCompletion);
6969
} else {
7070
span.finish();
7171
}
@@ -99,8 +99,8 @@ public static void exit(
9999
}
100100
if (future != null) {
101101
future =
102-
HttpResponseWrappers.wrapFutureHttpResponseStream(
103-
future, span, CompletionDecorator.DECORATE::withCompletions, DECORATE::onError);
102+
HttpStreamResponseWrapper.wrapFuture(
103+
future, span, CompletionDecorator.DECORATE::withCompletions);
104104
} else {
105105
span.finish();
106106
}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,8 @@ public static void exit(
7474
}
7575
if (response != null) {
7676
response =
77-
HttpResponseWrappers.wrapHttpResponse(
78-
response, span, CompletionDecorator.DECORATE::withCompletion, DECORATE::onError);
77+
HttpResponseWrapper.wrap(
78+
response, span, CompletionDecorator.DECORATE::withCompletion);
7979
}
8080
DECORATE.beforeFinish(span);
8181
} finally {
@@ -109,8 +109,8 @@ public static void exit(
109109
}
110110
if (response != null) {
111111
response =
112-
HttpResponseWrappers.wrapHttpResponseStream(
113-
response, span, CompletionDecorator.DECORATE::withCompletions, DECORATE::onError);
112+
HttpStreamResponseWrapper.wrap(
113+
response, span, CompletionDecorator.DECORATE::withCompletions);
114114
} else {
115115
span.finish();
116116
}

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,9 @@ public String[] helperClassNames() {
1717
return new String[] {
1818
packageName + ".EmbeddingDecorator",
1919
packageName + ".OpenAiDecorator",
20-
packageName + ".HttpResponseWrappers",
21-
packageName + ".HttpResponseWrappers$DDHttpResponseFor",
22-
packageName + ".HttpResponseWrappers$1",
23-
packageName + ".HttpResponseWrappers$2",
24-
packageName + ".HttpResponseWrappers$2$1"
20+
packageName + ".HttpResponseWrapper",
21+
packageName + ".HttpStreamResponseWrapper",
22+
packageName + ".HttpStreamResponseStreamWrapper",
2523
};
2624
}
2725

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,8 @@ public static void exit(
5555
}
5656
if (response != null) {
5757
response =
58-
HttpResponseWrappers.wrapHttpResponse(
59-
response, span, EmbeddingDecorator.DECORATE::withCreateEmbeddingResponse, DECORATE::onError);
58+
HttpResponseWrapper.wrap(
59+
response, span, EmbeddingDecorator.DECORATE::withCreateEmbeddingResponse);
6060
}
6161
DECORATE.beforeFinish(span);
6262
} finally {

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

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package datadog.trace.instrumentation.openai_java;
2+
3+
import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.DECORATE;
4+
5+
import com.openai.core.http.Headers;
6+
import com.openai.core.http.HttpResponseFor;
7+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
8+
import java.io.InputStream;
9+
import java.util.concurrent.CompletableFuture;
10+
import java.util.function.BiConsumer;
11+
import org.jetbrains.annotations.NotNull;
12+
13+
public final class HttpResponseWrapper<T> implements HttpResponseFor<T> {
14+
15+
public static <T> HttpResponseFor<T> wrap(
16+
HttpResponseFor<T> response, AgentSpan span, BiConsumer<AgentSpan, T> afterParse) {
17+
DECORATE.withHttpResponse(span, response.headers());
18+
return new HttpResponseWrapper<>(response, span, afterParse);
19+
}
20+
21+
public static <T> CompletableFuture<HttpResponseFor<T>> wrapFuture(
22+
CompletableFuture<HttpResponseFor<T>> future,
23+
AgentSpan span,
24+
BiConsumer<AgentSpan, T> afterParse) {
25+
return future
26+
.thenApply(response -> wrap(response, span, afterParse))
27+
.whenComplete(
28+
(r, t) -> {
29+
try {
30+
if (t != null) {
31+
DECORATE.onError(span, t);
32+
}
33+
DECORATE.beforeFinish(span);
34+
} finally {
35+
span.finish();
36+
}
37+
});
38+
}
39+
40+
private final HttpResponseFor<T> delegate;
41+
private final AgentSpan span;
42+
private final BiConsumer<AgentSpan, T> afterParse;
43+
44+
private HttpResponseWrapper(
45+
HttpResponseFor<T> delegate, AgentSpan span, BiConsumer<AgentSpan, T> afterParse) {
46+
this.delegate = delegate;
47+
this.span = span;
48+
this.afterParse = afterParse;
49+
}
50+
51+
@Override
52+
public T parse() {
53+
try {
54+
T parsed = delegate.parse();
55+
afterParse.accept(span, parsed);
56+
return parsed;
57+
} catch (Throwable err) {
58+
DECORATE.onError(span, err);
59+
throw err;
60+
}
61+
}
62+
63+
@Override
64+
public int statusCode() {
65+
return delegate.statusCode();
66+
}
67+
68+
@NotNull
69+
@Override
70+
public Headers headers() {
71+
return delegate.headers();
72+
}
73+
74+
@NotNull
75+
@Override
76+
public InputStream body() {
77+
return delegate.body();
78+
}
79+
80+
@Override
81+
public void close() {
82+
// span finished after the response is available
83+
delegate.close();
84+
}
85+
}

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

Lines changed: 0 additions & 138 deletions
This file was deleted.

0 commit comments

Comments
 (0)