Skip to content

Commit 66e4b88

Browse files
tzolovtzolov
committed
Add ObservationRegistry support to ChatClient
 - Implement observable chat responses in DefaultChatClient - Add ChatClientObservationContext and related classes for metrics - Update ChatClient and builder methods to support ObservationRegistry - Enhance RequestResponseAdvisor with getName() method - Add ChatClient streaming observability support - Introduce ChatClientObservationDocumentation for metric key names - Create DefaultChatClientObservationConvention for implementing conventions - Add ChatClientInputContentObservationFilter for optional input content logging - Update ChatClientAutoConfiguration to include new observation components - Extend ChatClientBuilderProperties with observation configuration options - Add unit tests for new observation classes and configurations - Update AiOperationType and AiProvider enums with new values - Implement safeguards and warnings for sensitive data in observations Resolves #1206
1 parent bf84d59 commit 66e4b88

File tree

17 files changed

+1035
-23
lines changed

17 files changed

+1035
-23
lines changed

‎spring-ai-core/src/main/java/org/springframework/ai/chat/client/ChatClient.java

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,17 +21,19 @@
2121
import java.util.Map;
2222
import java.util.function.Consumer;
2323

24-
import org.springframework.ai.model.Media;
24+
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
2525
import org.springframework.ai.chat.messages.Message;
2626
import org.springframework.ai.chat.model.ChatModel;
2727
import org.springframework.ai.chat.model.ChatResponse;
2828
import org.springframework.ai.chat.prompt.ChatOptions;
2929
import org.springframework.ai.chat.prompt.Prompt;
3030
import org.springframework.ai.converter.StructuredOutputConverter;
31+
import org.springframework.ai.model.Media;
3132
import org.springframework.core.ParameterizedTypeReference;
3233
import org.springframework.core.io.Resource;
3334
import org.springframework.util.MimeType;
3435

36+
import io.micrometer.observation.ObservationRegistry;
3537
import reactor.core.publisher.Flux;
3638

