Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,11 @@ void addAndGet() {
memory.deleteByConversationId(sessionId);
assertThat(memory.findByConversationId(sessionId)).isEmpty();

AssistantMessage assistantMessage = new AssistantMessage("test answer", Map.of(),
List.of(new AssistantMessage.ToolCall("id", "type", "name", "arguments")));
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("test answer")
.properties(Map.of())
.toolCalls(List.of(new AssistantMessage.ToolCall("id", "type", "name", "arguments")))
.build();

memory.saveAll(sessionId, List.of(userMessage, assistantMessage));
messages = memory.findByConversationId(sessionId);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,10 @@ public void verifyThatChromaCanHandleComplexMetadataValues() {

var response = ChatClientResponse.builder()
.chatResponse(ChatResponse.builder()
.generations(List
.of(new Generation(new AssistantMessage("AssistantMessage", Map.of("annotations", List.of())))))
.generations(List.of(new Generation(AssistantMessage.builder()
.content("AssistantMessage")
.properties(Map.of("annotations", List.of()))
.build())))
.build())
.build();
var res2 = advisor.after(response, null);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -433,7 +433,7 @@ public static List<ToolCallback> getToolCallbacksFromAsyncClients(List<McpAsyncC
if (CollectionUtils.isEmpty(asyncMcpClients)) {
return List.of();
}
return List.of((new AsyncMcpToolCallbackProvider(asyncMcpClients).getToolCallbacks()));
return List.of((AsyncMcpToolCallbackProvider.builder().mcpClients(asyncMcpClients).build().getToolCallbacks()));
}

@JsonIgnoreProperties(ignoreUnknown = true)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ private Message getMessage(UdtValue udt) {
Map<String, Object> props = Map.of(CONVERSATION_TS, udt.getInstant(this.conf.messageUdtTimestampColumn));
switch (MessageType.valueOf(udt.getString(this.conf.messageUdtTypeColumn))) {
case ASSISTANT:
return new AssistantMessage(content, props);
return AssistantMessage.builder().content(content).properties(props).build();
case USER:
return UserMessage.builder().text(content).metadata(props).build();
case SYSTEM:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -183,14 +183,16 @@ private Message buildToolMessage(org.neo4j.driver.Record record) {

private Message buildAssistantMessage(org.neo4j.driver.Record record, Map<String, Object> messageMap,
List<Media> mediaList) {
Message message;
message = new AssistantMessage(messageMap.get(MessageAttributes.TEXT_CONTENT.getValue()).toString(),
record.get("metadata").asMap(Map.of()), record.get("toolCalls").asList(v -> {
var toolCallMap = v.asMap();
return new AssistantMessage.ToolCall((String) toolCallMap.get("id"),
(String) toolCallMap.get("type"), (String) toolCallMap.get("name"),
(String) toolCallMap.get("arguments"));
}), mediaList);
Message message = AssistantMessage.builder()
.content(messageMap.get(MessageAttributes.TEXT_CONTENT.getValue()).toString())
.properties(record.get("metadata").asMap(Map.of()))
.toolCalls(record.get("toolCalls").asList(v -> {
var toolCallMap = v.asMap();
return new AssistantMessage.ToolCall((String) toolCallMap.get("id"), (String) toolCallMap.get("type"),
(String) toolCallMap.get("name"), (String) toolCallMap.get("arguments"));
}))
.media(mediaList)
.build();
return message;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -263,9 +263,12 @@ void handleMediaContent() {
void handleAssistantMessageWithToolCalls() {
var conversationId = UUID.randomUUID().toString();

AssistantMessage assistantMessage = new AssistantMessage("Message with tool calls", Map.of(),
List.of(new AssistantMessage.ToolCall("id1", "type1", "name1", "arguments1"),
new AssistantMessage.ToolCall("id2", "type2", "name2", "arguments2")));
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("Message with tool calls")
.properties(Map.of())
.toolCalls(List.of(new AssistantMessage.ToolCall("id1", "type1", "name1", "arguments1"),
new AssistantMessage.ToolCall("id2", "type2", "name2", "arguments2")))
.build();

this.chatMemoryRepository.saveAll(conversationId, List.<Message>of(assistantMessage));

Expand Down Expand Up @@ -389,7 +392,8 @@ void saveAndFindMessagesWithEmptyContentOrMetadata() {
assertThat(retrievedEmptyContentMsg.getMetadata().keySet()).hasSize(1); // Only
// messageType

// Verify second message (empty metadata from input, should only have messageType
// Verify second message (empty metadata from input, should only have
// messageType
// after retrieval)
Message retrievedEmptyMetadataMsg = retrievedMessages.get(1);
assertThat(retrievedEmptyMetadataMsg).isInstanceOf(UserMessage.class);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -321,19 +321,24 @@ private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage
for (ContentBlock content : chatCompletion.content()) {
switch (content.type()) {
case TEXT, TEXT_DELTA:
generations.add(new Generation(new AssistantMessage(content.text(), Map.of()),
generations.add(new Generation(
AssistantMessage.builder().content(content.text()).properties(Map.of()).build(),
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
break;
case THINKING, THINKING_DELTA:
Map<String, Object> thinkingProperties = new HashMap<>();
thinkingProperties.put("signature", content.signature());
generations.add(new Generation(new AssistantMessage(content.thinking(), thinkingProperties),
generations.add(new Generation(
AssistantMessage.builder()
.content(content.thinking())
.properties(thinkingProperties)
.build(),
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
break;
case REDACTED_THINKING:
Map<String, Object> redactedProperties = new HashMap<>();
redactedProperties.put("data", content.data());
generations.add(new Generation(new AssistantMessage(null, redactedProperties),
generations.add(new Generation(AssistantMessage.builder().properties(redactedProperties).build(),
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
break;
case TOOL_USE:
Expand All @@ -347,13 +352,17 @@ private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage
}

if (chatCompletion.stopReason() != null && generations.isEmpty()) {
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
Generation generation = new Generation(AssistantMessage.builder().content("").properties(Map.of()).build(),
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
generations.add(generation);
}

if (!CollectionUtils.isEmpty(toolCalls)) {
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("")
.properties(Map.of())
.toolCalls(toolCalls)
.build();
Generation toolCallGeneration = new Generation(assistantMessage,
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
generations.add(toolCallGeneration);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -493,7 +493,11 @@ private Generation buildGeneration(ChatChoice choice, Map<String, Object> metada
}

var content = responseMessage == null ? "" : responseMessage.getContent();
var assistantMessage = new AssistantMessage(content, metadata, toolCalls);
var assistantMessage = AssistantMessage.builder()
.content(content)
.properties(metadata)
.toolCalls(toolCalls)
.build();
var generationMetadata = generateChoiceMetadata(choice);

return new Generation(assistantMessage, generationMetadata);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -572,14 +572,15 @@ private ChatResponse toChatResponse(ConverseResponse response, ChatResponse perv
List<Generation> generations = message.content()
.stream()
.filter(content -> content.type() != ContentBlock.Type.TOOL_USE)
.map(content -> new Generation(new AssistantMessage(content.text(), Map.of()),
.map(content -> new Generation(
AssistantMessage.builder().content(content.text()).properties(Map.of()).build(),
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build()))
.toList();

List<Generation> allGenerations = new ArrayList<>(generations);

if (response.stopReasonAsString() != null && generations.isEmpty()) {
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
Generation generation = new Generation(AssistantMessage.builder().properties(Map.of()).build(),
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build());
allGenerations.add(generation);
}
Expand All @@ -603,7 +604,11 @@ private ChatResponse toChatResponse(ConverseResponse response, ChatResponse perv
.add(new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments));
}

AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("")
.properties(Map.of())
.toolCalls(toolCalls)
.build();
Generation toolCallGeneration = new Generation(assistantMessage,
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build());
allGenerations.add(toolCallGeneration);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,11 @@ public static Flux<ChatResponse> toChatResponse(Flux<ConverseStreamOutput> respo
}
}

AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("")
.properties(Map.of())
.toolCalls(toolCalls)
.build();
Generation toolCallGeneration = new Generation(assistantMessage,
ChatGenerationMetadata.builder().finishReason("tool_use").build());

Expand Down Expand Up @@ -176,7 +180,10 @@ else if (nextEvent instanceof ContentBlockDeltaEvent contentBlockDeltaEvent) {
if (contentBlockDeltaEvent.delta().type().equals(ContentBlockDelta.Type.TEXT)) {

var generation = new Generation(
new AssistantMessage(contentBlockDeltaEvent.delta().text(), Map.of()),
AssistantMessage.builder()
.content(contentBlockDeltaEvent.delta().text())
.properties(Map.of())
.build(),
ChatGenerationMetadata.builder()
.finishReason(lastAggregation.metadataAggregation().stopReason())
.build());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -616,7 +616,11 @@ protected List<Generation> responseCandidateToGeneration(Candidate candidate) {
})
.toList();

AssistantMessage assistantMessage = new AssistantMessage("", messageMetadata, assistantToolCalls);
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("")
.properties(messageMetadata)
.toolCalls(assistantToolCalls)
.build();

return List.of(new Generation(assistantMessage, chatGenerationMetadata));
}
Expand All @@ -626,7 +630,10 @@ protected List<Generation> responseCandidateToGeneration(Candidate candidate) {
.parts()
.orElse(List.of())
.stream()
.map(part -> new AssistantMessage(part.text().orElse(""), messageMetadata))
.map(part -> AssistantMessage.builder()
.content(part.text().orElse(""))
.properties(messageMetadata)
.build())
.map(assistantMessage -> new Generation(assistantMessage, chatGenerationMetadata))
.toList();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,8 @@ public ChatResponse call(Prompt prompt) {
new TypeReference<>() {

});
Generation generation = new Generation(new AssistantMessage(generatedText, detailsMap));
Generation generation = new Generation(
AssistantMessage.builder().content(generatedText).properties(detailsMap).build());
generations.add(generation);
}
return new ChatResponse(generations);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -226,7 +226,11 @@ private static Generation buildGeneration(Choice choice, Map<String, Object> met
acc1.addAll(acc2);
return acc1;
});
var assistantMessage = new AssistantMessage(choice.message().content(), metadata, toolCalls);
var assistantMessage = AssistantMessage.builder()
.content(choice.message().content())
.properties(metadata)
.toolCalls(toolCalls)
.build();
String finishReason = (choice.finishReason() != null ? choice.finishReason().name() : "");
var generationMetadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
return new Generation(assistantMessage, generationMetadata);
Expand Down Expand Up @@ -432,7 +436,11 @@ private Generation buildGeneration(ChatCompletionMessage message, ChatCompletion
toolCall.function().name(), toolCall.function().arguments()))
.toList();

var assistantMessage = new AssistantMessage(message.content(), metadata, toolCalls);
var assistantMessage = AssistantMessage.builder()
.content(message.content())
.properties(metadata)
.toolCalls(toolCalls)
.build();
String finishReason = (completionFinishReason != null ? completionFinishReason.name() : "");
var generationMetadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
return new Generation(assistantMessage, generationMetadata);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -362,7 +362,11 @@ private Generation buildGeneration(Choice choice, Map<String, Object> metadata)
toolCall.function().name(), toolCall.function().arguments()))
.toList();

var assistantMessage = new AssistantMessage(choice.message().content(), metadata, toolCalls);
var assistantMessage = AssistantMessage.builder()
.content(choice.message().content())
.properties(metadata)
.toolCalls(toolCalls)
.build();
String finishReason = (choice.finishReason() != null ? choice.finishReason().name() : "");
var generationMetadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
return new Generation(assistantMessage, generationMetadata);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ private List<Generation> toGenerations(com.oracle.bmc.generativeaiinference.resp
ChatGenerationMetadata metadata = ChatGenerationMetadata.builder()
.finishReason(resp.getFinishReason().getValue())
.build();
AssistantMessage message = new AssistantMessage(resp.getText(), Map.of());
AssistantMessage message = AssistantMessage.builder().content(resp.getText()).properties(Map.of()).build();
generations.add(new Generation(message, metadata));
return generations;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -255,7 +255,11 @@ private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespon
ModelOptionsUtils.toJsonString(toolCall.function().arguments())))
.toList();

var assistantMessage = new AssistantMessage(ollamaResponse.message().content(), Map.of(), toolCalls);
var assistantMessage = AssistantMessage.builder()
.content(ollamaResponse.message().content())
.properties(Map.of())
.toolCalls(toolCalls)
.build();

ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
if (ollamaResponse.promptEvalCount() != null && ollamaResponse.evalCount() != null) {
Expand Down Expand Up @@ -333,7 +337,11 @@ private Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousCh
.toList();
}

var assistantMessage = new AssistantMessage(content, Map.of(), toolCalls);
var assistantMessage = AssistantMessage.builder()
.content(content)
.properties(Map.of())
.toolCalls(toolCalls)
.build();

ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
if (chunk.promptEvalCount() != null && chunk.evalCount() != null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -448,7 +448,12 @@ private Generation buildGeneration(Choice choice, Map<String, Object> metadata,
generationMetadataBuilder.metadata("logprobs", choice.logprobs());
}

var assistantMessage = new AssistantMessage(textContent, metadata, toolCalls, media);
var assistantMessage = AssistantMessage.builder()
.content(textContent)
.properties(metadata)
.toolCalls(toolCalls)
.media(media)
.build();
return new Generation(assistantMessage, generationMetadataBuilder.build());
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -631,15 +631,19 @@ protected List<Generation> responseCandidateToGeneration(Candidate candidate) {
})
.toList();

AssistantMessage assistantMessage = new AssistantMessage("", messageMetadata, assistantToolCalls);
AssistantMessage assistantMessage = AssistantMessage.builder()
.content("")
.properties(messageMetadata)
.toolCalls(assistantToolCalls)
.build();

return List.of(new Generation(assistantMessage, chatGenerationMetadata));
}
else {
List<Generation> generations = candidate.getContent()
.getPartsList()
.stream()
.map(part -> new AssistantMessage(part.getText(), messageMetadata))
.map(part -> AssistantMessage.builder().content(part.getText()).properties(messageMetadata).build())
.map(assistantMessage -> new Generation(assistantMessage, chatGenerationMetadata))
.toList();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import reactor.core.publisher.Flux;

import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.AssistantMessage.ToolCall;
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.EmptyRateLimit;
Expand Down Expand Up @@ -157,12 +158,17 @@ public Flux<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,

if (!CollectionUtils.isEmpty(collectedToolCalls)) {

finalAssistantMessage = new AssistantMessage(messageTextContentRef.get().toString(),
messageMetadataMapRef.get(), collectedToolCalls);
finalAssistantMessage = AssistantMessage.builder()
.content(messageTextContentRef.get().toString())
.properties(messageMetadataMapRef.get())
.toolCalls(collectedToolCalls)
.build();
}
else {
finalAssistantMessage = new AssistantMessage(messageTextContentRef.get().toString(),
messageMetadataMapRef.get());
finalAssistantMessage = AssistantMessage.builder()
.content(messageTextContentRef.get().toString())
.properties(messageMetadataMapRef.get())
.build();
}
onAggregationComplete.accept(new ChatResponse(List.of(new Generation(finalAssistantMessage,

Expand Down
Loading