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

GH-938: Add Couchbase vector store support

Browse Source 
Fixes: #938

Issue link: https://github.com/spring-projects/spring-ai/issues/938

This commit integrates Couchbase as a vector store option in Spring AI, providing:

- CouchbaseSearchVectorStore implementation with vector similarity search capabilities
- Support for metadata filtering with SQL++ expression conversion
- Spring Boot auto-configuration and starter module for easy integration
- Comprehensive documentation covering setup, configuration, and usage examples
- Integration tests using TestContainers with Couchbase 7.6

The implementation supports configuring dimensions, similarity functions (dot_product/l2_norm),
and optimization strategies (recall/latency). Schema initialization is now opt-in via
the initializeSchema property. Documentation includes both auto-configuration and
manual configuration instructions, along with property configuration details.

Signed-off-by: Abhiraj <abhiraj.official15@gmail.com>

co-authored-by: Laurent Doguin <laurent.doguin@gmail.com>
This commit is contained in:
Laurent Doguin
2024-06-23 19:19:03 +02:00
committed by Soby Chacko
parent 6b25b62771
commit d25d37ab12
22 changed files with 1777 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
pom.xml
View File

@@ -56,6 +56,7 @@
<module>vector-stores/spring-ai-cassandra-store</module>
<module>vector-stores/spring-ai-chroma-store</module>
<module>vector-stores/spring-ai-coherence-store</module>
<module>vector-stores/spring-ai-couchbase-store</module>
<module>vector-stores/spring-ai-elasticsearch-store</module>
<module>vector-stores/spring-ai-gemfire-store</module>
<module>vector-stores/spring-ai-hanadb-store</module>
@@ -78,6 +79,7 @@
<module>spring-ai-spring-boot-starters/spring-ai-starter-cassandra-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-chroma-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-coherence-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-couchbase-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-gemfire-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-hanadb-store</module>
@@ -235,6 +237,7 @@
<mariadb.version>3.5.1</mariadb.version>
<commonmark.version>0.22.0</commonmark.version>
<couchbase.version>3.7.8</couchbase.version>
<!-- testing dependencies -->
<okhttp3.version>4.12.0</okhttp3.version>

32
spring-ai-bom/pom.xml
View File

@@ -291,11 +291,17 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-opensearch-store</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-opensearch-store</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-couchbase-store</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
@@ -599,11 +605,17 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-qianfan-spring-boot-starter</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-qianfan-spring-boot-starter</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-couchbase-store-spring-boot-starter</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>

5
spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java
View File

