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

Fix Anthropic chat model functioncalling token usage

Browse Source 
- Accumulate the token usage when functioncalling is used
   - Fix both call() as well as stream() operations

 - Add/update tests
This commit is contained in:
Ilayaperumal Gopinathan
2024-12-12 14:52:44 +00:00
committed by Christian Tzolov
parent 81dfd3b6b5
commit cadea8b730
4 changed files with 68 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java
View File

@@ -47,6 +47,9 @@ 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.metadata.Usage;
import org.springframework.ai.chat.metadata.UsageUtils;
import org.springframework.ai.chat.model.AbstractToolCallSupport;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
@@ -210,6 +213,10 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
@Override
public ChatResponse call(Prompt prompt) {
return this.internalCall(prompt, null);
}
public ChatResponse internalCall(Prompt prompt, ChatResponse previousChatResponse) {
ChatCompletionRequest request = createRequest(prompt, false);
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
@@ -226,8 +233,14 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
ResponseEntity<ChatCompletionResponse> completionEntity = this.retryTemplate
.execute(ctx -> this.anthropicApi.chatCompletionEntity(request));
ChatResponse chatResponse = toChatResponse(completionEntity.getBody());
AnthropicApi.ChatCompletionResponse completionResponse = completionEntity.getBody();
AnthropicApi.Usage usage = completionResponse.usage();
Usage currentChatResponseUsage = usage != null ? AnthropicUsage.from(completionResponse.usage())
: new EmptyUsage();
Usage accumulatedUsage = UsageUtils.getCumulativeUsage(currentChatResponseUsage, previousChatResponse);
ChatResponse chatResponse = toChatResponse(completionEntity.getBody(), accumulatedUsage);
observationContext.setResponse(chatResponse);
return chatResponse;
@@ -236,7 +249,7 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
if (!isProxyToolCalls(prompt, this.defaultOptions) && response != null
&& this.isToolCall(response, Set.of("tool_use"))) {
var toolCallConversation = handleToolCalls(prompt, response);
return this.call(new Prompt(toolCallConversation, prompt.getOptions()));
return this.internalCall(new Prompt(toolCallConversation, prompt.getOptions()), response);
}
return response;
@@ -244,6 +257,10 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
return this.internalStream(prompt, null);
}
public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousChatResponse) {
return Flux.deferContextual(contextView -> {
ChatCompletionRequest request = createRequest(prompt, true);
@@ -263,11 +280,14 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
// @formatter:off
Flux<ChatResponse> chatResponseFlux = response.switchMap(chatCompletionResponse -> {
ChatResponse chatResponse = toChatResponse(chatCompletionResponse);
AnthropicApi.Usage usage = chatCompletionResponse.usage();
Usage currentChatResponseUsage = usage != null ? AnthropicUsage.from(chatCompletionResponse.usage()) : new EmptyUsage();
Usage accumulatedUsage = UsageUtils.getCumulativeUsage(currentChatResponseUsage, previousChatResponse);
ChatResponse chatResponse = toChatResponse(chatCompletionResponse, accumulatedUsage);
if (!isProxyToolCalls(prompt, this.defaultOptions) && this.isToolCall(chatResponse, Set.of("tool_use"))) {
var toolCallConversation = handleToolCalls(prompt, chatResponse);
return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
return this.internalStream(new Prompt(toolCallConversation, prompt.getOptions()), chatResponse);
}
return Mono.just(chatResponse);
@@ -281,7 +301,7 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
});
}
private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion) {
private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage usage) {
if (chatCompletion == null) {
logger.warn("Null chat completion returned");
@@ -327,12 +347,15 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
allGenerations.add(toolCallGeneration);
}
return new ChatResponse(allGenerations, this.from(chatCompletion));
return new ChatResponse(allGenerations, this.from(chatCompletion, usage));
}
private ChatResponseMetadata from(AnthropicApi.ChatCompletionResponse result) {
return from(result, AnthropicUsage.from(result.usage()));
}
private ChatResponseMetadata from(AnthropicApi.ChatCompletionResponse result, Usage usage) {
Assert.notNull(result, "Anthropic ChatCompletionResult must not be null");
AnthropicUsage usage = AnthropicUsage.from(result.usage());
return ChatResponseMetadata.builder()
.withId(result.id())
.withModel(result.model())

35
models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/AnthropicChatModelIT.java
View File

@@ -37,6 +37,7 @@ import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
@@ -288,7 +289,12 @@ class AnthropicChatModelIT {
logger.info("Response: {}", response);
Generation generation = response.getResult();
assertThat(generation).isNotNull();
assertThat(generation.getOutput()).isNotNull();
assertThat(generation.getOutput().getText()).contains("30", "10", "15");
assertThat(response.getMetadata()).isNotNull();
assertThat(response.getMetadata().getUsage()).isNotNull();
assertThat(response.getMetadata().getUsage().getTotalTokens()).isLessThan(4000).isGreaterThan(1800);
}
@Test
@@ -324,6 +330,35 @@ class AnthropicChatModelIT {
assertThat(content).contains("30", "10", "15");
}
@Test
void streamFunctionCallUsageTest() {
UserMessage userMessage = new UserMessage(
"What's the weather like in San Francisco, Tokyo and Paris? Return the result in Celsius.");
List<Message> messages = new ArrayList<>(List.of(userMessage));
var promptOptions = AnthropicChatOptions.builder()
.withModel(AnthropicApi.ChatModel.CLAUDE_3_5_SONNET.getName())
.withFunctionCallbacks(List.of(FunctionCallback.builder()
.function("getCurrentWeather", new MockWeatherService())
.description(
"Get the weather in location. Return temperature in 36°F or 36°C format. Use multi-turn if needed.")
.inputType(MockWeatherService.Request.class)
.build()))
.build();
Flux<ChatResponse> responseFlux = this.chatModel.stream(new Prompt(messages, promptOptions));
ChatResponse chatResponse = responseFlux.last().block();
logger.info("Response: {}", chatResponse);
Usage usage = chatResponse.getMetadata().getUsage();
assertThat(usage).isNotNull();
assertThat(usage.getTotalTokens()).isLessThan(4000).isGreaterThan(1800);
}
@Test
void validateCallResponseMetadata() {
String model = AnthropicApi.ChatModel.CLAUDE_2_1.getName();

4
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithFunctionBeanIT.java
View File

@@ -57,7 +57,7 @@ class FunctionCallWithFunctionBeanIT {
this.contextRunner
.withPropertyValues(
"spring.ai.anthropic.chat.options.model=" + AnthropicApi.ChatModel.CLAUDE_3_OPUS.getValue())
"spring.ai.anthropic.chat.options.model=" + AnthropicApi.ChatModel.CLAUDE_3_5_HAIKU.getValue())
.run(context -> {
AnthropicChatModel chatModel = context.getBean(AnthropicChatModel.class);
@@ -87,7 +87,7 @@ class FunctionCallWithFunctionBeanIT {
this.contextRunner
.withPropertyValues(
"spring.ai.anthropic.chat.options.model=" + AnthropicApi.ChatModel.CLAUDE_3_OPUS.getValue())
"spring.ai.anthropic.chat.options.model=" + AnthropicApi.ChatModel.CLAUDE_3_5_HAIKU.getValue())
.run(context -> {
AnthropicChatModel chatModel = context.getBean(AnthropicChatModel.class);

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithPromptFunctionIT.java
View File

@@ -49,7 +49,7 @@ public class FunctionCallWithPromptFunctionIT {
void functionCallTest() {
this.contextRunner
.withPropertyValues(
"spring.ai.anthropic.chat.options.model=" + AnthropicApi.ChatModel.CLAUDE_3_OPUS.getValue())
"spring.ai.anthropic.chat.options.model=" + AnthropicApi.ChatModel.CLAUDE_3_5_HAIKU.getValue())
.run(context -> {
AnthropicChatModel chatModel = context.getBean(AnthropicChatModel.class);