Skip to content

Commit 6720c2e

Browse files
rpanackalrpanackal
committed
Centralize MDC management
1 parent 2478a55 commit 6720c2e

File tree

5 files changed

+217
-59
lines changed

5 files changed

+217
-59
lines changed

‎core/src/main/java/com/sap/ai/sdk/core/common/ClientResponseHandler.java‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import com.fasterxml.jackson.core.JsonProcessingException;
66
import com.fasterxml.jackson.databind.ObjectMapper;
77
import com.google.common.annotations.Beta;
8+
import com.sap.ai.sdk.core.common.MdcHelper.RequestContext;
89
import io.vavr.control.Try;
910
import java.nio.charset.StandardCharsets;
1011
import java.util.Optional;
@@ -18,7 +19,6 @@
1819
import org.apache.hc.core5.http.HttpEntity;
1920
import org.apache.hc.core5.http.io.HttpClientResponseHandler;
2021
import org.apache.hc.core5.http.io.entity.EntityUtils;
21-
import org.slf4j.MDC;
2222

2323
/**
2424
* Parse incoming JSON responses and handles any errors. For internal use only.
@@ -87,7 +87,7 @@ private T parseSuccess(@Nonnull final ClassicHttpResponse response) throws E {
8787
val content =
8888
tryGetContent(responseEntity)
8989
.getOrElseThrow(e -> exceptionFactory.build(message, e).setHttpResponse(response));
90-
logResponseSuccess(response);
90+
RequestContext.logResponseSuccess(response);
9191

9292
try {
9393
return objectMapper.readValue(content, successType);
@@ -172,16 +172,4 @@ private static String getErrorMessage(
172172
val message = Optional.ofNullable(additionalMessage).orElse("");
173173
return message.isEmpty() ? baseErrorMessage : "%s: %s".formatted(baseErrorMessage, message);
174174
}
175-
176-
private static void logResponseSuccess(final @Nonnull ClassicHttpResponse response) {
177-
if (!log.isDebugEnabled()) {
178-
return;
179-
}
180-
val headerTime = Optional.ofNullable(response.getFirstHeader("x-upstream-service-time"));
181-
val duration = headerTime.map(h -> h.getValue() + "ms").orElseGet(() -> "unknown");
182-
val entityLength = response.getEntity().getContentLength();
183-
val sizeInfo = entityLength >= 0 ? String.format("%.1fKB", entityLength / 1024.0) : "unknown";
184-
val msg = "[reqId={}] {} request completed successfully with duration={}, size={}.";
185-
log.debug(msg, MDC.get("reqId"), MDC.get("service"), duration, sizeInfo);
186-
}
187175
}

‎core/src/main/java/com/sap/ai/sdk/core/common/MdcHelper.java‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
package com.sap.ai.sdk.core.common;
2+
3+
import com.google.common.annotations.Beta;
4+
import java.util.Optional;
5+
import java.util.UUID;
6+
import javax.annotation.Nonnull;
7+
import lombok.Getter;
8+
import lombok.RequiredArgsConstructor;
9+
import lombok.experimental.UtilityClass;
10+
import lombok.extern.slf4j.Slf4j;
11+
import lombok.val;
12+
import org.apache.hc.core5.http.ClassicHttpResponse;
13+
import org.apache.hc.core5.http.HttpEntity;
14+
import org.slf4j.MDC;
15+
16+
/**
17+
* Utility for managing MDC (Mapped Diagnostic Context) for logging of AI Core requests.
18+
*
19+
* <p>This class is intended for internal use only.
20+
*/
21+
@UtilityClass
22+
@Beta
23+
public class MdcHelper {
24+
25+
@UtilityClass
26+
private static class MdcKeys {
27+
private static final String REQUEST_ID = "reqId";
28+
private static final String ENDPOINT = "endpoint";
29+
private static final String DESTINATION = "destination";
30+
private static final String MODE = "mode";
31+
private static final String SERVICE = "service";
32+
}
33+
34+
/** Utility for managing request context information in MDC. */
35+
@Slf4j
36+
@UtilityClass
37+
public static class RequestContext {
38+
private static void setRequestId(@Nonnull final String requestId) {
39+
MDC.put(MdcKeys.REQUEST_ID, requestId);
40+
}
41+
42+
/**
43+
* Set the endpoint for the current request context.
44+
*
45+
* @param endpoint the endpoint URL
46+
*/
47+
public static void setEndpoint(@Nonnull final String endpoint) {
48+
MDC.put(MdcKeys.ENDPOINT, endpoint);
49+
}
50+
51+
/**
52+
* Set the destination for the current request context.
53+
*
54+
* @param destination the destination name
55+
*/
56+
public static void setDestination(@Nonnull final String destination) {
57+
MDC.put(MdcKeys.DESTINATION, destination);
58+
}
59+
60+
/**
61+
* Set the mode for the current request context.
62+
*
63+
* @param mode the request mode
64+
*/
65+
public static void setMode(@Nonnull final Mode mode) {
66+
MDC.put(MdcKeys.MODE, mode.getValue());
67+
}
68+
69+
/**
70+
* Set the service for the current request context.
71+
*
72+
* @param service the service type
73+
*/
74+
public static void setService(@Nonnull final Service service) {
75+
MDC.put(MdcKeys.SERVICE, service.getValue());
76+
}
77+
78+
/** Clear all MDC request context information. */
79+
public static void clear() {
80+
MDC.remove(MdcKeys.REQUEST_ID);
81+
MDC.remove(MdcKeys.ENDPOINT);
82+
MDC.remove(MdcKeys.DESTINATION);
83+
MDC.remove(MdcKeys.MODE);
84+
MDC.remove(MdcKeys.SERVICE);
85+
}
86+
87+
/** Log the start of a request with generated request ID. */
88+
public static void logRequestStart() {
89+
val reqId = UUID.randomUUID().toString().substring(0, 8);
90+
RequestContext.setRequestId(reqId);
91+
92+
val message = "[reqId={}] Starting {} {} request to {}, destination={}";
93+
log.debug(
94+
message,
95+
reqId,
96+
MDC.get(MdcKeys.SERVICE),
97+
MDC.get(MdcKeys.MODE),
98+
MDC.get(MdcKeys.ENDPOINT),
99+
MDC.get(MdcKeys.DESTINATION));
100+
}
101+
102+
/**
103+
* Log successful response with duration and size information.
104+
*
105+
* @param response the HTTP response
106+
*/
107+
@SuppressWarnings("PMD.CloseResource")
108+
public static void logResponseSuccess(@Nonnull final ClassicHttpResponse response) {
109+
if (!log.isDebugEnabled()) {
110+
return;
111+
}
112+
113+
val headerTime = Optional.ofNullable(response.getFirstHeader("x-upstream-service-time"));
114+
val duration = headerTime.map(h -> h.getValue() + "ms").orElse("unknown");
115+
val sizeInfo =
116+
Optional.ofNullable(response.getEntity())
117+
.map(HttpEntity::getContentLength)
118+
.filter(length -> length >= 0)
119+
.map(length -> "%.1fKB".formatted(length / 1024.0))
120+
.orElse("unknown");
121+
val message = "[reqId={}] {} request completed successfully with duration={}, size={}.";
122+
log.debug(message, MDC.get(MdcKeys.REQUEST_ID), MDC.get(MdcKeys.SERVICE), duration, sizeInfo);
123+
}
124+
}
125+
126+
/** Request execution modes. */
127+
@RequiredArgsConstructor
128+
public enum Mode {
129+
/** Synchronous request mode */
130+
SYNCHRONOUS("synchronous"),
131+
/** Streaming request mode */
132+
STREAMING("streaming");
133+
@Getter private final String value;
134+
}
135+
136+
/** AI service types. */
137+
@RequiredArgsConstructor
138+
public enum Service {
139+
/** OpenAI service */
140+
OPENAI("openai"),
141+
/** Orchestration service */
142+
ORCHESTRATION("orchestration");
143+
@Getter private final String value;
144+
}
145+
}

