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.
This commit is contained in:
Christian Tzolov
2024-02-06 18:19:00 +01:00
parent 7b58f426ec
commit ed6a464ba8
16 changed files with 759 additions and 255 deletions

View File

@@ -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.Arrays;
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.EmbeddingOptions;
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 @@ import org.springframework.util.StringUtils;
* <a href="https://postgresml.org">PostgresML</a> EmbeddingClient
*
* @author Toshiaki Maki
* @author Christian Tzolov
*/
public class PostgresMlEmbeddingClient extends AbstractEmbeddingClient implements InitializingBean {
public static final String DEFAULT_TRANSFORMER_MODEL = "distilbert-base-uncased";
private final PostgresMlEmbeddingOptions defaultOptions;
private final JdbcTemplate jdbcTemplate;
private final String transformer;
private final VectorType vectorType;
private final String kwargs;
private final MetadataMode metadataMode;
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,7 +84,24 @@ public class PostgresMlEmbeddingClient extends AbstractEmbeddingClient implement
* @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;
}
/**
@@ -80,27 +109,32 @@ public class PostgresMlEmbeddingClient extends AbstractEmbeddingClient implement
* @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 class PostgresMlEmbeddingClient extends AbstractEmbeddingClient implement
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));
}
@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;
});
}
@Override
public EmbeddingResponse embedForResponse(List<String> texts) {
return this.call(new EmbeddingRequest(texts, EmbeddingOptions.EMPTY));
return this.embed(document.getFormattedContent(this.defaultOptions.getMetadataMode()));
}
@SuppressWarnings("null")
@Override
public EmbeddingResponse call(EmbeddingRequest request) {
final PostgresMlEmbeddingOptions optionsToUse = this.mergeOptions(request.getOptions());
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);
}
}

View File

@@ -0,0 +1,134 @@
/*
* Copyright 2024-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.util.Map;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.postgresml.PostgresMlEmbeddingClient.VectorType;
/**
* @author Christian Tzolov
*/
@JsonInclude(Include.NON_NULL)
public class PostgresMlEmbeddingOptions implements EmbeddingOptions {
// @formatter:off
/**
* The Huggingface transformer model to use for the embedding.
*/
private @JsonProperty("transformer") String transformer = PostgresMlEmbeddingClient.DEFAULT_TRANSFORMER_MODEL;
/**
* PostgresML vector type to use for the embedding.
* Two options are supported: PG_ARRAY and PG_VECTOR.
*/
private @JsonProperty("vectorType") VectorType vectorType = VectorType.PG_ARRAY;
/**
* Additional transformer specific options.
*/
private @JsonProperty("kwargs") Map<String, Object> kwargs = Map.of();
/**
* The Document metadata aggregation mode.
*/
private @JsonProperty("metadataMode") MetadataMode metadataMode = MetadataMode.EMBED;
// @formatter:on
public static Builder builder() {
return new Builder();
}
public static class Builder {
protected PostgresMlEmbeddingOptions options;
public Builder() {
this.options = new PostgresMlEmbeddingOptions();
}
public Builder withTransformer(String transformer) {
this.options.setTransformer(transformer);
return this;
}
public Builder withVectorType(VectorType vectorType) {
this.options.setVectorType(vectorType);
return this;
}
public Builder withKwargs(String kwargs) {
this.options.setKwargs(ModelOptionsUtils.objectToMap(kwargs));
return this;
}
public Builder withKwargs(Map<String, Object> kwargs) {
this.options.setKwargs(kwargs);
return this;
}
public Builder withMetadataMode(MetadataMode metadataMode) {
this.options.setMetadataMode(metadataMode);
return this;
}
public PostgresMlEmbeddingOptions build() {
return this.options;
}
}
public String getTransformer() {
return this.transformer;
}
public void setTransformer(String transformer) {
this.transformer = transformer;
}
public VectorType getVectorType() {
return this.vectorType;
}
public void setVectorType(VectorType vectorType) {
this.vectorType = vectorType;
}
public Map<String, Object> getKwargs() {
return this.kwargs;
}
public void setKwargs(Map<String, Object> kwargs) {
this.kwargs = kwargs;
}
public MetadataMode getMetadataMode() {
return metadataMode;
}
public void setMetadataMode(MetadataMode metadataMode) {
this.metadataMode = metadataMode;
}
}

View File

