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

Add Anthropic Claude streaming Function Calling support

Browse Source 
Resolves #930
This commit is contained in:
Mariusz Bernacki
2024-06-25 21:11:49 +02:00
committed by Christian Tzolov
parent a22e62d1fa
commit f5b72242c0
4 changed files with 250 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<StreamResponse> response = this.anthropicApi.chatCompletionStream(request);
return this.retryTemplate.execute(ctx -> {
AtomicReference<ChatCompletionBuilder> chatCompletionReference = new AtomicReference<>();
Flux<StreamResponse> response = this.anthropicApi.chatCompletionStream(request);
AtomicReference<ChatCompletionBuilder> 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<ChatCompletionBuilder> 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<MediaContent> toolUses;
public ChatCompletionBuilder() {
}
@@ -381,6 +401,11 @@ public class AnthropicChatModel extends
return this;
}
public ChatCompletionBuilder withToolUses(List<MediaContent> 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<ResponseEntity<ChatCompletion>> 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<ChatCompletionBuilder> chatCompletionReference = new AtomicReference<>();
return this.anthropicApi.chatCompletionStream(request)
.map(chunk -> this.chunkToChatCompletion(chunk, chatCompletionReference))
.map(Optional::ofNullable)
.map(ResponseEntity::of);
}
@Override

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

@@ -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<String, Object> 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<String, Object> delta,
@JsonProperty("usage") OutputUsage usage) {
@JsonProperty("usage") OutputUsage usage,
List<MediaContent> 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<StreamResponse> monoChunk = window.reduce(this.chunkMerger.emptyChunk(),
this.chunkMerger::mergeChunks);
return List.of(monoChunk);
})
// Flux<Mono<StreamResponse>> -> Flux<StreamResponse>
.flatMap(mono -> mono);
}
}

118
models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicStreamFunctionCallingHelper.java Normal file
View File

@@ -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.
* <p>
* 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<MediaContent> 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<MediaContent> mergeToolUses(StreamResponse previous, StreamResponse current) {
List<MediaContent> 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<MediaContent> 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> 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;
}
}

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

@@ -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<String> 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();