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)(chat) Support integrating embeddingStore into system settings. (#1405)

Browse Source
This commit is contained in:
lexluo09
2024-07-14 21:50:12 +08:00
committed by GitHub
parent 4eb6193699
commit d5c78d87e7
9 changed files with 41 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
common/src/main/java/dev/langchain4j/chroma/spring/ChromaAutoConfig.java
View File

@@ -16,6 +16,6 @@ public class ChromaAutoConfig {
@Bean
@ConditionalOnProperty(PREFIX + ".embedding-store.base-url")
EmbeddingStoreFactory chromaChatModel(Properties properties) {
return new ChromaEmbeddingStoreFactory(properties);
return new ChromaEmbeddingStoreFactory(properties.getEmbeddingStore());
}
}

13
common/src/main/java/dev/langchain4j/chroma/spring/ChromaEmbeddingStoreFactory.java
View File

@@ -12,19 +12,18 @@ import java.time.Duration;
@Slf4j
public class ChromaEmbeddingStoreFactory extends BaseEmbeddingStoreFactory {
private Properties properties;
private EmbeddingStoreProperties storeProperties;
public ChromaEmbeddingStoreFactory(EmbeddingStoreConfig storeConfig) {
this(createPropertiesFromConfig(storeConfig));
}
public ChromaEmbeddingStoreFactory(Properties properties) {
this.properties = properties;
public ChromaEmbeddingStoreFactory(EmbeddingStoreProperties storeProperties) {
this.storeProperties = storeProperties;
}
@Override
public EmbeddingStore createEmbeddingStore(String collectionName) {
EmbeddingStoreProperties storeProperties = properties.getEmbeddingStore();
return ChromaEmbeddingStore.builder()
.baseUrl(storeProperties.getBaseUrl())
.collectionName(collectionName)
@@ -32,12 +31,10 @@ public class ChromaEmbeddingStoreFactory extends BaseEmbeddingStoreFactory {
.build();
}
private static Properties createPropertiesFromConfig(EmbeddingStoreConfig storeConfig) {
Properties properties = new Properties();
private static EmbeddingStoreProperties createPropertiesFromConfig(EmbeddingStoreConfig storeConfig) {
EmbeddingStoreProperties embeddingStore = new EmbeddingStoreProperties();
BeanUtils.copyProperties(storeConfig, embeddingStore);
embeddingStore.setTimeout(Duration.ofSeconds(storeConfig.getTimeOut()));
properties.setEmbeddingStore(embeddingStore);
return properties;
return embeddingStore;
}
}

2
common/src/main/java/dev/langchain4j/chroma/spring/EmbeddingStoreProperties.java
View File

@@ -2,11 +2,13 @@ package dev.langchain4j.chroma.spring;
import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
import java.time.Duration;
@Getter
@Setter
@ToString
public class EmbeddingStoreProperties {
private String baseUrl;

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

@@ -6,7 +6,7 @@ import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
public abstract class BaseEmbeddingStoreFactory implements EmbeddingStoreFactory {
protected static final Map<String, EmbeddingStore<TextSegment>> collectionNameToStore = new ConcurrentHashMap<>();
protected final Map<String, EmbeddingStore<TextSegment>> collectionNameToStore = new ConcurrentHashMap<>();
public EmbeddingStore<TextSegment> create(String collectionName) {
return collectionNameToStore.computeIfAbsent(collectionName, this::createEmbeddingStore);

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

@@ -1,5 +1,6 @@
package dev.langchain4j.store.embedding;
import com.tencent.supersonic.common.config.EmbeddingStoreParameterConfig;
import com.tencent.supersonic.common.pojo.EmbeddingStoreConfig;
import com.tencent.supersonic.common.util.ContextUtils;
import dev.langchain4j.chroma.spring.ChromaEmbeddingStoreFactory;
@@ -7,20 +8,33 @@ import dev.langchain4j.inmemory.spring.InMemoryEmbeddingStoreFactory;
import dev.langchain4j.milvus.spring.MilvusEmbeddingStoreFactory;
import org.apache.commons.lang3.StringUtils;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
public class EmbeddingStoreFactoryProvider {
public static EmbeddingStoreFactory getFactory(EmbeddingStoreConfig storeConfig) {
if (storeConfig == null || StringUtils.isBlank(storeConfig.getProvider())) {
protected static final Map<EmbeddingStoreConfig, EmbeddingStoreFactory> factoryMap = new ConcurrentHashMap<>();
public static EmbeddingStoreFactory getFactory() {
EmbeddingStoreParameterConfig parameterConfig = ContextUtils.getBean(EmbeddingStoreParameterConfig.class);
return getFactory(parameterConfig.convert());
}
public static EmbeddingStoreFactory getFactory(EmbeddingStoreConfig embeddingStoreConfig) {
if (embeddingStoreConfig == null || StringUtils.isBlank(embeddingStoreConfig.getProvider())) {
return ContextUtils.getBean(EmbeddingStoreFactory.class);
}
if (EmbeddingStoreType.CHROMA.name().equalsIgnoreCase(storeConfig.getProvider())) {
return new ChromaEmbeddingStoreFactory(storeConfig);
if (EmbeddingStoreType.CHROMA.name().equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
return factoryMap.computeIfAbsent(embeddingStoreConfig,
storeConfig -> new ChromaEmbeddingStoreFactory(storeConfig));
}
if (EmbeddingStoreType.MILVUS.name().equalsIgnoreCase(storeConfig.getProvider())) {
return new MilvusEmbeddingStoreFactory(storeConfig);
if (EmbeddingStoreType.MILVUS.name().equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
return factoryMap.computeIfAbsent(embeddingStoreConfig,
storeConfig -> new MilvusEmbeddingStoreFactory(storeConfig));
}
if (EmbeddingStoreType.IN_MEMORY.name().equalsIgnoreCase(storeConfig.getProvider())) {
return new InMemoryEmbeddingStoreFactory(storeConfig);
if (EmbeddingStoreType.IN_MEMORY.name().equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
return factoryMap.computeIfAbsent(embeddingStoreConfig,
storeConfig -> new InMemoryEmbeddingStoreFactory(storeConfig));
}
throw new RuntimeException("Unsupported EmbeddingStore provider: " + storeConfig.getProvider());
throw new RuntimeException("Unsupported EmbeddingStoreFactory provider: " + embeddingStoreConfig.getProvider());
}
}