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

Updgrade Vertex AI Gemini to 26.37.0

Browse Source 
- Adjust to the new, breaking, API changes.
   - Update the docs adding the transport endpoint and scopes to the connection properties.
This commit is contained in:
Christian Tzolov
2024-04-10 09:09:58 +02:00
parent 0b04c8928a
commit ae9d3f4168
8 changed files with 94 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -23,7 +23,6 @@ import java.util.stream.Collectors;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.google.cloud.vertexai.Transport;
import com.google.cloud.vertexai.VertexAI;
import com.google.cloud.vertexai.api.Content;
import com.google.cloud.vertexai.api.FunctionCall;
@@ -146,8 +145,6 @@ public class VertexAiGeminiChatClient
var geminiRequest = createGeminiRequest(prompt);
// GenerateContentResponse response =
// this.chatCompletionWithFunctionCallSupport(geminiRequest);
GenerateContentResponse response = this.callWithFunctionSupport(geminiRequest);
List<Generation> generations = response.getCandidatesList()
@@ -168,7 +165,7 @@ public class VertexAiGeminiChatClient
var request = createGeminiRequest(prompt);
ResponseStream<GenerateContentResponse> responseStream = request.model
.generateContentStream(request.contents, request.config);
.generateContentStream(request.contents);
return Flux.fromStream(responseStream.stream()).map(response -> {
response = handleFunctionCallOrReturn(request, response);
@@ -193,7 +190,7 @@ public class VertexAiGeminiChatClient
}
@JsonInclude(Include.NON_NULL)
public record GeminiRequest(List<Content> contents, GenerativeModel model, GenerationConfig config) {
public record GeminiRequest(List<Content> contents, GenerativeModel model) {
}
private GeminiRequest createGeminiRequest(Prompt prompt) {
@@ -202,7 +199,8 @@ public class VertexAiGeminiChatClient
GenerationConfig generationConfig = this.generationConfig;
GenerativeModel generativeModel = new GenerativeModel(this.defaultOptions.getModel(), this.vertexAI);
var generativeModelBuilder = new GenerativeModel.Builder().setModelName(this.defaultOptions.getModel())
.setVertexAi(this.vertexAI);
VertexAiGeminiChatOptions updatedRuntimeOptions = null;
@@ -237,14 +235,8 @@ public class VertexAiGeminiChatClient
if (StringUtils.hasText(updatedRuntimeOptions.getModel())
&& !updatedRuntimeOptions.getModel().equals(this.defaultOptions.getModel())) {
generativeModel = new GenerativeModel(updatedRuntimeOptions.getModel(), vertexAI);
}
if (updatedRuntimeOptions.getTransportType() != null) {
Transport transport = (updatedRuntimeOptions
.getTransportType() == VertexAiGeminiChatOptions.TransportType.GRPC) ? Transport.GRPC
: Transport.REST;
generativeModel.setTransport(transport);
// Override model name
generativeModelBuilder.setModelName(updatedRuntimeOptions.getModel());
}
generationConfig = toGenerationConfig(updatedRuntimeOptions);
@@ -253,10 +245,14 @@ public class VertexAiGeminiChatClient
// Add the enabled functions definitions to the request's tools parameter.
if (!CollectionUtils.isEmpty(functionsForThisRequest)) {
List<Tool> tools = this.getFunctionTools(functionsForThisRequest);
generativeModel.setTools(tools);
generativeModelBuilder.setTools(tools);
}
return new GeminiRequest(toGeminiContent(prompt), generativeModel, generationConfig);
generativeModelBuilder.setGenerationConfig(generationConfig);
GenerativeModel generativeModel = generativeModelBuilder.build();
return new GeminiRequest(toGeminiContent(prompt), generativeModel);
}
private GenerationConfig toGenerationConfig(VertexAiGeminiChatOptions options) {
@@ -429,7 +425,7 @@ public class VertexAiGeminiChatClient
conversationHistory.add(contentFnResp);
return new GeminiRequest(conversationHistory, previousRequest.model(), previousRequest.config());
return new GeminiRequest(conversationHistory, previousRequest.model());
}
@Override
@@ -445,7 +441,7 @@ public class VertexAiGeminiChatClient
@Override
protected GenerateContentResponse doChatCompletion(GeminiRequest request) {
try {
return request.model.generateContent(request.contents, request.config);
return request.model.generateContent(request.contents);
}
catch (Exception e) {
throw new RuntimeException("Failed to generate content", e);

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

@@ -99,10 +99,6 @@ public class VertexAiGeminiChatOptions implements FunctionCallingOptions, ChatOp
@JsonIgnore
private Set<String> functions = new HashSet<>();
/**
* The transport type to use for the Gemini Chat Client.
*/
private TransportType transportType = TransportType.GRPC;
// @formatter:on
public static Builder builder() {
@@ -165,11 +161,6 @@ public class VertexAiGeminiChatOptions implements FunctionCallingOptions, ChatOp
return this;
}
public Builder withTransportType(TransportType transportType) {
this.options.setTransportType(transportType);
return this;
}
public VertexAiGeminiChatOptions build() {
return this.options;
}
@@ -257,14 +248,6 @@ public class VertexAiGeminiChatOptions implements FunctionCallingOptions, ChatOp
this.functions = functions;
}
public TransportType getTransportType() {
return this.transportType;
}
public void setTransportType(TransportType transportType) {
this.transportType = transportType;
}
@Override
public int hashCode() {
final int prime = 31;
@@ -278,7 +261,6 @@ public class VertexAiGeminiChatOptions implements FunctionCallingOptions, ChatOp
result = prime * result + ((model == null) ? 0 : model.hashCode());
result = prime * result + ((functionCallbacks == null) ? 0 : functionCallbacks.hashCode());
result = prime * result + ((functions == null) ? 0 : functions.hashCode());
result = prime * result + ((transportType == null) ? 0 : transportType.hashCode());
return result;
}
@@ -345,8 +327,6 @@ public class VertexAiGeminiChatOptions implements FunctionCallingOptions, ChatOp
}
else if (!functions.equals(other.functions))
return false;
if (transportType != other.transportType)
return false;
return true;
}

8
models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClientIT.java
View File

@@ -21,6 +21,7 @@ import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import com.google.cloud.vertexai.Transport;
import com.google.cloud.vertexai.VertexAI;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -37,7 +38,6 @@ import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.parser.BeanOutputParser;
import org.springframework.ai.parser.ListOutputParser;
import org.springframework.ai.parser.MapOutputParser;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatOptions.TransportType;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
@@ -224,7 +224,10 @@ class VertexAiGeminiChatClientIT {
public VertexAI vertexAiApi() {
String projectId = System.getenv("VERTEX_AI_GEMINI_PROJECT_ID");
String location = System.getenv("VERTEX_AI_GEMINI_LOCATION");
return new VertexAI(projectId, location);
return new VertexAI.Builder().setProjectId(projectId)
.setLocation(location)
.setTransport(Transport.REST)
.build();
}
@Bean
@@ -232,7 +235,6 @@ class VertexAiGeminiChatClientIT {
return new VertexAiGeminiChatClient(vertexAi,
VertexAiGeminiChatOptions.builder()
.withModel(VertexAiGeminiChatClient.ChatModel.GEMINI_PRO_VISION.getValue())
.withTransportType(TransportType.REST)
.build());
}

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

@@ -19,6 +19,7 @@ import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import com.google.cloud.vertexai.Transport;
import com.google.cloud.vertexai.VertexAI;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
@@ -37,7 +38,6 @@ import org.springframework.ai.model.function.FunctionCallbackWrapper;
import org.springframework.ai.model.function.FunctionCallbackWrapper.Builder.SchemaType;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatClient;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatOptions;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatOptions.TransportType;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
@@ -185,7 +185,10 @@ public class VertexAiGeminiChatClientFunctionCallingIT {
public VertexAI vertexAiApi() {
String projectId = System.getenv("VERTEX_AI_GEMINI_PROJECT_ID");
String location = System.getenv("VERTEX_AI_GEMINI_LOCATION");
return new VertexAI(projectId, location);
return new VertexAI.Builder().setLocation(location)
.setProjectId(projectId)
.setTransport(Transport.REST)
.build();
}
@Bean
@@ -194,7 +197,6 @@ public class VertexAiGeminiChatClientFunctionCallingIT {
VertexAiGeminiChatOptions.builder()
.withModel(VertexAiGeminiChatClient.ChatModel.GEMINI_PRO.getValue())
.withTemperature(0.9f)
.withTransportType(TransportType.REST)
.build());
}

2
pom.xml
View File

@@ -127,7 +127,7 @@
<jackson.version>2.16.1</jackson.version>
<djl.version>0.26.0</djl.version>
<onnxruntime.version>1.17.0</onnxruntime.version>
<com.google.cloud.version>26.34.0</com.google.cloud.version>
<com.google.cloud.version>26.37.0</com.google.cloud.version>
<qdrant.version>1.7.1</qdrant.version>
<spring-retry.version>2.0.5</spring-retry.version>
<ibm.sdk.version>9.20.0</ibm.sdk.version>

3
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/vertexai-gemini-chat.adoc
View File

@@ -54,6 +54,9 @@ The prefix `spring.ai.vertex.ai.gemini` is used as the property prefix that lets
| spring.ai.vertex.ai.gemini.projectId | Google Cloud Platform project ID | -
| spring.ai.vertex.ai.gemini.location | Region | -
| spring.ai.vertex.ai.gemini.credentialsUri | URI to Vertex AI Gemini credentials. When provided it is used to create an a `GoogleCredentials` instance to authenticate the `VertexAI`. | -
| spring.ai.vertex.ai.gemini.apiEndpoint | Vertex AI Gemini API endpoint. | -
| spring.ai.vertex.ai.gemini.scopes | | -
| spring.ai.vertex.ai.gemini.transport | API transport. GRPC or REST. | GRPC
|====
The prefix `spring.ai.vertex.ai.gemini.chat` is the property prefix that lets you configure the chat client implementation for VertexAI Gemini Chat.

20
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiAutoConfiguration.java
View File

@@ -32,6 +32,7 @@ import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
/**
* Auto-configuration for Vertex AI Gemini Chat.
@@ -49,15 +50,26 @@ public class VertexAiGeminiAutoConfiguration {
Assert.hasText(connectionProperties.getProjectId(), "Vertex AI project-id must be set!");
Assert.hasText(connectionProperties.getLocation(), "Vertex AI location must be set!");
Assert.notNull(connectionProperties.getTransport(), "Vertex AI transport must be set!");
var vertexAIBuilder = new VertexAI.Builder().setProjectId(connectionProperties.getProjectId())
.setLocation(connectionProperties.getLocation())
.setTransport(com.google.cloud.vertexai.Transport.valueOf(connectionProperties.getTransport().name()));
if (StringUtils.hasText(connectionProperties.getApiEndpoint())) {
vertexAIBuilder.setApiEndpoint(connectionProperties.getApiEndpoint());
}
if (!CollectionUtils.isEmpty(connectionProperties.getScopes())) {
vertexAIBuilder.setScopes(connectionProperties.getScopes());
}
if (connectionProperties.getCredentialsUri() != null) {
GoogleCredentials credentials = GoogleCredentials
.fromStream(connectionProperties.getCredentialsUri().getInputStream());
return new VertexAI(connectionProperties.getProjectId(), connectionProperties.getLocation(), credentials);
}
else {
return new VertexAI(connectionProperties.getProjectId(), connectionProperties.getLocation());
vertexAIBuilder.setCredentials(credentials);
}
return vertexAIBuilder.build();
}
@Bean

50
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiConnectionProperties.java
View File

@@ -15,6 +15,8 @@
*/
package org.springframework.ai.autoconfigure.vertexai.gemini;
import java.util.List;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.core.io.Resource;
@@ -44,6 +46,30 @@ public class VertexAiGeminiConnectionProperties {
*/
private Resource credentialsUri;
/**
* Vertex AI Gemini API endpoint.
*/
private String apiEndpoint;
/**
*
*/
private List<String> scopes = List.of();
private Transport transport = Transport.GRPC;
public enum Transport {
/** When used, the clients will send REST requests to the backing service. */
REST,
/**
* When used, the clients will send gRPC to the backing service. This is usually
* more efficient and is the default transport.
*/
GRPC
}
public String getProjectId() {
return this.projectId;
}
@@ -68,4 +94,28 @@ public class VertexAiGeminiConnectionProperties {
this.credentialsUri = credentialsUri;
}
public String getApiEndpoint() {
return this.apiEndpoint;
}
public List<String> getScopes() {
return this.scopes;
}
public void setScopes(List<String> scopes) {
this.scopes = scopes;
}
public void setApiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
}
public Transport getTransport() {
return this.transport;
}
public void setTransport(Transport transport) {
this.transport = transport;
}
}