feat: add observability to zhipu chat model

Author: mxsl-gr

models/spring-ai-zhipuai/pom.xml

@@ -54,6 +54,12 @@
             <scope>test</scope>
         </dependency>
 
+		<dependency>
+			<groupId>io.micrometer</groupId>
+			<artifactId>micrometer-observation-test</artifactId>
+			<scope>test</scope>
+		</dependency>
+
     </dependencies>
 
 </project>

models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java

@@ -15,6 +15,9 @@
  */
 package org.springframework.ai.zhipuai;
 
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.ai.chat.messages.AssistantMessage;
@@ -23,12 +26,19 @@
 import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
 import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.EmptyUsage;
 import org.springframework.ai.chat.model.AbstractToolCallSupport;
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.chat.model.ChatResponse;
 import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.chat.model.MessageAggregator;
 import org.springframework.ai.chat.model.StreamingChatModel;
+import org.springframework.ai.chat.observation.ChatModelObservationContext;
+import org.springframework.ai.chat.observation.ChatModelObservationConvention;
+import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
+import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
 import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
 import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.model.ModelOptionsUtils;
 import org.springframework.ai.model.function.FunctionCallback;
@@ -47,6 +57,7 @@
 import org.springframework.ai.zhipuai.api.ZhiPuAiApi.ChatCompletionMessage.ToolCall;
 import org.springframework.ai.zhipuai.api.ZhiPuAiApi.ChatCompletionRequest;
 import org.springframework.ai.zhipuai.api.ZhiPuAiApi.FunctionTool;
+import org.springframework.ai.zhipuai.api.ZhiPuApiConstants;
 import org.springframework.ai.zhipuai.metadata.ZhiPuAiUsage;
 import org.springframework.http.ResponseEntity;
 import org.springframework.retry.support.RetryTemplate;
@@ -78,6 +89,8 @@ public class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod
 
 	private static final Logger logger = LoggerFactory.getLogger(ZhiPuAiChatModel.class);
 
+	private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
+
 	/**
 	 * The default options used for the chat completion requests.
 	 */
@@ -93,6 +106,16 @@ public class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod
 	 */
 	private final ZhiPuAiApi zhiPuAiApi;
 
+	/**
+	 * Observation registry used for instrumentation.
+	 */
+	private final ObservationRegistry observationRegistry;
+
+	/**
+	 * Conventions to use for generating observations.
+	 */
+	private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
+
 	/**
 	 * Creates an instance of the ZhiPuAiChatModel.
 	 * @param zhiPuAiApi The ZhiPuAiApi instance to be used for interacting with the
@@ -124,7 +147,7 @@ public ZhiPuAiChatModel(ZhiPuAiApi zhiPuAiApi, ZhiPuAiChatOptions options) {
 	 */
 	public ZhiPuAiChatModel(ZhiPuAiApi zhiPuAiApi, ZhiPuAiChatOptions options,
 			FunctionCallbackContext functionCallbackContext, RetryTemplate retryTemplate) {
-		this(zhiPuAiApi, options, functionCallbackContext, List.of(), retryTemplate);
+		this(zhiPuAiApi, options, functionCallbackContext, List.of(), retryTemplate, ObservationRegistry.NOOP);
 	}
 
 	/**
@@ -135,58 +158,77 @@ public ZhiPuAiChatModel(ZhiPuAiApi zhiPuAiApi, ZhiPuAiChatOptions options,
 	 * @param functionCallbackContext The function callback context.
 	 * @param toolFunctionCallbacks The tool function callbacks.
 	 * @param retryTemplate The retry template.
+	 * @param observationRegistry The ObservationRegistry used for instrumentation.
 	 */
 	public ZhiPuAiChatModel(ZhiPuAiApi zhiPuAiApi, ZhiPuAiChatOptions options,
 			FunctionCallbackContext functionCallbackContext, List<FunctionCallback> toolFunctionCallbacks,
-			RetryTemplate retryTemplate) {
+			RetryTemplate retryTemplate, ObservationRegistry observationRegistry) {
 		super(functionCallbackContext, options, toolFunctionCallbacks);
 		Assert.notNull(zhiPuAiApi, "ZhiPuAiApi must not be null");
 		Assert.notNull(options, "Options must not be null");
 		Assert.notNull(retryTemplate, "RetryTemplate must not be null");
 		Assert.isTrue(CollectionUtils.isEmpty(options.getFunctionCallbacks()),
 				"The default function callbacks must be set via the toolFunctionCallbacks constructor parameter");
+		Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
 		this.zhiPuAiApi = zhiPuAiApi;
 		this.defaultOptions = options;
 		this.retryTemplate = retryTemplate;
+		this.observationRegistry = observationRegistry;
 	}
 
 	@Override
 	public ChatResponse call(Prompt prompt) {
 		ChatCompletionRequest request = createRequest(prompt, false);
 
-		ResponseEntity<ChatCompletion> completionEntity = this.retryTemplate
-			.execute(ctx -> this.zhiPuAiApi.chatCompletionEntity(request));
+		ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
+			.prompt(prompt)
+			.provider(ZhiPuApiConstants.PROVIDER_NAME)
+			.requestOptions(buildRequestOptions(request))
+			.build();
 
-		var chatCompletion = completionEntity.getBody();
+		ChatResponse response = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION
+			.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
+					this.observationRegistry)
+			.observe(() -> {
 
-		if (chatCompletion == null) {
-			logger.warn("No chat completion returned for prompt: {}", prompt);
-			return new ChatResponse(List.of());
-		}
+				ResponseEntity<ChatCompletion> completionEntity = this.retryTemplate
+					.execute(ctx -> this.zhiPuAiApi.chatCompletionEntity(request));
 
-		List<Choice> choices = chatCompletion.choices();
+				var chatCompletion = completionEntity.getBody();
 
-		List<Generation> generations = choices.stream().map(choice -> {
-			// @formatter:off
-			Map<String, Object> metadata = Map.of(
-					"id", chatCompletion.id(),
-					"role", choice.message().role() != null ? choice.message().role().name() : "",
-					"finishReason", choice.finishReason() != null ? choice.finishReason().name() : "");
-			// @formatter:on
-			return buildGeneration(choice, metadata);
-		}).toList();
+				if (chatCompletion == null) {
+					logger.warn("No chat completion returned for prompt: {}", prompt);
+					return new ChatResponse(List.of());
+				}
 
-		ChatResponse chatResponse = new ChatResponse(generations, from(completionEntity.getBody()));
+				List<Choice> choices = chatCompletion.choices();
 
-		if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(chatResponse,
+				List<Generation> generations = choices.stream().map(choice -> {
+			// @formatter:off
+					Map<String, Object> metadata = Map.of(
+						"id", chatCompletion.id(),
+						"role", choice.message().role() != null ? choice.message().role().name() : "",
+						"finishReason", choice.finishReason() != null ? choice.finishReason().name() : ""
+					);
+					// @formatter:on
+					return buildGeneration(choice, metadata);
+				}).toList();
+
+				ChatResponse chatResponse = new ChatResponse(generations, from(completionEntity.getBody()));
+
+				observationContext.setResponse(chatResponse);
+
+				return chatResponse;
+			});
+		if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response,
 				Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) {
-			var toolCallConversation = handleToolCalls(prompt, chatResponse);
+			var toolCallConversation = handleToolCalls(prompt, response);
 			// Recursively call the call method with the tool call message
 			// conversation that contains the call responses.
 			return this.call(new Prompt(toolCallConversation, prompt.getOptions()));
 		}
 
