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

Add VertexAI Gemini system-instruction support

Browse Source 
Resolves #906
This commit is contained in:
Christian Tzolov
2024-06-20 15:50:12 +02:00
parent 75d68662ae
commit 400d0d156c
2 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java
View File

@@ -28,6 +28,7 @@ import com.google.cloud.vertexai.api.GenerationConfig;
import com.google.cloud.vertexai.api.Part;
import com.google.cloud.vertexai.api.Schema;
import com.google.cloud.vertexai.api.Tool;
import com.google.cloud.vertexai.generativeai.ContentMaker;
import com.google.cloud.vertexai.generativeai.GenerativeModel;
import com.google.cloud.vertexai.generativeai.PartMaker;
import com.google.cloud.vertexai.generativeai.ResponseStream;
@@ -258,6 +259,16 @@ public class VertexAiGeminiChatModel
GenerativeModel generativeModel = generativeModelBuilder.build();
String systemContext = prompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() == MessageType.SYSTEM)
.map(m -> m.getContent())
.collect(Collectors.joining(System.lineSeparator()));
if (StringUtils.hasText(systemContext)) {
generativeModel.withSystemInstruction(ContentMaker.fromString(systemContext));
}
return new GeminiRequest(toGeminiContent(prompt), generativeModel);
}
@@ -289,18 +300,12 @@ public class VertexAiGeminiChatModel
private List<Content> toGeminiContent(Prompt prompt) {
String systemContext = prompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() == MessageType.SYSTEM)
.map(m -> m.getContent())
.collect(Collectors.joining(System.lineSeparator()));
List<Content> contents = prompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() == MessageType.USER || m.getMessageType() == MessageType.ASSISTANT)
.map(message -> Content.newBuilder()
.setRole(toGeminiMessageType(message.getMessageType()).getValue())
.addAllParts(messageToGeminiParts(message, systemContext))
.addAllParts(messageToGeminiParts(message))
.build())
.toList();
@@ -321,14 +326,11 @@ public class VertexAiGeminiChatModel
}
}
static List<Part> messageToGeminiParts(Message message, String systemContext) {
static List<Part> messageToGeminiParts(Message message) {
if (message instanceof UserMessage userMessage) {
String messageTextContent = (userMessage.getContent() == null) ? "null" : userMessage.getContent();
if (StringUtils.hasText(systemContext)) {
messageTextContent = systemContext + "\n\n" + messageTextContent;
}
Part textPart = Part.newBuilder().setText(messageTextContent).build();
List<Part> parts = new ArrayList<>(List.of(textPart));

1
models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatModelFunctionCallingIT.java
View File

@@ -28,7 +28,6 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;