‎core/src/test/java/com/sap/ai/sdk/core/common/MdcHelperTest.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package com.sap.ai.sdk.core.common;
2+
3+
import static com.sap.ai.sdk.core.common.MdcHelper.Mode.STREAMING;
4+
import static com.sap.ai.sdk.core.common.MdcHelper.Service.OPENAI;
5+
import static org.assertj.core.api.Assertions.assertThat;
6+
7+
import com.sap.ai.sdk.core.common.MdcHelper.RequestContext;
8+
import org.junit.jupiter.api.Test;
9+
import org.slf4j.MDC;
10+
11+
class MdcHelperTest {
12+
13+
@Test
14+
void testRequestContextLifecycle() {
15+
// Setup customer MDC entries (to test clear() safety)
16+
MDC.put("consumer-key", "consumer-value");
17+
18+
RequestContext.setService(OPENAI);
19+
RequestContext.setMode(STREAMING);
20+
RequestContext.setEndpoint("/api/endpoint");
21+
RequestContext.setDestination("http://localhost:8000");
22+
23+
assertThat(MDC.get("service")).isEqualTo("openai");
24+
assertThat(MDC.get("mode")).isEqualTo("streaming");
25+
assertThat(MDC.get("endpoint")).isEqualTo("/api/endpoint");
26+
assertThat(MDC.get("destination")).isEqualTo("http://localhost:8000");
27+
28+
RequestContext.logRequestStart();
29+
assertThat(MDC.get("reqId")).isNotNull().hasSize(8);
30+
31+
RequestContext.clear();
32+
assertThat(MDC.get("service")).isNull();
33+
assertThat(MDC.get("mode")).isNull();
34+
assertThat(MDC.get("endpoint")).isNull();
35+
assertThat(MDC.get("destination")).isNull();
36+
assertThat(MDC.get("reqId")).isNull();
37+
assertThat(MDC.get("consumer-key")).isEqualTo("consumer-value");
38+
}
39+
}

