Fix Anthropic token usage handling in streaming

Author: didalgolab

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

@@ -59,6 +59,7 @@
  *
  * @author Christian Tzolov
  * @author luocongqiu
+ * @author Mariusz Bernacki
  * @since 1.0.0
  */
 public class AnthropicChatModel extends
@@ -192,6 +193,11 @@ else if (chunk.type().equals("message_delta")) {
 				ChatCompletion delta = ModelOptionsUtils.mapToClass(chunk.delta(), ChatCompletion.class);
 
 				chatCompletionReference.get().withType(chunk.type());
+				if (chunk.usage() != null) {
+					var totalUsage = new Usage(chatCompletionReference.get().usage.inputTokens(),
+							chunk.usage().outputTokens());
+					chatCompletionReference.get().withUsage(totalUsage);
+				}
 				if (delta.id() != null) {
 					chatCompletionReference.get().withId(delta.id());
 				}
@@ -201,9 +207,6 @@ else if (chunk.type().equals("message_delta")) {
 				if (delta.model() != null) {
 					chatCompletionReference.get().withModel(delta.model());
 				}
-				if (delta.usage() != null) {
-					chatCompletionReference.get().withUsage(delta.usage());
-				}
 				if (delta.content() != null) {
 					chatCompletionReference.get().withContent(delta.content());
 				}

models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java

@@ -40,6 +40,7 @@
 
 /**
  * @author Christian Tzolov
+ * @author Mariusz Bernacki
  * @since 1.0.0
  */
 public class AnthropicApi {
@@ -521,6 +522,17 @@ public record Usage( // @formatter:off
 		 // @formatter:off
 	}
 
+	/**
+	 * Usage statistics with output only tokens for streamed completions.
+	 *
+	 * @param outputTokens The number of output tokens which were used in a completion.
+	 */
+	@JsonInclude(Include.NON_NULL)
+	public record OutputUsage( // @formatter:off
+		@JsonProperty("output_tokens") Integer outputTokens) {
+		// @formatter:off
+	}
+
 	/**
 	 * The role of the author of this message.
 	 */
@@ -557,7 +569,8 @@ public record StreamResponse( // @formatter:off
 		@JsonProperty("index") Integer index,
 		@JsonProperty("message") ChatCompletion message,
 		@JsonProperty("content_block") MediaContent contentBlock,
-		@JsonProperty("delta") Map<String, Object> delta) {
+		@JsonProperty("delta") Map<String, Object> delta,
+		@JsonProperty("usage") OutputUsage usage) {
 		// @formatter:on
 	}
 

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

@@ -94,6 +94,24 @@ void roleTest(String modelName) {
 		logger.info(response.toString());
 	}
 
+	@Test
+	void streamingWithTokenUsage() {
+		var promptOptions = AnthropicChatOptions.builder().withTemperature(0f).build();
+
+		var prompt = new Prompt("List two colors of the Polish flag. Be brief.", promptOptions);
+		var streamingTokenUsage = this.chatModel.stream(prompt).blockLast().getMetadata().getUsage();
+		var referenceTokenUsage = this.chatModel.call(prompt).getMetadata().getUsage();
+
+		assertThat(streamingTokenUsage.getPromptTokens()).isGreaterThan(0);
+		assertThat(streamingTokenUsage.getGenerationTokens()).isGreaterThan(0);
+		assertThat(streamingTokenUsage.getTotalTokens()).isGreaterThan(0);
+
+		assertThat(streamingTokenUsage.getPromptTokens()).isEqualTo(referenceTokenUsage.getPromptTokens());
+		assertThat(streamingTokenUsage.getGenerationTokens()).isEqualTo(referenceTokenUsage.getGenerationTokens());
+		assertThat(streamingTokenUsage.getTotalTokens()).isEqualTo(referenceTokenUsage.getTotalTokens());
+
+	}
+
 	@Test
 	void listOutputConverter() {
 		DefaultConversionService conversionService = new DefaultConversionService();
@@ -195,7 +213,7 @@ void multiModalityTest() throws IOException {
 		var response = chatModel.call(new Prompt(List.of(userMessage)));
 
 		logger.info(response.getResult().getOutput().getContent());
-		assertThat(response.getResult().getOutput().getContent()).contains("bananas", "apple", "basket");
+		assertThat(response.getResult().getOutput().getContent()).contains("banan", "apple", "basket");
 	}
 
 	@Test

models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/client/AnthropicChatClientIT.java

@@ -230,8 +230,10 @@ void defaultFunctionCallTest() {
 		String response = ChatClient.builder(chatModel)
 				.defaultFunction("getCurrentWeather", "Get the weather in location", new MockWeatherService())
 				.defaultUser(u -> u.text("What's the weather like in San Francisco, Tokyo, and Paris? Use Celsius."))
-			.build()
-			.prompt().call().content();
+				.build()
+			.prompt()
+			.call()
+			.content();
 		// @formatter:on
 
 		logger.info("Response: {}", response);
@@ -304,7 +306,7 @@ void streamingMultiModality() throws IOException {
 
 		// @formatter:off
 		Flux<String> response = ChatClient.create(chatModel).prompt()
-				.options(AnthropicChatOptions.builder().withModel(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
+				.options(AnthropicChatOptions.builder().withModel(AnthropicApi.ChatModel.CLAUDE_3_5_SONNET)
 						.build())
 				.user(u -> u.text("Explain what do you see on this picture?")
 						.media(MimeTypeUtils.IMAGE_PNG, new ClassPathResource("/test.png")))