@@ -51,7 +51,10 @@ public enum VectorStoreProvider {
* Vector store provided by CosmosDB.
*/
COSMOSDB("cosmosdb"),
/**
* Vector store provided by CosmosDB.
*/
COUCHBASE("couchbase"),
/**
* Vector store provided by Elasticsearch.
*/

1
spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc
View File

@@ -73,6 +73,7 @@
** xref:api/vectordbs/azure-cosmos-db.adoc[]
** xref:api/vectordbs/apache-cassandra.adoc[]
** xref:api/vectordbs/chroma.adoc[]
** xref:api/vectordbs/couchbase.adoc[]
** xref:api/vectordbs/elasticsearch.adoc[]
** xref:api/vectordbs/gemfire.adoc[GemFire]
** xref:api/vectordbs/mariadb.adoc[]

249
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/couchbase.adoc Normal file
View File

@@ -0,0 +1,249 @@
= Couchbase
This section will walk you through setting up the `CouchbaseSearchVectorStore` to store document embeddings and perform similarity searches using Couchbase.
link:https://docs.couchbase.com/server/current/vector-search/vector-search.html[Couchbase] is a distributed, JSON document database, with all the desired capabilities of a relational DBMS. Among other features, it allows users to query information using vector-based storage and retrieval.
== Prerequisites
A running Couchbase instance. The following options are available:
Couchbase
* link:https://hub.docker.com/_/couchbase/[Docker]
* link:https://cloud.couchbase.com/[Capella - Couchbase as a Service]
* link:https://www.couchbase.com/downloads/?family=couchbase-server[Install Couchbase locally]
* link:https://www.couchbase.com/downloads/?family=open-source-kubernetes[Couchbase Kubernetes Operator]
== Auto-configuration
Spring AI provides Spring Boot auto-configuration for the Couchbase Vector Store.
To enable it, add the following dependency to your project's Maven `pom.xml` file:
[source,xml]
----
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-couchbase-store-spring-boot-starter</artifactId>
</dependency>
----
or to your Gradle `build.gradle` build file.
[source,groovy]
----
dependencies {
implementation 'org.springframework.ai:spring-ai-couchbase-store-spring-boot-starter'
}
----
NOTE: Couchbase Vector search is only available in starting version 7.6 and Java SDK version 3.6.0"
TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file.
TIP: Refer to the xref:getting-started.adoc#repositories[Repositories] section to add Milestone and/or Snapshot Repositories to your build file.
The vector store implementation can initialize the configured bucket, scope, collection and search index for you, with default options, but you must opt-in by specifying the `initializeSchema` boolean in the appropriate constructor.
NOTE: This is a breaking change! In earlier versions of Spring AI, this schema initialization happened by default.
Please have a look at the list of <<couchbasevector-properties,configuration parameters>> for the vector store to learn about the default values and configuration options.
Additionally, you will need a configured `EmbeddingModel` bean. Refer to the xref:api/embeddings.adoc#available-implementations[EmbeddingModel] section for more information.
Now you can auto-wire the `CouchbaseSearchVectorStore` as a vector store in your application.
[source,java]
----
@Autowired VectorStore vectorStore;
// ...
List <Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
new Document("The World is Big and Salvation Lurks Around the Corner"),
new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));
// Add the documents to Qdrant
vectorStore.add(documents);
// Retrieve documents similar to a query
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
----
[[couchbasevector-properties]]
=== Configuration Properties
To connect to Couchbase and use the `CouchbaseSearchVectorStore`, you need to provide access details for your instance.
A simple configuration can either be provided via Spring Boot's `application.properties`,
[application,properties]
----
spring.ai.openai.api-key=<key>
spring.couchbase.connection-string=<conn_string>
spring.couchbase.username=<username>
spring.couchbase.password=<password>
----
environment variables,
[source,bash]
----
export SPRING_COUCHBASE_CONNECTION_STRINGS=<couchbase connection string like couchbase://localhost>
export SPRING_COUCHBASE_USERNAME=<couchbase username>
export SPRING_COUCHBASE_PASSWORD=<couchbase password>
# API key if needed, e.g. OpenAI
export SPRING_AI_OPENAI_API_KEY=<api-key>
----
or can be a mix of those.
For example, if you want to store your password as an environment variable but keep the rest in the plain `application.yml` file.
NOTE: If you choose to create a shell script for ease in future work, be sure to run it prior to starting your application by "sourcing" the file, i.e. `source <your_script_name>.sh`.
Spring Boot's auto-configuration feature for the Couchbase Cluster will create a bean instance that will be used by the `CouchbaseSearchVectorStore`.
The Spring Boot properties starting with `spring.couchbase.*` are used to configure the Couchbase cluster instance:
|===
|Property | Description | Default Value
| `spring.couchbase.connection-string` | A couchbase connection string | `couchbase://localhost`
| `spring.couchbase.password` | Password for authentication with Couchbase. | -
| `spring.couchbase.username` | Username for authentication with Couchbase.| -
| `spring.couchbase.env.io.minEndpoints` | Minimum number of sockets per node.| 1
| `spring.couchbase.env.io.maxEndpoints` | Maximum number of sockets per node.| 12
| `spring.couchbase.env.io.idleHttpConnectionTimeout` | Length of time an HTTP connection may remain idle before it is closed and removed from the pool.| 1s
| `spring.couchbase.env.ssl.enabled` | Whether to enable SSL support. Enabled automatically if a "bundle" is provided unless specified otherwise.| -
| `spring.couchbase.env.ssl.bundle` | SSL bundle name.| -
| `spring.couchbase.env.timeouts.connect` | Bucket connect timeout.| 10s
| `spring.couchbase.env.timeouts.disconnect` | Bucket disconnect timeout.| 10s
| `spring.couchbase.env.timeouts.key-value` | Timeout for operations on a specific key-value.| 2500ms
| `spring.couchbase.env.timeouts.key-value` | Timeout for operations on a specific key-value with a durability level.| 10s
| `spring.couchbase.env.timeouts.key-value-durable` | Timeout for operations on a specific key-value with a durability level.| 10s
| `spring.couchbase.env.timeouts.query` | SQL++ query operations timeout.| 75s
| `spring.couchbase.env.timeouts.view` | Regular and geospatial view operations timeout.| 75s
| `spring.couchbase.env.timeouts.search` | Timeout for the search service.| 75s
| `spring.couchbase.env.timeouts.analytics` | Timeout for the analytics service.| 75s
| `spring.couchbase.env.timeouts.management` | Timeout for the management operations.| 75s
|===
Properties starting with the `spring.ai.vectorstore.couchbase.*` prefix are used to configure `CouchbaseSearchVectorStore`.
|===
|Property | Description | Default Value
|`spring.ai.vectorstore.couchbase.index-name` | The name of the index to store the vectors. | spring-ai-document-index
|`spring.ai.vectorstore.couchbase.bucket-name` | The name of the Couchbase Bucket, parent of the scope. | default
|`spring.ai.vectorstore.couchbase.scope-name` |The name of the Couchbase scope, parent of the collection. Search queries will be executed in the scope context.| _default_
|`spring.ai.vectorstore.couchbase.collection-name` | The name of the Couchbase collection to store the Documents. | _default_
|`spring.ai.vectorstore.couchbase.dimensions` | The number of dimensions in the vector. | 1536
|`spring.ai.vectorstore.couchbase.similarity` | The similarity function to use. | `dot_product`
|`spring.ai.vectorstore.couchbase.optimization` | The similarity function to use. | `recall`
|`spring.ai.vectorstore.couchbase.initialize-schema`| whether to initialize the required schema | `false`
|===
The following similarity functions are available:
* l2_norm
* dot_product
The following index optimizations are available:
* recall
* latency
More details about each in the https://docs.couchbase.com/server/current/search/child-field-options-reference.html[Couchbase Documentation] on vector searches.
== Metadata Filtering
You can leverage the generic, portable link:https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_metadata_filters[metadata filters] with the Couchbase store.
For example, you can use either the text expression language:
[source,java]
----
vectorStore.similaritySearch(
SearchRequest.defaults()
.query("The World")
.topK(TOP_K)
.filterExpression("author in ['john', 'jill'] && article_type == 'blog'"));
----
or programmatically using the `Filter.Expression` DSL:
[source,java]
----
FilterExpressionBuilder b = new FilterExpressionBuilder();
vectorStore.similaritySearch(SearchRequest.defaults()
.query("The World")
.topK(TOP_K)
.filterExpression(b.and(
b.in("author","john", "jill"),
b.eq("article_type", "blog")).build()));
----
NOTE: These filter expressions are converted into the equivalent Couchbase SQL++ filters.
== Manual Configuration
Instead of using the Spring Boot auto-configuration, you can manually configure the Couchbase vector store. For this you need to add the `spring-ai-couchbase-store` to your project:
[source,xml]
----
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-couchbase-store</artifactId>
</dependency>
----
or to your Gradle `build.gradle` build file.
[source,groovy]
----
dependencies {
implementation 'org.springframework.ai:spring-ai-couchbase-store'
}
----
Create a Couchbase `Cluster` bean.
Read the link:https://docs.couchbase.com/java-sdk/current/hello-world/start-using-sdk.html[Couchbase Documentation] for more in-depth information about the configuration of a custom Cluster instance.
[source,java]
----
@Bean
public Cluster cluster() {
Cluster cluster = Cluster.connect("couchbase://localhost",
"username", "password");
}
----
and then create the `CouchbaseSearchVectorStore` bean using the builder pattern:
[source,java]
----
@Bean
public VectorStore couchbaseSearchVectorStore(Cluster cluster,
EmbeddingModel embeddingModel,
Boolean initializeSchema) {
return CouchbaseSearchVectorStore
.builder(cluster, embeddingModel)
.bucketName("test")
.scopeName("test")
.collectionName("test")
.initializeSchema(initializeSchema)
.build();
}
// This can be any EmbeddingModel implementation.
@Bean
public EmbeddingModel embeddingModel() {
return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(this.openaiKey).build());
}
----
== Limitations
NOTE: It is mandatory to have the following Couchbase services activated: Data, Query, Index, Search. While Data and Search could be enough, Query and Index are necessary to support the complete metadata filtering mechanism.

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

