From d25d37ab12c0fd4eb77773d157eb15002b4ff897 Mon Sep 17 00:00:00 2001 From: Laurent Doguin Date: Sun, 23 Jun 2024 19:19:03 +0200 Subject: [PATCH] GH-938: Add Couchbase vector store support 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 co-authored-by: Laurent Doguin --- pom.xml | 3 + spring-ai-bom/pom.xml | 32 +- .../conventions/VectorStoreProvider.java | 5 +- .../src/main/antora/modules/ROOT/nav.adoc | 1 + .../ROOT/pages/api/vectordbs/couchbase.adoc | 249 +++++++++ spring-ai-spring-boot-autoconfigure/pom.xml | 15 +- ...aseSearchVectorStoreAutoConfiguration.java | 69 +++ .../CouchbaseSearchVectorStoreProperties.java | 127 +++++ ...ot.autoconfigure.AutoConfiguration.imports | 1 + .../couchbase/CouchbaseContainerMetadata.java | 39 ++ ...eSearchVectorStoreAutoConfigurationIT.java | 144 ++++++ .../spring-ai-starter-couchbase-store/pom.xml | 42 ++ .../spring-ai-couchbase-store/README.md | 1 + .../spring-ai-couchbase-store/pom.xml | 74 +++ ...baseAiSearchFilterExpressionConverter.java | 82 +++ .../CouchbaseIndexOptimization.java | 40 ++ .../CouchbaseSearchVectorStore.java | 481 ++++++++++++++++++ .../CouchbaseSimilarityFunction.java | 42 ++ .../src/main/resources/application.properties | 1 + .../CouchbaseSearchVectorStoreIT.java | 301 +++++++++++ .../CouchbaseContainerMetadata.java | 39 ++ .../src/test/resources/application.properties | 1 + 22 files changed, 1777 insertions(+), 12 deletions(-) create mode 100644 spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/couchbase.adoc create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfiguration.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreProperties.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseContainerMetadata.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfigurationIT.java create mode 100644 spring-ai-spring-boot-starters/spring-ai-starter-couchbase-store/pom.xml create mode 100644 vector-stores/spring-ai-couchbase-store/README.md create mode 100644 vector-stores/spring-ai-couchbase-store/pom.xml create mode 100644 vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseAiSearchFilterExpressionConverter.java create mode 100644 vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseIndexOptimization.java create mode 100644 vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStore.java create mode 100644 vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSimilarityFunction.java create mode 100644 vector-stores/spring-ai-couchbase-store/src/main/resources/application.properties create mode 100644 vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStoreIT.java create mode 100644 vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/testcontainer/CouchbaseContainerMetadata.java create mode 100644 vector-stores/spring-ai-couchbase-store/src/test/resources/application.properties diff --git a/pom.xml b/pom.xml index c08ed21d4..1e1b90440 100644 --- a/pom.xml +++ b/pom.xml @@ -56,6 +56,7 @@ vector-stores/spring-ai-cassandra-store vector-stores/spring-ai-chroma-store vector-stores/spring-ai-coherence-store + vector-stores/spring-ai-couchbase-store vector-stores/spring-ai-elasticsearch-store vector-stores/spring-ai-gemfire-store vector-stores/spring-ai-hanadb-store @@ -78,6 +79,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-cassandra-store spring-ai-spring-boot-starters/spring-ai-starter-chroma-store spring-ai-spring-boot-starters/spring-ai-starter-coherence-store + spring-ai-spring-boot-starters/spring-ai-starter-couchbase-store spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store spring-ai-spring-boot-starters/spring-ai-starter-gemfire-store spring-ai-spring-boot-starters/spring-ai-starter-hanadb-store @@ -235,6 +237,7 @@ 3.5.1 0.22.0 + 3.7.8 4.12.0 diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index afba258b6..300ee79c4 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -291,11 +291,17 @@ ${project.version} - - org.springframework.ai - spring-ai-opensearch-store - ${project.version} - + + org.springframework.ai + spring-ai-opensearch-store + ${project.version} + + + + org.springframework.ai + spring-ai-couchbase-store + ${project.version} + org.springframework.ai @@ -599,11 +605,17 @@ ${project.version} - - org.springframework.ai - spring-ai-qianfan-spring-boot-starter - ${project.version} - + + org.springframework.ai + spring-ai-qianfan-spring-boot-starter + ${project.version} + + + + org.springframework.ai + spring-ai-couchbase-store-spring-boot-starter + ${project.version} + org.springframework.ai diff --git a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java index a65bb57c3..a9c3e3f2b 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java @@ -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. */ diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc index 012f8441e..04f495f9b 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc @@ -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[] diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/couchbase.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/couchbase.adoc new file mode 100644 index 000000000..f6fbb34fd --- /dev/null +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/couchbase.adoc @@ -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] +---- + + org.springframework.ai + spring-ai-couchbase-store-spring-boot-starter + +---- + +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 <> 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 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 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= +spring.couchbase.connection-string= +spring.couchbase.username= +spring.couchbase.password= +---- + +environment variables, + +[source,bash] +---- +export SPRING_COUCHBASE_CONNECTION_STRINGS= +export SPRING_COUCHBASE_USERNAME= +export SPRING_COUCHBASE_PASSWORD= +# API key if needed, e.g. OpenAI +export SPRING_AI_OPENAI_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 .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] +---- + + org.springframework.ai + spring-ai-couchbase-store + +---- + +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. diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index 402abd895..4ce105a7f 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -427,6 +427,13 @@ ${project.parent.version} true + + + org.springframework.ai + spring-ai-couchbase-store + ${project.parent.version} + true + @@ -606,6 +613,12 @@ test - + + org.testcontainers + couchbase + test + + + diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfiguration.java new file mode 100644 index 000000000..41cb99fdb --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfiguration.java @@ -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(); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreProperties.java new file mode 100644 index 000000000..288df510c --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreProperties.java @@ -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; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index 286f6b36a..132c474a9 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -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 diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseContainerMetadata.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseContainerMetadata.java new file mode 100644 index 000000000..ae3dbbf65 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseContainerMetadata.java @@ -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"); + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfigurationIT.java new file mode 100644 index 000000000..a53e02862 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/couchbase/CouchbaseSearchVectorStoreAutoConfigurationIT.java @@ -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 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(); + }); + } + +} diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-couchbase-store/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-couchbase-store/pom.xml new file mode 100644 index 000000000..632ec455a --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-couchbase-store/pom.xml @@ -0,0 +1,42 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-couchbase-store-spring-boot-starter + jar + Spring AI Starter - Couchbase Store + Spring AI Couchbase Store Auto Configuration + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + + org.springframework.boot + spring-boot-starter + + + + org.springframework.ai + spring-ai-spring-boot-autoconfigure + ${project.parent.version} + + + + org.springframework.ai + spring-ai-couchbase-store + ${project.parent.version} + + + + diff --git a/vector-stores/spring-ai-couchbase-store/README.md b/vector-stores/spring-ai-couchbase-store/README.md new file mode 100644 index 000000000..7e2f69702 --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/README.md @@ -0,0 +1 @@ +[Couchbase Vector Store Documentation](https://docs.spring.io/spring-ai/reference/1.0-SNAPSHOT/api/vectordbs/couchbase.html) \ No newline at end of file diff --git a/vector-stores/spring-ai-couchbase-store/pom.xml b/vector-stores/spring-ai-couchbase-store/pom.xml new file mode 100644 index 000000000..3d9127739 --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/pom.xml @@ -0,0 +1,74 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-couchbase-store + jar + Spring AI Vector Store - Couchbase + Spring AI Couchbase Vector Store + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + com.couchbase.client + java-client + ${couchbase.version} + + + org.springframework.ai + spring-ai-core + ${parent.version} + + + + + + org.springframework.ai + spring-ai-openai + ${parent.version} + test + + + + org.springframework.ai + spring-ai-test + ${parent.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + org.springframework.boot + spring-boot-testcontainers + test + + + org.testcontainers + couchbase + test + + + org.testcontainers + junit-jupiter + test + + + + diff --git a/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseAiSearchFilterExpressionConverter.java b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseAiSearchFilterExpressionConverter.java new file mode 100644 index 000000000..415247daf --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseAiSearchFilterExpressionConverter.java @@ -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(")"); + } + +} diff --git a/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseIndexOptimization.java b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseIndexOptimization.java new file mode 100644 index 000000000..c36bd3d72 --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseIndexOptimization.java @@ -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 + +} diff --git a/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStore.java b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStore.java new file mode 100644 index 000000000..58fabf5ec --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStore.java @@ -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 documents) { + logger.info("Trying Add"); + logger.info(this.bucketName); + logger.info(this.scopeName); + List 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 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 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 Optional 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 { + + 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 metadata, float[] embedding) { + } + +} diff --git a/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSimilarityFunction.java b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSimilarityFunction.java new file mode 100644 index 000000000..e2ac036da --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/main/java/org/springframework/ai/vectorstore/CouchbaseSimilarityFunction.java @@ -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 + +} diff --git a/vector-stores/spring-ai-couchbase-store/src/main/resources/application.properties b/vector-stores/spring-ai-couchbase-store/src/main/resources/application.properties new file mode 100644 index 000000000..33239c9cd --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/main/resources/application.properties @@ -0,0 +1 @@ +spring.application.name=demo \ No newline at end of file diff --git a/vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStoreIT.java b/vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStoreIT.java new file mode 100644 index 000000000..b670b9e60 --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/CouchbaseSearchVectorStoreIT.java @@ -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 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 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 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 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 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 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 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 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 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()); + } + + } + +} diff --git a/vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/testcontainer/CouchbaseContainerMetadata.java b/vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/testcontainer/CouchbaseContainerMetadata.java new file mode 100644 index 000000000..0b9a2e621 --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/test/java/org/springframework/ai/vectorstore/testcontainer/CouchbaseContainerMetadata.java @@ -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"); + +} diff --git a/vector-stores/spring-ai-couchbase-store/src/test/resources/application.properties b/vector-stores/spring-ai-couchbase-store/src/test/resources/application.properties new file mode 100644 index 000000000..33239c9cd --- /dev/null +++ b/vector-stores/spring-ai-couchbase-store/src/test/resources/application.properties @@ -0,0 +1 @@ +spring.application.name=demo \ No newline at end of file