diff --git a/models/spring-ai-zhipuai/pom.xml b/models/spring-ai-zhipuai/pom.xml index 0fea2ed64..c0e5128ec 100644 --- a/models/spring-ai-zhipuai/pom.xml +++ b/models/spring-ai-zhipuai/pom.xml @@ -54,6 +54,12 @@ test + + io.micrometer + micrometer-observation-test + test + + diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java index 2b2ed1ea0..0980dc928 100644 --- a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java +++ b/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.ToolResponseMessage; 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.Role; 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 class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod */ 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 class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod * @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 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 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 completionEntity = this.retryTemplate + .execute(ctx -> this.zhiPuAiApi.chatCompletionEntity(request)); - List choices = chatCompletion.choices(); + var chatCompletion = completionEntity.getBody(); - List generations = choices.stream().map(choice -> { + if (chatCompletion == null) { + logger.warn("No chat completion returned for prompt: {}", prompt); + return new ChatResponse(List.of()); + } + + List choices = chatCompletion.choices(); + + List generations = choices.stream().map(choice -> { // @formatter:off - Map 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(); + Map 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())); + ChatResponse chatResponse = new ChatResponse(generations, from(completionEntity.getBody())); - if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(chatResponse, + 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 class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod @Override public Flux stream(Prompt prompt) { - ChatCompletionRequest request = createRequest(prompt, true); + return Flux.deferContextual(contextView -> { + ChatCompletionRequest request = createRequest(prompt, true); - Flux completionChunks = this.retryTemplate - .execute(ctx -> this.zhiPuAiApi.chatCompletionStream(request)); + Flux 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 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 roleMap = new ConcurrentHashMap<>(); - // Convert the ChatCompletionChunk into a ChatCompletion to be able to reuse - // the function call handling logic. - Flux 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 = completionChunks.map(this::chunkToChatCompletion) + .switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> { + try { + String id = chatCompletion2.id(); + + // @formatter:off List generations = chatCompletion2.choices().stream().map(choice -> { if (choice.message().role() != null) { roleMap.putIfAbsent(id, choice.message().role().name()); } Map 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()); } + + })); + + // @formatter:off + Flux 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())); } - 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 { 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 @@ public class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod } } + 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 getFunctionTools(Set functionNames) { return this.resolveFunctionCallbacks(functionNames).stream().map(functionCallback -> { var function = new FunctionTool.Function(functionCallback.getDescription(), functionCallback.getName(), @@ -414,4 +481,8 @@ public class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod }).toList(); } + public void setObservationConvention(ChatModelObservationConvention observationConvention) { + this.observationConvention = observationConvention; + } + } diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuApiConstants.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuApiConstants.java index 28dcd229d..36d0c4292 100644 --- a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuApiConstants.java +++ b/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(); + } diff --git a/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java index 8b81ce131..3f2dfe312 100644 --- a/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java +++ b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java @@ -164,7 +164,7 @@ public class ZhiPuAiRetryTests { 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 diff --git a/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/chat/ZhiPuAiChatModelObservationIT.java b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/chat/ZhiPuAiChatModelObservationIT.java new file mode 100644 index 000000000..162a56b4f --- /dev/null +++ b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/chat/ZhiPuAiChatModelObservationIT.java @@ -0,0 +1,172 @@ +/* + * Copyright 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.zhipuai.chat; + +import io.micrometer.observation.tck.TestObservationRegistry; +import io.micrometer.observation.tck.TestObservationRegistryAssert; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.model.function.FunctionCallbackContext; +import org.springframework.ai.observation.conventions.AiOperationType; +import org.springframework.ai.observation.conventions.AiProvider; +import org.springframework.ai.zhipuai.ZhiPuAiChatModel; +import org.springframework.ai.zhipuai.ZhiPuAiChatOptions; +import org.springframework.ai.zhipuai.api.ZhiPuAiApi; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.context.annotation.Bean; +import org.springframework.retry.support.RetryTemplate; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.HighCardinalityKeyNames; +import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.LowCardinalityKeyNames; + +/** + * Integration tests for observation instrumentation in {@link ZhiPuAiChatModel}. + * + * @author Geng Rong + */ +@SpringBootTest(classes = ZhiPuAiChatModelObservationIT.Config.class) +@EnabledIfEnvironmentVariable(named = "ZHIPU_AI_API_KEY", matches = ".+") +public class ZhiPuAiChatModelObservationIT { + + @Autowired + TestObservationRegistry observationRegistry; + + @Autowired + ZhiPuAiChatModel chatModel; + + @BeforeEach + void beforeEach() { + observationRegistry.clear(); + } + + @Test + void observationForChatOperation() { + + var options = ZhiPuAiChatOptions.builder() + .withModel(ZhiPuAiApi.ChatModel.GLM_4_Air.getValue()) + .withMaxTokens(2048) + .withStop(List.of("this-is-the-end")) + .withTemperature(0.7) + .withTopP(1.0) + .build(); + + Prompt prompt = new Prompt("Why does a raven look like a desk?", options); + + ChatResponse chatResponse = chatModel.call(prompt); + assertThat(chatResponse.getResult().getOutput().getContent()).isNotEmpty(); + + ChatResponseMetadata responseMetadata = chatResponse.getMetadata(); + assertThat(responseMetadata).isNotNull(); + + validate(responseMetadata); + } + + @Test + void observationForStreamingChatOperation() { + var options = ZhiPuAiChatOptions.builder() + .withModel(ZhiPuAiApi.ChatModel.GLM_4_Air.getValue()) + .withMaxTokens(2048) + .withStop(List.of("this-is-the-end")) + .withTemperature(0.7) + .withTopP(1.0) + .build(); + + Prompt prompt = new Prompt("Why does a raven look like a desk?", options); + + Flux chatResponseFlux = chatModel.stream(prompt); + + List responses = chatResponseFlux.collectList().block(); + assertThat(responses).isNotEmpty(); + assertThat(responses).hasSizeGreaterThan(10); + + String aggregatedResponse = responses.subList(0, responses.size() - 1) + .stream() + .map(r -> r.getResult().getOutput().getContent()) + .collect(Collectors.joining()); + assertThat(aggregatedResponse).isNotEmpty(); + + ChatResponse lastChatResponse = responses.get(responses.size() - 1); + + ChatResponseMetadata responseMetadata = lastChatResponse.getMetadata(); + assertThat(responseMetadata).isNotNull(); + + validate(responseMetadata); + } + + private void validate(ChatResponseMetadata responseMetadata) { + TestObservationRegistryAssert.assertThat(observationRegistry) + .doesNotHaveAnyRemainingCurrentObservation() + .hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME) + .that() + .hasContextualNameEqualTo("chat " + ZhiPuAiApi.ChatModel.GLM_4_Air.getValue()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(), + AiOperationType.CHAT.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.ZHIPUAI.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(), + ZhiPuAiApi.ChatModel.GLM_4_Air.getValue()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_STOP_SEQUENCES.asString(), + "[\"this-is-the-end\"]") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TEMPERATURE.asString(), "0.7") + .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.REQUEST_TOP_K.asString()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_P.asString(), "1.0") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_ID.asString(), responseMetadata.getId()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_FINISH_REASONS.asString(), "[\"STOP\"]") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getPromptTokens())) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_OUTPUT_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getGenerationTokens())) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getTotalTokens())) + .hasBeenStarted() + .hasBeenStopped(); + } + + @SpringBootConfiguration + static class Config { + + @Bean + public TestObservationRegistry observationRegistry() { + return TestObservationRegistry.create(); + } + + @Bean + public ZhiPuAiApi zhiPuAiApi() { + return new ZhiPuAiApi(System.getenv("ZHIPU_AI_API_KEY")); + } + + @Bean + public ZhiPuAiChatModel zhiPuAiChatModel(ZhiPuAiApi zhiPuAiApi, TestObservationRegistry observationRegistry) { + return new ZhiPuAiChatModel(zhiPuAiApi, ZhiPuAiChatOptions.builder().build(), new FunctionCallbackContext(), + List.of(), RetryTemplate.defaultInstance(), observationRegistry); + } + + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiProvider.java b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiProvider.java index d88d72b7f..8ff242721 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiProvider.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiProvider.java @@ -39,6 +39,7 @@ public enum AiProvider { MINIMAX("minimax"), MOONSHOT("moonshot"), QIANFAN("qianfan"), + ZHIPUAI("zhipuai"), SPRING_AI("spring_ai"), VERTEX_AI("vertex_ai"); diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/zhipuai/ZhiPuAiAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/zhipuai/ZhiPuAiAutoConfiguration.java index 036a1b806..0831e2b2c 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/zhipuai/ZhiPuAiAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/zhipuai/ZhiPuAiAutoConfiguration.java @@ -15,7 +15,9 @@ */ package org.springframework.ai.autoconfigure.zhipuai; +import io.micrometer.observation.ObservationRegistry; import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration; +import org.springframework.ai.chat.observation.ChatModelObservationConvention; import org.springframework.ai.model.function.FunctionCallback; import org.springframework.ai.model.function.FunctionCallbackContext; import org.springframework.ai.zhipuai.ZhiPuAiChatModel; @@ -23,6 +25,7 @@ import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingModel; import org.springframework.ai.zhipuai.ZhiPuAiImageModel; import org.springframework.ai.zhipuai.api.ZhiPuAiApi; import org.springframework.ai.zhipuai.api.ZhiPuAiImageApi; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; @@ -55,13 +58,19 @@ public class ZhiPuAiAutoConfiguration { public ZhiPuAiChatModel zhiPuAiChatModel(ZhiPuAiConnectionProperties commonProperties, ZhiPuAiChatProperties chatProperties, RestClient.Builder restClientBuilder, List toolFunctionCallbacks, FunctionCallbackContext functionCallbackContext, - RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) { + RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler, + ObjectProvider observationRegistry, + ObjectProvider observationConvention) { var zhiPuAiApi = zhiPuAiApi(chatProperties.getBaseUrl(), commonProperties.getBaseUrl(), chatProperties.getApiKey(), commonProperties.getApiKey(), restClientBuilder, responseErrorHandler); - return new ZhiPuAiChatModel(zhiPuAiApi, chatProperties.getOptions(), functionCallbackContext, - toolFunctionCallbacks, retryTemplate); + var chatModel = new ZhiPuAiChatModel(zhiPuAiApi, chatProperties.getOptions(), functionCallbackContext, + toolFunctionCallbacks, retryTemplate, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)); + + observationConvention.ifAvailable(chatModel::setObservationConvention); + + return chatModel; } @Bean