‎foundation-models/openai/src/main/java/com/sap/ai/sdk/foundationmodels/openai/OpenAiClient.java‎

Lines changed: 16 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
package com.sap.ai.sdk.foundationmodels.openai;
22

3+
import static com.sap.ai.sdk.core.common.MdcHelper.Mode.STREAMING;
4+
import static com.sap.ai.sdk.core.common.MdcHelper.Mode.SYNCHRONOUS;
5+
import static com.sap.ai.sdk.core.common.MdcHelper.Service.OPENAI;
36
import static com.sap.ai.sdk.foundationmodels.openai.OpenAiClientException.FACTORY;
47
import static com.sap.ai.sdk.foundationmodels.openai.OpenAiUtils.getOpenAiObjectMapper;
58

@@ -10,6 +13,7 @@
1013
import com.sap.ai.sdk.core.DeploymentResolutionException;
1114
import com.sap.ai.sdk.core.common.ClientResponseHandler;
1215
import com.sap.ai.sdk.core.common.ClientStreamingHandler;
16+
import com.sap.ai.sdk.core.common.MdcHelper.RequestContext;
1317
import com.sap.ai.sdk.core.common.StreamedDelta;
1418
import com.sap.ai.sdk.foundationmodels.openai.generated.model.ChatCompletionStreamOptions;
1519
import com.sap.ai.sdk.foundationmodels.openai.generated.model.CreateChatCompletionRequest;
@@ -30,19 +34,16 @@
3034
import java.io.IOException;
3135
import java.util.ArrayList;
3236
import java.util.List;
33-
import java.util.UUID;
3437
import java.util.stream.Stream;
3538
import javax.annotation.Nonnull;
3639
import javax.annotation.Nullable;
3740
import lombok.AccessLevel;
3841
import lombok.RequiredArgsConstructor;
3942
import lombok.extern.slf4j.Slf4j;
40-
import lombok.val;
4143
import org.apache.hc.client5.http.classic.methods.HttpPost;
4244
import org.apache.hc.core5.http.ContentType;
4345
import org.apache.hc.core5.http.io.entity.StringEntity;
4446
import org.apache.hc.core5.http.message.BasicClassicHttpRequest;
45-
import org.slf4j.MDC;
4647