@@ -427,6 +427,13 @@
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<!-- Couchbase Vector Search Store -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-couchbase-store</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<!-- test dependencies -->
@@ -606,6 +613,12 @@
<scope>test</scope>
</dependency>
</dependencies>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>couchbase</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

69
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfiguration.java Normal file
View File

@@ -0,0 +1,69 @@
/*
* Copyright 2023 - 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.vectorstore.couchbase;
import com.couchbase.client.java.Cluster;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.CouchbaseSearchVectorStore;
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.couchbase.CouchbaseAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.util.StringUtils;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
@AutoConfiguration(after = CouchbaseAutoConfiguration.class)
@ConditionalOnClass({ CouchbaseSearchVectorStore.class, EmbeddingModel.class, Cluster.class })
@EnableConfigurationProperties(CouchbaseSearchVectorStoreProperties.class)
public class CouchbaseSearchVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public CouchbaseSearchVectorStore vectorStore(CouchbaseSearchVectorStoreProperties properties, Cluster cluster,
EmbeddingModel embeddingModel) {
var builder = CouchbaseSearchVectorStore.builder(cluster, embeddingModel);
if (StringUtils.hasText(properties.getIndexName())) {
builder.vectorIndexName(properties.getIndexName());
}
if (StringUtils.hasText(properties.getBucketName())) {
builder.bucketName(properties.getBucketName());
}
if (StringUtils.hasText(properties.getScopeName())) {
builder.scopeName(properties.getScopeName());
}
if (StringUtils.hasText(properties.getCollectionName())) {
builder.collectionName(properties.getCollectionName());
}
if (properties.getDimensions() != null) {
builder.dimensions(properties.getDimensions());
}
if (properties.getSimilarity() != null) {
builder.similarityFunction(properties.getSimilarity());
}
if (properties.getOptimization() != null) {
builder.indexOptimization(properties.getOptimization());
}
return builder.initializeSchema(properties.isInitializeSchema()).build();
}
}

127
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreProperties.java Normal file
View File

@@ -0,0 +1,127 @@
/*
* Copyright 2023 - 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.vectorstore.couchbase;
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
import org.springframework.ai.vectorstore.CouchbaseIndexOptimization;
import org.springframework.ai.vectorstore.CouchbaseSimilarityFunction;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
@ConfigurationProperties(prefix = CouchbaseSearchVectorStoreProperties.CONFIG_PREFIX)
public class CouchbaseSearchVectorStoreProperties extends CommonVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.couchbase";
/**
* The name of the index to store the vectors.
*/
private String indexName;
/**
* The name of the Couchbase collection to store the Documents.
*/
private String collectionName;
/**
* The name of the Couchbase scope, parent of the collection. Search queries will be
* executed in the scope context.
*/
private String scopeName;
/**
* The name of the Couchbase Bucket, parent of the scope.
*/
private String bucketName;
/**
* The total number of elements in the vector embedding array, up to 2048 elements.
* Arrays can be an array of arrays.
*/
private Integer dimensions;
/**
* The method to calculate the similarity between the vector embedding in a Vector
* Search index and the vector embedding in a Vector Search query.
*/
private CouchbaseSimilarityFunction similarity;
/**
* Choose whether the Search Service should prioritize recall or latency when
* returning similar vectors in search results.
*/
private CouchbaseIndexOptimization optimization;
public String getIndexName() {
return this.indexName;
}
public void setIndexName(String indexName) {
this.indexName = indexName;
}
public String getCollectionName() {
return collectionName;
}
public void setCollectionName(String collectionName) {
this.collectionName = collectionName;
}
public String getScopeName() {
return scopeName;
}
public void setScopeName(String scopeName) {
this.scopeName = scopeName;
}
public String getBucketName() {
return bucketName;
}
public void setBucketName(String bucketName) {
this.bucketName = bucketName;
}
public Integer getDimensions() {
return dimensions;
}
public void setDimensions(Integer dimensions) {
this.dimensions = dimensions;
}
public CouchbaseSimilarityFunction getSimilarity() {
return similarity;
}
public void setSimilarity(CouchbaseSimilarityFunction similarity) {
this.similarity = similarity;
}
public CouchbaseIndexOptimization getOptimization() {
return optimization;
}
public void setOptimization(CouchbaseIndexOptimization optimization) {
this.optimization = optimization;
}
}

1
spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
View File

@@ -59,3 +59,4 @@ org.springframework.ai.autoconfigure.minimax.MiniMaxAutoConfiguration
org.springframework.ai.autoconfigure.vertexai.embedding.VertexAiEmbeddingAutoConfiguration
org.springframework.ai.autoconfigure.chat.memory.cassandra.CassandraChatMemoryAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.observation.VectorStoreObservationAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.couchbase.CouchbaseSearchVectorStoreAutoConfiguration

39
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseContainerMetadata.java Normal file
View File

