Add options support to PostgresMlEmbeddingClient

- Add postgremaddmbedding adoc page.
- Auto-configuration:
  - add missing  boot-starter.
  - refactor autoconf class and properties to accomodate the PostgresMlEmbeddingOptions.
- PostgesMlEmbeddingClient
  - Add the (default) options field and remove old fields.
  - Implement default and request options merging.
  - Add tests for options and merging.
- Remove redundant code.
- Code style fixes.

Author: tzolov

models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClient.java

@@ -1,3 +1,18 @@
+/*
+ * Copyright 2023-2024 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.postgresml;
 
 import java.sql.Array;
@@ -7,9 +22,6 @@
 import java.util.List;
 import java.util.Map;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.ObjectMapper;
-
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.MetadataMode;
 import org.springframework.ai.embedding.AbstractEmbeddingClient;
@@ -18,6 +30,7 @@
 import org.springframework.ai.embedding.EmbeddingRequest;
 import org.springframework.ai.embedding.EmbeddingResponse;
 import org.springframework.ai.embedding.EmbeddingResponseMetadata;
+import org.springframework.ai.model.ModelOptionsUtils;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.jdbc.core.RowMapper;
@@ -29,25 +42,24 @@
  * <a href="https://postgresml.org">PostgresML</a> EmbeddingClient
  *
  * @author Toshiaki Maki
+ * @author Christian Tzolov
  */
 public class PostgresMlEmbeddingClient extends AbstractEmbeddingClient implements InitializingBean {
 
-	private final JdbcTemplate jdbcTemplate;
+	public static final String DEFAULT_TRANSFORMER_MODEL = "distilbert-base-uncased";
 
-	private final String transformer;
+	private final PostgresMlEmbeddingOptions defaultOptions;
 
-	private final VectorType vectorType;
-
-	private final String kwargs;
-
-	private final MetadataMode metadataMode;
+	private final JdbcTemplate jdbcTemplate;
 
 	public enum VectorType {
 
 		PG_ARRAY("", null, (rs, i) -> {
 			Array embedding = rs.getArray("embedding");
 			return Arrays.stream((Float[]) embedding.getArray()).map(Float::doubleValue).toList();
-		}), PG_VECTOR("::vector", "vector", (rs, i) -> {
+		}),
+
+		PG_VECTOR("::vector", "vector", (rs, i) -> {
 			String embedding = rs.getString("embedding");
 			return Arrays.stream((embedding.substring(1, embedding.length() - 1)
 				/* remove leading '[' and trailing ']' */.split(","))).map(Double::parseDouble).toList();
@@ -72,35 +84,57 @@ public enum VectorType {
 	 * @param jdbcTemplate JdbcTemplate
 	 */
 	public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate) {
-		this(jdbcTemplate, "distilbert-base-uncased");
+		this(jdbcTemplate, PostgresMlEmbeddingOptions.builder().build());
+	}
+
+	/**
+	 * a PostgresMlEmbeddingClient constructor
+	 * @param jdbcTemplate JdbcTemplate to use to interact with the database.
+	 * @param options PostgresMlEmbeddingOptions to configure the client.
+	 */
+	public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, PostgresMlEmbeddingOptions options) {
+		Assert.notNull(jdbcTemplate, "jdbc template must not be null.");
+		Assert.notNull(options, "options must not be null.");
+		Assert.notNull(options.getTransformer(), "transformer must not be null.");
+		Assert.notNull(options.getVectorType(), "vectorType must not be null.");
+		Assert.notNull(options.getKwargs(), "kwargs must not be null.");
+		Assert.notNull(options.getMetadataMode(), "metadataMode must not be null.");
+
+		this.jdbcTemplate = jdbcTemplate;
+		this.defaultOptions = options;
 	}
 
 	/**
 	 * a constructor
 	 * @param jdbcTemplate JdbcTemplate
 	 * @param transformer huggingface sentence-transformer name
 	 */
+	@Deprecated(since = "0.8.0", forRemoval = true)
 	public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, String transformer) {
 		this(jdbcTemplate, transformer, VectorType.PG_ARRAY);
 	}
 
 	/**
 	 * a constructor
+	 * @deprecated Use the constructor with {@link PostgresMlEmbeddingOptions} instead.
 	 * @param jdbcTemplate JdbcTemplate
 	 * @param transformer huggingface sentence-transformer name
 	 * @param vectorType vector type in PostgreSQL
 	 */
+	@Deprecated(since = "0.8.0", forRemoval = true)
 	public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, String transformer, VectorType vectorType) {
 		this(jdbcTemplate, transformer, vectorType, Map.of(), MetadataMode.EMBED);
 	}
 
 	/**
-	 * a constructor
+	 * a constructor * @deprecated Use the constructor with
+	 * {@link PostgresMlEmbeddingOptions} instead.
 	 * @param jdbcTemplate JdbcTemplate
 	 * @param transformer huggingface sentence-transformer name
 	 * @param vectorType vector type in PostgreSQL
 	 * @param kwargs optional arguments
 	 */
+	@Deprecated(since = "0.8.0", forRemoval = true)
 	public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, String transformer, VectorType vectorType,
 			Map<String, Object> kwargs, MetadataMode metadataMode) {
 		Assert.notNull(jdbcTemplate, "jdbc template must not be null.");
@@ -110,73 +144,93 @@ public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, String transformer,
 		Assert.notNull(metadataMode, "metadataMode must not be null.");
 
 		this.jdbcTemplate = jdbcTemplate;
-		this.transformer = transformer;
-		this.vectorType = vectorType;
-		this.metadataMode = metadataMode;
-		try {
-			this.kwargs = new ObjectMapper().writeValueAsString(kwargs);
-		}
-		catch (JsonProcessingException e) {
-			throw new IllegalArgumentException(e);
-		}
+
+		this.defaultOptions = PostgresMlEmbeddingOptions.builder()
+			.withTransformer(transformer)
+			.withVectorType(vectorType)
+			.withMetadataMode(metadataMode)
+			.withKwargs(ModelOptionsUtils.toJsonString(kwargs))
+			.build();
 	}
 
