Adding Vertex AI Imagen models support

Closes spring-projectsgh-2133 Signed-off-by: Marzouki-Sami samymarzouki502@gmail.com
Marzouki-Sami · Jan 29, 2025 · 735bd7d · 735bd7d
1 parent ab022fa
commit 735bd7d
Show file tree

Hide file tree

Showing 19 changed files with 2,170 additions and 0 deletions.
diff --git a/models/spring-ai-vertex-ai-imagen/pom.xml b/models/spring-ai-vertex-ai-imagen/pom.xml
@@ -0,0 +1,109 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Copyright 2025-2026 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.
+  -->
+
+<project xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+		 xmlns="http://maven.apache.org/POM/4.0.0"
+		 xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+	<modelVersion>4.0.0</modelVersion>
+	<parent>
+		<groupId>org.springframework.ai</groupId>
+		<artifactId>spring-ai</artifactId>
+		<version>1.0.0-M5</version>
+		<relativePath>../../pom.xml</relativePath>
+	</parent>
+
+	<artifactId>spring-ai-vertex-ai-imagen</artifactId>
+	<packaging>jar</packaging>
+	<name>Spring AI Model - Vertex AI Imagen</name>
+	<description>Vertex AI Imagen models support</description>
+	<url>https://github.com/spring-projects/spring-ai</url>
+
+	<scm>
+		<url>https://github.com/spring-projects/spring-ai</url>
+		<connection>git://github.com/spring-projects/spring-ai.git</connection>
+		<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
+	</scm>
+
+	<properties>
+	</properties>
+
+	<dependencyManagement>
+		<dependencies>
+			<dependency>
+				<groupId>com.google.cloud</groupId>
+				<artifactId>libraries-bom</artifactId>
+				<version>${com.google.cloud.version}</version>
+				<type>pom</type>
+				<scope>import</scope>
+			</dependency>
+		</dependencies>
+	</dependencyManagement>
+
+	<dependencies>
+
+		<dependency>
+			<groupId>com.google.cloud</groupId>
+			<artifactId>google-cloud-aiplatform</artifactId>
+			<exclusions>
+				<exclusion>
+					<groupId>commons-logging</groupId>
+					<artifactId>commons-logging</artifactId>
+				</exclusion>
+			</exclusions>
+		</dependency>
+
+		<!-- production dependencies -->
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-core</artifactId>
+			<version>${project.parent.version}</version>
+		</dependency>
+
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-retry</artifactId>
+			<version>${project.parent.version}</version>
+		</dependency>
+
+		<!-- Spring Framework -->
+		<dependency>
+			<groupId>org.springframework</groupId>
+			<artifactId>spring-context-support</artifactId>
+		</dependency>
+
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-starter-logging</artifactId>
+		</dependency>
+
+		<dependency>
+			<groupId>io.micrometer</groupId>
+			<artifactId>micrometer-observation-test</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+		<!-- test dependencies -->
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-test</artifactId>
+			<version>${project.version}</version>
+			<scope>test</scope>
+		</dependency>
+
+	</dependencies>
+
+</project>
diff --git a/...src/main/java/org/springframework/ai/vertexai/imagen/VertexAiImagenConnectionDetails.java b/...src/main/java/org/springframework/ai/vertexai/imagen/VertexAiImagenConnectionDetails.java
@@ -0,0 +1,181 @@
+/*
+ * Copyright 2025-2026 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.vertexai.imagen;
+
+import java.io.IOException;
+
+import com.google.cloud.aiplatform.v1.EndpointName;
+import com.google.cloud.aiplatform.v1.PredictionServiceSettings;
+
+import org.springframework.util.StringUtils;
+
+
+/**
+ * VertexAiImagenConnectionDetails represents the details of a connection to the Vertex AI imagen service.
+ * It provides methods to access the project ID, location, publisher, and PredictionServiceSettings.
+ *
+ * @author Sami Marzouki
+ */
+public class VertexAiImagenConnectionDetails {
+
+	public static final String DEFAULT_ENDPOINT = "us-central1-aiplatform.googleapis.com:443";
+
+	public static final String DEFAULT_ENDPOINT_SUFFIX = "-aiplatform.googleapis.com:443";
+
+	public static final String DEFAULT_PUBLISHER = "google";
+
+	private static final String DEFAULT_LOCATION = "us-central1";
+
+	/**
+	 * Your project ID.
+	 */
+	private final String projectId;
+
+	/**
+	 * A location is a <a href="https://cloud.google.com/about/locations?hl=en">region</a>
+	 * you can specify in a request to control where data is stored at rest. For a list of
+	 * available regions, see <a href=
+	 * "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations?hl=en">Generative
+	 * AI on Vertex AI locations</a>.
+	 */
+	private final String location;
+
+	private final String publisher;
+
+	private final PredictionServiceSettings predictionServiceSettings;
+
+	public VertexAiImagenConnectionDetails(String projectId, String location, String publisher,
+										   PredictionServiceSettings predictionServiceSettings) {
+		this.projectId = projectId;
+		this.location = location;
+		this.publisher = publisher;
+		this.predictionServiceSettings = predictionServiceSettings;
+	}
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	public String getProjectId() {
+		return this.projectId;
+	}
+
+	public String getLocation() {
+		return this.location;
+	}
+
+	public String getPublisher() {
+		return this.publisher;
+	}
+
+	public EndpointName getEndpointName(String modelName) {
+		return EndpointName.ofProjectLocationPublisherModelName(this.projectId, this.location, this.publisher,
+				modelName);
+	}
+
+	public com.google.cloud.aiplatform.v1.PredictionServiceSettings getPredictionServiceSettings() {
+		return this.predictionServiceSettings;
+	}
+
+	public static class Builder {
+
+		/**
+		 * The Vertex AI embedding endpoint.
+		 */
+		private String endpoint;
+
+		/**
+		 * Your project ID.
+		 */
+		private String projectId;
+
+		/**
+		 * A location is a
+		 * <a href="https://cloud.google.com/about/locations?hl=en">region</a> you can
+		 * specify in a request to control where data is stored at rest. For a list of
+		 * available regions, see <a href=
+		 * "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations?hl=en">Generative
+		 * AI on Vertex AI locations</a>.
+		 */
+		private String location;
+
+		/**
+		 *
+		 */
+		private String publisher;
+
+		/**
+		 * Allows the connection settings to be customised
+		 */
+		private PredictionServiceSettings predictionServiceSettings;
+
+		public Builder apiEndpoint(String endpoint) {
+			this.endpoint = endpoint;
+			return this;
+		}
+
+		public Builder projectId(String projectId) {
+			this.projectId = projectId;
+			return this;
+		}
+
+		public Builder location(String location) {
+			this.location = location;
+			return this;
+		}
+
+		public Builder publisher(String publisher) {
+			this.publisher = publisher;
+			return this;
+		}
+
+		public Builder predictionServiceSettings(PredictionServiceSettings predictionServiceSettings) {
+			this.predictionServiceSettings = predictionServiceSettings;
+			return this;
+		}
+
+		public VertexAiImagenConnectionDetails build() {
+			if (!StringUtils.hasText(this.endpoint)) {
+				if (!StringUtils.hasText(this.location)) {
+					this.endpoint = DEFAULT_ENDPOINT;
+					this.location = DEFAULT_LOCATION;
+				} else {
+					this.endpoint = this.location + DEFAULT_ENDPOINT_SUFFIX;
+				}
+			}
+
+			if (!StringUtils.hasText(this.publisher)) {
+				this.publisher = DEFAULT_PUBLISHER;
+			}
+
+			if (this.predictionServiceSettings == null) {
+				try {
+					this.predictionServiceSettings = PredictionServiceSettings.newBuilder()
+							.setEndpoint(this.endpoint)
+							.build();
+				} catch (IOException e) {
+					throw new RuntimeException(e);
+				}
+			}
+
+			return new VertexAiImagenConnectionDetails(this.projectId, this.location, this.publisher,
+					this.predictionServiceSettings);
+		}
+
+	}
+
+}