Add PostgresML support for EmbeddingClient
- Implement dimensions method - Add MetadataMode support. Defaults to EMBED - Drop the pgml extension between tests. - Disable the PostgresMlEmbeddingClientIT by default. Resolves #33
This commit is contained in:
Toshiaki Maki
committed by
Christian Tzolov
Christian Tzolov
parent
7c3c8445f0
commit
caa8cb2021
@@ -0,0 +1,74 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>org.springframework.experimental.ai</groupId>
|
||||
<artifactId>spring-ai</artifactId>
|
||||
<version>0.7.0-SNAPSHOT</version>
|
||||
<relativePath>../../pom.xml</relativePath>
|
||||
</parent>
|
||||
<artifactId>spring-ai-postgresml-embedding-client</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
<name>Spring AI Embedding Client - PostgresML</name>
|
||||
<description>Spring AI PostgresML Embedding Client</description>
|
||||
<url>https://github.com/spring-projects-experimental/spring-ai</url>
|
||||
|
||||
<scm>
|
||||
<url>https://github.com/spring-projects-experimental/spring-ai</url>
|
||||
<connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
|
||||
<developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
|
||||
</scm>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.springframework.experimental.ai</groupId>
|
||||
<artifactId>spring-ai-core</artifactId>
|
||||
<version>${parent.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.postgresql</groupId>
|
||||
<artifactId>postgresql</artifactId>
|
||||
<scope>runtime</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework</groupId>
|
||||
<artifactId>spring-jdbc</artifactId>
|
||||
</dependency>
|
||||
|
||||
<!-- TESTING -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-test</artifactId>
|
||||
<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>junit-jupiter</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>postgresql</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.zaxxer</groupId>
|
||||
<artifactId>HikariCP</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
@@ -0,0 +1,181 @@
|
||||
package org.springframework.ai.embedding;
|
||||
|
||||
import java.sql.Array;
|
||||
import java.sql.PreparedStatement;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
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.beans.factory.InitializingBean;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
import org.springframework.jdbc.core.RowMapper;
|
||||
import org.springframework.util.Assert;
|
||||
import org.springframework.util.CollectionUtils;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
/**
|
||||
* <a href="https://postgresml.org">PostgresML</a> EmbeddingClient
|
||||
*
|
||||
* @author Toshiaki Maki
|
||||
*/
|
||||
public class PostgresMlEmbeddingClient implements EmbeddingClient, InitializingBean {
|
||||
|
||||
private final JdbcTemplate jdbcTemplate;
|
||||
|
||||
private final String transformer;
|
||||
|
||||
private final VectorType vectorType;
|
||||
|
||||
private final String kwargs;
|
||||
|
||||
private final AtomicInteger embeddingDimensions = new AtomicInteger(-1);
|
||||
|
||||
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) -> {
|
||||
String embedding = rs.getString("embedding");
|
||||
return Arrays.stream((embedding.substring(1, embedding.length() - 1)
|
||||
/* remove leading '[' and trailing ']' */.split(","))).map(Double::parseDouble).toList();
|
||||
});
|
||||
|
||||
private final String cast;
|
||||
|
||||
private final String extensionName;
|
||||
|
||||
private final RowMapper<List<Double>> rowMapper;
|
||||
|
||||
VectorType(String cast, String extensionName, RowMapper<List<Double>> rowMapper) {
|
||||
this.cast = cast;
|
||||
this.extensionName = extensionName;
|
||||
this.rowMapper = rowMapper;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* a constructor
|
||||
* @param jdbcTemplate JdbcTemplate
|
||||
*/
|
||||
public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate) {
|
||||
this(jdbcTemplate, "distilbert-base-uncased");
|
||||
}
|
||||
|
||||
/**
|
||||
* a constructor
|
||||
* @param jdbcTemplate JdbcTemplate
|
||||
* @param transformer huggingface sentence-transformer name
|
||||
*/
|
||||
public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, String transformer) {
|
||||
this(jdbcTemplate, transformer, VectorType.PG_ARRAY);
|
||||
}
|
||||
|
||||
/**
|
||||
* a constructor
|
||||
* @param jdbcTemplate JdbcTemplate
|
||||
* @param transformer huggingface sentence-transformer name
|
||||
* @param vectorType vector type in PostgreSQL
|
||||
*/
|
||||
public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, String transformer, VectorType vectorType) {
|
||||
this(jdbcTemplate, transformer, vectorType, Map.of(), MetadataMode.EMBED);
|
||||
}
|
||||
|
||||
/**
|
||||
* a constructor
|
||||
* @param jdbcTemplate JdbcTemplate
|
||||
* @param transformer huggingface sentence-transformer name
|
||||
* @param vectorType vector type in PostgreSQL
|
||||
* @param kwargs optional arguments
|
||||
*/
|
||||
public PostgresMlEmbeddingClient(JdbcTemplate jdbcTemplate, String transformer, VectorType vectorType,
|
||||
Map<String, Object> kwargs, MetadataMode metadataMode) {
|
||||
Assert.notNull(jdbcTemplate, "jdbc template must not be null.");
|
||||
Assert.notNull(transformer, "transformer must not be null.");
|
||||
Assert.notNull(vectorType, "vectorType must not be null.");
|
||||
Assert.notNull(kwargs, "kwargs must not be null.");
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Double> embed(String text) {
|
||||
return this.jdbcTemplate.queryForObject(
|
||||
"SELECT pgml.embed(?, ?, ?::JSONB)" + this.vectorType.cast + " AS embedding", this.vectorType.rowMapper,
|
||||
this.transformer, text, this.kwargs);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Double> embed(Document document) {
|
||||
return this.embed(document.getFormattedContent(this.metadataMode));
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<List<Double>> embed(List<String> texts) {
|
||||
if (CollectionUtils.isEmpty(texts)) {
|
||||
return List.of();
|
||||
}
|
||||
return this.jdbcTemplate.query(connection -> {
|
||||
PreparedStatement preparedStatement = connection.prepareStatement("SELECT pgml.embed(?, text, ?::JSONB)"
|
||||
+ vectorType.cast + " AS embedding FROM (SELECT unnest(?) AS text) AS texts");
|
||||
preparedStatement.setString(1, transformer);
|
||||
preparedStatement.setString(2, kwargs);
|
||||
preparedStatement.setArray(3, connection.createArrayOf("TEXT", texts.toArray(Object[]::new)));
|
||||
return preparedStatement;
|
||||
}, rs -> {
|
||||
List<List<Double>> result = new ArrayList<>();
|
||||
while (rs.next()) {
|
||||
result.add(vectorType.rowMapper.mapRow(rs, -1));
|
||||
}
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public EmbeddingResponse embedForResponse(List<String> texts) {
|
||||
List<Embedding> data = new ArrayList<>();
|
||||
List<List<Double>> embed = this.embed(texts);
|
||||
for (int i = 0; i < embed.size(); i++) {
|
||||
data.add(new Embedding(embed.get(i), i));
|
||||
}
|
||||
return new EmbeddingResponse(data,
|
||||
Map.of("transformer", this.transformer, "vector-type", this.vectorType.name(), "kwargs", this.kwargs));
|
||||
}
|
||||
|
||||
@Override
|
||||
public int dimensions() {
|
||||
if (this.embeddingDimensions.get() < 0) {
|
||||
this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, this.transformer));
|
||||
}
|
||||
return this.embeddingDimensions.get();
|
||||
}
|
||||
|
||||
@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);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package org.springframework.ai.embedding;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.time.temporal.ChronoUnit;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
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.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.document.Document;
|
||||
import org.springframework.ai.document.MetadataMode;
|
||||
import org.springframework.ai.embedding.PostgresMlEmbeddingClient.VectorType;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
||||
import org.springframework.boot.test.autoconfigure.jdbc.AutoConfigureTestDatabase;
|
||||
import org.springframework.boot.test.autoconfigure.jdbc.JdbcTest;
|
||||
import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.springframework.ai.embedding.PostgresMlEmbeddingClient.VectorType.PG_ARRAY;
|
||||
import static org.springframework.ai.embedding.PostgresMlEmbeddingClient.VectorType.PG_VECTOR;
|
||||
|
||||
/**
|
||||
* @author Toshiaki Maki
|
||||
*/
|
||||
|
||||
@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)!")
|
||||
class PostgresMlEmbeddingClientIT {
|
||||
|
||||
@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;
|
||||
|
||||
@AfterEach
|
||||
void dropPgmlExtension() {
|
||||
this.jdbcTemplate.execute("DROP EXTENSION IF EXISTS pgml");
|
||||
}
|
||||
|
||||
@Test
|
||||
void embed() {
|
||||
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate);
|
||||
embeddingClient.afterPropertiesSet();
|
||||
List<Double> embed = embeddingClient.embed("Hello World!");
|
||||
assertThat(embed).hasSize(768);
|
||||
// embeddingClient.dropPgmlExtension();
|
||||
}
|
||||
|
||||
@Test
|
||||
void embedWithPgVector() {
|
||||
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
|
||||
"distilbert-base-uncased", PG_VECTOR);
|
||||
embeddingClient.afterPropertiesSet();
|
||||
List<Double> embed = embeddingClient.embed(new Document("Hello World!"));
|
||||
assertThat(embed).hasSize(768);
|
||||
// embeddingClient.dropPgmlExtension();
|
||||
}
|
||||
|
||||
@Test
|
||||
void embedWithDifferentModel() {
|
||||
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
|
||||
"intfloat/e5-small");
|
||||
embeddingClient.afterPropertiesSet();
|
||||
List<Double> embed = embeddingClient.embed(new Document("Hello World!"));
|
||||
assertThat(embed).hasSize(384);
|
||||
// embeddingClient.dropPgmlExtension();
|
||||
}
|
||||
|
||||
@Test
|
||||
void embedWithKwargs() {
|
||||
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
|
||||
"distilbert-base-uncased", PG_ARRAY, Map.of("device", "cpu"), MetadataMode.EMBED);
|
||||
embeddingClient.afterPropertiesSet();
|
||||
List<Double> embed = embeddingClient.embed(new Document("Hello World!"));
|
||||
assertThat(embed).hasSize(768);
|
||||
// embeddingClient.dropPgmlExtension();
|
||||
}
|
||||
|
||||
@ParameterizedTest
|
||||
@ValueSource(strings = { "PG_ARRAY", "PG_VECTOR" })
|
||||
void embedForResponse(String vectorType) {
|
||||
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate,
|
||||
"distilbert-base-uncased", VectorType.valueOf(vectorType));
|
||||
embeddingClient.afterPropertiesSet();
|
||||
EmbeddingResponse embeddingResponse = embeddingClient
|
||||
.embedForResponse(List.of("Hello World!", "Spring AI!", "LLM!"));
|
||||
assertThat(embeddingResponse).isNotNull();
|
||||
assertThat(embeddingResponse.getData()).hasSize(3);
|
||||
assertThat(embeddingResponse.getMetadata()).containsExactlyEntriesOf(
|
||||
Map.of("transformer", "distilbert-base-uncased", "vector-type", vectorType, "kwargs", "{}"));
|
||||
assertThat(embeddingResponse.getData().get(0).getIndex()).isEqualTo(0);
|
||||
assertThat(embeddingResponse.getData().get(0).getEmbedding()).hasSize(768);
|
||||
assertThat(embeddingResponse.getData().get(1).getIndex()).isEqualTo(1);
|
||||
assertThat(embeddingResponse.getData().get(1).getEmbedding()).hasSize(768);
|
||||
assertThat(embeddingResponse.getData().get(2).getIndex()).isEqualTo(2);
|
||||
assertThat(embeddingResponse.getData().get(2).getEmbedding()).hasSize(768);
|
||||
// embeddingClient.dropPgmlExtension();
|
||||
}
|
||||
|
||||
@Test
|
||||
void dimensions() {
|
||||
PostgresMlEmbeddingClient embeddingClient = new PostgresMlEmbeddingClient(this.jdbcTemplate);
|
||||
embeddingClient.afterPropertiesSet();
|
||||
assertThat(embeddingClient.dimensions()).isEqualTo(768);
|
||||
// cached
|
||||
assertThat(embeddingClient.dimensions()).isEqualTo(768);
|
||||
}
|
||||
|
||||
@SpringBootApplication
|
||||
public static class TestApplication {
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user