Add vertex ai autoconfig and spring boot

- Add Vertex AI Autoconfigurations for chat and embedding clients.
 - Factor out the embeding client dimensions() computation into an abstract parent AbstractEmbeddingClient.
 - Add ITs
 - Vertex dos.

embedding-clients/spring-ai-postgresml-embedding-client/src/main/java/org/springframework/ai/embedding/PostgresMlEmbeddingClient.java

@@ -6,7 +6,6 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -25,7 +24,7 @@ import org.springframework.util.StringUtils;
  *
  * @author Toshiaki Maki
  */
-public class PostgresMlEmbeddingClient implements EmbeddingClient, InitializingBean {
+public class PostgresMlEmbeddingClient extends AbstractEmbeddingClient implements InitializingBean {
 
 	private final JdbcTemplate jdbcTemplate;
 
@@ -35,8 +34,6 @@ public class PostgresMlEmbeddingClient implements EmbeddingClient, InitializingB
 
 	private final String kwargs;
 
-	private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
-
 	private final MetadataMode metadataMode;
 
 	public enum VectorType {
@@ -162,14 +159,6 @@ public class PostgresMlEmbeddingClient implements EmbeddingClient, InitializingB
 				Map.of("transformer", this.transformer, "vector-type", this.vectorType.name(), "kwargs", this.kwargs));
 	}
 
-	@Override
-	public int dimensions() {
-		if (this.embeddingDimensions.get() < 0) {
-			this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, this.transformer));
-		}
-		return this.embeddingDimensions.get();
-	}
-
 	@Override
 	public void afterPropertiesSet() {
 		this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS pgml");

embedding-clients/transformers-embedding/src/main/java/org/springframework/ai/embedding/TransformersEmbeddingClient.java

@@ -5,7 +5,6 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
 import ai.djl.huggingface.tokenizers.Encoding;
@@ -35,7 +34,7 @@ import org.springframework.util.StringUtils;
  *
  * @author Christian Tzolov
  */
-public class TransformersEmbeddingClient implements EmbeddingClient, InitializingBean {
+public class TransformersEmbeddingClient extends AbstractEmbeddingClient implements InitializingBean {
 
 	private static final Log logger = LogFactory.getLog(TransformersEmbeddingClient.class);
 
@@ -72,8 +71,6 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin
 	 */
 	private OrtSession session;
 
-	private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
-
 	/**
 	 * Specifies what parts of the {@link Document}'s content and metadata will be used
 	 * for computing the embeddings. Applicable for the {@link #embed(Document)} method
@@ -148,10 +145,6 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin
 		this.modelResource = toResource(modelResourceUri);
 	}
 
-	public void setEmbeddingDimensions(int dimension) {
-		this.embeddingDimensions.set(dimension);
-	}
-
 	public void setModelOutputName(String modelOutputName) {
 		this.modelOutputName = modelOutputName;
 	}
@@ -325,14 +318,6 @@ public class TransformersEmbeddingClient implements EmbeddingClient, Initializin
 		return result;
 	}
 
-	@Override
-	public int dimensions() {
-		if (this.embeddingDimensions.get() < 0) {
-			this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, "Test"));
-		}
-		return this.embeddingDimensions.get();
-	}
-
 	private static Resource toResource(String uri) {
 		return new DefaultResourceLoader().getResource(uri);
 	}

pom.xml

@@ -37,6 +37,7 @@
 		<module>vector-stores/spring-ai-azure</module>
 		<module>vector-stores/spring-ai-weaviate</module>
 		<module>spring-ai-vertex-ai</module>
+		<module>spring-ai-spring-boot-starters/spring-ai-starter-vertex-ai</module>
 
 	</modules>
 

spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/embedding/AzureOpenAiEmbeddingClient.java

@@ -4,7 +4,6 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import com.azure.ai.openai.OpenAIClient;
 import com.azure.ai.openai.models.EmbeddingItem;
@@ -16,13 +15,12 @@ import org.slf4j.LoggerFactory;
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.MetadataMode;
+import org.springframework.ai.embedding.AbstractEmbeddingClient;
 import org.springframework.ai.embedding.Embedding;
-import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.embedding.EmbeddingResponse;
-import org.springframework.ai.embedding.EmbeddingUtil;
 import org.springframework.util.Assert;
 
-public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
+public class AzureOpenAiEmbeddingClient extends AbstractEmbeddingClient {
 
 	private static final Logger logger = LoggerFactory.getLogger(AzureOpenAiEmbeddingClient.class);
 
@@ -30,8 +28,6 @@ public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
 
 	private final String model;
 
-	private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
-
 	private final MetadataMode metadataMode;
 
 	public AzureOpenAiEmbeddingClient(OpenAIClient azureOpenAiClient) {
@@ -113,12 +109,4 @@ public class AzureOpenAiEmbeddingClient implements EmbeddingClient {
 		return data;
 	}
 
-	@Override
-	public int dimensions() {
-		if (this.embeddingDimensions.get() < 0) {
-			this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, this.model));
-		}
-		return this.embeddingDimensions.get();
-	}
-
 }

spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingClient.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.embedding;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Abstract implementation of the {@link EmbeddingClient} interface that provides
+ * dimensions calculation caching.
+ *
+ * @author Christian Tzolov
+ */
+public abstract class AbstractEmbeddingClient implements EmbeddingClient {
+
+	private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
+
+	@Override
+	public int dimensions() {
+		if (this.embeddingDimensions.get() < 0) {
+			this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, "Test"));
+		}
+		return this.embeddingDimensions.get();
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingClient.java

@@ -4,16 +4,42 @@ import org.springframework.ai.document.Document;
 
 import java.util.List;
 
+/**
+ * EmbeddingClient is a generic interface for embedding clients.
+ */
 public interface EmbeddingClient {
 
+	/**
+	 * Embeds the given text into a vector.
+	 * @param text the text to embed.
+	 * @return the embedded vector.
+	 */
 	List<Double> embed(String text);
 
+	/**
+	 * Embeds the given document's content into a vector.
+	 * @param document the document to embed.
+	 * @return the embedded vector.
+	 */
 	List<Double> embed(Document document);
 
+	/**
+	 * Embeds a batch of texts into vectors.
+	 * @param texts list of texts to embed.
+	 * @return list of list of embedded vectors.
+	 */
 	List<List<Double>> embed(List<String> texts);
 
+	/**
+	 * Embeds a batch of texts into vectors and returns the {@link EmbeddingResponse}.
+	 * @param texts list of texts to embed.
+	 * @return the embedding response.
+	 */
 	EmbeddingResponse embedForResponse(List<String> texts);
 
+	/**
+	 * @return the number of dimensions of the embedded vectors. It is model specific.
+	 */
 	default int dimensions() {
 		return embed("Test String").size();
 	}

spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java

@@ -4,7 +4,6 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import com.theokanning.openai.Usage;
 import com.theokanning.openai.embedding.EmbeddingRequest;
@@ -14,13 +13,12 @@ import org.slf4j.LoggerFactory;
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.MetadataMode;
+import org.springframework.ai.embedding.AbstractEmbeddingClient;
 import org.springframework.ai.embedding.Embedding;
-import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.embedding.EmbeddingResponse;
-import org.springframework.ai.embedding.EmbeddingUtil;
 import org.springframework.util.Assert;
 
-public class OpenAiEmbeddingClient implements EmbeddingClient {
+public class OpenAiEmbeddingClient extends AbstractEmbeddingClient {
 
 	private static final Logger logger = LoggerFactory.getLogger(OpenAiEmbeddingClient.class);
 
@@ -28,8 +26,6 @@ public class OpenAiEmbeddingClient implements EmbeddingClient {
 
 	private final String model;
 
-	private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
-
 	private final MetadataMode metadataMode;
 
 	public OpenAiEmbeddingClient(OpenAiService openAiService) {
@@ -108,12 +104,4 @@ public class OpenAiEmbeddingClient implements EmbeddingClient {
 		return metadata;
 	}
 
-	@Override
-	public int dimensions() {
-		if (this.embeddingDimensions.get() < 0) {
-			this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, this.model));
-		}
-		return this.embeddingDimensions.get();
-	}
-
 }

spring-ai-spring-boot-autoconfigure/pom.xml

@@ -125,6 +125,14 @@
 			<optional>true</optional>
 		</dependency>
 
+		<!-- Vertex AI LLM -->
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-vertex-ai</artifactId>
+			<version>${project.parent.version}</version>
+			<optional>true</optional>
+		</dependency>
+
 
 		<dependency>
 			<groupId>org.springframework.boot</groupId>

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.autoconfigure.vertexai;
+
+import org.springframework.ai.vertex.api.VertexAiApi;
+import org.springframework.ai.vertex.embedding.VertexAiEmbeddingClient;
+import org.springframework.ai.vertex.generation.VertexAiChatClient;
+import org.springframework.boot.autoconfigure.AutoConfiguration;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.context.annotation.Bean;
+import org.springframework.web.client.RestClient;
+
+@AutoConfiguration
+@ConditionalOnClass(VertexAiApi.class)
+@EnableConfigurationProperties({ VertexAiConnectionProperties.class, VertexAiChatProperties.class,
+		VertexAiEmbeddingProperties.class })
+public class VertexAiAutoConfiguration {
+
+	@Bean
+	@ConditionalOnMissingBean
+	public VertexAiChatClient vertexAiClient(VertexAiApi vertexAiApi, VertexAiChatProperties chatProperties) {
+
+		VertexAiChatClient client = new VertexAiChatClient(vertexAiApi);
+
+		client.setTemperature(chatProperties.getTemperature());
+		client.setTopP(chatProperties.getTopP());
+		client.setTopK(chatProperties.getTopK());
+		client.setCandidateCount(chatProperties.getCandidateCount());
+
+		return client;
+	}
+
+	@Bean
+	@ConditionalOnMissingBean
+	public VertexAiEmbeddingClient vertexAiEmbeddingClient(VertexAiApi vertexAiApi) {
+		return new VertexAiEmbeddingClient(vertexAiApi);
+	}
+
+	@Bean
+	@ConditionalOnMissingBean
+	public VertexAiApi vertexAiApi(VertexAiConnectionProperties connectionProperties,
+			VertexAiEmbeddingProperties embeddingAiProperties, VertexAiChatProperties chatProperties) {
+
+		return new VertexAiApi(connectionProperties.getBaseUrl(), connectionProperties.getApiKey(),
+				chatProperties.getModel(), embeddingAiProperties.getModel(), RestClient.builder());
+	}
+
+}

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiChatProperties.java

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.autoconfigure.vertexai;
+
+import org.springframework.ai.vertex.api.VertexAiApi;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+
+@ConfigurationProperties(VertexAiChatProperties.CONFIG_PREFIX)
+public class VertexAiChatProperties {
+
+	public static final String CONFIG_PREFIX = "spring.ai.vertex.ai.chat";
+
+	/**
+	 * Controls the randomness of the output. Values can range over [0.0,1.0], inclusive.
+	 * A value closer to 1.0 will produce responses that are more varied, while a value
+	 * closer to 0.0 will typically result in less surprising responses from the model.
+	 * This value specifies default to be used by the backend while making the call to the
+	 * model.
+	 */
+	private Float temperature = 0.7f;
+
+	/**
+	 * The maximum cumulative probability of tokens to consider when sampling. The model
+	 * uses combined Top-k and nucleus sampling. Nucleus sampling considers the smallest
+	 * set of tokens whose probability sum is at least topP.
+	 */
+	private Float topP = null;
+
+	/**
+	 * The number of generated response messages to return. This value must be between [1,
+	 * 8], inclusive. Defaults to 1.
+	 */
+	private Integer candidateCount = 1;
+
+	/**
+	 * The maximum number of tokens to consider when sampling. The model uses combined
+	 * Top-k and nucleus sampling. Top-k sampling considers the set of topK most probable
+	 * tokens.
+	 */
+	private Integer topK = 20;
+
+	/**
+	 * Vertex AI PaLM API model name. Defaults to chat-bison-001
+	 */
+	private String model = VertexAiApi.DEFAULT_GENERATE_MODEL;
+
+	public String getModel() {
+		return this.model;
+	}
+
+	public void setModel(String model) {
+		this.model = model;
+	}
+
+	public Float getTemperature() {
+		return this.temperature;
+	}
+
+	public void setTemperature(Float temperature) {
+		this.temperature = temperature;
+	}
+
+	public Float getTopP() {
+		return this.topP;
+	}
+
+	public void setTopP(Float topP) {
+		this.topP = topP;
+	}
+
+	public Integer getCandidateCount() {
+		return this.candidateCount;
+	}
+
+	public void setCandidateCount(Integer candidateCount) {
+		this.candidateCount = candidateCount;
+	}
+
+	public Integer getTopK() {
+		return this.topK;
+	}
+
+	public void setTopK(Integer topK) {
+		this.topK = topK;
+	}
+
+}

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiConnectionProperties.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.autoconfigure.vertexai;
+
+import org.springframework.ai.vertex.api.VertexAiApi;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+
+@ConfigurationProperties(VertexAiConnectionProperties.CONFIG_PREFIX)
+public class VertexAiConnectionProperties {
+
+	public static final String CONFIG_PREFIX = "spring.ai.vertex.ai";
+
+	/**
+	 * Vertex AI PaLM API access key.
+	 */
+	private String apiKey;
+
+	/**
+	 * Vertex AI PaLM API base URL. Defaults to
+	 * https://generativelanguage.googleapis.com/v1beta3
+	 */
+	private String baseUrl = VertexAiApi.DEFAULT_BASE_URL;
+
+	public String getApiKey() {
+		return this.apiKey;
+	}
+
+	public void setApiKey(String apiKey) {
+		this.apiKey = apiKey;
+	}
+
+	public String getBaseUrl() {
+		return this.baseUrl;
+	}
+
+	public void setBaseUrl(String baseUrl) {
+		this.baseUrl = baseUrl;
+	}
+
+}

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiEmbeddingProperties.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.autoconfigure.vertexai;
+
+import org.springframework.ai.vertex.api.VertexAiApi;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+
+@ConfigurationProperties(VertexAiEmbeddingProperties.CONFIG_PREFIX)
+public class VertexAiEmbeddingProperties {
+
+	public static final String CONFIG_PREFIX = "spring.ai.vertex.ai.embedding";
+
+	/**
+	 * Vertex AI PaLM API embedding model name. Defaults to embedding-gecko-001.
+	 */
+	private String model = VertexAiApi.DEFAULT_EMBEDDING_MODEL;
+
+	public String getModel() {
+		return this.model;
+	}
+
+	public void setModel(String model) {
+		this.model = model;
+	}
+
+}

spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports

@@ -9,3 +9,4 @@ org.springframework.ai.autoconfigure.huggingface.HuggingfaceAutoConfiguration
 org.springframework.ai.autoconfigure.vectorstore.chroma.ChromaVectorStoreAutoConfiguration
 org.springframework.ai.autoconfigure.vectorstore.azure.AzureVectorStoreAutoConfiguration
 org.springframework.ai.autoconfigure.vectorstore.weaviate.WeaviateVectorStoreAutoConfiguration
+org.springframework.ai.autoconfigure.vertexai.VertexAiAutoConfiguration

spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfigurationIT.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.autoconfigure.vertexai;
+
+import java.util.List;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+
+import org.springframework.ai.embedding.EmbeddingResponse;
+import org.springframework.ai.vertex.embedding.VertexAiEmbeddingClient;
+import org.springframework.ai.vertex.generation.VertexAiChatClient;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@EnabledIfEnvironmentVariable(named = "PALM_API_KEY", matches = ".*")
+public class VertexAiAutoConfigurationIT {
+
+	private static final Log logger = LogFactory.getLog(VertexAiAutoConfigurationIT.class);
+
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+		.withPropertyValues("spring.ai.vertex.ai.baseUrl=https://generativelanguage.googleapis.com/v1beta3",
+				"spring.ai.vertex.ai.apiKey=" + System.getenv("PALM_API_KEY"),
+				"spring.ai.vertex.ai.chat.model=chat-bison-001",
+				"spring.ai.vertex.ai.embedding.model=embedding-gecko-001")
+		.withConfiguration(AutoConfigurations.of(VertexAiAutoConfiguration.class));
+
+	@Test
+	void generate() {
+		contextRunner.run(context -> {
+			VertexAiChatClient client = context.getBean(VertexAiChatClient.class);
+
+			String response = client.generate("Hello");
+
+			assertThat(response).isNotEmpty();
+			logger.info("Response: " + response);
+		});
+	}
+
+	@Test
+	void embedding() {
+		contextRunner.run(context -> {
+			VertexAiEmbeddingClient embeddingClient = context.getBean(VertexAiEmbeddingClient.class);
+
+			EmbeddingResponse embeddingResponse = embeddingClient
+				.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
+			assertThat(embeddingResponse.getData()).hasSize(2);
+			assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
+			assertThat(embeddingResponse.getData().get(0).getIndex()).isEqualTo(0);
+			assertThat(embeddingResponse.getData().get(1).getEmbedding()).isNotEmpty();
+			assertThat(embeddingResponse.getData().get(1).getIndex()).isEqualTo(1);
+
+			assertThat(embeddingClient.dimensions()).isEqualTo(768);
+		});
+	}
+
+}

spring-ai-spring-boot-starters/spring-ai-starter-vertex-ai/pom.xml

@@ -0,0 +1,51 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <groupId>org.springframework.ai</groupId>
+        <artifactId>spring-ai</artifactId>
+        <version>0.8.0-SNAPSHOT</version>
+        <relativePath>../../pom.xml</relativePath>
+    </parent>
+    <artifactId>spring-ai-vertex-ai-spring-boot-starter</artifactId>
+    <packaging>jar</packaging>
+    <name>Spring AI Starter - VertexAI</name>
+    <description>Spring AI VertexAI Auto Configuration</description>
+    <url>https://github.com/spring-projects-experimental/spring-ai</url>
+
+    <scm>
+        <url>https://github.com/spring-projects-experimental/spring-ai</url>
+        <connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
+        <developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
+    </scm>
+
+    <dependencies>
+
+        <!-- production dependencies -->
+
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-vertex-ai</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+
+        <!-- test dependencies -->
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-test</artifactId>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+</project>

spring-ai-vertex-ai/README.md

@@ -0,0 +1,37 @@
+
+# Vertex AI API client for the Generative Language model
+
+The [Generative Language](https://developers.generativeai.google/api/rest/generativelanguage) PaLM API allows developers to build generative AI applications using the PaLM model. Large Language Models (LLMs) are a powerful, versatile type of machine learning model that enables computers to comprehend and generate natural language through a series of prompts. The PaLM API is based on Google's next generation LLM, PaLM. It excels at a variety of different tasks like code generation, reasoning, and writing. You can use the PaLM API to build generative AI applications for use cases like content generation, dialogue agents, summarization and classification systems, and more.
+
+Based on the [Models REST API](https://developers.generativeai.google/api/rest/generativelanguage/models).
+
+## Prerequisite
+
+To access the PaLM2 REST API you need to obtain an access API KEY form [makersuite](https://makersuite.google.com/app/apikey).
+Note: Currently it is not available outside US, but you can use VPN for testing.
+
+## PaLM API
+
+The VertexAI, AiClient and EmbeddingClient are built on top the [VertexAiApi.java](./src/main/java/org/springframework/ai/vertex/api/VertexAiApi.java) client library:
+
+![PaLM API](./src/test/resources/Google%20Generative%20AI%20-%20PaLM2%20REST%20API.jpg)
+
+Following snippets show how to use the `VertexAiApi` client directly:
+
+```java
+
+VertexAiApi vertexAiApi = new VertexAiApi(< YOUR PALM_API_KEY>);
+
+// Generate
+var prompt = new MessagePrompt(List.of(new Message("0", "Hello, how are you?")));
+
+GenerateMessageRequest request = new GenerateMessageRequest(prompt);
+
+GenerateMessageResponse response = vertexAiApi.generateMessage(request);
+
+// Embed text
+Embedding embedding = vertexAiApi.embedText("Hello, how are you?");
+
+// Batch embedding
+List<Embedding> embeddings = vertexAiApi.batchEmbedText(List.of("Hello, how are you?", "I am fine, thank you!"));
+```

spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/api/VertexAiApi.java

@@ -49,7 +49,7 @@ import org.springframework.web.client.RestClient;
  * Supported models:
  *
  * <pre>
- * name=models/chat-bison-001,
+ * 		name=models/chat-bison-001,
  * 		version=001,
  * 		displayName=Chat Bison,
  * 		description=Chat-optimized generative language model.,
@@ -60,7 +60,7 @@ import org.springframework.web.client.RestClient;
  * 		topP=0.95,
  *		topK=40
  *
- * name=models/text-bison-001,
+ * 		name=models/text-bison-001,
  *		version=001,
  *		displayName=Text Bison,
  *		description=Model targeted for text generation.,
@@ -71,7 +71,7 @@ import org.springframework.web.client.RestClient;
  *		topP=0.95,
  *		topK=40
  *
- * name=models/embedding-gecko-001,
+ * 		name=models/embedding-gecko-001,
  * 		version=001,
  * 		displayName=Embedding Gecko, description=Obtain a distributed representation of a text.,
  * 		inputTokenLimit=1024,
@@ -98,7 +98,10 @@ public class VertexAiApi {
 	 */
 	public static final String DEFAULT_EMBEDDING_MODEL = "embedding-gecko-001";
 
-	private static final String DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta3";
+	/**
+	 * The default base URL for accessing the Vertex AI API.
+	 */
+	public static final String DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta3";
 
 	private final RestClient restClient;
 
@@ -582,33 +585,5 @@ public class VertexAiApi {
 			}
 		}
 	}
-
-	/**
-	 * Main method to test the VertexAiApi.
-	 * @param args blank.
-	 */
-	public static void main(String[] args) {
-		VertexAiApi vertexAiApi = new VertexAiApi(System.getenv("PALM_API_KEY"));
-
-		var prompt = new MessagePrompt(List.of(new Message("0", "Hello, how are you?")));
-
-		GenerateMessageRequest request = new GenerateMessageRequest(prompt);
-
-		GenerateMessageResponse response = vertexAiApi.generateMessage(request);
-
-		System.out.println(response);
-
-		System.out.println(vertexAiApi.embedText("Hello, how are you?"));
-
-		System.out.println(vertexAiApi.batchEmbedText(List.of("Hello, how are you?", "I am fine, thank you!")));
-
-		System.out.println(vertexAiApi.countMessageTokens(prompt));
-
-		System.out.println(vertexAiApi.listModels());
-
-		System.out.println(vertexAiApi.listModels().stream().map(vertexAiApi::getModel).toList());
-
-	}
-
 }
 // @formatter:on

spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/embedding/VertexAiEmbeddingClient.java

@@ -16,20 +16,20 @@
 
 package org.springframework.ai.vertex.embedding;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.AbstractEmbeddingClient;
 import org.springframework.ai.embedding.Embedding;
-import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.embedding.EmbeddingResponse;
 import org.springframework.ai.vertex.api.VertexAiApi;
 
 /**
  * @author Christian Tzolov
  */
-public class VertexAiEmbeddingClient implements EmbeddingClient {
+public class VertexAiEmbeddingClient extends AbstractEmbeddingClient {
 
 	private final VertexAiApi vertexAiApi;
 
@@ -56,11 +56,10 @@ public class VertexAiEmbeddingClient implements EmbeddingClient {
 	@Override
 	public EmbeddingResponse embedForResponse(List<String> texts) {
 		List<VertexAiApi.Embedding> vertexEmbeddings = this.vertexAiApi.batchEmbedText(texts);
-		int index = 0;
-		List<Embedding> embeddings = new ArrayList<>();
-		for (VertexAiApi.Embedding vertexEmbedding : vertexEmbeddings) {
-			embeddings.add(new Embedding(vertexEmbedding.value(), index++));
-		}
+		AtomicInteger indexCounter = new AtomicInteger(0);
+		List<Embedding> embeddings = vertexEmbeddings.stream()
+			.map(vm -> new Embedding(vm.value(), indexCounter.getAndIncrement()))
+			.toList();
 		return new EmbeddingResponse(embeddings, Map.of());
 	}
 

spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/generation/VertexAiChatClient.java

@@ -34,7 +34,7 @@ import org.springframework.util.CollectionUtils;
 /**
  * @author Christian Tzolov
  */
-public class VertexAiChatGenerationClient implements AiClient {
+public class VertexAiChatClient implements AiClient {
 
 	private final VertexAiApi vertexAiApi;
 
@@ -42,14 +42,30 @@ public class VertexAiChatGenerationClient implements AiClient {
 
 	private Float topP;
 
+	private Integer topK;
+
 	private Integer candidateCount;
 
-	private Integer maxTokens;
-
-	public VertexAiChatGenerationClient(VertexAiApi vertexAiApi) {
+	public VertexAiChatClient(VertexAiApi vertexAiApi) {
 		this.vertexAiApi = vertexAiApi;
 	}
 
+	public void setTemperature(Float temperature) {
+		this.temperature = temperature;
+	}
+
+	public void setTopK(Integer candidateCount) {
+		this.topK = candidateCount;
+	}
+
+	public void setTopP(Float topP) {
+		this.topP = topP;
+	}
+
+	public void setCandidateCount(Integer maxTokens) {
+		this.candidateCount = maxTokens;
+	}
+
 	@Override
 	public AiResponse generate(Prompt prompt) {
 
@@ -70,7 +86,7 @@ public class VertexAiChatGenerationClient implements AiClient {
 		var vertexPrompt = new MessagePrompt(vertexContext, vertexMessages);
 
 		GenerateMessageRequest request = new GenerateMessageRequest(vertexPrompt, this.temperature, this.candidateCount,
-				this.topP, this.maxTokens);
+				this.topP, this.topK);
 
 		GenerateMessageResponse response = this.vertexAiApi.generateMessage(request);
 

spring-ai-vertex-ai/src/test/java/org/springframework/ai/vertex/embedding/VertexAiEmbeddingClientIT.java

@@ -30,6 +30,20 @@ class VertexAiEmbeddingClientIT {
 		assertThat(embeddingClient.dimensions()).isEqualTo(768);
 	}
 
+	@Test
+	void batchEmbedding() {
+		assertThat(embeddingClient).isNotNull();
+		EmbeddingResponse embeddingResponse = embeddingClient
+			.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
+		assertThat(embeddingResponse.getData()).hasSize(2);
+		assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
+		assertThat(embeddingResponse.getData().get(0).getIndex()).isEqualTo(0);
+		assertThat(embeddingResponse.getData().get(1).getEmbedding()).isNotEmpty();
+		assertThat(embeddingResponse.getData().get(1).getIndex()).isEqualTo(1);
+
+		assertThat(embeddingClient.dimensions()).isEqualTo(768);
+	}
+
 	@SpringBootConfiguration
 	public static class TestConfiguration {
 

spring-ai-vertex-ai/src/test/java/org/springframework/ai/vertex/generation/VertexAiChatGenerationClientIT.java

@@ -33,7 +33,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 class VertexAiChatGenerationClientIT {
 
 	@Autowired
-	private VertexAiChatGenerationClient client;
+	private VertexAiChatClient client;
 
 	@Value("classpath:/prompts/system-message.st")
 	private Resource systemResource;
@@ -121,8 +121,8 @@ class VertexAiChatGenerationClientIT {
 		}
 
 		@Bean
-		public VertexAiChatGenerationClient vertexAiEmbedding(VertexAiApi vertexAiApi) {
-			return new VertexAiChatGenerationClient(vertexAiApi);
+		public VertexAiChatClient vertexAiEmbedding(VertexAiApi vertexAiApi) {
+			return new VertexAiChatClient(vertexAiApi);
 		}
 
 	}