From f5b72242c03fd52688dcc3d741a188019e07170b Mon Sep 17 00:00:00 2001 From: Mariusz Bernacki Date: Tue, 25 Jun 2024 21:11:49 +0200 Subject: [PATCH] Add Anthropic Claude streaming Function Calling support Resolves #930 --- .../ai/anthropic/AnthropicChatModel.java | 149 +++++++++++------- .../ai/anthropic/api/AnthropicApi.java | 49 +++++- .../AnthropicStreamFunctionCallingHelper.java | 118 ++++++++++++++ .../client/AnthropicChatClientIT.java | 3 +- 4 files changed, 250 insertions(+), 69 deletions(-) create mode 100644 models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicStreamFunctionCallingHelper.java 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 8ba821e48..8aedc14f1 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 @@ -20,15 +20,13 @@ import java.util.Base64; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import org.springframework.ai.chat.model.ChatModel; -import reactor.core.publisher.Flux; - import org.springframework.ai.anthropic.api.AnthropicApi; import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletion; import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionRequest; @@ -39,6 +37,7 @@ import org.springframework.ai.anthropic.api.AnthropicApi.Role; import org.springframework.ai.anthropic.api.AnthropicApi.StreamResponse; import org.springframework.ai.anthropic.api.AnthropicApi.Usage; import org.springframework.ai.anthropic.metadata.AnthropicChatResponseMetadata; +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.messages.MessageType; @@ -54,6 +53,8 @@ import org.springframework.retry.support.RetryTemplate; import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; +import reactor.core.publisher.Flux; + /** * The {@link ChatModel} implementation for the Anthropic service. * @@ -160,69 +161,86 @@ public class AnthropicChatModel extends ChatCompletionRequest request = createRequest(prompt, true); - Flux response = this.anthropicApi.chatCompletionStream(request); + return this.retryTemplate.execute(ctx -> { - AtomicReference chatCompletionReference = new AtomicReference<>(); + Flux response = this.anthropicApi.chatCompletionStream(request); + + AtomicReference chatCompletionReference = new AtomicReference<>(); + + return response.map(chunk -> chunkToChatCompletion(chunk, chatCompletionReference)) + .switchMap( + cc -> handleFunctionCallOrReturnStream(request, Flux.just(ResponseEntity.of(Optional.of(cc))))) + .map(ResponseEntity::getBody) + .map(this::toChatResponse); + }); + } + + private ChatCompletion chunkToChatCompletion(StreamResponse chunk, + AtomicReference chatCompletionReference) { // https://docs.anthropic.com/claude/reference/messages-streaming - return response.map(chunk -> { + if (chunk.type().equals("message_start")) { + chatCompletionReference.set(new ChatCompletionBuilder()); + chatCompletionReference.get() + .withType(chunk.type()) + .withId(chunk.message().id()) + .withRole(chunk.message().role()) + .withModel(chunk.message().model()) + .withUsage(chunk.message().usage()) + .withContent(new ArrayList<>()); + } + else if (chunk.type().equals("content_block_start")) { + var content = new MediaContent(chunk.contentBlock().type(), null, chunk.contentBlock().text(), + chunk.index()); + chatCompletionReference.get().withType(chunk.type()).withContent(List.of(content)); + } + else if (chunk.type().equals("content_block_delta")) { + var content = new MediaContent(Type.TEXT_DELTA, null, (String) chunk.delta().get("text"), chunk.index()); + chatCompletionReference.get().withType(chunk.type()).withContent(List.of(content)); + } + else if (chunk.type().equals("message_delta")) { - if (chunk.type().equals("message_start")) { - chatCompletionReference.set(new ChatCompletionBuilder()); - chatCompletionReference.get() - .withType(chunk.type()) - .withId(chunk.message().id()) - .withRole(chunk.message().role()) - .withModel(chunk.message().model()) - .withUsage(chunk.message().usage()) - .withContent(new ArrayList<>()); - } - else if (chunk.type().equals("content_block_start")) { - var content = new MediaContent(chunk.contentBlock().type(), null, chunk.contentBlock().text(), - chunk.index()); - chatCompletionReference.get().withType(chunk.type()).withContent(List.of(content)); - } - else if (chunk.type().equals("content_block_delta")) { - var content = new MediaContent(Type.TEXT_DELTA, null, (String) chunk.delta().get("text"), - chunk.index()); - chatCompletionReference.get().withType(chunk.type()).withContent(List.of(content)); - } - else if (chunk.type().equals("message_delta")) { + ChatCompletion delta = ModelOptionsUtils.mapToClass(chunk.delta(), ChatCompletion.class); - 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()); - } - if (delta.role() != null) { - chatCompletionReference.get().withRole(delta.role()); - } - if (delta.model() != null) { - chatCompletionReference.get().withModel(delta.model()); - } - if (delta.content() != null) { - chatCompletionReference.get().withContent(delta.content()); - } - if (delta.stopReason() != null) { - chatCompletionReference.get().withStopReason(delta.stopReason()); - } - if (delta.stopSequence() != null) { - chatCompletionReference.get().withStopSequence(delta.stopSequence()); - } + chatCompletionReference.get().withType(chunk.type()); + if (chunk.usage() != null) { + var totalUsage = new Usage(chatCompletionReference.get().usage.inputTokens(), + chunk.usage().outputTokens()); + chatCompletionReference.get().withUsage(totalUsage); } - else { - chatCompletionReference.get().withType(chunk.type()).withContent(List.of()); + if (chunk.mergedToolUses() != null) { + chatCompletionReference.get().withToolUses(chunk.mergedToolUses()); } - return chatCompletionReference.get().build(); + if (delta.id() != null) { + chatCompletionReference.get().withId(delta.id()); + } + if (delta.role() != null) { + chatCompletionReference.get().withRole(delta.role()); + } + if (delta.model() != null) { + chatCompletionReference.get().withModel(delta.model()); + } + if (delta.content() != null) { + chatCompletionReference.get().withContent(delta.content()); + } + if (delta.stopReason() != null) { + chatCompletionReference.get().withStopReason(delta.stopReason()); + } + if (delta.stopSequence() != null) { + chatCompletionReference.get().withStopSequence(delta.stopSequence()); + } + } + else if (chunk.type().equals("message_stop")) { + if (chatCompletionReference.get().toolUses != null) { + chatCompletionReference.get().withContent(chatCompletionReference.get().toolUses); + } + } + else { + chatCompletionReference.get().withType(chunk.type()).withContent(List.of()); + } - }).map(this::toChatResponse); + return chatCompletionReference.get().build(); } private ChatResponse toChatResponse(ChatCompletion chatCompletion) { @@ -338,6 +356,8 @@ public class AnthropicChatModel extends private Usage usage; + private List toolUses; + public ChatCompletionBuilder() { } @@ -381,6 +401,11 @@ public class AnthropicChatModel extends return this; } + public ChatCompletionBuilder withToolUses(List toolUses) { + this.toolUses = toolUses; + return this; + } + public ChatCompletion build() { return new ChatCompletion(this.id, this.type, this.role, this.content, this.model, this.stopReason, this.stopSequence, this.usage); @@ -449,9 +474,13 @@ public class AnthropicChatModel extends @Override protected Flux> doChatCompletionStream(ChatCompletionRequest request) { - // https://docs.anthropic.com/en/docs/tool-use - throw new UnsupportedOperationException( - "Streaming (stream=true) is not yet supported. We plan to add streaming support in a future beta version."); + + AtomicReference chatCompletionReference = new AtomicReference<>(); + + return this.anthropicApi.chatCompletionStream(request) + .map(chunk -> this.chunkToChatCompletion(chunk, chatCompletionReference)) + .map(Optional::ofNullable) + .map(ResponseEntity::of); } @Override diff --git a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java index bcb93e749..3c8e14ec3 100644 --- a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java +++ b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java @@ -17,6 +17,7 @@ package org.springframework.ai.anthropic.api; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import java.util.function.Predicate; @@ -388,6 +389,7 @@ public class AnthropicApi { @JsonProperty("id") String id, @JsonProperty("name") String name, @JsonProperty("input") Map input, + @JsonProperty("input_json") String inputJson, // tool_result response only @JsonProperty("tool_use_id") String toolUseId, @@ -400,20 +402,25 @@ public class AnthropicApi { } public MediaContent(Source source) { - this(Type.IMAGE, source, null, null, null, null, null, null, null); + this(Type.IMAGE, source, null, null, null, null, null, null, null, null); } public MediaContent(String text) { - this(Type.TEXT, null, text, null, null, null, null, null, null); + this(Type.TEXT, null, text, null, null, null, null, null, null, null); } // Tool result public MediaContent(Type type, String toolUseId, String content) { - this(type, null, null, null, null, null, null, toolUseId, content); + this(type, null, null, null, null, null, null, null, toolUseId, content); } public MediaContent(Type type, Source source, String text, Integer index) { - this(type, source, text, index, null, null, null, null, null); + this(type, source, text, index, null, null, null, null, null, null); + } + + // Tool use input JSON delta streaming + public MediaContent(Type type, Integer index, String id, String name, String inputJson) { + this(type, null, null, index, id, name, null, inputJson, null, null); } /** @@ -492,7 +499,7 @@ public class AnthropicApi { * @param content Content generated by the model. This is an array of content blocks. * @param model The model that handled the request. * @param stopReason The reason the model stopped generating tokens. This will be one - * of "end_turn", "max_tokens", "stop_sequence", or "timeout". + * of "end_turn", "max_tokens", "stop_sequence", "tool_use", or "timeout". * @param stopSequence Which custom stop sequence was generated, if any. * @param usage Input and output token usage. */ @@ -570,7 +577,8 @@ public class AnthropicApi { @JsonProperty("message") ChatCompletion message, @JsonProperty("content_block") MediaContent contentBlock, @JsonProperty("delta") Map delta, - @JsonProperty("usage") OutputUsage usage) { + @JsonProperty("usage") OutputUsage usage, + List mergedToolUses) { // @formatter:on } @@ -588,6 +596,8 @@ public class AnthropicApi { return this.restClient.post().uri("/v1/messages").body(chatRequest).retrieve().toEntity(ChatCompletion.class); } + private final AnthropicStreamFunctionCallingHelper chunkMerger = new AnthropicStreamFunctionCallingHelper(); + /** * Creates a streaming chat response for the given chat conversation. * @param chatRequest The chat completion request. Must have the stream property set @@ -599,6 +609,8 @@ public class AnthropicApi { Assert.notNull(chatRequest, "The request body can not be null."); Assert.isTrue(chatRequest.stream(), "Request must set the steam property to true."); + AtomicBoolean isInsideTool = new AtomicBoolean(false); + return this.webClient.post() .uri("/v1/messages") .body(Mono.just(chatRequest), ChatCompletionRequest.class) @@ -606,7 +618,30 @@ public class AnthropicApi { .bodyToFlux(String.class) .takeUntil(SSE_DONE_PREDICATE) .filter(SSE_DONE_PREDICATE.negate()) - .map(content -> ModelOptionsUtils.jsonToObject(content, StreamResponse.class)); + .map(content -> ModelOptionsUtils.jsonToObject(content, StreamResponse.class)) + // Detect if the chunk is part of a streaming function call. + .map(chunk -> { + if (this.chunkMerger.isStreamingToolFunctionCall(chunk)) { + isInsideTool.set(true); + } + return chunk; + }) + // Group all chunks belonging to the same function call. + .windowUntil(chunk -> { + if (isInsideTool.get() && this.chunkMerger.isStreamingToolFunctionCallFinish(chunk)) { + isInsideTool.set(false); + return true; + } + return !isInsideTool.get(); + }) + // Merging the window chunks into a single chunk. + .concatMapIterable(window -> { + Mono monoChunk = window.reduce(this.chunkMerger.emptyChunk(), + this.chunkMerger::mergeChunks); + return List.of(monoChunk); + }) + // Flux> -> Flux + .flatMap(mono -> mono); } } diff --git a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicStreamFunctionCallingHelper.java b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicStreamFunctionCallingHelper.java new file mode 100644 index 000000000..a9700ea0a --- /dev/null +++ b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicStreamFunctionCallingHelper.java @@ -0,0 +1,118 @@ +/* + * Copyright 2023 - 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.anthropic.api; + +import org.springframework.ai.anthropic.api.AnthropicApi.MediaContent; +import org.springframework.ai.anthropic.api.AnthropicApi.StreamResponse; +import org.springframework.ai.model.ModelOptionsUtils; + +import java.util.ArrayList; +import java.util.List; + +/** + * Helper class to support streaming function calling. + *