@@ -1,3 +1,18 @@
/*
* Copyright 2024-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.time.Duration;
@@ -11,7 +26,12 @@ import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.postgresml.PostgresMlEmbeddingClient.VectorType;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.containers.wait.strategy.LogMessageWaitStrategy;
import org.testcontainers.junit.jupiter.Container;
@@ -32,7 +52,6 @@ import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Toshiaki Maki
*/
@JdbcTest(properties = "logging.level.sql=TRACE")
@AutoConfigureTestDatabase(replace = AutoConfigureTestDatabase.Replace.NONE)
@Testcontainers
@@ -42,7 +61,7 @@ class PostgresMlEmbeddingClientIT {
@Container
@ServiceConnection
static PostgreSQLContainer<?> postgres = new PostgreSQLContainer<>(
DockerImageName.parse("ghcr.io/postgresml/postgresml:2.7.3").asCompatibleSubstituteFor("postgres"))
DockerImageName.parse("ghcr.io/postgresml/postgresml:2.8.1").asCompatibleSubstituteFor("postgres"))
.withCommand("sleep", "infinity")
.withLabel("org.springframework.boot.service-connection", "postgres")
.withUsername("postgresml")
@@ -63,53 +82,69 @@ class PostgresMlEmbeddingClientIT {
void embed() {
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate);
embeddingClient.afterPropertiesSet();
List<Double> embed = embeddingClient.embed("Hello World!");
assertThat(embed).hasSize(768);
// embeddingClient.dropPgmlExtension();
}
@Test
void embedWithPgVector() {
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
"distilbert-base-uncased", PostgresMlEmbeddingClient.VectorType.PG_VECTOR);
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(PostgresMlEmbeddingClient.VectorType.PG_VECTOR)
.build());
embeddingClient.afterPropertiesSet();
List<Double> embed = embeddingClient.embed(new Document("Hello World!"));
assertThat(embed).hasSize(768);
// embeddingClient.dropPgmlExtension();
}
@Test
void embedWithDifferentModel() {
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
"intfloat/e5-small");
PostgresMlEmbeddingOptions.builder().withTransformer("intfloat/e5-small").build());
embeddingClient.afterPropertiesSet();
List<Double> embed = embeddingClient.embed(new Document("Hello World!"));
assertThat(embed).hasSize(384);
// embeddingClient.dropPgmlExtension();
}
@Test
void embedWithKwargs() {
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
"distilbert-base-uncased", PostgresMlEmbeddingClient.VectorType.PG_ARRAY, Map.of("device", "cpu"),
MetadataMode.EMBED);
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(PostgresMlEmbeddingClient.VectorType.PG_ARRAY)
.withKwargs(Map.of("device", "cpu"))
.withMetadataMode(MetadataMode.EMBED)
.build());
embeddingClient.afterPropertiesSet();
List<Double> embed = embeddingClient.embed(new Document("Hello World!"));
assertThat(embed).hasSize(768);
// embeddingClient.dropPgmlExtension();
}
@ParameterizedTest
@ValueSource(strings = { "PG_ARRAY", "PG_VECTOR" })
void embedForResponse(String vectorType) {
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
"distilbert-base-uncased", PostgresMlEmbeddingClient.VectorType.valueOf(vectorType));
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(VectorType.valueOf(vectorType))
.build());
embeddingClient.afterPropertiesSet();
EmbeddingResponse embeddingResponse = embeddingClient
.embedForResponse(List.of("Hello World!", "Spring AI!", "LLM!"));
assertThat(embeddingResponse).isNotNull();
assertThat(embeddingResponse.getResults()).hasSize(3);
assertThat(embeddingResponse.getMetadata()).containsExactlyEntriesOf(
assertThat(embeddingResponse.getMetadata()).containsExactlyInAnyOrderEntriesOf(
Map.of("transformer", "distilbert-base-uncased", "vector-type", vectorType, "kwargs", "{}"));
assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(768);
@@ -117,7 +152,55 @@ class PostgresMlEmbeddingClientIT {
assertThat(embeddingResponse.getResults().get(1).getOutput()).hasSize(768);
assertThat(embeddingResponse.getResults().get(2).getIndex()).isEqualTo(2);
assertThat(embeddingResponse.getResults().get(2).getOutput()).hasSize(768);
// embeddingClient.dropPgmlExtension();
}
@Test
void embedCallWithRequestOptionsOverride() {
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(VectorType.PG_VECTOR)
.build());
embeddingClient.afterPropertiesSet();
var request1 = new EmbeddingRequest(List.of("Hello World!", "Spring AI!", "LLM!"), EmbeddingOptions.EMPTY);
EmbeddingResponse embeddingResponse = embeddingClient.call(request1);
assertThat(embeddingResponse).isNotNull();
assertThat(embeddingResponse.getResults()).hasSize(3);
assertThat(embeddingResponse.getMetadata()).containsExactlyInAnyOrderEntriesOf(Map.of("transformer",
"distilbert-base-uncased", "vector-type", VectorType.PG_VECTOR.name(), "kwargs", "{}"));
assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(768);
assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1);
assertThat(embeddingResponse.getResults().get(1).getOutput()).hasSize(768);
assertThat(embeddingResponse.getResults().get(2).getIndex()).isEqualTo(2);
assertThat(embeddingResponse.getResults().get(2).getOutput()).hasSize(768);
// Override the default options in the request
var request2 = new EmbeddingRequest(List.of("Hello World!", "Spring AI!", "LLM!"),
PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withVectorType(VectorType.PG_ARRAY)
.withMetadataMode(MetadataMode.EMBED)
.withKwargs(Map.of("device", "cpu"))
.build());
embeddingResponse = embeddingClient.call(request2);
assertThat(embeddingResponse).isNotNull();
assertThat(embeddingResponse.getResults()).hasSize(3);
assertThat(embeddingResponse.getMetadata()).containsExactlyInAnyOrderEntriesOf(Map.of("transformer",
"intfloat/e5-small", "vector-type", VectorType.PG_ARRAY.name(), "kwargs", "{\"device\":\"cpu\"}"));
assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(384);
assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1);
assertThat(embeddingResponse.getResults().get(1).getOutput()).hasSize(384);
assertThat(embeddingResponse.getResults().get(2).getIndex()).isEqualTo(2);
assertThat(embeddingResponse.getResults().get(2).getOutput()).hasSize(384);
}
@Test

