topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2025-12-11 12:07:42 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

[improvement][common]Add pgvector vector library adaptation. (#1800)

Browse Source
This commit is contained in:
Zhengyang Jia
2024-10-15 09:16:32 +08:00
committed by GitHub
parent 0b71390fde
commit 1ef642d0dd
10 changed files with 169 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
common/src/main/java/dev/langchain4j/pgvector/spring/EmbeddingStoreProperties.java Normal file
View File

@@ -0,0 +1,23 @@
package dev.langchain4j.pgvector.spring;
import dev.langchain4j.store.embedding.pgvector.MetadataStorageConfig;
import lombok.Getter;
import lombok.Setter;
@Getter
@Setter
class EmbeddingStoreProperties {
private String host;
private Integer port;
private String user;
private String password;
private String database;
private String table;
private Integer dimension;
private Boolean useIndex;
private Integer indexListSize;
private Boolean createTable;
private Boolean dropTableFirst;
private MetadataStorageConfig metadataStorageConfig;
}

20
common/src/main/java/dev/langchain4j/pgvector/spring/PgvectorAutoConfig.java Normal file
View File

@@ -0,0 +1,20 @@
package dev.langchain4j.pgvector.spring;
import dev.langchain4j.store.embedding.EmbeddingStoreFactory;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import static dev.langchain4j.pgvector.spring.Properties.PREFIX;
@Configuration
@EnableConfigurationProperties(Properties.class)
public class PgvectorAutoConfig {
@Bean
@ConditionalOnProperty(PREFIX + ".embedding-store.host")
EmbeddingStoreFactory pgvectorChatModel(Properties properties) {
return new PgvectorEmbeddingStoreFactory(properties.getEmbeddingStore());
}
}

46
common/src/main/java/dev/langchain4j/pgvector/spring/PgvectorEmbeddingStoreFactory.java Normal file
View File

@@ -0,0 +1,46 @@
package dev.langchain4j.pgvector.spring;
import com.tencent.supersonic.common.pojo.EmbeddingStoreConfig;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.BaseEmbeddingStoreFactory;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import org.springframework.beans.BeanUtils;
public class PgvectorEmbeddingStoreFactory extends BaseEmbeddingStoreFactory {
private final EmbeddingStoreProperties storeProperties;
public PgvectorEmbeddingStoreFactory(EmbeddingStoreConfig storeConfig) {
this(createPropertiesFromConfig(storeConfig));
}
public PgvectorEmbeddingStoreFactory(EmbeddingStoreProperties storeProperties) {
this.storeProperties = storeProperties;
}
private static EmbeddingStoreProperties createPropertiesFromConfig(
EmbeddingStoreConfig storeConfig) {
EmbeddingStoreProperties embeddingStore = new EmbeddingStoreProperties();
BeanUtils.copyProperties(storeConfig, embeddingStore);
embeddingStore.setHost(storeConfig.getBaseUrl());
embeddingStore.setPort(storeConfig.getPost());
embeddingStore.setDatabase(storeConfig.getDatabaseName());
embeddingStore.setUser(storeConfig.getUser());
embeddingStore.setPassword(storeConfig.getApiKey());
return embeddingStore;
}
@Override
public EmbeddingStore<TextSegment> createEmbeddingStore(String collectionName) {
return PgVectorEmbeddingStore.builder()
.host(storeProperties.getHost())
.port(storeProperties.getPort())
.database(storeProperties.getDatabase())
.user(storeProperties.getUser())
.password(storeProperties.getPassword())
.table(collectionName)
.dimension(storeProperties.getDimension())
.build();
}
}

16
common/src/main/java/dev/langchain4j/pgvector/spring/Properties.java Normal file
View File

@@ -0,0 +1,16 @@
package dev.langchain4j.pgvector.spring;
import lombok.Getter;
import lombok.Setter;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.context.properties.NestedConfigurationProperty;
@Getter
@Setter
@ConfigurationProperties(prefix = Properties.PREFIX)
public class Properties {
static final String PREFIX = "langchain4j.pgvector";
@NestedConfigurationProperty EmbeddingStoreProperties embeddingStore;
}

6
common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java
View File

@@ -6,6 +6,7 @@ import com.tencent.supersonic.common.util.ContextUtils;
import dev.langchain4j.chroma.spring.ChromaEmbeddingStoreFactory;
import dev.langchain4j.inmemory.spring.InMemoryEmbeddingStoreFactory;
import dev.langchain4j.milvus.spring.MilvusEmbeddingStoreFactory;
import dev.langchain4j.pgvector.spring.PgvectorEmbeddingStoreFactory;
import org.apache.commons.lang3.StringUtils;
import java.util.Map;
@@ -34,6 +35,11 @@ public class EmbeddingStoreFactoryProvider {
return factoryMap.computeIfAbsent(embeddingStoreConfig,
storeConfig -> new MilvusEmbeddingStoreFactory(storeConfig));
}
if (EmbeddingStoreType.PGVECTOR.name().equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
return factoryMap.computeIfAbsent(
embeddingStoreConfig,
storeConfig -> new PgvectorEmbeddingStoreFactory(storeConfig));
}
if (EmbeddingStoreType.IN_MEMORY.name()
.equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
return factoryMap.computeIfAbsent(embeddingStoreConfig,

2
common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreType.java
View File

@@ -1,5 +1,5 @@
package dev.langchain4j.store.embedding;
public enum EmbeddingStoreType {
IN_MEMORY, MILVUS, CHROMA
IN_MEMORY, MILVUS, CHROMA, PGVECTOR
}