+ * It can merge the streamed {@link StreamResponse} chunks in case of function calling + * message. + * + * @author Mariusz Bernacki + * @since 1.0.0 + */ +public class AnthropicStreamFunctionCallingHelper { + + public boolean isStreamingToolFunctionCall(StreamResponse response) { + if (response == null || response.contentBlock() == null) { + return false; + } + + return response.contentBlock().type() == MediaContent.Type.TOOL_USE; + } + + public boolean isStreamingToolFunctionCallFinish(StreamResponse response) { + return response.delta() != null && "tool_use".equals(response.delta().get("stop_reason")); + } + + public StreamResponse emptyChunk() { + return new StreamResponse(null, null, null, null, null, null, new ArrayList<>()); + } + + public StreamResponse mergeChunks(StreamResponse previous, StreamResponse current) { + if (previous == null) { + return current; + } + if (current == null) { + return previous; + } + + List mergedContent = mergeToolUses(previous, current); + + if (isStreamingToolFunctionCallFinish(current)) { + finalizeToolUsesAggregation(mergedContent); + } + + return new StreamResponse(lastElement(previous.type(), current.type()), + lastElement(previous.index(), current.index()), lastElement(previous.message(), current.message()), + lastElement(previous.contentBlock(), current.contentBlock()), + lastElement(previous.delta(), current.delta()), lastElement(previous.usage(), current.usage()), + mergedContent); + } + + private List mergeToolUses(StreamResponse previous, StreamResponse current) { + List mergedContent = new ArrayList<>(previous.mergedToolUses()); + + if (current.contentBlock() != null) { + mergedContent.add(current.contentBlock()); + } + else if (!mergedContent.isEmpty() && current.delta() != null && current.delta().containsKey("partial_json")) { + int lastIndex = mergedContent.size() - 1; + MediaContent previousMedia = mergedContent.get(lastIndex); + MediaContent currentMedia = new MediaContent(previousMedia.type(), previousMedia.index(), + previousMedia.id(), previousMedia.name(), + concat(previousMedia.inputJson(), (String) current.delta().get("partial_json"))); + + mergedContent.set(lastIndex, currentMedia); + } + + return mergedContent; + } + + public void finalizeToolUsesAggregation(List mergedToolUses) { + mergedToolUses.replaceAll(media -> { + if (media.inputJson() == null) { + return media; + } + + return new MediaContent(media.type(), media.source(), media.text(), null, media.id(), media.name(), + ModelOptionsUtils.jsonToMap(media.inputJson()), null, media.toolUseId(), media.content()); + }); + } + + private static T lastElement(T left, T right) { + return right != null ? right : left; + } + + private static String concat(String left, String right) { + if (left == null) { + return right; + } + if (right == null) { + return left; + } + + return left + right; + } + +} diff --git a/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/client/AnthropicChatClientIT.java b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/client/AnthropicChatClientIT.java index df3fc935f..ee8ccb30c 100644 --- a/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/client/AnthropicChatClientIT.java +++ b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/client/AnthropicChatClientIT.java @@ -241,13 +241,12 @@ class AnthropicChatClientIT { assertThat(response).contains("30", "10", "15"); } - @Disabled("SpringAI has not implemented streaming for function calls for Anthropic yet.") @Test void streamFunctionCallTest() { // @formatter:off Flux response = ChatClient.create(chatModel).prompt() - .user("What's the weather like in San Francisco, Tokyo, and Paris?") + .user("What's the weather like in San Francisco, Tokyo, and Paris? Use Celsius.") .function("getCurrentWeather", "Get the weather in location", new MockWeatherService()) .stream() .content();