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

Make consistent sync/stream AssistantMessage properties for OpenAI and Mistral AI

Browse Source
This commit is contained in:
Christian Tzolov
2024-03-12 11:34:56 +01:00
parent 7f1570d722
commit 4209d5a65d
2 changed files with 31 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatClient.java
View File

@@ -15,6 +15,7 @@
*/
package org.springframework.ai.mistralai;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
@@ -113,8 +114,7 @@ public class MistralAiChatClient extends
List<Generation> generations = chatCompletion.choices()
.stream()
.map(choice -> new Generation(choice.message().content(),
Map.of("role", choice.message().role().name()))
.map(choice -> new Generation(choice.message().content(), toMap(chatCompletion.id(), choice))
.withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null)))
.toList();
@@ -122,6 +122,20 @@ public class MistralAiChatClient extends
});
}
private Map<String, Object> toMap(String id, ChatCompletion.Choice choice) {
Map<String, Object> map = new HashMap<>();
var message = choice.message();
if (message.role() != null) {
map.put("role", message.role().name());
}
if (choice.finishReason() != null) {
map.put("finishReason", choice.finishReason().name());
}
map.put("id", id);
return map;
}
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
var request = createRequest(prompt, true);

25
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatClient.java
View File

@@ -147,7 +147,7 @@ public class OpenAiChatClient extends
RateLimit rateLimits = OpenAiResponseHeaderExtractor.extractAiResponseHeaders(completionEntity);
List<Generation> generations = chatCompletion.choices().stream().map(choice -> {
return new Generation(choice.message().content(), toMap(choice.message()))
return new Generation(choice.message().content(), toMap(chatCompletion.id(), choice))
.withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null));
}).toList();
@@ -156,6 +156,20 @@ public class OpenAiChatClient extends
});
}
private Map<String, Object> toMap(String id, ChatCompletion.Choice choice) {
Map<String, Object> map = new HashMap<>();
var message = choice.message();
if (message.role() != null) {
map.put("role", message.role().name());
}
if (choice.finishReason() != null) {
map.put("finishReason", choice.finishReason().name());
}
map.put("id", id);
return map;
}
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
@@ -280,15 +294,6 @@ public class OpenAiChatClient extends
}).toList();
}
private Map<String, Object> toMap(ChatCompletionMessage message) {
Map<String, Object> map = new HashMap<>();
if (message.role() != null) {
map.put("role", message.role().name());
}
return map;
}
@Override
protected ChatCompletionRequest doCreateToolResponseRequest(ChatCompletionRequest previousRequest,
ChatCompletionMessage responseMessage, List<ChatCompletionMessage> conversationHistory) {