@@ -0,0 +1,39 @@
/*
* Copyright 2023 - 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.vectorstore.couchbase;
import org.testcontainers.couchbase.BucketDefinition;
import org.testcontainers.utility.DockerImageName;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
public class CouchbaseContainerMetadata {
public static final String BUCKET_NAME = "example";
public static final String USERNAME = "Administrator";
public static final String PASSWORD = "password";
public static final BucketDefinition bucketDefinition = new BucketDefinition(BUCKET_NAME);
public static final DockerImageName COUCHBASE_IMAGE_ENTERPRISE = DockerImageName.parse("couchbase:enterprise")
.asCompatibleSubstituteFor("couchbase/server")
.withTag("enterprise-7.6.1");
}

144
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,144 @@
/*
* Copyright 2023 - 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.vectorstore.couchbase;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.CouchbaseIndexOptimization;
import org.springframework.ai.vectorstore.CouchbaseSimilarityFunction;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.couchbase.CouchbaseAutoConfiguration;
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.couchbase.CouchbaseContainer;
import org.testcontainers.couchbase.CouchbaseService;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.autoconfigure.vectorstore.couchbase.CouchbaseContainerMetadata.*;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
class CouchbaseSearchVectorStoreAutoConfigurationIT {
// Define the couchbase container.
@Container
final static CouchbaseContainer couchbaseContainer = new CouchbaseContainer(COUCHBASE_IMAGE_ENTERPRISE)
.withCredentials(USERNAME, PASSWORD)
.withEnabledServices(CouchbaseService.KV, CouchbaseService.QUERY, CouchbaseService.INDEX,
CouchbaseService.SEARCH)
.withBucket(bucketDefinition)
.withStartupAttempts(4)
.withStartupTimeout(Duration.ofSeconds(90))
.waitingFor(Wait.forHealthcheck());
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(CouchbaseAutoConfiguration.class,
CouchbaseSearchVectorStoreAutoConfiguration.class, RestClientAutoConfiguration.class,
SpringAiRetryAutoConfiguration.class, OpenAiAutoConfiguration.class))
.withPropertyValues("spring.couchbase.connection-string=" + couchbaseContainer.getConnectionString(),
"spring.couchbase.username=" + couchbaseContainer.getUsername(),
"spring.couchbase.password=" + couchbaseContainer.getPassword(),
"spring.ai.vectorstore.couchbase.initialize-schema=true",
"spring.ai.vectorstore.couchbase.index-name=example",
"spring.ai.vectorstore.couchbase.collection-name=example",
"spring.ai.vectorstore.couchbase.scope-name=example",
"spring.ai.vectorstore.couchbase.bucket-name=example",
"spring.ai.openai.api-key=" + System.getenv("OPENAI_API_KEY"));
@Test
public void addAndSearchWithFilters() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "Bulgaria"));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "Netherlands"));
vectorStore.add(List.of(bgDocument, nlDocument));
var requestBuilder = SearchRequest.builder().query("The World").topK(5);
List<Document> results = vectorStore.similaritySearch(requestBuilder.build());
assertThat(results).hasSize(2);
results = vectorStore.similaritySearch(
requestBuilder.similarityThresholdAll().filterExpression("country == 'Bulgaria'").build());
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
results = vectorStore.similaritySearch(
requestBuilder.similarityThresholdAll().filterExpression("country == 'Netherlands'").build());
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
// Remove all documents from the store
vectorStore.delete(List.of(bgDocument, nlDocument).stream().map(doc -> doc.getId()).toList());
});
}
@Test
public void propertiesTest() {
new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(CouchbaseAutoConfiguration.class,
CouchbaseSearchVectorStoreAutoConfiguration.class, RestClientAutoConfiguration.class,
SpringAiRetryAutoConfiguration.class, OpenAiAutoConfiguration.class))
.withPropertyValues("spring.couchbase.connection-string=" + couchbaseContainer.getConnectionString(),
"spring.couchbase.username=" + couchbaseContainer.getUsername(),
"spring.couchbase.password=" + couchbaseContainer.getPassword(),
"spring.ai.openai.api-key=" + System.getenv("OPENAI_API_KEY"),
"spring.ai.vectorstore.couchbase.index-name=example",
"spring.ai.vectorstore.couchbase.collection-name=example",
"spring.ai.vectorstore.couchbase.scope-name=example",
"spring.ai.vectorstore.couchbase.bucket-name=example",
"spring.ai.vectorstore.couchbase.dimensions=1024",
"spring.ai.vectorstore.couchbase.optimization=latency",
"spring.ai.vectorstore.couchbase.similarity=l2_norm")
.run(context -> {
var properties = context.getBean(CouchbaseSearchVectorStoreProperties.class);
var vectorStore = context.getBean(VectorStore.class);
assertThat(properties).isNotNull();
assertThat(properties.getIndexName()).isEqualTo("example");
assertThat(properties.getCollectionName()).isEqualTo("example");
assertThat(properties.getScopeName()).isEqualTo("example");
assertThat(properties.getBucketName()).isEqualTo("example");
assertThat(properties.getDimensions()).isEqualTo(1024);
assertThat(properties.getOptimization()).isEqualTo(CouchbaseIndexOptimization.latency);
assertThat(properties.getSimilarity()).isEqualTo(CouchbaseSimilarityFunction.l2_norm);
assertThat(vectorStore).isNotNull();
});
}
}

42
spring-ai-spring-boot-starters/spring-ai-starter-couchbase-store/pom.xml Normal file
View File

@@ -0,0 +1,42 @@
<?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.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>1.0.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-couchbase-store-spring-boot-starter</artifactId>
<packaging>jar</packaging>
<name>Spring AI Starter - Couchbase Store</name>
<description>Spring AI Couchbase Store Auto Configuration</description>
<url>https://github.com/spring-projects/spring-ai</url>
<scm>
<url>https://github.com/spring-projects/spring-ai</url>
<connection>git://github.com/spring-projects/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-couchbase-store</artifactId>
<version>${project.parent.version}</version>
</dependency>
</dependencies>
</project>

1
vector-stores/spring-ai-couchbase-store/README.md Normal file
View File

@@ -0,0 +1 @@
[Couchbase Vector Store Documentation](https://docs.spring.io/spring-ai/reference/1.0-SNAPSHOT/api/vectordbs/couchbase.html)

74
vector-stores/spring-ai-couchbase-store/pom.xml Normal file
View File

@@ -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.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>1.0.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-couchbase-store</artifactId>
<packaging>jar</packaging>
<name>Spring AI Vector Store - Couchbase</name>
<description>Spring AI Couchbase Vector Store</description>
<url>https://github.com/spring-projects/spring-ai</url>
<scm>
<url>https://github.com/spring-projects/spring-ai</url>
<connection>git://github.com/spring-projects/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>com.couchbase.client</groupId>
<artifactId>java-client</artifactId>
<version>${couchbase.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-core</artifactId>
<version>${parent.version}</version>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-openai</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<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>couchbase</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

82
vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseAiSearchFilterExpressionConverter.java Normal file
View File

@@ -0,0 +1,82 @@
/*
* Copyright 2023 - 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.vectorstore;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
public class CouchbaseAiSearchFilterExpressionConverter extends AbstractFilterExpressionConverter {
public CouchbaseAiSearchFilterExpressionConverter() {
}
@Override
protected void doExpression(Expression expression, StringBuilder context) {
this.convertOperand(expression.left(), context);
context.append(getOperationSymbol(expression));
this.convertOperand(expression.right(), context);
}
private String getOperationSymbol(Expression exp) {
switch (exp.type()) {
case AND:
return " AND ";
case OR:
return " OR ";
case EQ:
return " == ";
case NE:
return " != ";
case LT:
return " < ";
case LTE:
return " <= ";
case GT:
return " > ";
case GTE:
return " >= ";
case IN:
return " IN ";
case NIN:
return " NOT IN ";
default:
throw new RuntimeException("Not supported expression type: " + exp.type());
}
}
@Override
protected void doKey(Key key, StringBuilder context) {
context.append("metadata.");
context.append(key.key());
}
@Override
protected void doStartGroup(Group group, StringBuilder context) {
context.append("(");
}
@Override
protected void doEndGroup(Group group, StringBuilder context) {
context.append(")");
}
}

40
vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseIndexOptimization.java Normal file
View File

@@ -0,0 +1,40 @@
/*
* Copyright 2023 - 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.vectorstore;
/**
* Choose whether the Vector store should prioritize recall or latency when returning
* similar vectors in search results. See
* https://docs.couchbase.com/server/current/search/child-field-options-reference.html for
* more details.
*
* @author Laurent Doguin
* @since 1.0.0
*/
public enum CouchbaseIndexOptimization {
/**
* recall: The Search Service prioritizes returning the most accurate result. This may
* increase resource usage for Search queries.
*/
recall,
/**
* latency: The Search Service prioritizes returning results with lower latency. This
* may reduce the accuracy of results.
*/
latency
}

