Add Spring Boot auto-configuration for Neo4j vector store
- resolve Neo4j auto-configuraion property expossing external API. - move neo4j auto-conf under the vectorstore parent package.
This commit is contained in:
oujingzhou
committed by
Christian Tzolov
Christian Tzolov
parent
8282ae2b74
commit
b625ab1c15
@@ -170,6 +170,14 @@
|
||||
<optional>true</optional>
|
||||
</dependency>
|
||||
|
||||
<!-- Neo4j Vector Store-->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-neo4j-store</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
<optional>true</optional>
|
||||
</dependency>
|
||||
|
||||
<!-- test dependencies -->
|
||||
|
||||
<dependency>
|
||||
@@ -223,6 +231,13 @@
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>neo4j</artifactId>
|
||||
<version>${testcontainers.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
|
||||
@@ -0,0 +1,289 @@
|
||||
/*
|
||||
* Copyright 2023-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.vectorstore.neo4j;
|
||||
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
|
||||
import java.io.File;
|
||||
import java.net.URI;
|
||||
import java.time.Duration;
|
||||
|
||||
/**
|
||||
* Properties for Neo4j driver
|
||||
*
|
||||
* @author Jingzhou Ou
|
||||
*/
|
||||
@ConfigurationProperties(Neo4jDriverProperties.CONFIG_PREFIX)
|
||||
public class Neo4jDriverProperties {
|
||||
|
||||
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.neo4j.driver";
|
||||
|
||||
/**
|
||||
* supports bolt or neo4j as schemes.
|
||||
*/
|
||||
private URI uri;
|
||||
|
||||
/**
|
||||
* optional
|
||||
*/
|
||||
private Authentication authentication = new Authentication();
|
||||
|
||||
/**
|
||||
* connection pool configuration
|
||||
*/
|
||||
private PoolSettings pool = new PoolSettings();
|
||||
|
||||
/**
|
||||
* Detailed driver configuration of the driver
|
||||
*/
|
||||
private DriverSettings config = new DriverSettings();
|
||||
|
||||
public URI getUri() {
|
||||
return this.uri;
|
||||
}
|
||||
|
||||
public void setUri(URI uri) {
|
||||
this.uri = uri;
|
||||
}
|
||||
|
||||
public Authentication getAuthentication() {
|
||||
return this.authentication;
|
||||
}
|
||||
|
||||
public void setAuthentication(Authentication authentication) {
|
||||
this.authentication = authentication;
|
||||
}
|
||||
|
||||
public PoolSettings getPool() {
|
||||
return this.pool;
|
||||
}
|
||||
|
||||
public void setPool(PoolSettings pool) {
|
||||
this.pool = pool;
|
||||
}
|
||||
|
||||
public DriverSettings getConfig() {
|
||||
return this.config;
|
||||
}
|
||||
|
||||
public void setConfig(DriverSettings config) {
|
||||
this.config = config;
|
||||
}
|
||||
|
||||
public static class Authentication {
|
||||
|
||||
private String username;
|
||||
|
||||
private String password;
|
||||
|
||||
private String realm;
|
||||
|
||||
/**
|
||||
* kerberos authentication
|
||||
*/
|
||||
private String kerberosTicket;
|
||||
|
||||
public String getUsername() {
|
||||
return this.username;
|
||||
}
|
||||
|
||||
public void setUsername(String username) {
|
||||
this.username = username;
|
||||
}
|
||||
|
||||
public String getPassword() {
|
||||
return this.password;
|
||||
}
|
||||
|
||||
public void setPassword(String password) {
|
||||
this.password = password;
|
||||
}
|
||||
|
||||
public String getRealm() {
|
||||
return this.realm;
|
||||
}
|
||||
|
||||
public void setRealm(String realm) {
|
||||
this.realm = realm;
|
||||
}
|
||||
|
||||
public String getKerberosTicket() {
|
||||
return this.kerberosTicket;
|
||||
}
|
||||
|
||||
public void setKerberosTicket(String kerberosTicket) {
|
||||
this.kerberosTicket = kerberosTicket;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
public static class PoolSettings {
|
||||
|
||||
private boolean metricsEnabled = false;
|
||||
|
||||
private boolean logLeakedSessions = false;
|
||||
|
||||
private int maxConnectionPoolSize = org.neo4j.driver.internal.async.pool.PoolSettings.DEFAULT_MAX_CONNECTION_POOL_SIZE;
|
||||
|
||||
private Duration idleTimeBeforeConnectionTest;
|
||||
|
||||
private Duration maxConnectionLifetime = Duration
|
||||
.ofMillis(org.neo4j.driver.internal.async.pool.PoolSettings.DEFAULT_MAX_CONNECTION_LIFETIME);
|
||||
|
||||
private Duration connectionAcquisitionTimeout = Duration
|
||||
.ofMillis(org.neo4j.driver.internal.async.pool.PoolSettings.DEFAULT_CONNECTION_ACQUISITION_TIMEOUT);
|
||||
|
||||
public boolean isLogLeakedSessions() {
|
||||
return this.logLeakedSessions;
|
||||
}
|
||||
|
||||
public void setLogLeakedSessions(boolean logLeakedSessions) {
|
||||
this.logLeakedSessions = logLeakedSessions;
|
||||
}
|
||||
|
||||
public int getMaxConnectionPoolSize() {
|
||||
return this.maxConnectionPoolSize;
|
||||
}
|
||||
|
||||
public void setMaxConnectionPoolSize(int maxConnectionPoolSize) {
|
||||
this.maxConnectionPoolSize = maxConnectionPoolSize;
|
||||
}
|
||||
|
||||
public Duration getIdleTimeBeforeConnectionTest() {
|
||||
return this.idleTimeBeforeConnectionTest;
|
||||
}
|
||||
|
||||
public void setIdleTimeBeforeConnectionTest(Duration idleTimeBeforeConnectionTest) {
|
||||
this.idleTimeBeforeConnectionTest = idleTimeBeforeConnectionTest;
|
||||
}
|
||||
|
||||
public Duration getMaxConnectionLifetime() {
|
||||
return this.maxConnectionLifetime;
|
||||
}
|
||||
|
||||
public void setMaxConnectionLifetime(Duration maxConnectionLifetime) {
|
||||
this.maxConnectionLifetime = maxConnectionLifetime;
|
||||
}
|
||||
|
||||
public Duration getConnectionAcquisitionTimeout() {
|
||||
return this.connectionAcquisitionTimeout;
|
||||
}
|
||||
|
||||
public void setConnectionAcquisitionTimeout(Duration connectionAcquisitionTimeout) {
|
||||
this.connectionAcquisitionTimeout = connectionAcquisitionTimeout;
|
||||
}
|
||||
|
||||
public boolean isMetricsEnabled() {
|
||||
return this.metricsEnabled;
|
||||
}
|
||||
|
||||
public void setMetricsEnabled(boolean metricsEnabled) {
|
||||
this.metricsEnabled = metricsEnabled;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
public static class DriverSettings {
|
||||
|
||||
private boolean encrypted = false;
|
||||
|
||||
private TrustSettings trustSettings = new TrustSettings();
|
||||
|
||||
private Duration connectionTimeout = Duration.ofSeconds(30);
|
||||
|
||||
private Duration maxTransactionRetryTime = Duration
|
||||
.ofMillis(org.neo4j.driver.internal.retry.RetrySettings.DEFAULT.maxRetryTimeMs());
|
||||
|
||||
public boolean isEncrypted() {
|
||||
return this.encrypted;
|
||||
}
|
||||
|
||||
public void setEncrypted(boolean encrypted) {
|
||||
this.encrypted = encrypted;
|
||||
}
|
||||
|
||||
public TrustSettings getTrustSettings() {
|
||||
return this.trustSettings;
|
||||
}
|
||||
|
||||
public void setTrustSettings(TrustSettings trustSettings) {
|
||||
this.trustSettings = trustSettings;
|
||||
}
|
||||
|
||||
public Duration getConnectionTimeout() {
|
||||
return this.connectionTimeout;
|
||||
}
|
||||
|
||||
public void setConnectionTimeout(Duration connectionTimeout) {
|
||||
this.connectionTimeout = connectionTimeout;
|
||||
}
|
||||
|
||||
public Duration getMaxTransactionRetryTime() {
|
||||
return this.maxTransactionRetryTime;
|
||||
}
|
||||
|
||||
public void setMaxTransactionRetryTime(Duration maxTransactionRetryTime) {
|
||||
this.maxTransactionRetryTime = maxTransactionRetryTime;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
public static class TrustSettings {
|
||||
|
||||
public enum Strategy {
|
||||
|
||||
TRUST_ALL_CERTIFICATES,
|
||||
|
||||
TRUST_CUSTOM_CA_SIGNED_CERTIFICATES,
|
||||
|
||||
TRUST_SYSTEM_CA_SIGNED_CERTIFICATES
|
||||
|
||||
}
|
||||
|
||||
private TrustSettings.Strategy strategy = Strategy.TRUST_SYSTEM_CA_SIGNED_CERTIFICATES;
|
||||
|
||||
private File certFile;
|
||||
|
||||
private boolean hostnameVerificationEnabled = false;
|
||||
|
||||
public TrustSettings.Strategy getStrategy() {
|
||||
return this.strategy;
|
||||
}
|
||||
|
||||
public void setStrategy(TrustSettings.Strategy strategy) {
|
||||
this.strategy = strategy;
|
||||
}
|
||||
|
||||
public File getCertFile() {
|
||||
return this.certFile;
|
||||
}
|
||||
|
||||
public void setCertFile(File certFile) {
|
||||
this.certFile = certFile;
|
||||
}
|
||||
|
||||
public boolean isHostnameVerificationEnabled() {
|
||||
return this.hostnameVerificationEnabled;
|
||||
}
|
||||
|
||||
public void setHostnameVerificationEnabled(boolean hostnameVerificationEnabled) {
|
||||
this.hostnameVerificationEnabled = hostnameVerificationEnabled;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
/*
|
||||
* Copyright 2023-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.vectorstore.neo4j;
|
||||
|
||||
import org.neo4j.driver.AuthToken;
|
||||
import org.neo4j.driver.AuthTokens;
|
||||
import org.neo4j.driver.Config;
|
||||
import org.neo4j.driver.Driver;
|
||||
import org.neo4j.driver.GraphDatabase;
|
||||
import org.neo4j.driver.internal.Scheme;
|
||||
import org.springframework.ai.embedding.EmbeddingClient;
|
||||
import org.springframework.ai.vectorstore.Neo4jVectorStore;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.boot.autoconfigure.AutoConfiguration;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
|
||||
import org.springframework.boot.context.properties.EnableConfigurationProperties;
|
||||
import org.springframework.boot.context.properties.source.InvalidConfigurationPropertyValueException;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
import java.io.File;
|
||||
import java.net.URI;
|
||||
import java.util.Locale;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
/**
|
||||
* @author Jingzhou Ou
|
||||
*/
|
||||
@AutoConfiguration
|
||||
@ConditionalOnClass({ Neo4jVectorStore.class, EmbeddingClient.class })
|
||||
@EnableConfigurationProperties({ Neo4jVectorStoreProperties.class, Neo4jDriverProperties.class })
|
||||
public class Neo4jVectorStoreAutoConfiguration {
|
||||
|
||||
@Bean
|
||||
@ConditionalOnMissingBean
|
||||
public VectorStore vectorStore(Driver driver, EmbeddingClient embeddingClient,
|
||||
Neo4jVectorStoreProperties properties) {
|
||||
Neo4jVectorStore.Neo4jVectorStoreConfig config = Neo4jVectorStore.Neo4jVectorStoreConfig.builder()
|
||||
.withDatabaseName(properties.getDatabaseName())
|
||||
.withEmbeddingDimension(properties.getEmbeddingDimension())
|
||||
.withDistanceType(properties.getDistanceType())
|
||||
.withLabel(properties.getLabel())
|
||||
.withEmbeddingProperty(properties.getEmbeddingProperty())
|
||||
.build();
|
||||
|
||||
return new Neo4jVectorStore(driver, embeddingClient, config);
|
||||
}
|
||||
|
||||
@Bean
|
||||
@ConditionalOnMissingBean(Driver.class)
|
||||
Driver neo4jDriver(Neo4jDriverProperties driverProperties) {
|
||||
AuthToken authToken = getAuthToken(driverProperties);
|
||||
Config config = getDriverConfig(driverProperties);
|
||||
return GraphDatabase.driver(driverProperties.getUri(), authToken, config);
|
||||
}
|
||||
|
||||
private Config getDriverConfig(Neo4jDriverProperties driverProperties) {
|
||||
Config.ConfigBuilder builder = Config.builder();
|
||||
buildWithPoolSettings(builder, driverProperties.getPool());
|
||||
URI uri = driverProperties.getUri();
|
||||
String scheme = uri == null ? "bolt" : uri.getScheme();
|
||||
buildWithDriverSettings(builder, driverProperties.getConfig(), isSimpleScheme(scheme));
|
||||
return builder.build();
|
||||
}
|
||||
|
||||
private AuthToken getAuthToken(Neo4jDriverProperties driverProperties) {
|
||||
String username = driverProperties.getAuthentication().getUsername();
|
||||
String password = driverProperties.getAuthentication().getPassword();
|
||||
String kerberosTicket = driverProperties.getAuthentication().getKerberosTicket();
|
||||
String realm = driverProperties.getAuthentication().getRealm();
|
||||
|
||||
boolean hasUsername = StringUtils.hasText(username);
|
||||
boolean hasPassword = StringUtils.hasText(password);
|
||||
boolean hasKerberosTicket = StringUtils.hasText(kerberosTicket);
|
||||
|
||||
if (hasUsername && hasKerberosTicket) {
|
||||
throw new InvalidConfigurationPropertyValueException("spring.ai.vectorstore.neo4j.driver.authentication",
|
||||
"username=" + username + ",kerberos-ticket=" + kerberosTicket,
|
||||
"Cannot specify both username and kerberos ticket.");
|
||||
}
|
||||
|
||||
if (hasUsername && hasPassword) {
|
||||
return AuthTokens.basic(username, password, realm);
|
||||
}
|
||||
|
||||
if (hasKerberosTicket) {
|
||||
return AuthTokens.kerberos(kerberosTicket);
|
||||
}
|
||||
|
||||
return AuthTokens.none();
|
||||
}
|
||||
|
||||
private void buildWithPoolSettings(Config.ConfigBuilder builder, Neo4jDriverProperties.PoolSettings poolSettings) {
|
||||
if (poolSettings.isLogLeakedSessions()) {
|
||||
builder.withLeakedSessionsLogging();
|
||||
}
|
||||
builder.withMaxConnectionPoolSize(poolSettings.getMaxConnectionPoolSize());
|
||||
if (poolSettings.getIdleTimeBeforeConnectionTest() != null) {
|
||||
builder.withConnectionLivenessCheckTimeout(poolSettings.getIdleTimeBeforeConnectionTest().toMillis(),
|
||||
TimeUnit.MILLISECONDS);
|
||||
}
|
||||
builder.withMaxConnectionLifetime(poolSettings.getMaxConnectionLifetime().toMillis(), TimeUnit.MILLISECONDS);
|
||||
builder.withConnectionAcquisitionTimeout(poolSettings.getConnectionAcquisitionTimeout().toMillis(),
|
||||
TimeUnit.MILLISECONDS);
|
||||
|
||||
if (poolSettings.isMetricsEnabled()) {
|
||||
builder.withDriverMetrics();
|
||||
}
|
||||
else {
|
||||
builder.withoutDriverMetrics();
|
||||
}
|
||||
}
|
||||
|
||||
private void buildWithDriverSettings(Config.ConfigBuilder builder,
|
||||
Neo4jDriverProperties.DriverSettings driverSettings, boolean withEncryptionAndTrustSettings) {
|
||||
if (withEncryptionAndTrustSettings) {
|
||||
if (driverSettings.isEncrypted()) {
|
||||
builder.withEncryption();
|
||||
}
|
||||
else {
|
||||
builder.withoutEncryption();
|
||||
}
|
||||
builder.withTrustStrategy(getTrustStrategy(driverSettings.getTrustSettings()));
|
||||
}
|
||||
builder.withConnectionTimeout(driverSettings.getConnectionTimeout().toMillis(), TimeUnit.MILLISECONDS);
|
||||
builder.withMaxTransactionRetryTime(driverSettings.getMaxTransactionRetryTime().toMillis(),
|
||||
TimeUnit.MILLISECONDS);
|
||||
}
|
||||
|
||||
private boolean isSimpleScheme(String scheme) {
|
||||
String lowerCaseScheme = scheme.toLowerCase(Locale.ENGLISH);
|
||||
try {
|
||||
Scheme.validateScheme(lowerCaseScheme);
|
||||
}
|
||||
catch (IllegalArgumentException ex) {
|
||||
throw new IllegalArgumentException(String.format("'%s' is not a supported scheme.", scheme));
|
||||
}
|
||||
return lowerCaseScheme.equals("bolt") || lowerCaseScheme.equals("neo4j");
|
||||
}
|
||||
|
||||
private Config.TrustStrategy getTrustStrategy(Neo4jDriverProperties.TrustSettings trustSettings) {
|
||||
String propertyName = "spring.ai.vectorstore.neo4j.driver.config.trust-settings";
|
||||
Config.TrustStrategy internalRepresentation;
|
||||
File certFile = trustSettings.getCertFile();
|
||||
switch (trustSettings.getStrategy()) {
|
||||
case TRUST_ALL_CERTIFICATES:
|
||||
internalRepresentation = Config.TrustStrategy.trustAllCertificates();
|
||||
break;
|
||||
case TRUST_SYSTEM_CA_SIGNED_CERTIFICATES:
|
||||
internalRepresentation = Config.TrustStrategy.trustSystemCertificates();
|
||||
break;
|
||||
case TRUST_CUSTOM_CA_SIGNED_CERTIFICATES:
|
||||
if (certFile == null || !certFile.isFile()) {
|
||||
throw new InvalidConfigurationPropertyValueException(propertyName,
|
||||
trustSettings.getStrategy().name(),
|
||||
"Configured trust strategy requires a certificate file.");
|
||||
}
|
||||
internalRepresentation = Config.TrustStrategy.trustCustomCertificateSignedBy(certFile);
|
||||
break;
|
||||
default:
|
||||
throw new InvalidConfigurationPropertyValueException(propertyName, trustSettings.getStrategy().name(),
|
||||
"Unknown strategy.");
|
||||
}
|
||||
if (trustSettings.isHostnameVerificationEnabled()) {
|
||||
internalRepresentation.withHostnameVerification();
|
||||
}
|
||||
else {
|
||||
internalRepresentation.withoutHostnameVerification();
|
||||
}
|
||||
return internalRepresentation;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
/*
|
||||
* Copyright 2023-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.vectorstore.neo4j;
|
||||
|
||||
import org.springframework.ai.vectorstore.Neo4jVectorStore;
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
|
||||
/**
|
||||
* @author Jingzhou Ou
|
||||
*/
|
||||
@ConfigurationProperties(Neo4jVectorStoreProperties.CONFIG_PREFIX)
|
||||
public class Neo4jVectorStoreProperties {
|
||||
|
||||
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.neo4j";
|
||||
|
||||
private String databaseName;
|
||||
|
||||
private int embeddingDimension = Neo4jVectorStore.DEFAULT_EMBEDDING_DIMENSION;
|
||||
|
||||
private Neo4jVectorStore.Neo4jDistanceType distanceType = Neo4jVectorStore.Neo4jDistanceType.COSINE;
|
||||
|
||||
private String label = Neo4jVectorStore.DEFAULT_LABEL;
|
||||
|
||||
private String embeddingProperty = Neo4jVectorStore.DEFAULT_EMBEDDING_PROPERTY;
|
||||
|
||||
public String getDatabaseName() {
|
||||
return databaseName;
|
||||
}
|
||||
|
||||
public void setDatabaseName(String databaseName) {
|
||||
this.databaseName = databaseName;
|
||||
}
|
||||
|
||||
public int getEmbeddingDimension() {
|
||||
return embeddingDimension;
|
||||
}
|
||||
|
||||
public void setEmbeddingDimension(int embeddingDimension) {
|
||||
this.embeddingDimension = embeddingDimension;
|
||||
}
|
||||
|
||||
public Neo4jVectorStore.Neo4jDistanceType getDistanceType() {
|
||||
return distanceType;
|
||||
}
|
||||
|
||||
public void setDistanceType(Neo4jVectorStore.Neo4jDistanceType distanceType) {
|
||||
this.distanceType = distanceType;
|
||||
}
|
||||
|
||||
public String getLabel() {
|
||||
return label;
|
||||
}
|
||||
|
||||
public void setLabel(String label) {
|
||||
this.label = label;
|
||||
}
|
||||
|
||||
public String getEmbeddingProperty() {
|
||||
return embeddingProperty;
|
||||
}
|
||||
|
||||
public void setEmbeddingProperty(String embeddingProperty) {
|
||||
this.embeddingProperty = embeddingProperty;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
/*
|
||||
* Copyright 2023-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.vectorstore.neo4j;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.testcontainers.containers.Neo4jContainer;
|
||||
import org.testcontainers.junit.jupiter.Container;
|
||||
import org.testcontainers.junit.jupiter.Testcontainers;
|
||||
import org.testcontainers.utility.DockerImageName;
|
||||
|
||||
import org.springframework.ai.ResourceUtils;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.embedding.EmbeddingClient;
|
||||
import org.springframework.ai.embedding.TransformersEmbeddingClient;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Jingzhou Ou
|
||||
*/
|
||||
@Testcontainers
|
||||
public class Neo4jVectorStoreAutoConfigurationIT {
|
||||
|
||||
// Needs to be Neo4j 5.13+, because Neo4j 5.13 deprecated the used embedding storing
|
||||
// function.
|
||||
@Container
|
||||
static Neo4jContainer<?> neo4jContainer = new Neo4jContainer<>(DockerImageName.parse("neo4j:5.14"))
|
||||
.withRandomPassword();
|
||||
|
||||
List<Document> documents = List.of(
|
||||
new Document(ResourceUtils.getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")),
|
||||
new Document(ResourceUtils.getText("classpath:/test/data/time.shelter.txt")), new Document(
|
||||
ResourceUtils.getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad")));
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
.withConfiguration(AutoConfigurations.of(Neo4jVectorStoreAutoConfiguration.class))
|
||||
.withUserConfiguration(Config.class)
|
||||
.withPropertyValues("spring.ai.vectorstore.neo4j.driver.uri=" + neo4jContainer.getBoltUrl(),
|
||||
"spring.ai.vectorstore.neo4j.driver.authentication.username=" + "neo4j",
|
||||
"spring.ai.vectorstore.neo4j.driver.authentication.password=" + neo4jContainer.getAdminPassword());
|
||||
|
||||
@Test
|
||||
void addAndSearch() {
|
||||
contextRunner
|
||||
.withPropertyValues("spring.ai.vectorstore.neo4j.label=my_test_label",
|
||||
"spring.ai.vectorstore.neo4j.embeddingDimension=384")
|
||||
.run(context -> {
|
||||
var properties = context.getBean(Neo4jVectorStoreProperties.class);
|
||||
assertThat(properties.getLabel()).isEqualTo("my_test_label");
|
||||
assertThat(properties.getEmbeddingDimension()).isEqualTo(384);
|
||||
|
||||
VectorStore vectorStore = context.getBean(VectorStore.class);
|
||||
vectorStore.add(documents);
|
||||
|
||||
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
|
||||
|
||||
assertThat(results).hasSize(1);
|
||||
Document resultDoc = results.get(0);
|
||||
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
|
||||
assertThat(resultDoc.getContent()).contains(
|
||||
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
|
||||
|
||||
// Remove all documents from the store
|
||||
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
|
||||
|
||||
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
|
||||
assertThat(results).isEmpty();
|
||||
});
|
||||
}
|
||||
|
||||
@Configuration(proxyBeanMethods = false)
|
||||
static class Config {
|
||||
|
||||
@Bean
|
||||
public EmbeddingClient embeddingClient() {
|
||||
return new TransformersEmbeddingClient();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -194,13 +194,13 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
|
||||
|
||||
}
|
||||
|
||||
private static final int DEFAULT_EMBEDDING_DIMENSION = 1536;
|
||||
public static final int DEFAULT_EMBEDDING_DIMENSION = 1536;
|
||||
|
||||
private static final String DEFAULT_LABEL = "Document";
|
||||
public static final String DEFAULT_LABEL = "Document";
|
||||
|
||||
private static final String INDEX_NAME = "spring-ai-document-index";
|
||||
|
||||
private static final String DEFAULT_EMBEDDING_PROPERTY = "embedding";
|
||||
public static final String DEFAULT_EMBEDDING_PROPERTY = "embedding";
|
||||
|
||||
private final Driver driver;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user