4748
/** Client for interacting with OpenAI models. */
4849
@Slf4j
@@ -420,7 +421,7 @@ private <T> T execute(
420421
@Nonnull final Class<T> responseType) {
421422
final var request = new HttpPost(path);
422423
serializeAndSetHttpEntity(request, payload, this.customHeaders);
423-
MDC.put("endpoint", path);
424+
RequestContext.setEndpoint(path);
424425
return executeRequest(request, responseType);
425426
}
426427

@@ -431,7 +432,7 @@ private <D extends StreamedDelta> Stream<D> executeStream(
431432
@Nonnull final Class<D> deltaType) {
432433
final var request = new HttpPost(path);
433434
serializeAndSetHttpEntity(request, payload, this.customHeaders);
434-
MDC.put("endpoint", path);
435+
RequestContext.setEndpoint(path);
435436
return streamRequest(request, deltaType);
436437
}
437438

@@ -454,15 +455,16 @@ private <T> T executeRequest(
454455
final BasicClassicHttpRequest request, @Nonnull final Class<T> responseType) {
455456
try {
456457
final var client = ApacheHttpClient5Accessor.getHttpClient(destination);
457-
MDC.put("destination", ((HttpDestination) destination).getUri().toASCIIString());
458-
MDC.put("mode", "synchronous");
459-
logRequestStart();
458+
RequestContext.setDestination(((HttpDestination) destination).getUri().toString());
459+
RequestContext.setMode(SYNCHRONOUS);
460+
RequestContext.setService(OPENAI);
461+
RequestContext.logRequestStart();
460462
return client.execute(
461463
request, new ClientResponseHandler<>(responseType, OpenAiError.class, FACTORY));
462464
} catch (final IOException e) {
463465
throw new OpenAiClientException("Request to OpenAI model failed", e).setHttpRequest(request);
464466
} finally {
465-
MDC.clear();
467+
RequestContext.clear();
466468
}
467469
}
468470

@@ -471,28 +473,17 @@ private <D extends StreamedDelta> Stream<D> streamRequest(
471473
final BasicClassicHttpRequest request, @Nonnull final Class<D> deltaType) {
472474
try {
473475
final var client = ApacheHttpClient5Accessor.getHttpClient(destination);
474-
MDC.put("destination", ((HttpDestination) destination).getUri().toASCIIString());
475-
MDC.put("mode", "streaming");
476-
logRequestStart();
476+
RequestContext.setDestination(((HttpDestination) destination).getUri().toASCIIString());
477+
RequestContext.setMode(STREAMING);
478+
RequestContext.setService(OPENAI);
479+
RequestContext.logRequestStart();
477480
return new ClientStreamingHandler<>(deltaType, OpenAiError.class, FACTORY)
478481
.objectMapper(JACKSON)
479482
.handleStreamingResponse(client.executeOpen(null, request, null));
480483
} catch (final IOException e) {
481484
throw new OpenAiClientException("Request to OpenAI model failed", e).setHttpRequest(request);
482485
} finally {
483-
MDC.clear();
486+
RequestContext.clear();
484487
}
485488
}
486-
487-
private static void logRequestStart() {
488-
val reqId = UUID.randomUUID().toString().substring(0, 8);
489-
MDC.put("reqId", reqId);
490-
MDC.put("service", "OpenAI");
491-
log.debug(
492-
"[reqId={}] Starting OpenAI {} request to {}, destination={}",
493-
reqId,
494-
MDC.get("mode"),
495-
MDC.get("endpoint"),
496-
MDC.get("destination"));
497-
}
498489
}

0 commit comments

Comments
 (0)