481
vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStore.java Normal file
View File

@@ -0,0 +1,481 @@
/*
* Copyright 2023 - 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.vectorstore;
import com.couchbase.client.core.util.ConsistencyUtil;
import com.couchbase.client.java.Bucket;
import com.couchbase.client.java.Cluster;
import com.couchbase.client.java.Collection;
import com.couchbase.client.java.Scope;
import com.couchbase.client.java.manager.bucket.BucketSettings;
import com.couchbase.client.java.manager.collection.CollectionSpec;
import com.couchbase.client.java.manager.collection.ScopeSpec;
import com.couchbase.client.java.manager.query.CreatePrimaryQueryIndexOptions;
import com.couchbase.client.java.manager.search.SearchIndex;
import com.couchbase.client.java.query.QueryOptions;
import com.couchbase.client.java.query.QueryResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import reactor.core.publisher.Mono;
import reactor.util.retry.RetrySpec;
import java.time.Duration;
import java.util.*;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
public class CouchbaseSearchVectorStore extends AbstractObservationVectorStore
implements InitializingBean, AutoCloseable {
private static final Logger logger = LoggerFactory.getLogger(CouchbaseSearchVectorStore.class);
private static final String DEFAULT_INDEX_NAME = "spring-ai-document-index";
private static final String DEFAULT_COLLECTION_NAME = "_default";
private static final String DEFAULT_SCOPE_NAME = "_default";
private static final String DEFAULT_BUCKET_NAME = "default";
private final EmbeddingModel embeddingModel;
private final String collectionName;
private final String scopeName;
private final String bucketName;
private final String vectorIndexName;
private final Integer dimensions;
private final CouchbaseSimilarityFunction similarityFunction;
private final CouchbaseIndexOptimization indexOptimization;
private final Cluster cluster;
private final CouchbaseAiSearchFilterExpressionConverter filterExpressionConverter;
private final boolean initializeSchema;
private final Collection collection;
private final Scope scope;
private final Bucket bucket;
protected CouchbaseSearchVectorStore(Builder builder) {
super(builder);
Objects.requireNonNull(builder.cluster, "CouchbaseCluster must not be null");
Objects.requireNonNull(builder.embeddingModel, "embeddingModel must not be null");
this.initializeSchema = builder.initializeSchema;
this.embeddingModel = builder.embeddingModel;
this.filterExpressionConverter = builder.filterExpressionConverter;
this.cluster = builder.cluster;
this.bucket = cluster.bucket(builder.bucketName);
this.scope = bucket.scope(builder.scopeName);
this.collection = scope.collection(builder.collectionName);
this.vectorIndexName = builder.vectorIndexName;
this.collectionName = builder.collectionName;
this.bucketName = builder.bucketName;
this.scopeName = builder.scopeName;
this.dimensions = builder.dimensions;
this.similarityFunction = builder.similarityFunction;
this.indexOptimization = builder.indexOptimization;
}
@Override
public void afterPropertiesSet() {
if (!this.initializeSchema) {
return;
}
try {
logger.info("Init Cluster Called");
initCluster();
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
@Override
public void doAdd(List<Document> documents) {
logger.info("Trying Add");
logger.info(this.bucketName);
logger.info(this.scopeName);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
for (Document document : documents) {
CouchbaseDocument cbDoc = new CouchbaseDocument(document.getId(), document.getText(),
document.getMetadata(), embeddings.get(documents.indexOf(document)));
collection.upsert(document.getId(), cbDoc);
}
}
@Override
public void doDelete(List<String> idList) {
for (String id : idList) {
collection.remove(id);
}
}
@Override
public void doDelete(Filter.Expression filterExpression) {
Assert.notNull(filterExpression, "Filter expression must not be null");
try {
String nativeFilter = this.filterExpressionConverter.convertExpression(filterExpression);
String sql = String.format("DELETE FROM %s WHERE %s", collection.name(), nativeFilter);
scope.query(sql, QueryOptions.queryOptions().metrics(true));
}
catch (Exception e) {
logger.error("Failed to delete documents by filter: {}", e.getMessage(), e);
throw new IllegalStateException("Failed to delete documents by filter", e);
}
}
@Override
public List<Document> doSimilaritySearch(org.springframework.ai.vectorstore.SearchRequest springAiRequest) {
float[] embeddings = this.embeddingModel.embed(springAiRequest.getQuery());
int topK = springAiRequest.getTopK();
double similarityThreshold = springAiRequest.getSimilarityThreshold();
Filter.Expression fe = springAiRequest.getFilterExpression();
String nativeFilterExpression = (fe != null) ? " AND " + this.filterExpressionConverter.convertExpression(fe)
: "";
String statement = String.format(
"""
SELECT c.* FROM `%s` AS c
WHERE SEARCH_SCORE() > %s AND SEARCH(`c`, {"query": {"match_none": {}}, "knn": [{"field": "embedding", "k": %s, "vector": %s } ] }, {"index": "%s.%s.%s"} )
%s
""",
this.collectionName, similarityThreshold, topK, Arrays.toString(embeddings), this.bucketName,
this.scopeName, this.vectorIndexName, nativeFilterExpression);
QueryResult result = scope.query(statement, QueryOptions.queryOptions());
return result.rowsAs(Document.class);
}
@Override
public <T> Optional<T> getNativeClient() {
@SuppressWarnings("unchecked")
T client = (T) this;
return Optional.of(client);
}
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.COUCHBASE.value(), operationName)
.collectionName(this.collection.name())
.dimensions(this.embeddingModel.dimensions());
}
public static Builder builder(Cluster cluster, EmbeddingModel embeddingModel) {
return new Builder(cluster, embeddingModel);
}
public static class Builder extends AbstractVectorStoreBuilder<Builder> {
private String collectionName = DEFAULT_COLLECTION_NAME;
private String scopeName = DEFAULT_SCOPE_NAME;
private String bucketName = DEFAULT_BUCKET_NAME;
private String vectorIndexName = DEFAULT_INDEX_NAME;
private Integer dimensions = 1536;
private CouchbaseSimilarityFunction similarityFunction = CouchbaseSimilarityFunction.dot_product;
private CouchbaseIndexOptimization indexOptimization = CouchbaseIndexOptimization.recall;
private final Cluster cluster;
private final CouchbaseAiSearchFilterExpressionConverter filterExpressionConverter = new CouchbaseAiSearchFilterExpressionConverter();
private boolean initializeSchema = false;
/**
* @throws IllegalArgumentException if couchbaseSearchVectorConfig or cluster is
* null
*/
private Builder(Cluster cluster, EmbeddingModel embeddingModel) {
super(embeddingModel);
Assert.notNull(cluster, "Cluster must not be null");
this.cluster = cluster;
}
/**
* Sets whether to initialize the schema.
* @param initializeSchema true to initialize schema, false otherwise
* @return the builder instance
*/
public Builder initializeSchema(boolean initializeSchema) {
this.initializeSchema = initializeSchema;
return this;
}
/**
* Configures the Couchbase collection storing {@link Document}.
* @param collectionName
* @return this builder
*/
public CouchbaseSearchVectorStore.Builder collectionName(String collectionName) {
Assert.notNull(collectionName, "Collection Name must not be null");
Assert.notNull(collectionName, "Collection Name must not be empty");
this.collectionName = collectionName;
return this;
}
/**
* Configures the Couchbase scope, parent of the selected collection. Search will
* be executed in this scope context.
* @param scopeName
* @return this builder
*/
public CouchbaseSearchVectorStore.Builder scopeName(String scopeName) {
Assert.notNull(scopeName, "Scope Name must not be null");
Assert.notNull(scopeName, "Scope Name must not be empty");
this.scopeName = scopeName;
return this;
}
/**
* Configures the Couchbase bucket, parent of the selected Scope.
* @param bucketName
* @return this builder
*/
public CouchbaseSearchVectorStore.Builder bucketName(String bucketName) {
Assert.notNull(bucketName, "Bucket Name must not be null");
Assert.notNull(bucketName, "Bucket Name must not be empty");
this.bucketName = bucketName;
return this;
}
/**
* Configures the vector index name. This must match the name of the Vector Search
* Index Name in Atlas
* @param vectorIndexName
* @return this builder
*/
public CouchbaseSearchVectorStore.Builder vectorIndexName(String vectorIndexName) {
Assert.notNull(vectorIndexName, "Vector Index Name must not be null");
Assert.notNull(vectorIndexName, "Vector Index Name must not be empty");
this.vectorIndexName = vectorIndexName;
return this;
}
/**
* The number of dimensions in the vector.
* @param dimensions
* @return this builder
*/
public CouchbaseSearchVectorStore.Builder dimensions(Integer dimensions) {
Assert.notNull(dimensions, "Dimensions must not be null");
Assert.notNull(dimensions, "Dimensions must not be empty");
this.dimensions = dimensions;
return this;
}
/**
* Choose the method to calculate the similarity between the vector embedding in a
* Vector Search index and the vector embedding in a Vector Search query.
* @param similarityFunction
* @return this builder
*/
public CouchbaseSearchVectorStore.Builder similarityFunction(CouchbaseSimilarityFunction similarityFunction) {
Assert.notNull(similarityFunction, "Couchbase Similarity Function must not be null");
Assert.notNull(similarityFunction, "Couchbase Similarity Function must not be empty");
this.similarityFunction = similarityFunction;
return this;
}
/**
* Choose to prioritize accuracy or latency.
* @param indexOptimization
* @return this builder
*/
public CouchbaseSearchVectorStore.Builder indexOptimization(CouchbaseIndexOptimization indexOptimization) {
Assert.notNull(indexOptimization, "Index Optimization must not be null");
Assert.notNull(indexOptimization, "Index Optimization must not be empty");
this.indexOptimization = indexOptimization;
return this;
}
public CouchbaseSearchVectorStore build() {
return new CouchbaseSearchVectorStore(this);
}
}
public void initCluster() throws InterruptedException {
// init scope, collection, indexes
BucketSettings bs = cluster.buckets().getAllBuckets().get(this.bucketName);
if (bs == null) {
cluster.buckets().createBucket(BucketSettings.create(this.bucketName));
}
logger.info("Created bucket");
Bucket b = cluster.bucket(this.bucketName);
b.waitUntilReady(Duration.ofSeconds(20));
logger.info("Opened Bucket");
boolean scopeExist = b.collections().getAllScopes().stream().anyMatch(sc -> sc.name().equals(this.scopeName));
if (!scopeExist) {
b.collections().createScope(this.scopeName);
}
ConsistencyUtil.waitUntilScopePresent(cluster.core(), this.bucketName, this.scopeName);
Scope s = b.scope(this.scopeName);
boolean collectionExist = bucket.collections()
.getAllScopes()
.stream()
.map(ScopeSpec::collections)
.flatMap(java.util.Collection::stream)
.filter(it -> it.scopeName().equals(this.scopeName))
.map(CollectionSpec::name)
.anyMatch(this.collectionName::equals);
if (!collectionExist) {
b.collections().createCollection(this.scopeName, this.collectionName);
ConsistencyUtil.waitUntilCollectionPresent(cluster.core(), this.bucketName, this.scopeName,
this.collectionName);
Collection c = s.collection(this.collectionName);
Mono.empty()
.then(Mono.fromRunnable(
() -> c.async()
.queryIndexes()
.createPrimaryIndex(CreatePrimaryQueryIndexOptions.createPrimaryQueryIndexOptions()
.ignoreIfExists(true))))
.retryWhen(RetrySpec.backoff(3, Duration.ofMillis(1000)));
}
boolean indexExist = s.searchIndexes()
.getAllIndexes()
.stream()
.anyMatch(idx -> this.vectorIndexName.equals(idx.name()));
if (!indexExist) {
String jsonIndexTemplate = """
{
"type": "fulltext-index",
"name": "%s",
"sourceType": "gocbcore",
"sourceName": "%s",
"planParams": {
"maxPartitionsPerPIndex": 1024,
"indexPartitions": 1
},
"params": {
"doc_config": {
"docid_prefix_delim": "",
"docid_regexp": "",
"mode": "scope.collection.type_field",
"type_field": "type"
},
"mapping": {
"analysis": {},
"default_analyzer": "standard",
"default_datetime_parser": "dateTimeOptional",
"default_field": "_all",
"default_mapping": {
"dynamic": false,
"enabled": false
},
"default_type": "%s",
"docvalues_dynamic": false,
"index_dynamic": false,
"store_dynamic": false,
"type_field": "_type",
"types": {
"%s.%s": {
"dynamic": false,
"enabled": true,
"properties": {
"embedding": {
"dynamic": false,
"enabled": true,
"fields": [
{
"dims": %s,
"index": true,
"name": "embedding",
"similarity": "%s",
"type": "vector",
"vector_index_optimized_for": "%s"
}
]
},
"content": {
"dynamic": false,
"enabled": true,
"fields": [
{
"analyzer": "keyword",
"docvalues": true,
"include_in_all": true,
"include_term_vectors": true,
"index": true,
"name": "text",
"store": true,
"type": "text"
}
]
}
}
}
}
},
"store": {
"indexType": "scorch",
"segmentVersion": 16
}
},
"sourceParams": {}
}
""";
String jsonIndexValue = String.format(jsonIndexTemplate, this.vectorIndexName, this.bucketName,
this.collectionName, this.scopeName, this.collectionName, this.dimensions, this.similarityFunction,
this.indexOptimization);
SearchIndex si = SearchIndex.fromJson(jsonIndexValue);
s.searchIndexes().upsertIndex(si);
}
}
public void close() throws Exception {
if (this.cluster != null) {
this.cluster.close();
logger.info("Connection with cluster closed");
}
}
public record CouchbaseDocument(String id, String content, Map<String, Object> metadata, float[] embedding) {
}
}