3739
/**
@@ -48,11 +50,25 @@
4850
public interface ChatClient {
4951

5052
static ChatClient create(ChatModel chatModel) {
51-
return builder(chatModel).build();
53+
return create(chatModel, ObservationRegistry.NOOP);
54+
}
55+
56+
static ChatClient create(ChatModel chatModel, ObservationRegistry observationRegistry) {
57+
return create(chatModel, observationRegistry, null);
58+
}
59+
60+
static ChatClient create(ChatModel chatModel, ObservationRegistry observationRegistry,
61+
ChatClientObservationConvention observationConvention) {
62+
return builder(chatModel, observationRegistry, observationConvention).build();
5263
}
5364

5465
static Builder builder(ChatModel chatModel) {
55-
return new DefaultChatClientBuilder(chatModel);
66+
return builder(chatModel, ObservationRegistry.NOOP, null);
67+
}
68+
69+
static Builder builder(ChatModel chatModel, ObservationRegistry observationRegistry,
70+
ChatClientObservationConvention customObservationConvention) {
71+
return new DefaultChatClientBuilder(chatModel, observationRegistry, customObservationConvention);
5672
}
5773

5874
ChatClientRequestSpec prompt();

‎spring-ai-core/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java

Lines changed: 77 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,10 @@
2828
import java.util.concurrent.ConcurrentHashMap;
2929
import java.util.function.Consumer;
3030

31-
import reactor.core.publisher.Flux;
32-
33-
import org.springframework.ai.model.Media;
31+
import org.springframework.ai.chat.client.observation.ChatClientObservationContext;
32+
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
33+
import org.springframework.ai.chat.client.observation.ChatClientObservationDocumentation;
34+
import org.springframework.ai.chat.client.observation.DefaultChatClientObservationConvention;
3435
import org.springframework.ai.chat.messages.Message;
3536
import org.springframework.ai.chat.messages.SystemMessage;
3637
import org.springframework.ai.chat.messages.UserMessage;
@@ -42,6 +43,7 @@
4243
import org.springframework.ai.chat.prompt.PromptTemplate;
4344
import org.springframework.ai.converter.BeanOutputConverter;
4445
import org.springframework.ai.converter.StructuredOutputConverter;
46+
import org.springframework.ai.model.Media;
4547
import org.springframework.ai.model.function.FunctionCallback;
4648
import org.springframework.ai.model.function.FunctionCallbackWrapper;
4749
import org.springframework.ai.model.function.FunctionCallingOptions;
@@ -52,6 +54,11 @@
5254
import org.springframework.util.MimeType;
5355
import org.springframework.util.StringUtils;
5456

57+
import io.micrometer.observation.Observation;
58+
import io.micrometer.observation.ObservationRegistry;
59+
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
60+
import reactor.core.publisher.Flux;
61+
5562
/**
5663
* The default implementation of {@link ChatClient} as created by the
5764
* {@link Builder#build()} } method.
@@ -65,6 +72,8 @@
6572
*/
6673
public class DefaultChatClient implements ChatClient {
6774

75+
private static final ChatClientObservationConvention DEFAULT_CHAT_CLIENT_OBSERVATION_CONVENTION = new DefaultChatClientObservationConvention();
76+
6877
private final ChatModel chatModel;
6978

7079
private final DefaultChatClientRequestSpec defaultChatClientRequest;
@@ -281,7 +290,7 @@ public <T> ResponseEntity<ChatResponse, T> responseEntity(
281290
}
282291

283292
protected <T> ResponseEntity<ChatResponse, T> doResponseEntity(StructuredOutputConverter<T> boc) {
284-
var chatResponse = doGetChatResponse(this.request, boc.getFormat());
293+
var chatResponse = doGetObservableChatResponse(this.request, boc.getFormat());
285294
var responseContent = chatResponse.getResult().getOutput().getContent();
286295
T entity = boc.convert(responseContent);
287296

@@ -297,7 +306,7 @@ public <T> T entity(StructuredOutputConverter<T> structuredOutputConverter) {
297306
}
298307

299308
private <T> T doSingleWithBeanOutputConverter(StructuredOutputConverter<T> boc) {
300-
var chatResponse = doGetChatResponse(this.request, boc.getFormat());
309+
var chatResponse = doGetObservableChatResponse(this.request, boc.getFormat());
301310
var stringResponse = chatResponse.getResult().getOutput().getContent();
302311
return boc.convert(stringResponse);
303312
}
@@ -309,7 +318,23 @@ public <T> T entity(Class<T> type) {
309318
}
310319

311320
private ChatResponse doGetChatResponse() {
312-
return this.doGetChatResponse(this.request, "");
321+
return this.doGetObservableChatResponse(this.request, "");
322+
}
323+
324+
private ChatResponse doGetObservableChatResponse(DefaultChatClientRequestSpec inputRequest,
325+
String formatParam) {
326+
327+
ChatClientObservationContext observationContext = new ChatClientObservationContext(inputRequest,
328+
formatParam, false);
329+
330+
return ChatClientObservationDocumentation.AI_CHAT_CLIENT
331+
.observation(inputRequest.customObservationConvention, DEFAULT_CHAT_CLIENT_OBSERVATION_CONVENTION,
332+
() -> observationContext, inputRequest.observationRegistry)
333+
.observe(() -> {
334+
ChatResponse chatResponse = doGetChatResponse(inputRequest, formatParam);
335+
return chatResponse;
336+
});
337+
313338
}
314339

315340
private ChatResponse doGetChatResponse(DefaultChatClientRequestSpec inputRequest, String formatParam) {
@@ -395,6 +420,29 @@ public DefaultStreamResponseSpec(ChatModel chatModel, DefaultChatClientRequestSp
395420
}
396421

397422
private Flux<ChatResponse> doGetFluxChatResponse(DefaultChatClientRequestSpec inputRequest) {
423+
return Flux.deferContextual(contextView -> {
424+
ChatClientObservationContext observationContext = new ChatClientObservationContext(inputRequest, "",
425+
true);
426+
427+
Observation observation = ChatClientObservationDocumentation.AI_CHAT_CLIENT.observation(
428+
inputRequest.customObservationConvention, DEFAULT_CHAT_CLIENT_OBSERVATION_CONVENTION,
429+
() -> observationContext, inputRequest.observationRegistry);
430+
431+
observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null))
432+
.start();
433+
434+
// @formatter:off
435+
return doGetFluxChatResponse2(inputRequest)
436+
.doOnError(observation::error)
437+
.doFinally(s -> {
438+
observation.stop();
439+
})
440+
.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
441+
// @formatter:on
442+
});
443+
}
444+
445+
private Flux<ChatResponse> doGetFluxChatResponse2(DefaultChatClientRequestSpec inputRequest) {
398446

399447
Map<String, Object> context = new ConcurrentHashMap<>();
400448
context.putAll(inputRequest.getAdvisorParams());
@@ -426,9 +474,7 @@ private Flux<ChatResponse> doGetFluxChatResponse(DefaultChatClientRequestSpec in
426474
messages.add(userMessage);
427475
}
428476

429-
if (advisedRequest.getChatOptions() instanceof
430-
431-
FunctionCallingOptions functionCallingOptions) {
477+
if (advisedRequest.getChatOptions() instanceof FunctionCallingOptions functionCallingOptions) {
432478
if (!advisedRequest.getFunctionNames().isEmpty()) {
433479
functionCallingOptions.setFunctions(new HashSet<>(advisedRequest.getFunctionNames()));
434480
}
@@ -470,6 +516,10 @@ public Flux<String> content() {
470516

471517
public static class DefaultChatClientRequestSpec implements ChatClientRequestSpec {
472518

519+
private final ObservationRegistry observationRegistry;
520+
521+
private final ChatClientObservationConvention customObservationConvention;
522+
473523
private final ChatModel chatModel;
474524

475525
private String userText = "";
@@ -494,6 +544,14 @@ public static class DefaultChatClientRequestSpec implements ChatClientRequestSpe
494544

495545
private final Map<String, Object> advisorParams = new HashMap<>();
496546

547+
private ObservationRegistry getObservationRegistry() {
548+
return observationRegistry;
549+
}
550+
551+
private ChatClientObservationConvention getCustomObservationConvention() {
552+
return customObservationConvention;
553+
}
554+
497555
public String getUserText() {
498556
return userText;
499557
}
@@ -541,13 +599,15 @@ public List<FunctionCallback> getFunctionCallbacks() {
541599
/* copy constructor */
542600
DefaultChatClientRequestSpec(DefaultChatClientRequestSpec ccr) {
543601
this(ccr.chatModel, ccr.userText, ccr.userParams, ccr.systemText, ccr.systemParams, ccr.functionCallbacks,
544-
ccr.messages, ccr.functionNames, ccr.media, ccr.chatOptions, ccr.advisors, ccr.advisorParams);
602+
ccr.messages, ccr.functionNames, ccr.media, ccr.chatOptions, ccr.advisors, ccr.advisorParams,
603+
ccr.observationRegistry, ccr.customObservationConvention);
545604
}
546605

547606
public DefaultChatClientRequestSpec(ChatModel chatModel, String userText, Map<String, Object> userParams,
548607
String systemText, Map<String, Object> systemParams, List<FunctionCallback> functionCallbacks,
549608
List<Message> messages, List<String> functionNames, List<Media> media, ChatOptions chatOptions,
550-
List<RequestResponseAdvisor> advisors, Map<String, Object> advisorParams) {
609+
List<RequestResponseAdvisor> advisors, Map<String, Object> advisorParams,
610+
ObservationRegistry observationRegistry, ChatClientObservationConvention customObservationConvention) {
551611

552612
this.chatModel = chatModel;
553613
this.chatOptions = chatOptions != null ? chatOptions.copy()
@@ -564,14 +624,17 @@ public DefaultChatClientRequestSpec(ChatModel chatModel, String userText, Map<St
564624
this.media.addAll(media);
565625
this.advisors.addAll(advisors);
566626
this.advisorParams.putAll(advisorParams);
627+
this.observationRegistry = observationRegistry;
628+
this.customObservationConvention = customObservationConvention;
567629
}
568630

569631
/**
570632
* Return a {@code ChatClient2Builder} to create a new {@code ChatClient2} whose
571633
* settings are replicated from this {@code ChatClientRequest}.
572634
*/
573635
public Builder mutate() {
574-
DefaultChatClientBuilder builder = (DefaultChatClientBuilder) ChatClient.builder(chatModel)
636+
DefaultChatClientBuilder builder = (DefaultChatClientBuilder) ChatClient
637+
.builder(chatModel, this.observationRegistry, this.customObservationConvention)
575638
.defaultSystem(s -> s.text(this.systemText).params(this.systemParams))
576639
.defaultUser(u -> u.text(this.userText)
577640
.params(this.userParams)
@@ -756,7 +819,8 @@ public static DefaultChatClientRequestSpec adviseOnRequest(DefaultChatClientRequ
756819
adviseRequest.userParams(), adviseRequest.systemText(), adviseRequest.systemParams(),
757820
adviseRequest.functionCallbacks(), adviseRequest.messages(), adviseRequest.functionNames(),
758821
adviseRequest.media(), adviseRequest.chatOptions(), adviseRequest.advisors(),
759-
adviseRequest.advisorParams());
822+
adviseRequest.advisorParams(), inputRequest.getObservationRegistry(),
823+
inputRequest.getCustomObservationConvention());
760824
}
761825

762826
return advisedRequest;

‎spring-ai-core/src/main/java/org/springframework/ai/chat/client/DefaultChatClientBuilder.java

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,14 @@
2626
import org.springframework.ai.chat.client.ChatClient.PromptSystemSpec;
2727
import org.springframework.ai.chat.client.ChatClient.PromptUserSpec;
2828
import org.springframework.ai.chat.client.DefaultChatClient.DefaultChatClientRequestSpec;
29+
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
2930
import org.springframework.ai.chat.model.ChatModel;
3031
import org.springframework.ai.chat.prompt.ChatOptions;
3132
import org.springframework.core.io.Resource;
3233
import org.springframework.util.Assert;
3334

35+
import io.micrometer.observation.ObservationRegistry;
36+
3437
/**
3538
* DefaultChatClientBuilder is a builder class for creating a ChatClient.
3639
*
@@ -48,11 +51,18 @@ public class DefaultChatClientBuilder implements Builder {
4851

4952
private final ChatModel chatModel;
5053

51-
public DefaultChatClientBuilder(ChatModel chatModel) {
54+
DefaultChatClientBuilder(ChatModel chatModel) {
55+
this(chatModel, ObservationRegistry.NOOP, null);
56+
}
57+
58+
public DefaultChatClientBuilder(ChatModel chatModel, ObservationRegistry observationRegistry,
59+
ChatClientObservationConvention customObservationConvention) {
5260
Assert.notNull(chatModel, "the " + ChatModel.class.getName() + " must be non-null");
61+
Assert.notNull(observationRegistry, "the " + ObservationRegistry.class.getName() + " must be non-null");
5362
this.chatModel = chatModel;
5463
this.defaultRequest = new DefaultChatClientRequestSpec(chatModel, "", Map.of(), "", Map.of(), List.of(),
55-
List.of(), List.of(), List.of(), null, List.of(), Map.of());
64+
List.of(), List.of(), List.of(), null, List.of(), Map.of(), observationRegistry,
65+
customObservationConvention);
5666
}
5767

5868
public ChatClient build() {

‎spring-ai-core/src/main/java/org/springframework/ai/chat/client/RequestResponseAdvisor.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,13 @@
3434
*/
3535
public interface RequestResponseAdvisor {
3636

37+
/**
38+
* @return the advisor name.
39+
*/
40+
default String getName() {
41+
return this.getClass().getSimpleName();
42+
}
43+
3744
/**
3845
* @param request the {@link AdvisedRequest} data to be advised. Represents the row
3946
* {@link ChatClient.ChatClientRequestSpec} data before sealed into a {@link Prompt}.

0 commit comments

Comments
 (0)