Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Added PostgresMlAutoConfiguration (#242)

Browse Source 
Fixes issue #242
This commit is contained in:
utkarsh
2024-01-12 17:21:38 +05:30
committed by Mark Pollack
parent 5e38536b6c
commit 5718f9683d
5 changed files with 373 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
spring-ai-spring-boot-autoconfigure/pom.xml
View File

@@ -35,6 +35,13 @@
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-postgresml</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-azure-openai</artifactId>
@@ -212,6 +219,12 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-testcontainers</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers</artifactId>
@@ -219,6 +232,12 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>postgresql</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>

45
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfiguration.java Normal file
View File

@@ -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());
}
}

155
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlProperties.java Normal file
View File

@@ -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<String, Object> 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<String, Object> getKwargs() {
return kwargs;
}
public void setKwargs(Map<String, Object> 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<String, Object> 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<String, Object> getKwargs() {
return this.kwargs != null ? this.kwargs : getPostgresMlProperties().getKwargs();
}
public void setKwargs(Map<String, Object> 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;
}
}
}

87
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlAutoConfigurationIT.java Normal file
View File

@@ -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);
});
}
}

67
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/postgresml/PostgresMlPropertiesTests.java Normal file
View File

@@ -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 {
}
}