diff --git a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClient.java b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClient.java index bd6b0d09a..b687ac97c 100644 --- a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClient.java +++ b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClient.java @@ -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 generations = response.getCandidatesList() @@ -168,7 +165,7 @@ public class VertexAiGeminiChatClient var request = createGeminiRequest(prompt); ResponseStream 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 contents, GenerativeModel model, GenerationConfig config) { + public record GeminiRequest(List 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 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); diff --git a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatOptions.java b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatOptions.java index 3e7e6fe5e..7d4e9875a 100644 --- a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatOptions.java +++ b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatOptions.java @@ -99,10 +99,6 @@ public class VertexAiGeminiChatOptions implements FunctionCallingOptions, ChatOp @JsonIgnore private Set 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; } diff --git a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClientIT.java b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClientIT.java index 8860dc9f9..3f5783c6e 100644 --- a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClientIT.java +++ b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClientIT.java @@ -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()); } diff --git a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatClientFunctionCallingIT.java b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatClientFunctionCallingIT.java index cc9a153bd..19d4b4e9c 100644 --- a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatClientFunctionCallingIT.java +++ b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatClientFunctionCallingIT.java @@ -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()); } diff --git a/pom.xml b/pom.xml index 7c0d9c0db..6bf7fb6b9 100644 --- a/pom.xml +++ b/pom.xml @@ -127,7 +127,7 @@ 2.16.1 0.26.0 1.17.0 - 26.34.0 + 26.37.0 1.7.1 2.0.5 9.20.0 diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/vertexai-gemini-chat.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/vertexai-gemini-chat.adoc index 5c275f42b..79c143501 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/vertexai-gemini-chat.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/vertexai-gemini-chat.adoc @@ -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. diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiAutoConfiguration.java index 93b93beac..12dfb8134 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiAutoConfiguration.java @@ -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 diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiConnectionProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiConnectionProperties.java index 35dcef985..ef65327b5 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiConnectionProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/gemini/VertexAiGeminiConnectionProperties.java @@ -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 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 getScopes() { + return this.scopes; + } + + public void setScopes(List 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; + } + }