From 6363352b820d434661a44afe26a334bd6a070296 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Thu, 25 Jul 2024 18:13:22 +0200 Subject: [PATCH] Improve Anthropic function calling - factor out the common funciton calling logic form AnthropicChatModel to the AbstractToolCallSupport. - improve the AbstractToolCallSupport isToolCall to handle OpenAi and Anthropic. - fix an issue with the function calling streaming aggreagation leading to lost usage statistics. - small code improvements for OpenAiChatModel. --- .../ai/anthropic/AnthropicChatModel.java | 196 ++++++++---------- .../ai/anthropic/api/AnthropicApi.java | 43 ++-- .../ai/anthropic/api/StreamHelper.java | 24 +-- .../api/tool/AnthropicApiToolIT.java | 6 +- .../ai/openai/OpenAiChatModel.java | 42 ++-- .../ai/chat/messages/AssistantMessage.java | 6 + .../chat/metadata/ChatGenerationMetadata.java | 6 + .../chat/model/AbstractToolCallSupport.java | 19 +- 8 files changed, 178 insertions(+), 164 deletions(-) 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 ff17b1ca0..b89cfa1d8 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 @@ -15,40 +15,6 @@ */ package org.springframework.ai.anthropic; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.springframework.ai.anthropic.api.AnthropicApi; -import org.springframework.ai.anthropic.api.AnthropicApi.AnthropicMessage; -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.ContentBlock.ContentBlockType; -import org.springframework.ai.anthropic.api.AnthropicApi.Role; -import org.springframework.ai.anthropic.metadata.AnthropicUsage; -import org.springframework.ai.chat.messages.AssistantMessage; -import org.springframework.ai.chat.messages.Message; -import org.springframework.ai.chat.messages.MessageType; -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.model.ChatModel; -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.chat.model.Generation; -import org.springframework.ai.chat.prompt.ChatOptions; -import org.springframework.ai.chat.prompt.Prompt; -import org.springframework.ai.model.ModelOptionsUtils; -import org.springframework.ai.chat.model.AbstractToolCallSupport; -import org.springframework.ai.model.function.FunctionCallbackContext; -import org.springframework.ai.retry.RetryUtils; -import org.springframework.http.ResponseEntity; -import org.springframework.retry.support.RetryTemplate; -import org.springframework.util.Assert; -import org.springframework.util.CollectionUtils; -import org.springframework.util.StringUtils; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - import java.util.ArrayList; import java.util.Base64; import java.util.HashSet; @@ -57,6 +23,40 @@ import java.util.Map; import java.util.Set; import java.util.stream.Collectors; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.anthropic.api.AnthropicApi; +import org.springframework.ai.anthropic.api.AnthropicApi.AnthropicMessage; +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.ContentBlock.Type; +import org.springframework.ai.anthropic.api.AnthropicApi.Role; +import org.springframework.ai.anthropic.metadata.AnthropicUsage; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.MessageType; +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.model.AbstractToolCallSupport; +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.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.model.ModelOptionsUtils; +import org.springframework.ai.model.function.FunctionCallbackContext; +import org.springframework.ai.retry.RetryUtils; +import org.springframework.http.ResponseEntity; +import org.springframework.retry.support.RetryTemplate; +import org.springframework.util.Assert; +import org.springframework.util.CollectionUtils; +import org.springframework.util.StringUtils; + +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + /** * The {@link ChatModel} implementation for the Anthropic service. * @@ -150,17 +150,17 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM ChatCompletionRequest request = createRequest(prompt, false); - return this.retryTemplate.execute(ctx -> { - ResponseEntity completionEntity = this.anthropicApi.chatCompletionEntity(request); + ResponseEntity completionEntity = this.retryTemplate + .execute(ctx -> this.anthropicApi.chatCompletionEntity(request)); - if (this.isToolFunctionCall(completionEntity.getBody())) { - List toolCallMessageConversation = this.handleToolCallRequests(prompt.getInstructions(), - completionEntity.getBody()); - return this.call(new Prompt(toolCallMessageConversation, prompt.getOptions())); - } + ChatResponse chatResponse = toChatResponse(completionEntity.getBody()); - return toChatResponse(completionEntity.getBody()); - }); + if (this.isToolCall(chatResponse, Set.of("tool_use"))) { + var toolCallConversation = handleToolCalls(prompt, chatResponse); + return this.call(new Prompt(toolCallConversation, prompt.getOptions())); + } + + return chatResponse; } @Override @@ -168,68 +168,66 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM ChatCompletionRequest request = createRequest(prompt, true); - return this.retryTemplate.execute(ctx -> { + Flux response = this.retryTemplate + .execute(ctx -> this.anthropicApi.chatCompletionStream(request)); - Flux response = this.anthropicApi.chatCompletionStream(request); + return response.switchMap(chatCompletionResponse -> { - return response.switchMap(chatCompletionResponse -> { + ChatResponse chatResponse = toChatResponse(chatCompletionResponse); - if (this.isToolFunctionCall(chatCompletionResponse)) { - List toolCallMessageConversation = this.handleToolCallRequests(prompt.getInstructions(), - chatCompletionResponse); - return this.stream(new Prompt(toolCallMessageConversation, prompt.getOptions())); - } + if (this.isToolCall(chatResponse, Set.of("tool_use"))) { + var toolCallConversation = handleToolCalls(prompt, chatResponse); + return this.stream(new Prompt(toolCallConversation, prompt.getOptions())); + } - return Mono.just(chatCompletionResponse).map(this::toChatResponse); - }); + return Mono.just(chatResponse); }); } - private List handleToolCallRequests(List previousMessages, - ChatCompletionResponse chatCompletionResponse) { - - AnthropicMessage anthropicAssistantMessage = new AnthropicMessage(chatCompletionResponse.content(), - Role.ASSISTANT); - - List toolToUseList = anthropicAssistantMessage.content() - .stream() - .filter(c -> c.type() == ContentBlock.ContentBlockType.TOOL_USE) - .toList(); - - List toolCalls = new ArrayList<>(); - - for (ContentBlock toolToUse : toolToUseList) { - - var functionCallId = toolToUse.id(); - var functionName = toolToUse.name(); - var functionArguments = ModelOptionsUtils.toJsonString(toolToUse.input()); - - toolCalls.add(new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments)); - } - - AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls); - ToolResponseMessage toolResponseMessage = this.executeFunctions(assistantMessage); - - // History - List toolCallMessageConversation = new ArrayList<>(previousMessages); - toolCallMessageConversation.add(assistantMessage); - toolCallMessageConversation.add(toolResponseMessage); - - return toolCallMessageConversation; - } - private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion) { + if (chatCompletion == null) { logger.warn("Null chat completion returned"); return new ChatResponse(List.of()); } - List generations = chatCompletion.content().stream().map(content -> { - return new Generation(content.text(), Map.of()) - .withGenerationMetadata(ChatGenerationMetadata.from(chatCompletion.stopReason(), null)); - }).toList(); + List generations = chatCompletion.content() + .stream() + .filter(content -> content.type() != ContentBlock.Type.TOOL_USE) + .map(content -> { + new AssistantMessage(content.text(), Map.of()); + return new Generation(new AssistantMessage(content.text(), Map.of()), + ChatGenerationMetadata.from(chatCompletion.stopReason(), null)); + }) + .toList(); - return new ChatResponse(generations, from(chatCompletion)); + List allGenerations = new ArrayList<>(generations); + + List toolToUseList = chatCompletion.content() + .stream() + .filter(c -> c.type() == ContentBlock.Type.TOOL_USE) + .toList(); + + if (!CollectionUtils.isEmpty(toolToUseList)) { + List toolCalls = new ArrayList<>(); + + for (ContentBlock toolToUse : toolToUseList) { + + var functionCallId = toolToUse.id(); + var functionName = toolToUse.name(); + var functionArguments = ModelOptionsUtils.toJsonString(toolToUse.input()); + + toolCalls + .add(new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments)); + } + + AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls); + Generation toolCallGeneration = new Generation(assistantMessage, + ChatGenerationMetadata.from(chatCompletion.stopReason(), null)); + allGenerations.add(toolCallGeneration); + } + + return new ChatResponse(allGenerations, this.from(chatCompletion)); } private ChatResponseMetadata from(AnthropicApi.ChatCompletionResponse result) { @@ -288,8 +286,8 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM } if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) { for (AssistantMessage.ToolCall toolCall : assistantMessage.getToolCalls()) { - contentBlocks.add(new ContentBlock(ContentBlockType.TOOL_USE, toolCall.id(), - toolCall.name(), ModelOptionsUtils.jsonToMap(toolCall.arguments()))); + contentBlocks.add(new ContentBlock(Type.TOOL_USE, toolCall.id(), toolCall.name(), + ModelOptionsUtils.jsonToMap(toolCall.arguments()))); } } return new AnthropicMessage(contentBlocks, Role.ASSISTANT); @@ -297,7 +295,7 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM else if (message.getMessageType() == MessageType.TOOL) { List toolResponses = ((ToolResponseMessage) message).getResponses() .stream() - .map(toolResponse -> new ContentBlock(ContentBlockType.TOOL_RESULT, toolResponse.id(), + .map(toolResponse -> new ContentBlock(Type.TOOL_RESULT, toolResponse.id(), toolResponse.responseData())) .toList(); return new AnthropicMessage(toolResponses, Role.USER); @@ -355,16 +353,6 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM }).toList(); } - @SuppressWarnings("null") - protected boolean isToolFunctionCall(ChatCompletionResponse response) { - if (response == null || CollectionUtils.isEmpty(response.content())) { - return false; - } - return response.content() - .stream() - .anyMatch(content -> content.type() == ContentBlock.ContentBlockType.TOOL_USE); - } - @Override public ChatOptions getDefaultOptions() { return AnthropicChatOptions.fromOptions(this.defaultOptions); 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 ec4e897a7..51eb70ab7 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 @@ -23,6 +23,8 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Predicate; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.springframework.ai.anthropic.api.StreamHelper.ChatCompletionResponseBuilder; import org.springframework.ai.model.ChatModelDescription; import org.springframework.ai.model.ModelOptionsUtils; @@ -53,6 +55,8 @@ import reactor.core.publisher.Mono; */ public class AnthropicApi { + private static final Logger logger = LoggerFactory.getLogger(AnthropicApi.class); + private static final String HEADER_X_API_KEY = "x-api-key"; private static final String HEADER_ANTHROPIC_VERSION = "anthropic-version"; @@ -415,7 +419,7 @@ public class AnthropicApi { */ @JsonInclude(Include.NON_NULL) public record ContentBlock( // @formatter:off - @JsonProperty("type") ContentBlockType type, + @JsonProperty("type") Type type, @JsonProperty("source") Source source, @JsonProperty("text") String text, @@ -438,67 +442,77 @@ public class AnthropicApi { } public ContentBlock(Source source) { - this(ContentBlockType.IMAGE, source, null, null, null, null, null, null, null); + this(Type.IMAGE, source, null, null, null, null, null, null, null); } public ContentBlock(String text) { - this(ContentBlockType.TEXT, null, text, null, null, null, null, null, null); + this(Type.TEXT, null, text, null, null, null, null, null, null); } // Tool result - public ContentBlock(ContentBlockType type, String toolUseId, String content) { + public ContentBlock(Type type, String toolUseId, String content) { this(type, null, null, null, null, null, null, toolUseId, content); } - public ContentBlock(ContentBlockType type, Source source, String text, Integer index) { + public ContentBlock(Type type, Source source, String text, Integer index) { this(type, source, text, index, null, null, null, null, null); } // Tool use input JSON delta streaming - public ContentBlock(ContentBlockType type, String id, String name, Map input) { + public ContentBlock(Type type, String id, String name, Map input) { this(type, null, null, null, id, name, input, null, null); } /** - * The type of this message. + * The ContentBlock type. */ - public enum ContentBlockType { + public enum Type { /** * Tool request */ @JsonProperty("tool_use") - TOOL_USE, + TOOL_USE("tool_use"), /** * Send tool result back to LLM. */ @JsonProperty("tool_result") - TOOL_RESULT, + TOOL_RESULT("tool_result"), /** * Text message. */ @JsonProperty("text") - TEXT, + TEXT("text"), /** * Text delta message. Returned from the streaming response. */ @JsonProperty("text_delta") - TEXT_DELTA, + TEXT_DELTA("text_delta"), /** * Tool use input partial JSON delta streaming. */ @JsonProperty("input_json_delta") - INPUT_JSON_DELTA, + INPUT_JSON_DELTA("input_json_delta"), /** * Image message. */ @JsonProperty("image") - IMAGE; + IMAGE("image"); + + public final String value; + + Type(String value) { + this.value = value; + } + + public String getValue() { + return this.value; + } } @@ -902,6 +916,7 @@ public class AnthropicApi { .takeUntil(SSE_DONE_PREDICATE) .filter(SSE_DONE_PREDICATE.negate()) .map(content -> ModelOptionsUtils.jsonToObject(content, StreamEvent.class)) + .filter(event -> event.type() != EventType.PING) // Detect if the chunk is part of a streaming function call. .map(event -> { if (this.streamHelper.isToolUseStart(event)) { 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 a699b0638..054bf023b 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 @@ -21,7 +21,7 @@ import java.util.concurrent.atomic.AtomicReference; import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionResponse; import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock; -import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock.ContentBlockType; +import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock.Type; import org.springframework.ai.anthropic.api.AnthropicApi.Role; import org.springframework.ai.anthropic.api.AnthropicApi.Usage; import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlockDeltaEvent; @@ -55,15 +55,15 @@ public class StreamHelper { if (event == null || event.type() == null || event.type() != EventType.CONTENT_BLOCK_START) { return false; } - return "tool_use".equals(((ContentBlockStartEvent) event).contentBlock().type()); + return ContentBlock.Type.TOOL_USE.getValue().equals(((ContentBlockStartEvent) event).contentBlock().type()); } public boolean isToolUseFinish(StreamEvent event) { - if (event == null || event.type() == null || event.type() != EventType.MESSAGE_DELTA) { + if (event == null || event.type() == null || event.type() != EventType.CONTENT_BLOCK_STOP) { return false; } - return "tool_use".equals(((MessageDeltaEvent) event).delta().stopReason()); + return true; } public StreamEvent mergeToolUseEvents(StreamEvent previousEvent, StreamEvent event) { @@ -73,7 +73,7 @@ public class StreamHelper { if (event.type() == EventType.CONTENT_BLOCK_START) { ContentBlockStartEvent contentBlockStart = (ContentBlockStartEvent) event; - if ("tool_use".equals(contentBlockStart.contentBlock().type())) { + if (ContentBlock.Type.TOOL_USE.getValue().equals(contentBlockStart.contentBlock().type())) { ContentBlockStartEvent.ContentBlockToolUse cbToolUse = (ContentBlockToolUse) contentBlockStart .contentBlock(); @@ -85,7 +85,7 @@ public class StreamHelper { } else if (event.type() == EventType.CONTENT_BLOCK_DELTA) { ContentBlockDeltaEvent contentBolckDelta = (ContentBlockDeltaEvent) event; - if ("input_json_delta".equals(contentBolckDelta.delta().type())) { + if (ContentBlock.Type.INPUT_JSON_DELTA.getValue().equals(contentBolckDelta.delta().type())) { return eventAggregator .appendPartialJson(((ContentBlockDeltaJson) contentBolckDelta.delta()).partialJson()); } @@ -96,9 +96,6 @@ public class StreamHelper { return eventAggregator; } } - else if (isToolUseFinish(event)) { - return eventAggregator; - } return event; } @@ -128,8 +125,7 @@ public class StreamHelper { List content = eventToolUseBuilder.getToolContentBlocks() .stream() - .map(tooToUse -> new ContentBlock(ContentBlockType.TOOL_USE, tooToUse.id(), tooToUse.name(), - tooToUse.input())) + .map(tooToUse -> new ContentBlock(Type.TOOL_USE, tooToUse.id(), tooToUse.name(), tooToUse.input())) .toList(); contentBlockReference.get().withContent(content); } @@ -142,7 +138,7 @@ public class StreamHelper { + contentBlockStartEvent.contentBlock().type()); ContentBlockText contentBlockText = (ContentBlockText) contentBlockStartEvent.contentBlock(); - ContentBlock contentBlock = new ContentBlock(ContentBlockType.TEXT, null, contentBlockText.text(), + ContentBlock contentBlock = new ContentBlock(Type.TEXT, null, contentBlockText.text(), contentBlockStartEvent.index()); contentBlockReference.get().withType(event.type().name()).withContent(List.of(contentBlock)); } @@ -156,8 +152,7 @@ public class StreamHelper { ContentBlockDeltaText deltaTxt = (ContentBlockDeltaText) contentBlockDeltaEvent.delta(); - var contentBlock = new ContentBlock(ContentBlockType.TEXT_DELTA, null, deltaTxt.text(), - contentBlockDeltaEvent.index()); + var contentBlock = new ContentBlock(Type.TEXT_DELTA, null, deltaTxt.text(), contentBlockDeltaEvent.index()); contentBlockReference.get().withType(event.type().name()).withContent(List.of(contentBlock)); } @@ -182,7 +177,6 @@ public class StreamHelper { } } else if (event.type().equals(EventType.MESSAGE_STOP)) { - } 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/tool/AnthropicApiToolIT.java b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/tool/AnthropicApiToolIT.java index c447c6a29..a4e9b9ed4 100644 --- a/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/tool/AnthropicApiToolIT.java +++ b/models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/tool/AnthropicApiToolIT.java @@ -29,7 +29,7 @@ import org.springframework.ai.anthropic.api.AnthropicApi; import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionResponse; import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionRequest; import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock; -import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock.ContentBlockType; +import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock.Type; import org.springframework.ai.anthropic.api.AnthropicApi.AnthropicMessage; import org.springframework.ai.anthropic.api.AnthropicApi.Role; import org.springframework.ai.anthropic.api.AnthropicApi.Tool; @@ -117,7 +117,7 @@ public class AnthropicApiToolIT { List toolToUseList = response.getBody() .content() .stream() - .filter(c -> c.type() == ContentBlock.ContentBlockType.TOOL_USE) + .filter(c -> c.type() == ContentBlock.Type.TOOL_USE) .toList(); if (CollectionUtils.isEmpty(toolToUseList)) { @@ -146,7 +146,7 @@ public class AnthropicApiToolIT { logger.info("Function response : " + content); - toolResults.add(new ContentBlock(ContentBlockType.TOOL_RESULT, id, content)); + toolResults.add(new ContentBlock(Type.TOOL_RESULT, id, content)); } // Add function response message to the conversation history diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java index 358847587..536420eec 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java @@ -186,18 +186,6 @@ public class OpenAiChatModel extends AbstractToolCallSupport implements ChatMode return chatResponse; } - public static ChatResponseMetadata from(OpenAiApi.ChatCompletion result, RateLimit rateLimit) { - Assert.notNull(result, "OpenAI ChatCompletionResult must not be null"); - return ChatResponseMetadata.builder() - .withId(result.id()) - .withUsage(OpenAiUsage.from(result.usage())) - .withModel(result.model()) - .withRateLimit(rateLimit) - .withKeyValue("created", result.created()) - .withKeyValue("system-fingerprint", result.systemFingerprint()) - .build(); - } - @Override public Flux stream(Prompt prompt) { @@ -232,7 +220,7 @@ public class OpenAiChatModel extends AbstractToolCallSupport implements ChatMode // @formatter:on if (chatCompletion2.usage() != null) { - return new ChatResponse(generations, from(chatCompletion2)); + return new ChatResponse(generations, from(chatCompletion2, null)); } else { return new ChatResponse(generations); @@ -259,18 +247,7 @@ public class OpenAiChatModel extends AbstractToolCallSupport implements ChatMode }); } - private ChatResponseMetadata from(OpenAiApi.ChatCompletion result) { - Assert.notNull(result, "OpenAI ChatCompletionResult must not be null"); - return ChatResponseMetadata.builder() - .withId(result.id()) - .withUsage(OpenAiUsage.from(result.usage())) - .withModel(result.model()) - .withKeyValue("created", result.created()) - .withKeyValue("system-fingerprint", result.systemFingerprint()) - .build(); - } - - private static Generation buildGeneration(Choice choice, Map metadata) { + private Generation buildGeneration(Choice choice, Map metadata) { List toolCalls = choice.message().toolCalls() == null ? List.of() : choice.message() .toolCalls() @@ -285,6 +262,21 @@ public class OpenAiChatModel extends AbstractToolCallSupport implements ChatMode return new Generation(assistantMessage, generationMetadata); } + private ChatResponseMetadata from(OpenAiApi.ChatCompletion result, RateLimit rateLimit) { + Assert.notNull(result, "OpenAI ChatCompletionResult must not be null"); + var builder = ChatResponseMetadata.builder() + .withId(result.id()) + .withUsage(OpenAiUsage.from(result.usage())) + .withModel(result.model()) + .withRateLimit(rateLimit) + .withKeyValue("created", result.created()) + .withKeyValue("system-fingerprint", result.systemFingerprint()); + if (rateLimit != null) { + builder.withRateLimit(rateLimit); + } + return builder.build(); + } + /** * Convert the ChatCompletionChunk into a ChatCompletion. The Usage is set to null. * @param chunk the ChatCompletionChunk to convert diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java index beac0344c..1fb46b580 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java @@ -77,4 +77,10 @@ public class AssistantMessage extends AbstractMessage { return Objects.hash(super.hashCode(), toolCalls); } + @Override + public String toString() { + return "AssistantMessage [messageType=" + messageType + ", toolCalls=" + toolCalls + ", textContent=" + + textContent + ", metadata=" + metadata + "]"; + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/ChatGenerationMetadata.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/ChatGenerationMetadata.java index 50848657d..744c3fdab 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/ChatGenerationMetadata.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/ChatGenerationMetadata.java @@ -51,6 +51,12 @@ public interface ChatGenerationMetadata extends ResultMetadata { public String getFinishReason() { return finishReason; } + + @Override + public String toString() { + return "ChatGenerationMetadata{finishReason=" + finishReason + "," + "contentFilterMetadata=" + + contentFilterMetadata + "}"; + } }; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/model/AbstractToolCallSupport.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/model/AbstractToolCallSupport.java index fcab56a4e..ba84329ab 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/model/AbstractToolCallSupport.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/model/AbstractToolCallSupport.java @@ -19,6 +19,7 @@ import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -98,7 +99,14 @@ public abstract class AbstractToolCallSupport { } protected List handleToolCalls(Prompt prompt, ChatResponse response) { - AssistantMessage assistantMessage = response.getResult().getOutput(); + Optional toolCallGeneration = response.getResults() + .stream() + .filter(g -> !CollectionUtils.isEmpty(g.getOutput().getToolCalls())) + .findFirst(); + if (toolCallGeneration.isEmpty()) { + throw new IllegalStateException("No tool call generation found in the response!"); + } + AssistantMessage assistantMessage = toolCallGeneration.get().getOutput(); ToolResponseMessage toolMessageResponse = this.executeFunctions(assistantMessage); return this.buildToolCallConversation(prompt.getInstructions(), assistantMessage, toolMessageResponse); } @@ -180,11 +188,16 @@ public abstract class AbstractToolCallSupport { return false; } - var generation = generations.get(0); + return generations.stream().anyMatch(g -> isToolCall(g, toolCallFinishReasons)); + } + + protected boolean isToolCall(Generation generation, Set toolCallFinishReasons) { + var finishReason = (generation.getMetadata().getFinishReason() != null) + ? generation.getMetadata().getFinishReason() : ""; return !CollectionUtils.isEmpty(generation.getOutput().getToolCalls()) && toolCallFinishReasons.stream() .map(s -> s.toLowerCase()) .toList() - .contains(generation.getMetadata().getFinishReason().toLowerCase()); + .contains(finishReason.toLowerCase()); } }