-		return chatResponse;
+		return response;
 	}
 
 	@Override
@@ -196,72 +238,87 @@ public ChatOptions getDefaultOptions() {
 
 	@Override
 	public Flux<ChatResponse> stream(Prompt prompt) {
-		ChatCompletionRequest request = createRequest(prompt, true);
+		return Flux.deferContextual(contextView -> {
+			ChatCompletionRequest request = createRequest(prompt, true);
 
-		Flux<ChatCompletionChunk> completionChunks = this.retryTemplate
-			.execute(ctx -> this.zhiPuAiApi.chatCompletionStream(request));
+			Flux<ChatCompletionChunk> completionChunks = this.retryTemplate
+				.execute(ctx -> this.zhiPuAiApi.chatCompletionStream(request));
 
-		// For chunked responses, only the first chunk contains the choice role.
-		// The rest of the chunks with same ID share the same role.
-		ConcurrentHashMap<String, String> roleMap = new ConcurrentHashMap<>();
+			// For chunked responses, only the first chunk contains the choice role.
+			// The rest of the chunks with same ID share the same role.
+			ConcurrentHashMap<String, String> roleMap = new ConcurrentHashMap<>();
 
-		// Convert the ChatCompletionChunk into a ChatCompletion to be able to reuse
-		// the function call handling logic.
-		Flux<ChatResponse> chatResponse = completionChunks.map(this::chunkToChatCompletion)
-			.switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> {
-				try {
-					@SuppressWarnings("null")
-					String id = chatCompletion2.id();
+			final ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
+				.prompt(prompt)
+				.provider(ZhiPuApiConstants.PROVIDER_NAME)
+				.requestOptions(buildRequestOptions(request))
+				.build();
 
-			// @formatter:off
+			Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(
+					this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
+					this.observationRegistry);
+
+			observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start();
+
+			Flux<ChatResponse> chatResponse = completionChunks.map(this::chunkToChatCompletion)
+				.switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> {
+					try {
+						String id = chatCompletion2.id();
+
+				// @formatter:off
 						List<Generation> generations = chatCompletion2.choices().stream().map(choice -> {
 							if (choice.message().role() != null) {
 								roleMap.putIfAbsent(id, choice.message().role().name());
 							}
 							Map<String, Object> metadata = Map.of(
-									"id", chatCompletion2.id(),
-									"role", roleMap.getOrDefault(id, ""),
-									"finishReason", choice.finishReason() != null ? choice.finishReason().name() : "");
+								"id", chatCompletion2.id(),
+								"role", roleMap.getOrDefault(id, ""),
+								"finishReason", choice.finishReason() != null ? choice.finishReason().name() : ""
+							);
 							return buildGeneration(choice, metadata);
 						}).toList();
 						// @formatter:on
 
-					if (chatCompletion2.usage() != null) {
 						return new ChatResponse(generations, from(chatCompletion2));
 					}
-					else {
-						return new ChatResponse(generations);
+					catch (Exception e) {
+						logger.error("Error processing chat completion", e);
+						return new ChatResponse(List.of());
 					}
-				}
-				catch (Exception e) {
-					logger.error("Error processing chat completion", e);
-					return new ChatResponse(List.of());
-				}
-
-			}));
 
-		return chatResponse.flatMap(response -> {
+				}));
 
-			if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response,
-					Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) {
-				var toolCallConversation = handleToolCalls(prompt, response);
-				// Recursively call the stream method with the tool call message
-				// conversation that contains the call responses.
-				return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
-			}
-			else {
+			// @formatter:off
+			Flux<ChatResponse> flux = chatResponse.flatMap(response -> {
+				if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response, Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) {
+					var toolCallConversation = handleToolCalls(prompt, response);
+					// Recursively call the stream method with the tool call message
+					// conversation that contains the call responses.
+					return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
+				}
 				return Flux.just(response);
-			}
+			}).doOnError(observation::error).doFinally(s -> {
+				// TODO: Consider a custom ObservationContext and
+				// include additional metadata
+				// if (s == SignalType.CANCEL) {
+				// observationContext.setAborted(true);
+				// }
+				observation.stop();
+			}).contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
+			// @formatter:on
+
+			return new MessageAggregator().aggregate(flux, observationContext::setResponse);
 		});
 	}
 
 	private ChatResponseMetadata from(ChatCompletion result) {
 		Assert.notNull(result, "ZhiPuAI ChatCompletionResult must not be null");
 		return ChatResponseMetadata.builder()
-			.withId(result.id())
-			.withUsage(ZhiPuAiUsage.from(result.usage()))
-			.withModel(result.model())
-			.withKeyValue("created", result.created())
+			.withId(result.id() != null ? result.id() : "")
+			.withUsage(result.usage() != null ? ZhiPuAiUsage.from(result.usage()) : new EmptyUsage())
+			.withModel(result.model() != null ? result.model() : "")
+			.withKeyValue("created", result.created() != null ? result.created() : 0L)
+			.withKeyValue("system-fingerprint", result.systemFingerprint() != null ? result.systemFingerprint() : "")
 			.build();
 	}
 
