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

Add observability to qianfan chat model

Browse Source
This commit is contained in:
GR
2024-10-08 10:26:01 +08:00
committed by Mark Pollack
parent 1606383704
commit 6203325f30
14 changed files with 363 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java
View File

@@ -256,14 +256,8 @@ public class AzureOpenAiChatModel extends AbstractToolCallSupport implements Cha
}
Flux<ChatResponse> flux = Flux.just(chatResponse).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);
});

5
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java
View File

@@ -347,11 +347,6 @@ public class OpenAiChatModel extends AbstractToolCallSupport implements ChatMode
})
.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));

6
models/spring-ai-qianfan/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>

158
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanChatModel.java
View File

@@ -15,13 +15,25 @@
*/
package org.springframework.ai.qianfan;
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;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.EmptyUsage;
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.qianfan.api.QianFanApi;
@@ -30,11 +42,14 @@ import org.springframework.ai.qianfan.api.QianFanApi.ChatCompletionChunk;
import org.springframework.ai.qianfan.api.QianFanApi.ChatCompletionMessage;
import org.springframework.ai.qianfan.api.QianFanApi.ChatCompletionMessage.Role;
import org.springframework.ai.qianfan.api.QianFanApi.ChatCompletionRequest;
import org.springframework.ai.qianfan.api.QianFanConstants;
import org.springframework.ai.qianfan.metadata.QianFanUsage;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import java.util.Collections;
import java.util.List;
@@ -45,15 +60,17 @@ import java.util.Map;
* backed by {@link QianFanApi}.
*
* @author Geng Rong
* @since 1.0
* @see ChatModel
* @see StreamingChatModel
* @see QianFanApi
* @since 1.0
*/
public class QianFanChatModel implements ChatModel, StreamingChatModel {
private static final Logger logger = LoggerFactory.getLogger(QianFanChatModel.class);
private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
/**
* The default options used for the chat completion requests.
*/
@@ -69,6 +86,16 @@ public class QianFanChatModel implements ChatModel, StreamingChatModel {
*/
private final QianFanApi qianFanApi;
/**
* 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 QianFanChatModel.
* @param qianFanApi The QianFanApi instance to be used for interacting with the
@@ -98,12 +125,27 @@ public class QianFanChatModel implements ChatModel, StreamingChatModel {
* @param retryTemplate The retry template.
*/
public QianFanChatModel(QianFanApi qianFanApi, QianFanChatOptions options, RetryTemplate retryTemplate) {
this(qianFanApi, options, retryTemplate, ObservationRegistry.NOOP);
}
/**
* Initializes a new instance of the QianFanChatModel.
* @param qianFanApi The QianFanApi instance to be used for interacting with the
* QianFan Chat API.
* @param options The QianFanChatOptions to configure the chat client.
* @param retryTemplate The retry template.
* @param observationRegistry The ObservationRegistry used for instrumentation.
*/
public QianFanChatModel(QianFanApi qianFanApi, QianFanChatOptions options, RetryTemplate retryTemplate,
ObservationRegistry observationRegistry) {
Assert.notNull(qianFanApi, "QianFanApi must not be null");
Assert.notNull(options, "Options must not be null");
Assert.notNull(retryTemplate, "RetryTemplate must not be null");
Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
this.qianFanApi = qianFanApi;
this.defaultOptions = options;
this.retryTemplate = retryTemplate;
this.observationRegistry = observationRegistry;
}
@Override
@@ -111,39 +153,80 @@ public class QianFanChatModel implements ChatModel, StreamingChatModel {
ChatCompletionRequest request = createRequest(prompt, false);
return this.retryTemplate.execute(ctx -> {
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
.prompt(prompt)
.provider(QianFanConstants.PROVIDER_NAME)
.requestOptions(buildRequestOptions(request))
.build();
ResponseEntity<ChatCompletion> completionEntity = this.doChatCompletion(request);
return ChatModelObservationDocumentation.CHAT_MODEL_OPERATION
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
this.observationRegistry)
.observe(() -> {
ResponseEntity<ChatCompletion> completionEntity = this.retryTemplate
.execute(ctx -> this.qianFanApi.chatCompletionEntity(request));
var chatCompletion = completionEntity.getBody();
if (chatCompletion == null) {
logger.warn("No chat completion returned for prompt: {}", prompt);
return new ChatResponse(List.of());
}
var chatCompletion = completionEntity.getBody();
if (chatCompletion == null) {
logger.warn("No chat completion returned for prompt: {}", prompt);
return new ChatResponse(List.of());
}
// if (chatCompletion.baseResponse() != null &&
// chatCompletion.baseResponse().statusCode() != 0) {
// throw new RuntimeException(chatCompletion.baseResponse().message());
// }
// @formatter:off
Map<String, Object> metadata = Map.of(
"id", chatCompletion.id(),
"role", Role.ASSISTANT
);
// @formatter:on
var generation = new Generation(chatCompletion.result(),
Map.of("id", chatCompletion.id(), "role", Role.ASSISTANT));
return new ChatResponse(Collections.singletonList(generation));
});
var assistantMessage = new AssistantMessage(chatCompletion.result(), metadata);
List<Generation> generations = Collections.singletonList(new Generation(assistantMessage));
ChatResponse chatResponse = new ChatResponse(generations, from(chatCompletion, request.model()));
observationContext.setResponse(chatResponse);
return chatResponse;
});
}
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
var request = createRequest(prompt, true);
return retryTemplate.execute(ctx -> {
return Flux.deferContextual(contextView -> {
ChatCompletionRequest request = createRequest(prompt, true);
var completionChunks = this.qianFanApi.chatCompletionStream(request);
return completionChunks.map(this::toChatCompletion).map(chatCompletion -> {
String id = chatCompletion.id();
var generation = new Generation(chatCompletion.result(), Map.of("id", id, "role", Role.ASSISTANT));
return new ChatResponse(Collections.singletonList(generation));
});
final ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
.prompt(prompt)
.provider(QianFanConstants.PROVIDER_NAME)
.requestOptions(buildRequestOptions(request))
.build();
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::toChatCompletion)
.switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> {
// @formatter:off
Map<String, Object> metadata = Map.of(
"id", chatCompletion.id(),
"role", Role.ASSISTANT
);
// @formatter:on
var assistantMessage = new AssistantMessage(chatCompletion.result(), metadata);
List<Generation> generations = Collections.singletonList(new Generation(assistantMessage));
return new ChatResponse(generations, from(chatCompletion, request.model()));
}))
.doOnError(observation::error)
.doFinally(s -> {
observation.stop();
})
.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
return new MessageAggregator().aggregate(chatResponse, observationContext::setResponse);
});
}
@@ -153,7 +236,8 @@ public class QianFanChatModel implements ChatModel, StreamingChatModel {
* @return the ChatCompletion
*/
private ChatCompletion toChatCompletion(ChatCompletionChunk chunk) {
return new ChatCompletion(chunk.id(), chunk.object(), chunk.created(), chunk.result(), chunk.usage());
return new ChatCompletion(chunk.id(), chunk.object(), chunk.created(), chunk.result(), chunk.finishReason(),
chunk.usage());
}
/**
@@ -193,8 +277,30 @@ public class QianFanChatModel implements ChatModel, StreamingChatModel {
return QianFanChatOptions.fromOptions(this.defaultOptions);
}
private ResponseEntity<ChatCompletion> doChatCompletion(ChatCompletionRequest request) {
return this.qianFanApi.chatCompletionEntity(request);
private ChatOptions buildRequestOptions(QianFanApi.ChatCompletionRequest request) {
return ChatOptionsBuilder.builder()
.withModel(request.model())
.withFrequencyPenalty(request.frequencyPenalty())
.withMaxTokens(request.maxTokens())
.withPresencePenalty(request.presencePenalty())
.withStopSequences(request.stop())
.withTemperature(request.temperature())
.withTopP(request.topP())
.build();
}
private ChatResponseMetadata from(QianFanApi.ChatCompletion result, String model) {
Assert.notNull(result, "QianFan ChatCompletionResult must not be null");
return ChatResponseMetadata.builder()
.withId(result.id() != null ? result.id() : "")
.withUsage(result.usage() != null ? QianFanUsage.from(result.usage()) : new EmptyUsage())
.withModel(model)
.withKeyValue("created", result.created() != null ? result.created() : 0L)
.build();
}
public void setObservationConvention(ChatModelObservationConvention observationConvention) {
this.observationConvention = observationConvention;
}
}

2
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanChatOptions.java
View File

@@ -51,7 +51,7 @@ public class QianFanChatOptions implements ChatOptions {
* The maximum number of tokens to generate in the chat completion. The total length of input
* tokens and generated tokens is limited by the model's context length.
*/
private @JsonProperty("max_tokens") Integer maxTokens;
private @JsonProperty("max_output_tokens") Integer maxTokens;
/**
* Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they
* appear in the text so far, increasing the model's likelihood to talk about new topics.

11
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanApi.java
View File

@@ -60,7 +60,7 @@ public class QianFanApi extends AuthApi {
* @param secretKey QianFan secret key.
*/
public QianFanApi(String apiKey, String secretKey) {
this(ApiUtils.DEFAULT_BASE_URL, apiKey, secretKey);
this(QianFanConstants.DEFAULT_BASE_URL, apiKey, secretKey);
}
/**
@@ -110,18 +110,18 @@ public class QianFanApi extends AuthApi {
* @param responseErrorHandler Response error handler.
*/
public QianFanApi(String baseUrl, String apiKey, String secretKey, RestClient.Builder restClientBuilder,
WebClient.Builder webClientBuilder,ResponseErrorHandler responseErrorHandler) {
WebClient.Builder webClientBuilder, ResponseErrorHandler responseErrorHandler) {
super(apiKey, secretKey);
this.restClient = restClientBuilder
.baseUrl(baseUrl)
.defaultHeaders(ApiUtils.getJsonContentHeaders())
.defaultHeaders(QianFanUtils.defaultHeaders())
.defaultStatusHandler(responseErrorHandler)
.build();
this.webClient = webClientBuilder
.baseUrl(baseUrl)
.defaultHeaders(ApiUtils.getJsonContentHeaders())
.defaultHeaders(QianFanUtils.defaultHeaders())
.build();
}
@@ -308,6 +308,7 @@ public class QianFanApi extends AuthApi {
@JsonProperty("object") String object,
@JsonProperty("created") Long created,
@JsonProperty("result") String result,
@JsonProperty("finish_reason") String finishReason,
@JsonProperty("usage") Usage usage) {
}
@@ -319,6 +320,7 @@ public class QianFanApi extends AuthApi {
*/
@JsonInclude(Include.NON_NULL)
public record Usage(
@JsonProperty("completion_tokens") Integer completionTokens,
@JsonProperty("prompt_tokens") Integer promptTokens,
@JsonProperty("total_tokens") Integer totalTokens) {
@@ -339,6 +341,7 @@ public class QianFanApi extends AuthApi {
@JsonProperty("object") String object,
@JsonProperty("created") Long created,
@JsonProperty("result") String result,
@JsonProperty("finish_reason") String finishReason,
@JsonProperty("is_end") Boolean end,
@JsonProperty("usage") Usage usage

11
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/ApiUtils.java → models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanConstants.java
View File

@@ -15,10 +15,7 @@
*/
package org.springframework.ai.qianfan.api;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import java.util.function.Consumer;
import org.springframework.ai.observation.conventions.AiProvider;
/**
* The ApiUtils class provides utility methods for working with API requests and
@@ -27,12 +24,10 @@ import java.util.function.Consumer;
* @author Geng Rong
* @since 1.0
*/
public class ApiUtils {
public class QianFanConstants {
public static final String DEFAULT_BASE_URL = "https://aip.baidubce.com/rpc/2.0/ai_custom";
public static Consumer<HttpHeaders> getJsonContentHeaders() {
return headers -> headers.setContentType(MediaType.APPLICATION_JSON);
}
public static final String PROVIDER_NAME = AiProvider.QIANFAN.value();
}

4
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanImageApi.java
View File

@@ -44,7 +44,7 @@ public class QianFanImageApi extends AuthApi {
* @param secretKey QianFan secret key.
*/
public QianFanImageApi(String apiKey, String secretKey) {
this(ApiUtils.DEFAULT_BASE_URL, apiKey, secretKey, RestClient.builder());
this(QianFanConstants.DEFAULT_BASE_URL, apiKey, secretKey, RestClient.builder());
}
/**
@@ -71,7 +71,7 @@ public class QianFanImageApi extends AuthApi {
super(apiKey, secretKey);
this.restClient = restClientBuilder.baseUrl(baseUrl)
.defaultHeaders(ApiUtils.getJsonContentHeaders())
.defaultHeaders(QianFanUtils.defaultHeaders())
.defaultStatusHandler(responseErrorHandler)
.build();
}

14
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanUtils.java Normal file
View File

@@ -0,0 +1,14 @@
package org.springframework.ai.qianfan.api;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import java.util.function.Consumer;
public class QianFanUtils {
public static Consumer<HttpHeaders> defaultHeaders() {
return headers -> headers.setContentType(MediaType.APPLICATION_JSON);
}
}

12
models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/api/QianFanRetryTests.java
View File

@@ -16,6 +16,7 @@
package org.springframework.ai.qianfan.api;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
@@ -108,8 +109,8 @@ public class QianFanRetryTests {
@Test
public void qianFanChatTransientError() {
ChatCompletion expectedChatCompletion = new ChatCompletion("id", "chat.completion", 666L, "Response",
new Usage(10, 10));
ChatCompletion expectedChatCompletion = new ChatCompletion("id", "chat.completion", 666L, "Response", "STOP",
new Usage(10, 10, 10));
when(qianFanApi.chatCompletionEntity(isA(ChatCompletionRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
@@ -132,9 +133,10 @@ public class QianFanRetryTests {
}
@Test
@Disabled("Currently stream() does not implmement retry")
public void qianFanChatStreamTransientError() {
ChatCompletionChunk expectedChatCompletion = new ChatCompletionChunk("id", "chat.completion", 666L, "Response",
true, null);
"", true, null);
when(qianFanApi.chatCompletionStream(isA(ChatCompletionRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
@@ -154,14 +156,14 @@ public class QianFanRetryTests {
public void qianFanChatStreamNonTransientError() {
when(qianFanApi.chatCompletionStream(isA(ChatCompletionRequest.class)))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> chatClient.stream(new Prompt("text")));
assertThrows(RuntimeException.class, () -> chatClient.stream(new Prompt("text")).collectList().block());
}
@Test
public void qianFanEmbeddingTransientError() {
QianFanApi.Embedding embedding = new QianFanApi.Embedding(1, new float[] { 9.9f, 8.8f });
EmbeddingList expectedEmbeddings = new EmbeddingList("embedding_list", List.of(embedding), "model", null, null,
new Usage(10, 10));
new Usage(10, 10, 10));
when(qianFanApi.embeddings(isA(EmbeddingRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))

177
models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/chat/QianFanChatModelObservationIT.java Normal file
View File

@@ -0,0 +1,177 @@
/*
* 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.qianfan.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.junit.jupiter.api.condition.EnabledIfEnvironmentVariables;
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.observation.conventions.AiOperationType;
import org.springframework.ai.observation.conventions.AiProvider;
import org.springframework.ai.qianfan.QianFanChatModel;
import org.springframework.ai.qianfan.QianFanChatOptions;
import org.springframework.ai.qianfan.api.QianFanApi;
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 QianFanChatModel}.
*
* @author Geng Rong
*/
@SpringBootTest(classes = QianFanChatModelObservationIT.Config.class)
@EnabledIfEnvironmentVariables(value = { @EnabledIfEnvironmentVariable(named = "QIANFAN_API_KEY", matches = ".+"),
@EnabledIfEnvironmentVariable(named = "QIANFAN_SECRET_KEY", matches = ".+") })
public class QianFanChatModelObservationIT {
@Autowired
TestObservationRegistry observationRegistry;
@Autowired
QianFanChatModel chatModel;
@BeforeEach
void beforeEach() {
observationRegistry.clear();
}
@Test
void observationForChatOperation() {
var options = QianFanChatOptions.builder()
.withModel(QianFanApi.ChatModel.ERNIE_Speed_8K.getValue())
.withFrequencyPenalty(0.0)
.withMaxTokens(2048)
.withPresencePenalty(0.0)
.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 = QianFanChatOptions.builder()
.withModel(QianFanApi.ChatModel.ERNIE_Speed_8K.getValue())
.withFrequencyPenalty(0.0)
.withMaxTokens(2048)
.withPresencePenalty(0.0)
.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();
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 " + QianFanApi.ChatModel.ERNIE_Speed_8K.getValue())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
AiOperationType.CHAT.value())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.QIANFAN.value())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(),
QianFanApi.ChatModel.ERNIE_Speed_8K.getValue())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_FREQUENCY_PENALTY.asString(), "0.0")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_PRESENCE_PENALTY.asString(), "0.0")
.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.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 QianFanApi qianFanApi() {
return new QianFanApi(System.getenv("QIANFAN_API_KEY"), System.getenv("QIANFAN_SECRET_KEY"));
}
@Bean
public QianFanChatModel qianFanChatModel(QianFanApi qianFanApi, TestObservationRegistry observationRegistry) {
return new QianFanChatModel(qianFanApi, QianFanChatOptions.builder().build(),
RetryTemplate.defaultInstance(), observationRegistry);
}
}
}

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

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

13
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/qianfan/QianFanAutoConfiguration.java
View File

@@ -15,13 +15,16 @@
*/
package org.springframework.ai.autoconfigure.qianfan;
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.FunctionCallbackContext;
import org.springframework.ai.qianfan.QianFanChatModel;
import org.springframework.ai.qianfan.QianFanEmbeddingModel;
import org.springframework.ai.qianfan.QianFanImageModel;
import org.springframework.ai.qianfan.api.QianFanApi;
import org.springframework.ai.qianfan.api.QianFanImageApi;
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;
@@ -51,13 +54,19 @@ public class QianFanAutoConfiguration {
matchIfMissing = true)
public QianFanChatModel qianFanChatModel(QianFanConnectionProperties commonProperties,
QianFanChatProperties chatProperties, RestClient.Builder restClientBuilder, RetryTemplate retryTemplate,
ResponseErrorHandler responseErrorHandler) {
ResponseErrorHandler responseErrorHandler, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<ChatModelObservationConvention> observationConvention) {
var qianFanApi = qianFanApi(chatProperties.getBaseUrl(), commonProperties.getBaseUrl(),
chatProperties.getApiKey(), commonProperties.getApiKey(), chatProperties.getSecretKey(),
commonProperties.getSecretKey(), restClientBuilder, responseErrorHandler);
return new QianFanChatModel(qianFanApi, chatProperties.getOptions(), retryTemplate);
var chatModel = new QianFanChatModel(qianFanApi, chatProperties.getOptions(), retryTemplate,
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
observationConvention.ifAvailable(chatModel::setObservationConvention);
return chatModel;
}
@Bean

4
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/qianfan/QianFanConnectionProperties.java
View File

@@ -15,7 +15,7 @@
*/
package org.springframework.ai.autoconfigure.qianfan;
import org.springframework.ai.qianfan.api.ApiUtils;
import org.springframework.ai.qianfan.api.QianFanConstants;
import org.springframework.boot.context.properties.ConfigurationProperties;
@ConfigurationProperties(QianFanConnectionProperties.CONFIG_PREFIX)
@@ -23,7 +23,7 @@ public class QianFanConnectionProperties extends QianFanParentProperties {
public static final String CONFIG_PREFIX = "spring.ai.qianfan";
public static final String DEFAULT_BASE_URL = ApiUtils.DEFAULT_BASE_URL;
public static final String DEFAULT_BASE_URL = QianFanConstants.DEFAULT_BASE_URL;
public QianFanConnectionProperties() {
super.setBaseUrl(DEFAULT_BASE_URL);