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 568a922fe..211a7029e 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 @@ -253,7 +253,7 @@ public class AnthropicChatModel implements ChatModel { this.getAdditionalHttpHeaders(prompt)); // @formatter:off - Flux chatResponseFlux = response.switchMap(chatCompletionResponse -> { + Flux chatResponseFlux = response.flatMap(chatCompletionResponse -> { AnthropicApi.Usage usage = chatCompletionResponse.usage(); Usage currentChatResponseUsage = usage != null ? this.getDefaultUsage(chatCompletionResponse.usage()) : new EmptyUsage(); Usage accumulatedUsage = UsageUtils.getCumulativeUsage(currentChatResponseUsage, previousChatResponse); diff --git a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/StreamHelper.java b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/StreamHelper.java index ae62eb074..e08a96690 100644 --- a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/StreamHelper.java +++ b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/StreamHelper.java @@ -179,7 +179,14 @@ public class StreamHelper { } } else if (event.type().equals(EventType.MESSAGE_STOP)) { - // pass through + // Don't return the latest Content block as it was before. Instead, return it + // with an updated event type and general information like: model, message + // type, id and usage + contentBlockReference.get() + .withType(event.type().name()) + .withContent(List.of()) + .withStopReason(null) + .withStopSequence(null); } else { contentBlockReference.get().withType(event.type().name()).withContent(List.of()); diff --git a/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/AnthropicApiIT.java b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/AnthropicApiIT.java index fe80f72a3..e2da41559 100644 --- a/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/AnthropicApiIT.java +++ b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/AnthropicApiIT.java @@ -16,6 +16,7 @@ package org.springframework.ai.anthropic.api; +import java.util.ArrayList; import java.util.List; import org.junit.jupiter.api.Test; @@ -27,6 +28,7 @@ import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionRequest; import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionResponse; import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock; import org.springframework.ai.anthropic.api.AnthropicApi.Role; +import org.springframework.ai.model.ModelOptionsUtils; import org.springframework.http.ResponseEntity; import static org.assertj.core.api.Assertions.assertThat; @@ -42,6 +44,24 @@ public class AnthropicApiIT { AnthropicApi anthropicApi = AnthropicApi.builder().apiKey(System.getenv("ANTHROPIC_API_KEY")).build(); + List tools = List.of(new AnthropicApi.Tool("getCurrentWeather", + "Get the weather in location. Return temperature in 30°F or 30°C format.", ModelOptionsUtils.jsonToMap(""" + { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state e.g. San Francisco, CA" + }, + "unit": { + "type": "string", + "enum": ["C", "F"] + } + }, + "required": ["location", "unit"] + } + """))); + @Test void chatCompletionEntity() { @@ -106,6 +126,47 @@ public class AnthropicApiIT { bla.stream().forEach(r -> System.out.println(r)); } + @Test + void chatCompletionStreamWithToolCall() { + List messageConversation = new ArrayList<>(); + + AnthropicMessage chatCompletionMessage = new AnthropicMessage( + List.of(new ContentBlock("What's the weather like in San Francisco? Show the temperature in Celsius.")), + Role.USER); + + messageConversation.add(chatCompletionMessage); + + ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder() + .model(AnthropicApi.ChatModel.CLAUDE_3_OPUS) + .messages(messageConversation) + .maxTokens(1500) + .stream(true) + .temperature(0.8) + .tools(tools) + .build(); + + List responses = this.anthropicApi.chatCompletionStream(chatCompletionRequest) + .collectList() + .block(); + + // Check that tool uses response returned only once + List toolCompletionResponses = responses.stream() + .filter(r -> r.stopReason() != null && r.stopReason().equals(ContentBlock.Type.TOOL_USE.value)) + .toList(); + assertThat(toolCompletionResponses).size().isEqualTo(1); + List toolContentBlocks = toolCompletionResponses.get(0).content(); + assertThat(toolContentBlocks).size().isEqualTo(1); + ContentBlock toolContentBlock = toolContentBlocks.get(0); + assertThat(toolContentBlock.type()).isEqualTo(ContentBlock.Type.TOOL_USE); + assertThat(toolContentBlock.name()).isEqualTo("getCurrentWeather"); + + // Check that message stop response also returned + List messageStopEvents = responses.stream() + .filter(r -> r.type().equals(AnthropicApi.EventType.MESSAGE_STOP.name())) + .toList(); + assertThat(messageStopEvents).size().isEqualTo(1); + } + @Test void chatCompletionStreamError() { AnthropicMessage chatCompletionMessage = new AnthropicMessage(List.of(new ContentBlock("Tell me a Joke?")),