View File

@@ -0,0 +1,98 @@
/*
* Copyright 2024-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.util.Map;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.jdbc.core.JdbcTemplate;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
public class PostgresMlEmbeddingOptionsTests {
@Test
public void defaultOptions() {
PostgresMlEmbeddingOptions options = PostgresMlEmbeddingOptions.builder().build();
assertThat(options.getTransformer()).isEqualTo(PostgresMlEmbeddingClient.DEFAULT_TRANSFORMER_MODEL);
assertThat(options.getVectorType()).isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_ARRAY);
assertThat(options.getKwargs()).isEqualTo(Map.of());
assertThat(options.getMetadataMode()).isEqualTo(org.springframework.ai.document.MetadataMode.EMBED);
}
@Test
public void newOptions() {
PostgresMlEmbeddingOptions options = PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withVectorType(PostgresMlEmbeddingClient.VectorType.PG_VECTOR)
.withMetadataMode(org.springframework.ai.document.MetadataMode.ALL)
.withKwargs(Map.of("device", "cpu"))
.build();
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
assertThat(options.getVectorType()).isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_VECTOR);
assertThat(options.getKwargs()).isEqualTo(Map.of("device", "cpu"));
assertThat(options.getMetadataMode()).isEqualTo(org.springframework.ai.document.MetadataMode.ALL);
}
@Test
public void mergeOptions() {
var jdbcTemplate = Mockito.mock(JdbcTemplate.class);
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(jdbcTemplate);
PostgresMlEmbeddingOptions options = embeddingClient.mergeOptions(EmbeddingOptions.EMPTY);
// Default options
assertThat(options.getTransformer()).isEqualTo(PostgresMlEmbeddingClient.DEFAULT_TRANSFORMER_MODEL);
assertThat(options.getVectorType()).isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_ARRAY);
assertThat(options.getKwargs()).isEqualTo(Map.of());
assertThat(options.getMetadataMode()).isEqualTo(org.springframework.ai.document.MetadataMode.EMBED);
// Partial override
options = embeddingClient.mergeOptions(PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withKwargs(Map.of("device", "cpu"))
.build());
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
assertThat(options.getVectorType()).isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_ARRAY); // Default
assertThat(options.getKwargs()).isEqualTo(Map.of("device", "cpu"));
assertThat(options.getMetadataMode()).isEqualTo(org.springframework.ai.document.MetadataMode.EMBED); // Default
// Complete override
options = embeddingClient.mergeOptions(PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withVectorType(PostgresMlEmbeddingClient.VectorType.PG_VECTOR)
.withMetadataMode(org.springframework.ai.document.MetadataMode.ALL)
.withKwargs(Map.of("device", "cpu"))
.build());
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
assertThat(options.getVectorType()).isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_VECTOR);
assertThat(options.getKwargs()).isEqualTo(Map.of("device", "cpu"));
assertThat(options.getMetadataMode()).isEqualTo(org.springframework.ai.document.MetadataMode.ALL);
}
}