From ed6a464ba8a54d19bb41676eeafc892d5f4e009c Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Tue, 6 Feb 2024 18:19:00 +0100 Subject: [PATCH] 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. --- .../postgresml/PostgresMlEmbeddingClient.java | 170 ++++++++++++------ .../PostgresMlEmbeddingOptions.java | 134 ++++++++++++++ .../PostgresMlEmbeddingClientIT.java | 109 +++++++++-- .../PostgresMlEmbeddingOptionsTests.java | 98 ++++++++++ .../TransformersEmbeddingClient.java | 2 +- pom.xml | 3 +- .../ai/model/ModelOptionsUtils.java | 14 ++ .../src/main/antora/modules/ROOT/nav.adoc | 1 + .../api/embeddings/postgresml-embeddings.adoc | 169 +++++++++++++++++ .../PostgresMlAutoConfiguration.java | 14 +- .../PostgresMlEmbeddingProperties.java | 60 +++++++ .../postgresml/PostgresMlProperties.java | 155 ---------------- ...TransformersEmbeddingClientProperties.java | 4 + .../PostgresMlAutoConfigurationIT.java | 4 +- ...> PostgresMlEmbeddingPropertiesTests.java} | 35 ++-- .../pom.xml | 42 +++++ 16 files changed, 759 insertions(+), 255 deletions(-) create mode 100644 models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptions.java create mode 100644 models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptionsTests.java create mode 100644 spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/postgresml-embeddings.adoc create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlEmbeddingProperties.java delete mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlProperties.java rename spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/{PostgresMlPropertiesTests.java => PostgresMlEmbeddingPropertiesTests.java} (50%) create mode 100644 spring-ai-spring-boot-starters/spring-ai-starter-postgresml-embedding/pom.xml diff --git a/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClient.java b/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClient.java index f14167b6b..9513577c6 100644 --- a/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClient.java +++ b/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.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; * PostgresML 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 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 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 embed(Document document) { - return this.embed(document.getFormattedContent(this.metadataMode)); - } - - @Override - public List> embed(List 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> result = new ArrayList<>(); - while (rs.next()) { - result.add(vectorType.rowMapper.mapRow(rs, -1)); - } - return result; - }); - } - - @Override - public EmbeddingResponse embedForResponse(List 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 data = new ArrayList<>(); - List> embed = this.embed(request.getInstructions()); - for (int i = 0; i < embed.size(); i++) { - data.add(new Embedding(embed.get(i), i)); + List> embed = List.of(); + + List 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> 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); } } diff --git a/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptions.java b/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptions.java new file mode 100644 index 000000000..129692730 --- /dev/null +++ b/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptions.java @@ -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 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 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 getKwargs() { + return this.kwargs; + } + + public void setKwargs(Map kwargs) { + this.kwargs = kwargs; + } + + public MetadataMode getMetadataMode() { + return metadataMode; + } + + public void setMetadataMode(MetadataMode metadataMode) { + this.metadataMode = metadataMode; + } + +} diff --git a/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClientIT.java b/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClientIT.java index 18a26372d..d68f52694 100644 --- a/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClientIT.java +++ b/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingClientIT.java @@ -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 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 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 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 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 diff --git a/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptionsTests.java b/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptionsTests.java new file mode 100644 index 000000000..62d9edeb6 --- /dev/null +++ b/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptionsTests.java @@ -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); + } + +} diff --git a/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingClient.java b/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingClient.java index 13e9acd34..954130d93 100644 --- a/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingClient.java +++ b/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingClient.java @@ -340,4 +340,4 @@ public class TransformersEmbeddingClient extends AbstractEmbeddingClient impleme return new DefaultResourceLoader().getResource(uri); } -} +} \ No newline at end of file diff --git a/pom.xml b/pom.xml index 5bd0790e7..d0336a39e 100644 --- a/pom.xml +++ b/pom.xml @@ -37,6 +37,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-weaviate-store spring-ai-spring-boot-starters/spring-ai-starter-redis spring-ai-spring-boot-starters/spring-ai-starter-neo4j-store + spring-ai-spring-boot-starters/spring-ai-starter-postgresml-embedding spring-ai-docs vector-stores/spring-ai-pgvector-store vector-stores/spring-ai-milvus-store @@ -104,7 +105,7 @@ 2.23.10 2.16.1 0.26.0 - 1.16.3 + 1.17.0 3.0.1 diff --git a/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java b/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java index 8fe0e8feb..6235e0439 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java @@ -56,6 +56,20 @@ public final class ModelOptionsUtils { } + /** + * Converts the given object to a JSON string. + * @param object the object to convert to a JSON string. + * @return the JSON string. + */ + public static String toJsonString(Object object) { + try { + return OBJECT_MAPPER.writeValueAsString(object); + } + catch (JsonProcessingException e) { + throw new RuntimeException(e); + } + } + /** * Merges the source object into the target object and returns an object represented * by the given class. The JSON property names are used to match the fields to merge. diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc index 452c7e019..f64c0ac68 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc @@ -18,6 +18,7 @@ *** xref:api/embeddings/openai-embeddings.adoc[] *** xref:api/embeddings/ollama-embeddings.adoc[] *** xref:api/embeddings/azure-openai-embeddings.adoc[] +*** xref:api/embeddings/postgresml-embeddings.adoc[] ** xref:api/vectordbs.adoc[] *** xref:api/vectordbs/azure.adoc[] *** xref:api/vectordbs/chroma.adoc[] diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/postgresml-embeddings.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/postgresml-embeddings.adoc new file mode 100644 index 000000000..67666d1da --- /dev/null +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/postgresml-embeddings.adoc @@ -0,0 +1,169 @@ += PostgresML Embeddings + +Spring AI supports the PostgresML text embeddings models. + +Embeddings are a numeric representation of text. +They are used to represent words and sentences as vectors, an array of numbers. +Embeddings can be used to find similar pieces of text, by comparing the similarity of the numeric vectors using a distance measure, or they can be used as input features for other machine learning models, since most algorithms can't use text directly. + +Many pretrained LLMs can be used to generate embeddings from text within PostgresML. +You can browse all the https://huggingface.co/models?library=sentence-transformers[models] available to find the best solution on Hugging Face. + +== Getting Started + +=== Configure the PostgresML Embeddings Client Manually + +Add the `spring-ai-postgresml` dependency to your project's Maven `pom.xml` file: +[source, xml] +---- + + org.springframework.ai + spring-ai-postgresml + 0.8.0-SNAPSHOT + +---- + +or to your Gradle `build.gradle` build file. + +[source,groovy] +---- +dependencies { + implementation 'org.springframework.ai:spring-ai-postgresml:0.8.0-SNAPSHOT' +} +---- + +Next, create an `PostgresMlEmbeddingClient` instance and use it to compute the similarity between two input texts: + +[source,java] +---- +var jdbcTemplate = new JdbcTemplate(dataSource); // your posgresml data source + +PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(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. + .build()); + +embeddingClient.afterPropertiesSet(); // initialize the jdbc template and database. + +EmbeddingResponse embeddingResponse = embeddingClient + .embedForResponse(List.of("Hello World", "World is big and salvation is near")); +---- + +NOTE: When created manually, you must call the `afterPropertiesSet()` after setting the properties and before using the client. +It is more convenient (and preferred) to create the PostgresMlEmbeddingClient as a `@Bean`. +Then you don’t have to call the `afterPropertiesSet()` manually: + +[source,java] +---- +@Bean +public EmbeddingClient embeddingClient(JdbcTemplate jdbcTemplate) { + + return new PostgresMlEmbeddingClient(jdbcTemplate, + PostgresMlEmbeddingOptions.builder() + .withTransformer("distilbert-base-uncased") + .withVectorType(VectorType.PG_VECTOR) + .withKwargs(Map.of("device", "cpu")) + .withMetadataMode(MetadataMode.EMBED) + .build()); +} +---- + +==== OpenAiEmbeddingOptions + +Use the https://github.com/spring-projects/spring-ai/blob/main/models/spring-ai-openai/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptions.java[PostgresMlEmbeddingOptions.java] to configure the `PostgresMlEmbeddingClient` with options, such as the model to use and etc. + + +On start you can pass a `PostgresMlEmbeddingOptions` to the `PostgresMlEmbeddingClient` constructor to configure the default options used for all embedding requests. + +At run-time you can override the default options, using a `PostgresMlEmbeddingOptions` in your `EmbeddingRequest`. + +For example to override the default model name for a specific request: + +[source,java] +---- + +EmbeddingResponse embeddingResponse = embeddingClient.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")) + .build())); +---- + +=== PostgresMlEmbeddingClient Auto-configuration + +Spring AI provides Spring Boot auto-configuration for the Azure PostgresML Embedding Client. +To enable it add the following dependency to your project's Maven `pom.xml` file: + +[source, xml] +---- + + org.springframework.ai + spring-ai-postgresml-spring-boot-starter + 0.8.0-SNAPSHOT + +---- + +or to your Gradle `build.gradle` build file. + +[source,groovy] +---- +dependencies { + implementation 'org.springframework.ai:spring-ai-postgresml-spring-boot-starter:0.8.0-SNAPSHOT' +} +---- + +NOTE: Refer to the xref:getting-started.adoc#_dependency_management[Dependency Management] section to add Milestone and/or Snapshot Repositories to your build file. + +Use the `spring.ai.postgresml.embedding.options.*` properties to configure your `PostgresMlEmbeddingClient`. links + +==== Sample Embedding Controller + +This will create a `EmbeddingClient` implementation that you can inject into your class. +Here is an example of a simple `@Controller` class that uses the `EmbeddingClient` implementation. + +[source,application.properties] +---- +spring.ai.postgres.embedding.options.transformer=distilbert-base-uncased +spring.ai.postgres.embedding.options.vectorType=PG_ARRAY +spring.ai.postgres.embedding.options.metadataMode=EMBED +spring.ai.postgres.embedding.options.kwargs.device=cpu +---- + +[source,java] +---- +@RestController +public class EmbeddingController { + + private final EmbeddingClient embeddingClient; + + @Autowired + public EmbeddingController(EmbeddingClient embeddingClient) { + this.embeddingClient = embeddingClient; + } + + @GetMapping("/ai/embedding") + public Map embed(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) { + EmbeddingResponse embeddingResponse = this.embeddingClient.embedForResponse(List.of(message)); + return Map.of("embedding", embeddingResponse); + } +} +---- + +== PostgresML Embedding Properties + +The prefix `spring.ai.postgres.embedding` is property prefix that configures the `EmbeddingClient` implementation for PostgresML embeddings. + +[cols="3,5,1"] +|==== +| Property | Description | Default +| spring.ai.postgres.embedding.options.transformer | The Huggingface transformer model to use for the embedding. | distilbert-base-uncased +| spring.ai.postgres.embedding.options.kwargs | Additional transformer specific options. | empty map +| spring.ai.postgres.embedding.options.vectorType | PostgresML vector type to use for the embedding. Two options are supported: `PG_ARRAY` and `PG_VECTOR`. | PG_ARRAY +| spring.ai.postgres.embedding.options.metadataMode | Document metadata aggregation mode | EMBED +|==== + diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfiguration.java index 003670cc9..2fe0277b1 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfiguration.java @@ -1,5 +1,5 @@ /* - * Copyright 2024 the original author or authors. + * 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. @@ -27,19 +27,21 @@ import org.springframework.jdbc.core.JdbcTemplate; /** * Auto-configuration class for PostgresMlEmbeddingClient. + * + * @author Utkarsh Srivastava + * @author Christian Tzolov */ @AutoConfiguration(after = JdbcTemplateAutoConfiguration.class) @ConditionalOnClass(PostgresMlEmbeddingClient.class) -@EnableConfigurationProperties(PostgresMlProperties.class) +@EnableConfigurationProperties(PostgresMlEmbeddingProperties.class) public class PostgresMlAutoConfiguration { @Bean @ConditionalOnMissingBean public EmbeddingClient postgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, - PostgresMlProperties postgresMlProperties) { - return new PostgresMlEmbeddingClient(jdbcTemplate, postgresMlProperties.getEmbedding().getTransformer(), - postgresMlProperties.getEmbedding().getVectorType(), postgresMlProperties.getEmbedding().getKwargs(), - postgresMlProperties.getEmbedding().getMetadataMode()); + PostgresMlEmbeddingProperties embeddingProperties) { + + return new PostgresMlEmbeddingClient(jdbcTemplate, embeddingProperties.getOptions()); } } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlEmbeddingProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlEmbeddingProperties.java new file mode 100644 index 000000000..bda214b02 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlEmbeddingProperties.java @@ -0,0 +1,60 @@ +/* + * 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.autoconfigure.postgresml; + +import java.util.Map; + +import org.springframework.ai.document.MetadataMode; +import org.springframework.ai.postgresml.PostgresMlEmbeddingClient; +import org.springframework.ai.postgresml.PostgresMlEmbeddingOptions; +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.boot.context.properties.NestedConfigurationProperty; +import org.springframework.util.Assert; + +/** + * Configuration properties for Postgres ML. + * + * @author Utkarsh Srivastava + * @author Christian Tzolov + */ +@ConfigurationProperties(PostgresMlEmbeddingProperties.CONFIG_PREFIX) +public class PostgresMlEmbeddingProperties { + + public static final String CONFIG_PREFIX = "spring.ai.postgresml.embedding"; + + @NestedConfigurationProperty + private PostgresMlEmbeddingOptions options = PostgresMlEmbeddingOptions.builder() + .withTransformer(PostgresMlEmbeddingClient.DEFAULT_TRANSFORMER_MODEL) + .withVectorType(PostgresMlEmbeddingClient.VectorType.PG_ARRAY) + .withKwargs(Map.of()) + .withMetadataMode(MetadataMode.EMBED) + .build(); + + public PostgresMlEmbeddingOptions getOptions() { + return this.options; + } + + public void setOptions(PostgresMlEmbeddingOptions options) { + 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.options = options; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlProperties.java deleted file mode 100644 index f1a3d5614..000000000 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlProperties.java +++ /dev/null @@ -1,155 +0,0 @@ -/* - * 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.postgresml; - -import java.util.Collections; -import java.util.Map; - -import org.springframework.ai.document.MetadataMode; -import org.springframework.ai.postgresml.PostgresMlEmbeddingClient; -import org.springframework.boot.context.properties.ConfigurationProperties; -import org.springframework.util.Assert; -import org.springframework.util.StringUtils; - -/** - * Configuration properties for Postgres ML. - */ -@ConfigurationProperties(PostgresMlProperties.CONFIG_PREFIX) -public class PostgresMlProperties { - - public static final String CONFIG_PREFIX = "spring.ai.postgresml"; - - private final PostgresMlProperties.Embedding embedding = new PostgresMlProperties.Embedding(this); - - private String transformer = "distilbert-base-uncased"; - - private PostgresMlEmbeddingClient.VectorType vectorType = PostgresMlEmbeddingClient.VectorType.PG_ARRAY; - - private Map kwargs = Collections.emptyMap(); - - private MetadataMode metadataMode = MetadataMode.EMBED; - - public PostgresMlProperties.Embedding getEmbedding() { - return this.embedding; - } - - public String getTransformer() { - return transformer; - } - - public void setTransformer(String transformer) { - this.transformer = transformer; - } - - public PostgresMlEmbeddingClient.VectorType getVectorType() { - return vectorType; - } - - public void setVectorType(PostgresMlEmbeddingClient.VectorType vectorType) { - this.vectorType = vectorType; - } - - public Map getKwargs() { - return kwargs; - } - - public void setKwargs(Map kwargs) { - this.kwargs = kwargs; - } - - public MetadataMode getMetadataMode() { - return metadataMode; - } - - public void setMetadataMode(MetadataMode metadataMode) { - this.metadataMode = metadataMode; - } - - public static class Embedding { - - private PostgresMlProperties postgresMlProperties; - - private String transformer; - - private PostgresMlEmbeddingClient.VectorType vectorType; - - private Map kwargs; - - private MetadataMode metadataMode; - - protected Embedding(PostgresMlProperties postgresMlProperties) { - Assert.notNull(postgresMlProperties, "PostgresMlProperties must not be null"); - this.postgresMlProperties = postgresMlProperties; - } - - public PostgresMlProperties getPostgresMlProperties() { - return postgresMlProperties; - } - - public String getTransformer() { - return StringUtils.hasText(this.transformer) ? this.transformer - : getPostgresMlProperties().getTransformer(); - } - - public void setTransformer(String transformer) { - this.transformer = transformer; - } - - public PostgresMlEmbeddingClient.VectorType getVectorType() { - return this.vectorType != null ? this.vectorType : getPostgresMlProperties().getVectorType(); - } - - public void setVectorType(PostgresMlEmbeddingClient.VectorType vectorType) { - this.vectorType = vectorType; - } - - public Map getKwargs() { - return this.kwargs != null ? this.kwargs : getPostgresMlProperties().getKwargs(); - } - - public void setKwargs(Map kwargs) { - this.kwargs = kwargs; - } - - public MetadataMode getMetadataMode() { - return this.metadataMode != null ? this.metadataMode : getPostgresMlProperties().getMetadataMode(); - } - - public void setMetadataMode(MetadataMode metadataMode) { - this.metadataMode = metadataMode; - } - - } - - public static class Metadata { - - private Boolean rateLimitMetricsEnabled; - - public boolean isRateLimitMetricsEnabled() { - return Boolean.TRUE.equals(getRateLimitMetricsEnabled()); - } - - public Boolean getRateLimitMetricsEnabled() { - return this.rateLimitMetricsEnabled; - } - - public void setRateLimitMetricsEnabled(Boolean rateLimitMetricsEnabled) { - this.rateLimitMetricsEnabled = rateLimitMetricsEnabled; - } - - } - -} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingClientProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingClientProperties.java index 6b1fa816d..1d24fe534 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingClientProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingClientProperties.java @@ -27,6 +27,7 @@ import org.springframework.ai.document.Document; import org.springframework.ai.document.MetadataMode; import org.springframework.ai.transformers.TransformersEmbeddingClient; import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.boot.context.properties.NestedConfigurationProperty; import static org.springframework.ai.autoconfigure.transformers.TransformersEmbeddingClientProperties.CONFIG_PREFIX; @@ -79,6 +80,7 @@ public class TransformersEmbeddingClientProperties { } + @NestedConfigurationProperty private final Tokenizer tokenizer = new Tokenizer(); public static class Cache { @@ -116,6 +118,7 @@ public class TransformersEmbeddingClientProperties { /** * Controls caching of remote, large resources to local file system. */ + @NestedConfigurationProperty private final Cache cache = new Cache(); public Cache getCache() { @@ -170,6 +173,7 @@ public class TransformersEmbeddingClientProperties { } + @NestedConfigurationProperty private final Onnx onnx = new Onnx(); public Onnx getOnnx() { diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfigurationIT.java index 62c6bc1a8..43bcb96dc 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfigurationIT.java @@ -1,5 +1,5 @@ /* - * Copyright 2024 the original author or authors. + * 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. @@ -52,7 +52,7 @@ public class PostgresMlAutoConfigurationIT { @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") diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlPropertiesTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlEmbeddingPropertiesTests.java similarity index 50% rename from spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlPropertiesTests.java rename to spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlEmbeddingPropertiesTests.java index eacbee536..5ca2a8d7d 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlPropertiesTests.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlEmbeddingPropertiesTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2024 the original author or authors. + * 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. @@ -30,36 +30,33 @@ import org.springframework.boot.test.context.SpringBootTest; import static org.assertj.core.api.Assertions.assertThat; /** - * Unit Tests for {@link PostgresMlProperties}. + * Unit Tests for {@link PostgresMlEmbeddingProperties}. * * @author Utkarsh Srivastava + * @author Christian Tzolov */ -@SpringBootTest(properties = { "spring.ai.postgresml.metadata-mode=all", "spring.ai.postgresml.kwargs.key1=value1", - "spring.ai.postgresml.kwargs.key2=value2", "spring.ai.postgresml.embedding.transformer=abc123" }) -class PostgresMlPropertiesTests { +@SpringBootTest(properties = { "spring.ai.postgresml.embedding.options.metadata-mode=all", + "spring.ai.postgresml.embedding.options.kwargs.key1=value1", + "spring.ai.postgresml.embedding.options.kwargs.key2=value2", + "spring.ai.postgresml.embedding.options.transformer=abc123" }) +class PostgresMlEmbeddingPropertiesTests { @Autowired - private PostgresMlProperties postgresMlProperties; + private PostgresMlEmbeddingProperties postgresMlProperties; @Test void postgresMlPropertiesAreCorrect() { assertThat(this.postgresMlProperties).isNotNull(); - assertThat(this.postgresMlProperties.getTransformer()).isEqualTo("distilbert-base-uncased"); - assertThat(this.postgresMlProperties.getVectorType()).isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_ARRAY); - assertThat(this.postgresMlProperties.getKwargs()).isEqualTo(Map.of("key1", "value1", "key2", "value2")); - assertThat(this.postgresMlProperties.getMetadataMode()).isEqualTo(MetadataMode.ALL); - - PostgresMlProperties.Embedding embedding = this.postgresMlProperties.getEmbedding(); - - assertThat(embedding).isNotNull(); - assertThat(embedding.getTransformer()).isEqualTo("abc123"); - assertThat(embedding.getVectorType()).isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_ARRAY); - assertThat(embedding.getKwargs()).isEqualTo(Map.of("key1", "value1", "key2", "value2")); - assertThat(embedding.getMetadataMode()).isEqualTo(MetadataMode.ALL); + assertThat(this.postgresMlProperties.getOptions().getTransformer()).isEqualTo("abc123"); + assertThat(this.postgresMlProperties.getOptions().getVectorType()) + .isEqualTo(PostgresMlEmbeddingClient.VectorType.PG_ARRAY); + assertThat(this.postgresMlProperties.getOptions().getKwargs()) + .isEqualTo(Map.of("key1", "value1", "key2", "value2")); + assertThat(this.postgresMlProperties.getOptions().getMetadataMode()).isEqualTo(MetadataMode.ALL); } @SpringBootConfiguration - @EnableConfigurationProperties(PostgresMlProperties.class) + @EnableConfigurationProperties(PostgresMlEmbeddingProperties.class) static class TestConfiguration { } diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-postgresml-embedding/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-postgresml-embedding/pom.xml new file mode 100644 index 000000000..d3b4e8f02 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-postgresml-embedding/pom.xml @@ -0,0 +1,42 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 0.8.0-SNAPSHOT + ../../pom.xml + + spring-ai-postgresml-spring-boot-starter + jar + Spring AI Starter - PostgresML Embedding + Spring PostgresML Embedding Auto Configuration + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + + org.springframework.boot + spring-boot-starter + + + + org.springframework.ai + spring-ai-spring-boot-autoconfigure + ${project.parent.version} + + + + org.springframework.ai + spring-ai-postgresml + ${project.parent.version} + + + +