Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add observability to zhipu chat model

Browse Source
This commit is contained in:
GR
2024-10-08 11:06:29 +08:00
committed by Mark Pollack
parent 6203325f30
commit 7bd8a17d33
7 changed files with 331 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
models/spring-ai-zhipuai/pom.xml
View File

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

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

@@ -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<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 -> {
if (chatCompletion == null) {
logger.warn("No chat completion returned for prompt: {}", prompt);
return new ChatResponse(List.of());
}
List<Choice> choices = chatCompletion.choices();
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();
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()));
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<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());
}
}));
// @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()));
}
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<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 @@ public class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod
}).toList();
}
public void setObservationConvention(ChatModelObservationConvention observationConvention) {
this.observationConvention = observationConvention;
}
}

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

@@ -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();
}

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

@@ -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

172
models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/chat/ZhiPuAiChatModelObservationIT.java Normal file
View File

@@ -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<ChatResponse> chatResponseFlux = chatModel.stream(prompt);
List<ChatResponse> 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);
}
}
}

1
spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiProvider.java
View File

@@ -39,6 +39,7 @@ public enum AiProvider {
MINIMAX("minimax"),
MOONSHOT("moonshot"),
QIANFAN("qianfan"),
ZHIPUAI("zhipuai"),
SPRING_AI("spring_ai"),
VERTEX_AI("vertex_ai");

15
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/zhipuai/ZhiPuAiAutoConfiguration.java
View File

@@ -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<FunctionCallback> toolFunctionCallbacks, FunctionCallbackContext functionCallbackContext,
RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) {
RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler,
ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<ChatModelObservationConvention> 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