+	@SuppressWarnings("null")
 	@Override
 	public List<Double> embed(String text) {
 		return this.jdbcTemplate.queryForObject(
-				"SELECT pgml.embed(?, ?, ?::JSONB)" + this.vectorType.cast + " AS embedding", this.vectorType.rowMapper,
-				this.transformer, text, this.kwargs);
+				"SELECT pgml.embed(?, ?, ?::JSONB)" + this.defaultOptions.getVectorType().cast + " AS embedding",
+				this.defaultOptions.getVectorType().rowMapper, this.defaultOptions.getTransformer(), text,
+				this.defaultOptions.getKwargs());
 	}
 
 	@Override
 	public List<Double> embed(Document document) {
-		return this.embed(document.getFormattedContent(this.metadataMode));
+		return this.embed(document.getFormattedContent(this.defaultOptions.getMetadataMode()));
 	}
 
+	@SuppressWarnings("null")
 	@Override
-	public List<List<Double>> embed(List<String> texts) {
-		if (CollectionUtils.isEmpty(texts)) {
-			return List.of();
-		}
-		return this.jdbcTemplate.query(connection -> {
-			PreparedStatement preparedStatement = connection.prepareStatement("SELECT pgml.embed(?, text, ?::JSONB)"
-					+ vectorType.cast + " AS embedding FROM (SELECT unnest(?) AS text) AS texts");
-			preparedStatement.setString(1, transformer);
-			preparedStatement.setString(2, kwargs);
-			preparedStatement.setArray(3, connection.createArrayOf("TEXT", texts.toArray(Object[]::new)));
-			return preparedStatement;
-		}, rs -> {
-			List<List<Double>> result = new ArrayList<>();
-			while (rs.next()) {
-				result.add(vectorType.rowMapper.mapRow(rs, -1));
-			}
-			return result;
-		});
-	}
+	public EmbeddingResponse call(EmbeddingRequest request) {
 
-	@Override
-	public EmbeddingResponse embedForResponse(List<String> texts) {
-		return this.call(new EmbeddingRequest(texts, EmbeddingOptions.EMPTY));
-	}
+		final PostgresMlEmbeddingOptions optionsToUse = this.mergeOptions(request.getOptions());
 
-	@Override
-	public EmbeddingResponse call(EmbeddingRequest request) {
 		List<Embedding> data = new ArrayList<>();
-		List<List<Double>> embed = this.embed(request.getInstructions());
-		for (int i = 0; i < embed.size(); i++) {
-			data.add(new Embedding(embed.get(i), i));
+		List<List<Double>> embed = List.of();
+
+		List<String> texts = request.getInstructions();
+		if (!CollectionUtils.isEmpty(texts)) {
+			embed = this.jdbcTemplate.query(connection -> {
+				PreparedStatement preparedStatement = connection.prepareStatement("SELECT pgml.embed(?, text, ?::JSONB)"
+						+ optionsToUse.getVectorType().cast + " AS embedding FROM (SELECT unnest(?) AS text) AS texts");
+				preparedStatement.setString(1, optionsToUse.getTransformer());
+				preparedStatement.setString(2, ModelOptionsUtils.toJsonString(optionsToUse.getKwargs()));
+				preparedStatement.setArray(3, connection.createArrayOf("TEXT", texts.toArray(Object[]::new)));
+				return preparedStatement;
+			}, rs -> {
+				List<List<Double>> result = new ArrayList<>();
+				while (rs.next()) {
+					result.add(optionsToUse.getVectorType().rowMapper.mapRow(rs, -1));
+				}
+				return result;
+			});
 		}
+
+		if (!CollectionUtils.isEmpty(embed)) {
+			for (int i = 0; i < embed.size(); i++) {
+				data.add(new Embedding(embed.get(i), i));
+			}
+		}
+
 		var metadata = new EmbeddingResponseMetadata(
-				Map.of("transformer", this.transformer, "vector-type", this.vectorType.name(), "kwargs", this.kwargs));
+				Map.of("transformer", optionsToUse.getTransformer(), "vector-type", optionsToUse.getVectorType().name(),
+						"kwargs", ModelOptionsUtils.toJsonString(optionsToUse.getKwargs())));
 
 		return new EmbeddingResponse(data, metadata);
 	}
 
+	/**
+	 * Merge the default and request options.
+	 * @param requestOptions request options to merge.
+	 * @return the merged options.
+	 */
+	PostgresMlEmbeddingOptions mergeOptions(EmbeddingOptions requestOptions) {
+
+		PostgresMlEmbeddingOptions options = (this.defaultOptions != null) ? this.defaultOptions
+				: PostgresMlEmbeddingOptions.builder().build();
+
+		if (requestOptions != null && !EmbeddingOptions.EMPTY.equals(requestOptions)) {
+			options = ModelOptionsUtils.merge(requestOptions, options, PostgresMlEmbeddingOptions.class);
+		}
+
+		return options;
+	}
+
 	@Override
 	public void afterPropertiesSet() {
 		this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS pgml");
-		if (StringUtils.hasText(this.vectorType.extensionName)) {
-			this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS " + this.vectorType.extensionName);
+		this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS hstore");
+		if (StringUtils.hasText(this.defaultOptions.getVectorType().extensionName)) {
+			this.jdbcTemplate
+				.execute("CREATE EXTENSION IF NOT EXISTS " + this.defaultOptions.getVectorType().extensionName);
 		}
 	}