42
vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSimilarityFunction.java Normal file
View File

@@ -0,0 +1,42 @@
/*
* Copyright 2023 - 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.vectorstore;
/**
* Choose the method to calculate the similarity between the vector embedding in a Vector
* Search index and the vector embedding in a Vector Search query. See
* https://docs.couchbase.com/server/current/search/child-field-options-reference.html for
* more details.
*
* @author Laurent Doguin
* @since 1.0.0
*/
public enum CouchbaseSimilarityFunction {
/**
* It’s best to use l2_norm similarity when your embeddings contain information about
* the count or measure of specific things, and your embedding model uses the same
* similarity metric.
*/
l2_norm,
/**
* Dot product similarity is commonly used by Large Language Models (LLMs). Use
* dot_product to get the best results with an embedding model that uses dot product
* similarity.
*/
dot_product
}

1
vector-stores/spring-ai-couchbase-store/src/main/resources/application.properties Normal file
View File

@@ -0,0 +1 @@
spring.application.name=demo

301
vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStoreIT.java Normal file
View File

@@ -0,0 +1,301 @@
/*
* Copyright 2023 - 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.vectorstore;
import com.couchbase.client.java.Cluster;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.couchbase.CouchbaseContainer;
import org.testcontainers.couchbase.CouchbaseService;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.vectorstore.testcontainer.CouchbaseContainerMetadata.*;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class CouchbaseSearchVectorStoreIT {
// Define the couchbase container.
@Container
final static CouchbaseContainer couchbaseContainer = new CouchbaseContainer(COUCHBASE_IMAGE_ENTERPRISE)
.withCredentials(USERNAME, PASSWORD)
.withEnabledServices(CouchbaseService.KV, CouchbaseService.QUERY, CouchbaseService.INDEX,
CouchbaseService.SEARCH)
.withBucket(bucketDefinition)
.withStartupAttempts(4)
.withStartupTimeout(Duration.ofSeconds(90))
.waitingFor(Wait.forHealthcheck());
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
Awaitility.setDefaultPollDelay(Duration.ZERO);
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
}
private ApplicationContextRunner getContextRunner() {
return new ApplicationContextRunner().withUserConfiguration(TestApplication.class);
}
@AfterAll
public static void stopContainers() {
couchbaseContainer.close();
}
@Test
void vectorStoreTest() {
getContextRunner().run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
List<Document> documents = List.of(
new Document(
"Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1")),
new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"),
new Document(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
vectorStore.add(documents);
Thread.sleep(5000); // wait for indexing
List<Document> results = vectorStore
.similaritySearch(SearchRequest.builder().query("Great").topK(1).build());
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getText()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsEntry("meta2", "meta2");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).collect(Collectors.toList()));
List<Document> results2 = vectorStore
.similaritySearch(SearchRequest.builder().query("Great").topK(1).build());
assertThat(results2).isEmpty();
});
}
@Test
void documentUpdateTest() {
getContextRunner().run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
vectorStore.add(List.of(document));
Thread.sleep(5000); // Await a second for the document to be indexed
List<Document> results = vectorStore
.similaritySearch(SearchRequest.builder().query("Spring").topK(5).build());
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsEntry("meta1", "meta1");
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
Collections.singletonMap("meta2", "meta2"));
vectorStore.add(List.of(sameIdDocument));
results = vectorStore.similaritySearch(SearchRequest.builder().query("FooBar").topK(5).build());
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsEntry("meta2", "meta2");
// Remove all documents from the store
vectorStore.delete(Collections.singletonList(document.getId()));
List<Document> results2 = vectorStore
.similaritySearch(SearchRequest.builder().query("Spring").topK(1).build());
assertThat(results2).isEmpty();
});
}
@Test
void searchWithFilters() {
getContextRunner().run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2020));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "NL"));
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2023));
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
Thread.sleep(5000); // Await a second for the document to be indexed
List<Document> results = vectorStore
.similaritySearch(SearchRequest.builder().query("The World").topK(5).build());
assertThat(results).hasSize(3);
results = vectorStore.similaritySearch(SearchRequest.builder()
.query("The World")
.topK(5)
.similarityThresholdAll()
.filterExpression("country == 'NL'")
.build());
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
results = vectorStore.similaritySearch(SearchRequest.builder()
.query("The World")
.topK(5)
.similarityThresholdAll()
.filterExpression("country == 'BG'")
.build());
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
results = vectorStore.similaritySearch(SearchRequest.builder()
.query("The World")
.topK(5)
.similarityThresholdAll()
.filterExpression("country == 'BG' && year == 2020")
.build());
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
results = vectorStore.similaritySearch(SearchRequest.builder()
.query("The World")
.topK(5)
.similarityThresholdAll()
.filterExpression("NOT(country == 'BG' && year == 2020)")
.build());
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(nlDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(nlDocument.getId(), bgDocument2.getId());
// Remove all documents from the store
vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId(), nlDocument.getId()));
List<Document> results2 = vectorStore
.similaritySearch(SearchRequest.builder().query("Spring").topK(1).build());
assertThat(results2).isEmpty();
});
}
@Test
void deleteWithComplexFilterExpression() {
getContextRunner().run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
var doc1 = new Document("Content 1", Map.of("type", "A", "priority", 1));
var doc2 = new Document("Content 2", Map.of("type", "A", "priority", 2));
var doc3 = new Document("Content 3", Map.of("type", "B", "priority", 1));
vectorStore.add(List.of(doc1, doc2, doc3));
Thread.sleep(5000); // Wait for indexing
// Complex filter expression: (type == 'A' AND priority > 1)
Filter.Expression priorityFilter = new Filter.Expression(Filter.ExpressionType.GT,
new Filter.Key("priority"), new Filter.Value(1));
Filter.Expression typeFilter = new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("type"),
new Filter.Value("A"));
Filter.Expression complexFilter = new Filter.Expression(Filter.ExpressionType.AND, typeFilter,
priorityFilter);
vectorStore.delete(complexFilter);
Thread.sleep(1000); // Wait for deletion to be processed
var results = vectorStore
.similaritySearch(SearchRequest.builder().query("Content").topK(5).similarityThresholdAll().build());
assertThat(results).hasSize(2);
assertThat(results.stream().map(doc -> doc.getMetadata().get("type")).collect(Collectors.toList()))
.containsExactlyInAnyOrder("A", "B");
assertThat(results.stream().map(doc -> doc.getMetadata().get("priority")).collect(Collectors.toList()))
.containsExactlyInAnyOrder(1, 1);
// Remove all documents from the store
vectorStore.delete(List.of(doc1.getId(), doc3.getId()));
List<Document> results2 = vectorStore
.similaritySearch(SearchRequest.builder().query("Content").topK(5).build());
assertThat(results2).isEmpty();
});
}
@Test
void getNativeClientTest() {
getContextRunner().run(context -> {
CouchbaseSearchVectorStore vectorStore = context.getBean(CouchbaseSearchVectorStore.class);
Optional<CouchbaseSearchVectorStore> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Bean
public CouchbaseSearchVectorStore vectorStore(EmbeddingModel embeddingModel) {
Cluster cluster = Cluster.connect(couchbaseContainer.getConnectionString(),
couchbaseContainer.getUsername(), couchbaseContainer.getPassword());
CouchbaseSearchVectorStore.Builder builder = CouchbaseSearchVectorStore.builder(cluster, embeddingModel)
.bucketName("springBucket")
.scopeName("springScope")
.collectionName("sprtingcollection");
return builder.initializeSchema(true).build();
}
@Bean
public EmbeddingModel embeddingModel() {
return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build());
}
}
}

39
vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/testcontainer/CouchbaseContainerMetadata.java Normal file
View File

@@ -0,0 +1,39 @@
/*
* Copyright 2023 - 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.vectorstore.testcontainer;
import org.testcontainers.couchbase.BucketDefinition;
import org.testcontainers.utility.DockerImageName;
/**
* @author Laurent Doguin
* @since 1.0.0
*/
public class CouchbaseContainerMetadata {
public static final String BUCKET_NAME = "springBucket";
public static final String USERNAME = "Administrator";
public static final String PASSWORD = "password";
public static final BucketDefinition bucketDefinition = new BucketDefinition(BUCKET_NAME);
public static final DockerImageName COUCHBASE_IMAGE_ENTERPRISE = DockerImageName.parse("couchbase:enterprise")
.asCompatibleSubstituteFor("couchbase/server")
.withTag("enterprise-7.6.1");
}

1
vector-stores/spring-ai-couchbase-store/src/test/resources/application.properties Normal file
View File

@@ -0,0 +1 @@
spring.application.name=demo