Skip to content

Commit 2d127f2

Browse files
ygreeygree
committed
Handle possible response parse errors
1 parent 854d7b0 commit 2d127f2

File tree

8 files changed

+39
-21
lines changed

8 files changed

+39
-21
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public static void exit(
7070
if (future != null) {
7171
future =
7272
HttpResponseWrappers.wrapFutureHttpResponse(
73-
future, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
73+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletion, DECORATE::onError);
7474
} else {
7575
span.finish();
7676
}
@@ -104,7 +104,7 @@ public static void exit(
104104
if (future != null) {
105105
future =
106106
HttpResponseWrappers.wrapFutureHttpResponseStream(
107-
future, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
107+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks, DECORATE::onError);
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: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -69,12 +69,12 @@ public static void exit(
6969
if (response != null) {
7070
response =
7171
HttpResponseWrappers.wrapHttpResponse(
72-
response, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
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
7373
}
7474
DECORATE.beforeFinish(span);
7575
} finally {
7676
scope.close();
77-
span.finish();
77+
span.finish(); // TODO so we finish here but decorate withChatCompletion later? should move
7878
}
7979
}
8080
}
@@ -104,7 +104,7 @@ public static void exit(
104104
if (response != null) {
105105
response =
106106
HttpResponseWrappers.wrapHttpResponseStream(
107-
response, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
107+
response, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks, DECORATE::onError);
108108
} else {
109109
span.finish();
110110
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ public static void exit(
6565
if (future != null) {
6666
future =
6767
HttpResponseWrappers.wrapFutureHttpResponse(
68-
future, span, CompletionDecorator.DECORATE::withCompletion);
68+
future, span, CompletionDecorator.DECORATE::withCompletion, DECORATE::onError);
6969
} else {
7070
span.finish();
7171
}
@@ -100,7 +100,7 @@ public static void exit(
100100
if (future != null) {
101101
future =
102102
HttpResponseWrappers.wrapFutureHttpResponseStream(
103-
future, span, CompletionDecorator.DECORATE::withCompletions);
103+
future, span, CompletionDecorator.DECORATE::withCompletions, DECORATE::onError);
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: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ public static void exit(
7575
if (response != null) {
7676
response =
7777
HttpResponseWrappers.wrapHttpResponse(
78-
response, span, CompletionDecorator.DECORATE::withCompletion);
78+
response, span, CompletionDecorator.DECORATE::withCompletion, DECORATE::onError);
7979
}
8080
DECORATE.beforeFinish(span);
8181
} finally {
@@ -110,7 +110,7 @@ public static void exit(
110110
if (response != null) {
111111
response =
112112
HttpResponseWrappers.wrapHttpResponseStream(
113-
response, span, CompletionDecorator.DECORATE::withCompletions);
113+
response, span, CompletionDecorator.DECORATE::withCompletions, DECORATE::onError);
114114
} else {
115115
span.finish();
116116
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ public static void exit(
5656
if (response != null) {
5757
response =
5858
HttpResponseWrappers.wrapHttpResponse(
59-
response, span, EmbeddingDecorator.DECORATE::withCreateEmbeddingResponse);
59+
response, span, EmbeddingDecorator.DECORATE::withCreateEmbeddingResponse, DECORATE::onError);
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/HttpResponseWrappers.java‎

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,16 @@ abstract static class DDHttpResponseFor<T> implements HttpResponseFor<T> {
2525

2626
abstract T afterParse(T resp);
2727

28+
abstract void parseFailed(Throwable ex);
29+
2830
@Override
2931
public T parse() {
30-
return afterParse(delegate.parse());
32+
try {
33+
return afterParse(delegate.parse());
34+
} catch (Throwable t) {
35+
parseFailed(t);
36+
throw t;
37+
}
3138
}
3239

3340
@Override
@@ -53,24 +60,30 @@ public void close() {
5360
}
5461
}
5562

63+
// TODO split to named classes grouped by streamed/single response wrappers
5664
public static <T> HttpResponseFor<T> wrapHttpResponse(
57-
HttpResponseFor<T> response, AgentSpan span, BiConsumer<AgentSpan, T> afterParse) {
65+
HttpResponseFor<T> response, AgentSpan span, BiConsumer<AgentSpan, T> afterParse, BiConsumer<AgentSpan, Throwable> parseFailed) {
5866
DECORATE.withHttpResponse(span, response.headers());
5967
return new DDHttpResponseFor<T>(response) {
6068
@Override
6169
public T afterParse(T t) {
6270
afterParse.accept(span, t);
6371
return t;
6472
}
73+
74+
@Override
75+
void parseFailed(Throwable ex) {
76+
parseFailed.accept(span, ex);
77+
}
6578
};
6679
}
6780

6881
public static <T> CompletableFuture<HttpResponseFor<T>> wrapFutureHttpResponse(
6982
CompletableFuture<HttpResponseFor<T>> future,
7083
AgentSpan span,
71-
BiConsumer<AgentSpan, T> afterParse) {
84+
BiConsumer<AgentSpan, T> afterParse, BiConsumer<AgentSpan, Throwable> parseFailed) {
7285
return future
73-
.thenApply(response -> wrapHttpResponse(response, span, afterParse))
86+
.thenApply(response -> wrapHttpResponse(response, span, afterParse, parseFailed))
7487
.whenComplete(
7588
(r, t) -> {
7689
DECORATE.beforeFinish(span);
@@ -81,7 +94,7 @@ public static <T> CompletableFuture<HttpResponseFor<T>> wrapFutureHttpResponse(
8194
public static <T> HttpResponseFor<StreamResponse<T>> wrapHttpResponseStream(
8295
HttpResponseFor<StreamResponse<T>> response,
8396
final AgentSpan span,
84-
BiConsumer<AgentSpan, List<T>> decorate) {
97+
BiConsumer<AgentSpan, List<T>> decorate, BiConsumer<AgentSpan, Throwable> parseFailed) {
8598
DECORATE.withHttpResponse(span, response.headers());
8699
return new DDHttpResponseFor<StreamResponse<T>>(response) {
87100
@Override
@@ -107,14 +120,19 @@ public void close() {
107120
}
108121
};
109122
}
123+
124+
@Override
125+
void parseFailed(Throwable ex) {
126+
parseFailed.accept(span, ex);
127+
}
110128
};
111129
}
112130

113131
public static <T>
114132
CompletableFuture<HttpResponseFor<StreamResponse<T>>> wrapFutureHttpResponseStream(
115133
CompletableFuture<HttpResponseFor<StreamResponse<T>>> future,
116134
AgentSpan span,
117-
BiConsumer<AgentSpan, List<T>> decorate) {
118-
return future.thenApply(r -> wrapHttpResponseStream(r, span, decorate));
135+
BiConsumer<AgentSpan, List<T>> decorate, BiConsumer<AgentSpan, Throwable> parseFailed) {
136+
return future.thenApply(r -> wrapHttpResponseStream(r, span, decorate, parseFailed));
119137
}
120138
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ public static void exit(
6666
if (future != null) {
6767
future =
6868
HttpResponseWrappers.wrapFutureHttpResponse(
69-
future, span, ResponseDecorator.DECORATE::withResponse);
69+
future, span, ResponseDecorator.DECORATE::withResponse, DECORATE::onError);
7070
} else {
7171
span.finish();
7272
}
@@ -101,7 +101,7 @@ public static void exit(
101101
if (future != null) {
102102
future =
103103
HttpResponseWrappers.wrapFutureHttpResponseStream(
104-
future, span, ResponseDecorator.DECORATE::withResponseStreamEvents);
104+
future, span, ResponseDecorator.DECORATE::withResponseStreamEvents, DECORATE::onError);
105105
} else {
106106
span.finish();
107107
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ public static void exit(
6868
if (response != null) {
6969
response =
7070
HttpResponseWrappers.wrapHttpResponse(
71-
response, span, ResponseDecorator.DECORATE::withResponse);
71+
response, span, ResponseDecorator.DECORATE::withResponse, DECORATE::onError);
7272
}
7373
DECORATE.beforeFinish(span);
7474
} finally {
@@ -103,7 +103,7 @@ public static void exit(
103103
if (response != null) {
104104
response =
105105
HttpResponseWrappers.wrapHttpResponseStream(
106-
response, span, ResponseDecorator.DECORATE::withResponseStreamEvents);
106+
response, span, ResponseDecorator.DECORATE::withResponseStreamEvents, DECORATE::onError);
107107
} else {
108108
span.finish();
109109
}

0 commit comments

Comments
 (0)