@@ -406,6 +463,16 @@ else if (mediaContentData instanceof String text) {
 		}
 	}
 
+	private ChatOptions buildRequestOptions(ZhiPuAiApi.ChatCompletionRequest request) {
+		return ChatOptionsBuilder.builder()
+			.withModel(request.model())
+			.withMaxTokens(request.maxTokens())
+			.withStopSequences(request.stop())
+			.withTemperature(request.temperature())
+			.withTopP(request.topP())
+			.build();
+	}
+
 	private List<FunctionTool> getFunctionTools(Set<String> functionNames) {
 		return this.resolveFunctionCallbacks(functionNames).stream().map(functionCallback -> {
 			var function = new FunctionTool.Function(functionCallback.getDescription(), functionCallback.getName(),
@@ -414,4 +481,8 @@ private List<FunctionTool> getFunctionTools(Set<String> functionNames) {
 		}).toList();
 	}
 
+	public void setObservationConvention(ChatModelObservationConvention observationConvention) {
+		this.observationConvention = observationConvention;
+	}
+
 }

models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuApiConstants.java

@@ -1,5 +1,7 @@
 package org.springframework.ai.zhipuai.api;
 
+import org.springframework.ai.observation.conventions.AiProvider;
+
 /**
  * Common value constants for ZhiPu api.
  *
@@ -10,4 +12,6 @@ public final class ZhiPuApiConstants {
 
 	public static final String DEFAULT_BASE_URL = "https://open.bigmodel.cn/api/paas";
 
+	public static final String PROVIDER_NAME = AiProvider.ZHIPUAI.value();
+
 }

models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java

@@ -164,7 +164,7 @@ public void zhiPuAiChatStreamTransientError() {
 	public void zhiPuAiChatStreamNonTransientError() {
 		when(zhiPuAiApi.chatCompletionStream(isA(ChatCompletionRequest.class)))
 				.thenThrow(new RuntimeException("Non Transient Error"));
-		assertThrows(RuntimeException.class, () -> chatModel.stream(new Prompt("text")));
+		assertThrows(RuntimeException.class, () -> chatModel.stream(new Prompt("text")).collectList().block());
 	}
 
 	@Test