Refactor PostgresML embedding options builder

- Deprecate the builder methods with the prefix `with`
 - Update the docs and references
This commit is contained in:
Ilayaperumal Gopinathan
2024-12-16 23:54:07 +00:00
committed by Mark Pollack
parent 95c32f05a3
commit 18eb4c3c61
5 changed files with 82 additions and 36 deletions

View File

@@ -33,6 +33,7 @@ import org.springframework.ai.postgresml.PostgresMlEmbeddingModel.VectorType;
*
* @author Christian Tzolov
* @author Thomas Vitale
* @author Ilayaperumal Gopinathan
*/
@JsonInclude(Include.NON_NULL)
public class PostgresMlEmbeddingOptions implements EmbeddingOptions {
@@ -116,26 +117,71 @@ public class PostgresMlEmbeddingOptions implements EmbeddingOptions {
this.options = new PostgresMlEmbeddingOptions();
}
public Builder transformer(String transformer) {
this.options.setTransformer(transformer);
return this;
}
public Builder vectorType(VectorType vectorType) {
this.options.setVectorType(vectorType);
return this;
}
public Builder kwargs(String kwargs) {
this.options.setKwargs(ModelOptionsUtils.objectToMap(kwargs));
return this;
}
public Builder kwargs(Map<String, Object> kwargs) {
this.options.setKwargs(kwargs);
return this;
}
public Builder metadataMode(MetadataMode metadataMode) {
this.options.setMetadataMode(metadataMode);
return this;
}
/**
* @deprecated use {@link #transformer(String)} )} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withTransformer(String transformer) {
this.options.setTransformer(transformer);
return this;
}
/**
* @deprecated use {@link #vectorType(VectorType)} )} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withVectorType(VectorType vectorType) {
this.options.setVectorType(vectorType);
return this;
}
/**
* @deprecated use {@link #kwargs(String)} )} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withKwargs(String kwargs) {
this.options.setKwargs(ModelOptionsUtils.objectToMap(kwargs));
return this;
}
/**
* @deprecated use {@link #kwargs(Map)} )} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withKwargs(Map<String, Object> kwargs) {
this.options.setKwargs(kwargs);
return this;
}
/**
* @deprecated use {@link #metadataMode(MetadataMode)} )} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withMetadataMode(MetadataMode metadataMode) {
this.options.setMetadataMode(metadataMode);
return this;

View File

@@ -90,8 +90,8 @@ class PostgresMlEmbeddingModelIT {
void embedWithPgVector() {
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
.transformer("distilbert-base-uncased")
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
.build(),
true);
embeddingModel.afterPropertiesSet();
@@ -104,7 +104,7 @@ class PostgresMlEmbeddingModelIT {
@Test
void embedWithDifferentModel() {
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
PostgresMlEmbeddingOptions.builder().withTransformer("intfloat/e5-small").build(), true);
PostgresMlEmbeddingOptions.builder().transformer("intfloat/e5-small").build(), true);
embeddingModel.afterPropertiesSet();
float[] embed = embeddingModel.embed(new Document("Hello World!"));
@@ -116,10 +116,10 @@ class PostgresMlEmbeddingModelIT {
void embedWithKwargs() {
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
.withKwargs(Map.of("device", "cpu"))
.withMetadataMode(MetadataMode.EMBED)
.transformer("distilbert-base-uncased")
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
.kwargs(Map.of("device", "cpu"))
.metadataMode(MetadataMode.EMBED)
.build(),
true);
embeddingModel.afterPropertiesSet();
@@ -134,8 +134,8 @@ class PostgresMlEmbeddingModelIT {
void embedForResponse(String vectorType) {
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(VectorType.valueOf(vectorType))
.transformer("distilbert-base-uncased")
.vectorType(VectorType.valueOf(vectorType))
.build(),
true);
embeddingModel.afterPropertiesSet();
@@ -173,8 +173,8 @@ class PostgresMlEmbeddingModelIT {
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased")
.withVectorType(VectorType.PG_VECTOR)
.transformer("distilbert-base-uncased")
.vectorType(VectorType.PG_VECTOR)
.build(),
true);
embeddingModel.afterPropertiesSet();
@@ -211,10 +211,10 @@ class PostgresMlEmbeddingModelIT {
// 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"))
.transformer("intfloat/e5-small")
.vectorType(VectorType.PG_ARRAY)
.metadataMode(MetadataMode.EMBED)
.kwargs(Map.of("device", "cpu"))
.build());
embeddingResponse = embeddingModel.call(request2);

View File

@@ -44,10 +44,10 @@ public class PostgresMlEmbeddingOptionsTests {
@Test
public void newOptions() {
PostgresMlEmbeddingOptions options = PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
.withMetadataMode(org.springframework.ai.document.MetadataMode.ALL)
.withKwargs(Map.of("device", "cpu"))
.transformer("intfloat/e5-small")
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
.metadataMode(org.springframework.ai.document.MetadataMode.ALL)
.kwargs(Map.of("device", "cpu"))
.build();
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
@@ -72,8 +72,8 @@ public class PostgresMlEmbeddingOptionsTests {
// Partial override
options = embeddingModel.mergeOptions(PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withKwargs(Map.of("device", "cpu"))
.transformer("intfloat/e5-small")
.kwargs(Map.of("device", "cpu"))
.build());
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
@@ -83,10 +83,10 @@ public class PostgresMlEmbeddingOptionsTests {
// Complete override
options = embeddingModel.mergeOptions(PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
.withMetadataMode(org.springframework.ai.document.MetadataMode.ALL)
.withKwargs(Map.of("device", "cpu"))
.transformer("intfloat/e5-small")
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
.metadataMode(org.springframework.ai.document.MetadataMode.ALL)
.kwargs(Map.of("device", "cpu"))
.build());
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");

View File

@@ -77,9 +77,9 @@ For example to override the default model name for a specific request:
EmbeddingResponse embeddingResponse = embeddingModel.call(
new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
PostgresMlEmbeddingOptions.builder()
.withTransformer("intfloat/e5-small")
.withVectorType(VectorType.PG_ARRAY)
.withKwargs(Map.of("device", "gpu"))
.transformer("intfloat/e5-small")
.vectorType(VectorType.PG_ARRAY)
.kwargs(Map.of("device", "gpu"))
.build()));
----
@@ -148,10 +148,10 @@ var jdbcTemplate = new JdbcTemplate(dataSource); // your posgresml data source
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
PostgresMlEmbeddingOptions.builder()
.withTransformer("distilbert-base-uncased") // huggingface transformer model name.
.withVectorType(VectorType.PG_VECTOR) //vector type in PostgreSQL.
.withKwargs(Map.of("device", "cpu")) // optional arguments.
.withMetadataMode(MetadataMode.EMBED) // Document metadata mode.
.transformer("distilbert-base-uncased") // huggingface transformer model name.
.vectorType(VectorType.PG_VECTOR) //vector type in PostgreSQL.
.kwargs(Map.of("device", "cpu")) // optional arguments.
.metadataMode(MetadataMode.EMBED) // Document metadata mode.
.build());
embeddingModel.afterPropertiesSet(); // initialize the jdbc template and database.

View File

@@ -48,10 +48,10 @@ public class PostgresMlEmbeddingProperties {
@NestedConfigurationProperty
private PostgresMlEmbeddingOptions options = PostgresMlEmbeddingOptions.builder()
.withTransformer(PostgresMlEmbeddingModel.DEFAULT_TRANSFORMER_MODEL)
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
.withKwargs(Map.of())
.withMetadataMode(MetadataMode.EMBED)
.transformer(PostgresMlEmbeddingModel.DEFAULT_TRANSFORMER_MODEL)
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
.kwargs(Map.of())
.metadataMode(MetadataMode.EMBED)
.build();
public PostgresMlEmbeddingOptions getOptions() {