diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index c188fa663..f4729ad61 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -35,6 +35,13 @@ true + + org.springframework.ai + spring-ai-postgresml + ${project.parent.version} + true + + org.springframework.ai spring-ai-azure-openai @@ -212,6 +219,12 @@ test + + org.springframework.boot + spring-boot-testcontainers + test + + org.testcontainers testcontainers @@ -219,6 +232,12 @@ test + + org.testcontainers + postgresql + test + + org.testcontainers junit-jupiter 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 new file mode 100644 index 000000000..003670cc9 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfiguration.java @@ -0,0 +1,45 @@ +/* + * Copyright 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 org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.postgresml.PostgresMlEmbeddingClient; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; +import org.springframework.jdbc.core.JdbcTemplate; + +/** + * Auto-configuration class for PostgresMlEmbeddingClient. + */ +@AutoConfiguration(after = JdbcTemplateAutoConfiguration.class) +@ConditionalOnClass(PostgresMlEmbeddingClient.class) +@EnableConfigurationProperties(PostgresMlProperties.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()); + } + +} 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 new file mode 100644 index 000000000..f1a3d5614 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlProperties.java @@ -0,0 +1,155 @@ +/* + * 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/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 new file mode 100644 index 000000000..62c6bc1a8 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfigurationIT.java @@ -0,0 +1,87 @@ +/* + * Copyright 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.time.Duration; +import java.time.temporal.ChronoUnit; +import java.util.List; + +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.containers.wait.strategy.LogMessageWaitStrategy; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.postgresml.PostgresMlEmbeddingClient; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.autoconfigure.jdbc.AutoConfigureTestDatabase; +import org.springframework.boot.test.autoconfigure.jdbc.JdbcTest; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.boot.testcontainers.service.connection.ServiceConnection; +import org.springframework.jdbc.core.JdbcTemplate; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Utkarsh Srivastava + */ +@JdbcTest(properties = "logging.level.sql=TRACE") +@AutoConfigureTestDatabase(replace = AutoConfigureTestDatabase.Replace.NONE) +@Testcontainers +@Disabled("Disabled from automatic execution, as it requires an excessive amount of memory (over 9GB)!") +public class PostgresMlAutoConfigurationIT { + + @Container + @ServiceConnection + static PostgreSQLContainer postgres = new PostgreSQLContainer<>( + DockerImageName.parse("ghcr.io/postgresml/postgresml:2.7.3").asCompatibleSubstituteFor("postgres")) + .withCommand("sleep", "infinity") + .withLabel("org.springframework.boot.service-connection", "postgres") + .withUsername("postgresml") + .withPassword("postgresml") + .withDatabaseName("postgresml") + .waitingFor(new LogMessageWaitStrategy().withRegEx(".*Starting dashboard.*\\s") + .withStartupTimeout(Duration.of(60, ChronoUnit.SECONDS))); + + @Autowired + JdbcTemplate jdbcTemplate; + + @Test + void embedding() { + ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withBean(JdbcTemplate.class, () -> jdbcTemplate) + .withConfiguration(AutoConfigurations.of(PostgresMlAutoConfiguration.class)); + contextRunner.run(context -> { + PostgresMlEmbeddingClient embeddingClient = context.getBean(PostgresMlEmbeddingClient.class); + + EmbeddingResponse embeddingResponse = embeddingClient + .embedForResponse(List.of("Hello World", "World is big and salvation is near")); + assertThat(embeddingResponse.getResults()).hasSize(2); + assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty(); + assertThat(embeddingResponse.getResults().get(0).getIndex()).isZero(); + assertThat(embeddingResponse.getResults().get(1).getOutput()).isNotEmpty(); + assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1); + + assertThat(embeddingClient.dimensions()).isEqualTo(768); + }); + } + +} 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/PostgresMlPropertiesTests.java new file mode 100644 index 000000000..eacbee536 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlPropertiesTests.java @@ -0,0 +1,67 @@ +/* + * Copyright 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.junit.jupiter.api.Test; + +import org.springframework.ai.document.MetadataMode; +import org.springframework.ai.postgresml.PostgresMlEmbeddingClient; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.boot.test.context.SpringBootTest; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit Tests for {@link PostgresMlProperties}. + * + * @author Utkarsh Srivastava + */ +@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 { + + @Autowired + private PostgresMlProperties 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); + } + + @SpringBootConfiguration + @EnableConfigurationProperties(PostgresMlProperties.class) + static class TestConfiguration { + + } + +}