diff --git a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java index 4f1c09df7..babf9a85c 100644 --- a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java +++ b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java @@ -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 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 stream(Prompt prompt) { + return this.internalStream(prompt, null); + } + + public Flux 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 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()) diff --git a/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/AnthropicChatModelIT.java b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/AnthropicChatModelIT.java index aa62746bc..824e1c42c 100644 --- a/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/AnthropicChatModelIT.java +++ b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/AnthropicChatModelIT.java @@ -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 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 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(); diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithFunctionBeanIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithFunctionBeanIT.java index 5a03f27ee..d42d2f2c2 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithFunctionBeanIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithFunctionBeanIT.java @@ -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); diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithPromptFunctionIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithPromptFunctionIT.java index 3c8867ac8..2884f9ee7 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithPromptFunctionIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/anthropic/tool/FunctionCallWithPromptFunctionIT.java @@ -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);