diff --git a/pom.xml b/pom.xml index 6bf7fb6b9..1aa983322 100644 --- a/pom.xml +++ b/pom.xml @@ -34,6 +34,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-azure-openai spring-ai-spring-boot-starters/spring-ai-starter-ollama spring-ai-spring-boot-starters/spring-ai-starter-transformers + spring-ai-spring-boot-starters/spring-ai-starter-cassandra spring-ai-spring-boot-starters/spring-ai-starter-chroma-store spring-ai-spring-boot-starters/spring-ai-starter-milvus-store spring-ai-spring-boot-starters/spring-ai-starter-pgvector-store @@ -47,6 +48,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-qdrant-store spring-ai-spring-boot-starters/spring-ai-starter-postgresml-embedding spring-ai-docs + vector-stores/spring-ai-cassandra vector-stores/spring-ai-pgvector-store vector-stores/spring-ai-hanadb-store vector-stores/spring-ai-milvus-store @@ -138,6 +140,7 @@ 3.25.2 + 4.18.0 3.0.1 0.1.4 2.20.11 diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index 1bc27de0c..d1c857f9d 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -132,6 +132,12 @@ ${project.version} + + org.springframework.ai + spring-ai-cassandra + ${project.version} + + org.springframework.ai spring-ai-chroma-store @@ -218,6 +224,12 @@ + + org.springframework.ai + spring-ai-apache-cassandra-store-spring-boot-starter + ${project.version} + + org.springframework.ai spring-ai-azure-openai-spring-boot-starter 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 cb34f6bbb..82793cc94 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc @@ -46,23 +46,25 @@ **** xref:api/audio/speech/openai-speech.adoc[OpenAI] ** xref:api/vectordbs.adoc[] *** xref:api/vectordbs/azure.adoc[] +*** xref:api/vectordbs/apache-cassandra.adoc[] *** xref:api/vectordbs/chroma.adoc[] +*** xref:api/vectordbs/gemfire.adoc[GemFire] *** xref:api/vectordbs/milvus.adoc[] *** xref:api/vectordbs/neo4j.adoc[] *** xref:api/vectordbs/pgvector.adoc[] -*** xref:api/vectordbs/weaviate.adoc[] -*** xref:api/vectordbs/redis.adoc[] *** xref:api/vectordbs/pinecone.adoc[] *** xref:api/vectordbs/qdrant.adoc[] -*** xref:api/vectordbs/gemfire.adoc[GemFire] +*** xref:api/vectordbs/redis.adoc[] *** xref:api/vectordbs/hana.adoc[SAP Hana] +*** xref:api/vectordbs/weaviate.adoc[] + ** xref:api/functions.adoc[Function Calling] ** xref:api/prompt.adoc[] ** xref:api/output-parser.adoc[] ** xref:api/etl-pipeline.adoc[] ** xref:api/testing.adoc[] ** xref:api/generic-model.adoc[] -* xref:api/testcontainers.adoc[Testcontainers] * xref:contribution-guidelines.adoc[Contribution Guidelines] * Appendices ** xref:upgrade-notes.adoc[] +** xref:api/testcontainers.adoc[Testcontainers] diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc index 0b4385c14..afc12ba2e 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc @@ -88,14 +88,18 @@ Find more information on the `Filter.Expression` in the <> sec These are the available implementations of the `VectorStore` interface: * xref:api/vectordbs/azure.adoc[ Azure Vector Search] - The https://learn.microsoft.com/en-us/azure/search/vector-search-overview[Azure] vector store. -* xref:api/vectordbs/chroma.adoc[ChromaVectorStore] - The https://www.trychroma.com/[Chroma] vector store. -* xref:api/vectordbs/milvus.adoc[MilvusVectorStore] - The https://milvus.io/[Milvus] vector store. -* xref:api/vectordbs/neo4j.adoc[Neo4jVectorStore] - The https://neo4j.com/[Neo4j] vector store. +* xref:api/vectordbs/apache-cassandra.adoc[Apache Cassandra] - The https://cassandra.apache.org/doc/latest/cassandra/vector-search/overview.html[Apache Cassandra] +* xref:api/vectordbs/chroma.adoc[Chroma Vector Store] - The https://www.trychroma.com/[Chroma] vector store. +* xref:api/vectordbs/gemfire.adoc[GemFire Vector Store] - The https://tanzu.vmware.com/content/blog/vmware-gemfire-vector-database-extension[GemFire] vector store. +* xref:api/vectordbs/milvus.adoc[Milvus Vector Store] - The https://milvus.io/[Milvus] vector store. +* xref:api/vectordbs/neo4j.adoc[Neo4j Vector Store] - The https://neo4j.com/[Neo4j] vector store. * xref:api/vectordbs/pgvector.adoc[PgVectorStore] - The https://github.com/pgvector/pgvector[PostgreSQL/PGVector] vector store. -* xref:api/vectordbs/pinecone.adoc[PineconeVectorStore] - https://www.pinecone.io/[PineCone] vector store. -* xref:api/vectordbs/qdrant.adoc[QdrantVectorStore] - https://www.qdrant.tech/[Qdrant] vector store. -* xref:api/vectordbs/redis.adoc[RedisVectorStore] - The https://redis.io/[Redis] vector store. -* xref:api/vectordbs/weaviate.adoc[WeaviateVectorStore] - The https://weaviate.io/[Weaviate] vector store. +* xref:api/vectordbs/pinecone.adoc[Pinecone Vector Store] - https://www.pinecone.io/[PineCone] vector store. +* xref:api/vectordbs/qdrant.adoc[Qdrant Vector Store] - https://www.qdrant.tech/[Qdrant] vector store. +* xref:api/vectordbs/redis.adoc[Redis Vector Store] - The https://redis.io/[Redis] vector store. +* xref:api/vectordbs/hana.adoc[SAP Hana Vector Store] - The https://news.sap.com/2024/04/sap-hana-cloud-vector-engine-ai-with-business-context/[SAP HANA] vector store. +* xref:api/vectordbs/weaviate.adoc[Weaviate Vector Store] - The https://weaviate.io/[Weaviate] vector store. +vector store. * link:https://github.com/spring-projects/spring-ai/blob/main/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java[SimpleVectorStore] - A simple implementation of persistent vector storage, good for educational purposes. More implementations may be supported in future releases. diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc new file mode 100644 index 000000000..a264c08c1 --- /dev/null +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc @@ -0,0 +1,260 @@ += Apache Cassandra + +This section walks you through setting up `CassandraVectorStore` to store document embeddings and perform similarity searches. + +== What is Apache Cassandra ? + +link:https://cassandra.apache.org[Apache Cassandra] is a true open source distributed database reknown for scalability and high availability without compromising performance. + +Linear scalability, proven fault-tolerance and low latency on commodity hardware makes it the perfect platform for mission-critical data. Its Vector Similarity Search (VSS) is based on the JVector library that ensures best-in-class performance and relevancy. + +A vector search in Apache Cassandra is done as simply as: +``` +SELECT content FROM table ORDER BY content_vector ANN OF query_embedding ; +``` + +More docs on this can be read https://cassandra.apache.org/doc/latest/cassandra/getting-started/vector-search-quickstart.html[here]. + +The Spring AI Cassandra Vector Store is designed to work for both brand new RAG applications as well as being able to be retrofitted on top of existing data and tables. This vector store may also equally be used for non-RAG non_AI use-cases, e.g. semantic searcing in an existing database. The Vector Store will automatically create, or enhance, the schema as needed according to its configuration. If you don't want the schema modifications, configure the store with `disallowSchemaChanges`. + +== What is JVector Vector Search ? + +link:https://github.com/jbellis/jvector[JVector] is a pure Java embedded vector search engine. + +It stands out from other HNSW Vector Similarity Search implementations by being + +* Algorithmic-fast. JVector uses state of the art graph algorithms inspired by DiskANN and related research that offer high recall and low latency. +* Implementation-fast. JVector uses the Panama SIMD API to accelerate index build and queries. +* Memory efficient. JVector compresses vectors using product quantization so they can stay in memory during searches. (As part of our PQ implementation, our SIMD-accelerated kmeans class is 5x faster than the one in Apache Commons Math.) +* Disk-aware. JVector’s disk layout is designed to do the minimum necessary iops at query time. +* Concurrent. Index builds scale linearly to at least 32 threads. Double the threads, half the build time. +* Incremental. Query your index as you build it. No delay between adding a vector and being able to find it in search results. +* Easy to embed. API designed for easy embedding, by people using it in production. + +== Prerequisites + +1. A `EmbeddingClient` instance to compute the document embeddings. This is usually configured as a Spring Bean. Several options are available: + +- `Transformers Embedding` - computes the embedding in your local environment. The default is via ONNX and the all-MiniLM-L6-v2 Sentence Transformers. This just works. +- If you want to use OpenAI's Embeddings` - uses the OpenAI embedding endpoint. You need to create an account at link:https://platform.openai.com/signup[OpenAI Signup] and generate the api-key token at link:https://platform.openai.com/account/api-keys[API Keys]. +- There are many more choices, see `Embeddings API` docs. + +2. An Apache Cassandra instance, from version 5.0-beta1 +a. link:https://cassandra.apache.org/_/quickstart.html[DIY Quick Start] +b. For a managed offering https://astra.datastax.com/[Astra DB] offers a healthy free tier offering. + +== Dependencies + +Add these dependencies to your project: + +* For just the Cassandra Vector Store + +[source,xml] +---- + + org.springframework.ai + spring-ai-cassandra + +---- + +* Or, for everything you need in a RAG application (using the default ONNX Embedding Client) + +[source,xml] +---- + + org.springframework.ai + spring-ai-cassandra-spring-boot-starter + +---- + + +TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file. + +* If for example you want to use the OpenAI modules, remember to provide your OpenAI API Key. Set it as an environment variable like so: + +[source,bash] +---- +export SPRING_AI_OPENAI_API_KEY='Your_OpenAI_API_Key' +---- + + +== Usage + +Create a CassandraVectorStore instance connected to your Apache Cassandra database: + +[source,java] +---- +@Bean +public VectorStore vectorStore(EmbeddingClient embeddingClient) { + + CassandraVectorStoreConfig config = CassandraVectorStoreConfig.builder().build(); + + return new CassandraVectorStore(config, embeddingClient); +} +---- + +NOTE: It is more convenient and preferred to create the `CassandraVectorStore` as a Bean. +But if you decide you can create it manually. + +[NOTE] +==== +The default configuration connects to Cassandra at localhost:9042 and will automatically create the default schema at `springframework_ai_vector.springframework_ai_vector_store`. + +Please see `CassandraVectorStoreConfig.Builder` for all the configuration options. +==== + +[NOTE] +==== +The Cassandra Java Driver is easiest configured via the `application.conf` file on the classpath. + +More info can be found link: https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration[here]. +==== + +Then in your main code, create some documents: + +[source,java] +---- +List documents = List.of( + new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("country", "UK", "year", 2020)), + new Document("The World is Big and Salvation Lurks Around the Corner", Map.of()), + new Document("You walk forward facing the past and you turn back toward the future.", Map.of("country", "NL", "year", 2023))); +---- + +Now add the documents to your vector store: + + +[source,java] +---- +vectorStore.add(documents); +---- + +And finally, retrieve documents similar to a query: + +[source,java] +---- +List results = vectorStore.similaritySearch( + SearchRequest.query("Spring").withTopK(5)); +---- + +If all goes well, you should retrieve the document containing the text "Spring AI rocks!!". + +You can also limit results based on a similarity threshold: +[source,java] +---- +List results = vectorStore.similaritySearch( + SearchRequest.query("Spring").withTopK(5) + .withSimilarityThreshold(0.5d)); +---- + +=== 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 CassandraVectorStore as well. Metadata fields must be configured in `CassandraVectorStoreConfig`. + +For example, you can use either the text expression language: + +[source,java] +---- +vectorStore.similaritySearch( + SearchRequest.query("The World").withTopK(TOP_K) + .withFilterExpression("country in ['UK', 'NL'] && year >= 2020")); +---- + +or programmatically using the expression DSL: + +[source,java] +---- +Filter.Expression f = new FilterExpressionBuilder() + .and(f.in("country", "UK", "NL"), f.gte("year", 2020)).build(); + +vectorStore.similaritySearch( + SearchRequest.query("The World").withTopK(TOP_K) + .withFilterExpression(f)); +---- + +The portable filter expressions get automatically converted into link:https://cassandra.apache.org/doc/latest/cassandra/developing/cql/index.html[CQL queries]. + +Metadata fields to be searchable need to be either primary key columns or SAI indexed. To do this configure the metadata field with the `SchemaColumnTags.INDEXED`. + + +== Advanced Example: Vector Store ontop full Wikipedia dataset + +The following example demonstrates how to use the store on an existing schema. Here we use the schema from the https://github.com/datastax-labs/colbert-wikipedia-data project which comes with the full wikipedia dataset ready vectorised for you. + + +== Usage + +Create the schema in the Cassandra database first: + +[source,bash] +---- +wget https://raw.githubusercontent.com/datastax-labs/colbert-wikipedia-data/main/schema.cql -O colbert-wikipedia-schema.cql +cqlsh -f colbert-wikipedia-schema.cql +---- + +Then configure the store like: + +[source,java] +---- +@Bean +public CassandraVectorStore store(EmbeddingClient embeddingClient) { + + List partitionColumns = List.of(new SchemaColumn("wiki", DataTypes.TEXT), + new SchemaColumn("language", DataTypes.TEXT), new SchemaColumn("title", DataTypes.TEXT)); + + List clusteringColumns = List.of(new SchemaColumn("chunk_no", DataTypes.INT), + new SchemaColumn("bert_embedding_no", DataTypes.INT)); + + List extraColumns = List.of(new SchemaColumn("revision", DataTypes.INT), + new SchemaColumn("id", DataTypes.INT)); + + CassandraVectorStoreConfig conf = CassandraVectorStoreConfig.builder() + .withKeyspaceName("wikidata") + .withTableName("articles") + .withPartitionKeys(partitionColumns) + .withClusteringKeys(clusteringColumns) + .withContentFieldName("body") + .withEmbeddingFieldName("all_minilm_l6_v2_embedding") + .withIndexName("all_minilm_l6_v2_ann") + .disallowSchemaChanges() + .addMetadataFields(extraColumns) + .withPrimaryKeyTranslator((List primaryKeys) -> { + // the deliminator used to join fields together into the document's id + // is arbitary, here "§¶" is used + if (primaryKeys.isEmpty()) { + return "test§¶0"; + } + return format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3)); + }) + .withDocumentIdTranslator((id) -> { + String[] parts = id.split("§¶"); + String title = parts[0]; + int chunk_no = 0 < parts.length ? Integer.parseInt(parts[1]) : 0; + return List.of("simplewiki", "en", title, chunk_no, 0); + }) + .build(); + + return new CassandraVectorStore(conf, embeddingClient()); +} + +@Bean +public EmbeddingClient embeddingClient() { + // default is ONNX all-MiniLM-L6-v2 which is what we want + return new TransformersEmbeddingClient(); +} +---- + +And, if you would like to load the full wikipedia dataset. +First download the `simplewiki-sstable.tar` from this link https://drive.google.com/file/d/1CcMMsj8jTKRVGep4A7hmOSvaPepsaKYP/view?usp=share_link . This will take a while, the file is tens of GBs. + +[source,bash] +---- +tar -xf simplewiki-sstable.tar -C ${CASSANDRA_DATA}/data/wikidata/articles-*/ + +nodetool import wikidata articles ${CASSANDRA_DATA}/data/wikidata/articles-*/ +---- + +NOTE: If you have existing data in this table you'll want to check the tarball's files don't clobber existing sstables when doing the `tar`. + +NOTE: An alternative to the `nodetool import` is to just restart Cassandra. + +NOTE: If there are any failures in the indexes they will be rebuilt automatically. diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/index.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/index.adoc index a1e95f488..a6761e612 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/index.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/index.adoc @@ -15,7 +15,7 @@ Spring AI provides the following features: * Supported Model types are Chat and Text to Image with more on the way. * Portable API across AI providers for Chat and for Embedding models. Both synchronous and stream API options are supported. Dropping down to access model specific features is also supported. * Mapping of AI Model output to POJOs. -* Support for all major Vector Database providers such as Azure Vector Search, Chroma, Milvus, Neo4j, PostgreSQL/PGVector, PineCone, Qdrant, Redis, and Weaviate +* Support for all major Vector Database providers such as Apache Cassandra, Azure Vector Search, Chroma, Milvus, Neo4j, PostgreSQL/PGVector, PineCone, Qdrant, Redis, and Weaviate * Portable API across Vector Store providers, including a novel SQL-like metadata filter API that is also portable. * Function calling * Spring Boot Auto Configuration and Starters for AI Models and Vector Stores. diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index 547d79b79..d48275f3d 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -146,6 +146,14 @@ true + + + org.springframework.ai + spring-ai-cassandra + ${project.parent.version} + true + + org.springframework.ai @@ -309,6 +317,12 @@ test + + org.testcontainers + cassandra + test + + com.redis testcontainers-redis diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraConnectionDetails.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraConnectionDetails.java new file mode 100644 index 000000000..b67f90f6a --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraConnectionDetails.java @@ -0,0 +1,37 @@ +/* + * Copyright 2024 - 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.cassandra; + +import java.net.InetSocketAddress; +import java.util.List; + +import org.springframework.boot.autoconfigure.service.connection.ConnectionDetails; + +/** + * @author Mick Semb Wever + * @since 1.0.0 + */ +public interface CassandraConnectionDetails extends ConnectionDetails { + + boolean hasCassandraContactPoints(); + + List getCassandraContactPoints(); + + boolean hasCassandraLocalDatacenter(); + + String getCassandraLocalDatacenter(); + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java new file mode 100644 index 000000000..1bb93f410 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java @@ -0,0 +1,118 @@ +/* + * Copyright 2024 - 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.cassandra; + +import java.net.InetSocketAddress; +import java.util.Arrays; +import java.util.List; + +import com.google.common.base.Preconditions; + +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.vectorstore.CassandraVectorStore; +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +/** + * @author Mick Semb Wever + * @since 1.0.0 + */ +@AutoConfiguration +@ConditionalOnClass({ CassandraVectorStore.class, EmbeddingClient.class }) +@EnableConfigurationProperties(CassandraVectorStoreProperties.class) +public class CassandraVectorStoreAutoConfiguration { + + @Bean + @ConditionalOnMissingBean(CassandraConnectionDetails.class) + public PropertiesCassandraConnectionDetails cassandraConnectionDetails(CassandraVectorStoreProperties properties) { + return new PropertiesCassandraConnectionDetails(properties); + } + + @Bean + @ConditionalOnMissingBean + public CassandraVectorStore vectorStore(EmbeddingClient embeddingClient, CassandraVectorStoreProperties properties, + CassandraConnectionDetails cassandraConnectionDetails) { + + var builder = CassandraVectorStoreConfig.builder(); + if (cassandraConnectionDetails.hasCassandraContactPoints()) { + for (InetSocketAddress contactPoint : cassandraConnectionDetails.getCassandraContactPoints()) { + builder = builder.addContactPoint(contactPoint); + } + } + if (cassandraConnectionDetails.hasCassandraLocalDatacenter()) { + builder = builder.withLocalDatacenter(cassandraConnectionDetails.getCassandraLocalDatacenter()); + } + + builder = builder.withKeyspaceName(properties.getKeyspace()) + .withTableName(properties.getTable()) + .withContentColumnName(properties.getContentFieldName()) + .withEmbeddingColumnName(properties.getEmbeddingFieldName()) + .withIndexName(properties.getIndexName()); + + if (properties.getDisallowSchemaCreation()) { + builder = builder.disallowSchemaChanges(); + } + + return new CassandraVectorStore(builder.build(), embeddingClient); + } + + private static class PropertiesCassandraConnectionDetails implements CassandraConnectionDetails { + + private final CassandraVectorStoreProperties properties; + + public PropertiesCassandraConnectionDetails(CassandraVectorStoreProperties properties) { + this.properties = properties; + } + + private String[] getCassandraContactPointHosts() { + return this.properties.getCassandraContactPointHosts().split("(,| )"); + } + + @Override + public List getCassandraContactPoints() { + + Preconditions.checkState(hasCassandraContactPoints(), "cassandraContactPointHosts has not been set"); + final int port = this.properties.getCassandraContactPointPort(); + + return Arrays.asList(getCassandraContactPointHosts()) + .stream() + .map((host) -> InetSocketAddress.createUnresolved(host, port)) + .toList(); + } + + @Override + public String getCassandraLocalDatacenter() { + Preconditions.checkState(hasCassandraLocalDatacenter(), "cassandraLocalDatacenter has not been set"); + return this.properties.getCassandraLocalDatacenter(); + } + + @Override + public boolean hasCassandraContactPoints() { + return null != this.properties.getCassandraContactPointHosts(); + } + + @Override + public boolean hasCassandraLocalDatacenter() { + return null != this.properties.getCassandraLocalDatacenter(); + } + + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java new file mode 100644 index 000000000..73b014ab6 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreProperties.java @@ -0,0 +1,121 @@ +/* + * Copyright 2024 - 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.cassandra; + +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * @author Mick Semb Wever + * @since 1.0.0 + */ +@ConfigurationProperties(CassandraVectorStoreProperties.CONFIG_PREFIX) +public class CassandraVectorStoreProperties { + + public static final String CONFIG_PREFIX = "spring.ai.vectorstore.cassandra"; + + private String cassandraContactPointHosts = null; + + private int cassandraContactPointPort = 9042; + + private String cassandraLocalDatacenter = null; + + private String keyspace = CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME; + + private String table = CassandraVectorStoreConfig.DEFAULT_TABLE_NAME; + + private String indexName = CassandraVectorStoreConfig.DEFAULT_INDEX_NAME; + + private String contentColumnName = CassandraVectorStoreConfig.DEFAULT_CONTENT_COLUMN_NAME; + + private String embeddingColumnName = CassandraVectorStoreConfig.DEFAULT_EMBEDDING_COLUMN_NAME; + + private boolean disallowSchemaChanges = false; + + public String getCassandraContactPointHosts() { + return this.cassandraContactPointHosts; + } + + /** comma or space separated */ + public void setCassandraContactPointHosts(String cassandraContactPointHosts) { + this.cassandraContactPointHosts = cassandraContactPointHosts; + } + + public int getCassandraContactPointPort() { + return this.cassandraContactPointPort; + } + + public void setCassandraContactPointPort(int cassandraContactPointPort) { + this.cassandraContactPointPort = cassandraContactPointPort; + } + + public String getCassandraLocalDatacenter() { + return this.cassandraLocalDatacenter; + } + + public void setCassandraLocalDatacenter(String cassandraLocalDatacenter) { + this.cassandraLocalDatacenter = cassandraLocalDatacenter; + } + + public String getKeyspace() { + return this.keyspace; + } + + public void setKeyspace(String keyspace) { + this.keyspace = keyspace; + } + + public String getTable() { + return this.table; + } + + public void setTable(String table) { + this.table = table; + } + + public String getIndexName() { + return this.indexName; + } + + public void setIndexName(String indexName) { + this.indexName = indexName; + } + + public String getContentFieldName() { + return this.contentColumnName; + } + + public void setContentFieldName(String contentFieldName) { + this.contentColumnName = contentFieldName; + } + + public String getEmbeddingFieldName() { + return this.embeddingColumnName; + } + + public void setEmbeddingFieldName(String embeddingFieldName) { + this.embeddingColumnName = embeddingFieldName; + } + + public Boolean getDisallowSchemaCreation() { + return this.disallowSchemaChanges; + } + + public void setDisallowSchemaCreation(boolean disallowSchemaCreation) { + this.disallowSchemaChanges = disallowSchemaCreation; + } + +} 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 3e3fd3574..59f104b68 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 @@ -31,3 +31,4 @@ org.springframework.ai.autoconfigure.vectorstore.mongo.MongoDBAtlasVectorStoreAu org.springframework.ai.autoconfigure.anthropic.AnthropicAutoConfiguration org.springframework.ai.autoconfigure.watsonxai.WatsonxAiAutoConfiguration org.springframework.ai.autoconfigure.vectorstore.elasticsearch.ElasticsearchVectorStoreAutoConfiguration +org.springframework.ai.autoconfigure.vectorstore.cassandra.CassandraVectorStoreAutoConfiguration diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java new file mode 100644 index 000000000..15530052f --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java @@ -0,0 +1,109 @@ +/* + * Copyright 2024 - 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.cassandra; + +import java.util.List; +import java.util.Map; + +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.containers.CassandraContainer; +import org.testcontainers.utility.DockerImageName; + +import org.springframework.ai.ResourceUtils; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.transformers.TransformersEmbeddingClient; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Mick Semb Wever + * @since 1.0.0 + */ +@Testcontainers +class CassandraVectorStoreAutoConfigurationIT { + + static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("cassandra"); + + @Container + static CassandraContainer cassandraContainer = new CassandraContainer(DEFAULT_IMAGE_NAME.withTag("5.0")); + + List documents = List.of( + new Document(ResourceUtils.getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")), + new Document(ResourceUtils.getText("classpath:/test/data/time.shelter.txt")), new Document( + ResourceUtils.getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad"))); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(CassandraVectorStoreAutoConfiguration.class)) + .withUserConfiguration(Config.class) + .withPropertyValues("spring.ai.vectorstore.cassandra.keyspace=test_autoconfigure") + .withPropertyValues("spring.ai.vectorstore.cassandra.contentFieldName=doc_chunk"); + + @Test + void addAndSearch() { + contextRunner + .withPropertyValues("spring.ai.vectorstore.cassandra.cassandraContactPointHosts=" + getContactPointHost()) + .withPropertyValues("spring.ai.vectorstore.cassandra.cassandraContactPointPort=" + getContactPointPort()) + .withPropertyValues("spring.ai.vectorstore.cassandra.cassandraLocalDatacenter=" + + cassandraContainer.getLocalDatacenter()) + + .run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + vectorStore.add(documents); + + List results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId()); + assertThat(resultDoc.getContent()).contains( + "Spring AI provides abstractions that serve as the foundation for developing AI applications."); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList()); + + results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + assertThat(results).isEmpty(); + }); + } + + @Configuration(proxyBeanMethods = false) + static class Config { + + @Bean + public EmbeddingClient embeddingClient() { + return new TransformersEmbeddingClient(); + } + + } + + private String getContactPointHost() { + return cassandraContainer.getContactPoint().getHostString(); + } + + private String getContactPointPort() { + return String.valueOf(cassandraContainer.getContactPoint().getPort()); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java new file mode 100644 index 000000000..c5a3e4d0f --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStorePropertiesTests.java @@ -0,0 +1,68 @@ +/* + * Copyright 2024 - 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.cassandra; + +import org.junit.jupiter.api.Test; + +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Mick Semb Wever + * @since 1.0.0 + */ +class CassandraVectorStorePropertiesTests { + + @Test + void defaultValues() { + var props = new CassandraVectorStoreProperties(); + assertThat(props.getCassandraContactPointHosts()).isNull(); + assertThat(props.getCassandraContactPointPort()).isEqualTo(9042); + assertThat(props.getCassandraLocalDatacenter()).isNull(); + assertThat(props.getKeyspace()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME); + assertThat(props.getTable()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_TABLE_NAME); + assertThat(props.getContentFieldName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_CONTENT_COLUMN_NAME); + assertThat(props.getEmbeddingFieldName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_EMBEDDING_COLUMN_NAME); + assertThat(props.getIndexName()).isEqualTo(CassandraVectorStoreConfig.DEFAULT_INDEX_NAME); + assertThat(props.getDisallowSchemaCreation()).isFalse(); + } + + @Test + void customValues() { + var props = new CassandraVectorStoreProperties(); + props.setCassandraContactPointHosts("127.0.0.1,127.0.0.2"); + props.setCassandraContactPointPort(9043); + props.setCassandraLocalDatacenter("dc1"); + props.setKeyspace("my_keyspace"); + props.setTable("my_table"); + props.setContentFieldName("my_content"); + props.setEmbeddingFieldName("my_vector"); + props.setIndexName("my_sai"); + props.setDisallowSchemaCreation(true); + + assertThat(props.getCassandraContactPointHosts()).isEqualTo("127.0.0.1,127.0.0.2"); + assertThat(props.getCassandraContactPointPort()).isEqualTo(9043); + assertThat(props.getCassandraLocalDatacenter()).isEqualTo("dc1"); + assertThat(props.getKeyspace()).isEqualTo("my_keyspace"); + assertThat(props.getTable()).isEqualTo("my_table"); + assertThat(props.getContentFieldName()).isEqualTo("my_content"); + assertThat(props.getEmbeddingFieldName()).isEqualTo("my_vector"); + assertThat(props.getIndexName()).isEqualTo("my_sai"); + assertThat(props.getDisallowSchemaCreation()).isTrue(); + } + +} diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-cassandra/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-cassandra/pom.xml new file mode 100644 index 000000000..a40111ee8 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-cassandra/pom.xml @@ -0,0 +1,42 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-cassandra-spring-boot-starter + jar + Spring AI Starter - Apache Cassandra Vector Store + Spring AI Apache Cassandra Vector 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-cassandra + ${project.parent.version} + + + + diff --git a/vector-stores/spring-ai-cassandra/README.md b/vector-stores/spring-ai-cassandra/README.md new file mode 100644 index 000000000..d35bcbfeb --- /dev/null +++ b/vector-stores/spring-ai-cassandra/README.md @@ -0,0 +1 @@ +[Apache Cassandra Vector Store Documentation](https://docs.spring.io/spring-ai/reference/api/vectordbs/apache-cassandra.html) \ No newline at end of file diff --git a/vector-stores/spring-ai-cassandra/pom.xml b/vector-stores/spring-ai-cassandra/pom.xml new file mode 100644 index 000000000..b01d200f8 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/pom.xml @@ -0,0 +1,76 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-cassandra + jar + Spring AI Vector Store – Apache Cassandra + Spring AI Vector Store for Apache Cassandra + 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 + + + + 17 + 17 + + + + + org.springframework.ai + spring-ai-core + ${project.parent.version} + + + + org.apache.cassandra + java-driver-query-builder + ${cassandra.java-driver.version} + + + + + org.springframework.ai + spring-ai-transformers + ${project.parent.version} + test + + + + org.springframework.ai + spring-ai-test + ${project.parent.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.testcontainers + junit-jupiter + test + + + + org.testcontainers + cassandra + test + + + + + diff --git a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java new file mode 100644 index 000000000..4a8b681a5 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java @@ -0,0 +1,139 @@ +/* + * Copyright 2024 - 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 java.util.Collection; +import java.util.Map; +import java.util.Optional; +import java.util.function.Function; +import java.util.stream.Collectors; + +import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata; +import com.datastax.oss.driver.api.core.type.DataTypes; +import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry; +import com.datastax.oss.driver.shaded.guava.common.base.Preconditions; + +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.Filter.ExpressionType; +import org.springframework.ai.vectorstore.filter.Filter.Key; +import org.springframework.ai.vectorstore.filter.Filter.Value; +import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter; + +/** + * Converts {@link Expression} into CQL where clauses. + * + * @author Mick Semb Wever + * @since 1.0.0 + */ +final class CassandraFilterExpressionConverter extends AbstractFilterExpressionConverter { + + private final Map columnsByName; + + public CassandraFilterExpressionConverter(Collection columns) { + + this.columnsByName = columns.stream() + .collect(Collectors.toMap((c) -> c.getName().asInternal(), Function.identity())); + } + + @Override + protected void doKey(Key key, StringBuilder context) { + String keyName = key.key(); + Optional column = getColumn(keyName); + Preconditions.checkArgument(column.isPresent(), "No metafield %s has been configured", keyName); + context.append(column.get().getName().asCql(false)); + } + + @Override + protected void doExpression(Filter.Expression expression, StringBuilder context) { + switch (expression.type()) { + case AND -> doBinaryOperation(" and ", expression, context); + case OR -> doBinaryOperation(" or ", expression, context); + case NIN, NOT -> throw new UnsupportedOperationException( + String.format("Expression type %s not yet implemented. Patches welcome.", expression.type())); + default -> doField(expression, context); + } + } + + private static void doOperand(ExpressionType type, StringBuilder context) { + switch (type) { + case EQ -> context.append(" = "); + case NE -> context.append(" != "); + case GT -> context.append(" > "); + case GTE -> context.append(" >= "); + case IN -> context.append(" IN "); + case LT -> context.append(" < "); + case LTE -> context.append(" <= "); + // TODO SAI supports collections + // reach out to mck@apache.org if you'd like these implemented + // case CONTAINS -> context.append(" CONTAINS "); + // case CONTAINS_KEY -> context.append(" CONTAINS KEY "); + default -> throw new UnsupportedOperationException( + String.format("Expression type %s not yet implemented. Patches welcome.", type)); + } + } + + private void doBinaryOperation(String operator, Filter.Expression expression, StringBuilder context) { + this.convertOperand(expression.left(), context); + context.append(operator); + this.convertOperand(expression.right(), context); + } + + private void doField(Filter.Expression expression, StringBuilder context) { + doKey((Key) expression.left(), context); + doOperand(expression.type(), context); + ColumnMetadata column = getColumn(((Key) expression.left()).key()).get(); + var v = ((Value) expression.right()).value(); + if (ExpressionType.IN.equals(expression.type())) { + Preconditions.checkArgument(v instanceof Collection); + doListValue(column, v, context); + } + else { + doValue(column, v, context); + } + } + + private void doListValue(ColumnMetadata column, Object v, StringBuilder context) { + context.append('('); + for (var e : (Collection) v) { + doValue(column, e, context); + context.append(','); + } + context.deleteCharAt(context.length() - 1); + context.append(')'); + } + + private void doValue(ColumnMetadata column, Object v, StringBuilder context) { + if (DataTypes.SMALLINT.equals(column.getType())) { + v = ((Number) v).shortValue(); + } + context.append(CodecRegistry.DEFAULT.codecFor(column.getType()).format(v)); + } + + private Optional getColumn(String name) { + Optional column = Optional.ofNullable(this.columnsByName.get(name)); + + // work around the need to escape filter keys the ANTLR parser doesn't like + // e.g. with underscores like chunk_no + if (column.isEmpty()) { + if (name.startsWith("\"") && name.endsWith("\"")) { + name = name.substring(1, name.length() - 1); + column = Optional.ofNullable(this.columnsByName.get(name)); + } + } + return column; + } + +} diff --git a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java new file mode 100644 index 000000000..6f651c944 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java @@ -0,0 +1,320 @@ +/* + * Copyright 2024 - 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 java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CompletableFuture; + +import com.datastax.oss.driver.api.core.cql.BoundStatement; +import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder; +import com.datastax.oss.driver.api.core.cql.PreparedStatement; +import com.datastax.oss.driver.api.core.cql.Row; +import com.datastax.oss.driver.api.core.data.CqlVector; +import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata; +import com.datastax.oss.driver.api.querybuilder.QueryBuilder; +import com.datastax.oss.driver.api.querybuilder.delete.Delete; +import com.datastax.oss.driver.api.querybuilder.delete.DeleteSelection; +import com.datastax.oss.driver.api.querybuilder.insert.InsertInto; +import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert; +import com.datastax.oss.driver.shaded.guava.common.base.Preconditions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn; +import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; +import org.springframework.beans.factory.InitializingBean; + +/** + * The CassandraVectorStore is for managing and querying vector data in an Apache + * Cassandra db. It offers functionalities like adding, deleting, and performing + * similarity searches on documents. + * + * The store utilizes CQL to index and search vector data. It allows for custom metadata + * fields in the documents to be stored alongside the vector and content data. + * + * This class requires a CassandraVectorStoreConfig configuration object for + * initialization, which includes settings like connection details, index name, field + * names, etc. It also requires an EmbeddingClient to convert documents into embeddings + * before storing them. + * + * A schema matching the configuration is automatically created if it doesn't exist. + * Missing columns and indexes in existing tables will also be automatically created. + * Disable this with the disallowSchemaCreation. + * + * This class is designed to work with brand new tables that it creates for you, or on top + * of existing Cassandra tables. The latter is appropriate when wanting to keep data in + * place, creating embeddings next to it, and performing vector similarity searches + * in-situ. + * + * Instances of this class are not dynamic against server-side schema changes. If you + * change the schema server-side you need a new CassandraVectorStore instance. + * + * @author Mick Semb Wever + * @see VectorStore + * @see CassandraVectorStoreConfig + * @see EmbeddingClient + * @since 1.0.0 + */ +public final class CassandraVectorStore implements VectorStore, InitializingBean, AutoCloseable { + + /** + * Indexes are automatically created with COSINE. This can be changed manually via + * cqlsh + */ + public enum Similarity { + + COSINE, DOT_PRODUCT, EUCLIDEAN; + + } + + private static final String QUERY_FORMAT = "select %s,%s,%s%s from %s.%s ? order by %s ann of ? limit ?"; + + public static final String SIMILARITY_FIELD_NAME = "similarity_score"; + + private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStore.class); + + private final CassandraVectorStoreConfig conf; + + private final EmbeddingClient embeddingClient; + + private final FilterExpressionConverter filterExpressionConverter; + + private final Map, PreparedStatement> addStmts = new HashMap<>(); + + private final PreparedStatement deleteStmt; + + private final String similarityStmt; + + private final Similarity similarity; + + public CassandraVectorStore(CassandraVectorStoreConfig conf, EmbeddingClient embeddingClient) { + + Preconditions.checkArgument(null != conf, "Config must not be null"); + Preconditions.checkArgument(null != embeddingClient, "Embedding client must not be null"); + + this.conf = conf; + this.embeddingClient = embeddingClient; + conf.ensureSchemaExists(embeddingClient.dimensions()); + prepareAddStatement(Set.of()); + this.deleteStmt = prepareDeleteStatement(); + + TableMetadata cassandraMetadata = conf.session.getMetadata() + .getKeyspace(conf.schema.keyspace()) + .get() + .getTable(conf.schema.table()) + .get(); + + this.similarity = getIndexSimilarity(cassandraMetadata); + this.similarityStmt = similaritySearchStatement(); + + this.filterExpressionConverter = new CassandraFilterExpressionConverter( + cassandraMetadata.getColumns().values()); + } + + @Override + public void add(List documents) { + CompletableFuture[] futures = new CompletableFuture[documents.size()]; + short i = 0; + for (Document d : documents) { + List primaryKeyValues = this.conf.documentIdTranslator.apply(d.getId()); + var embedding = this.embeddingClient.embed(d).stream().map(Double::floatValue).toList(); + + BoundStatementBuilder builder = prepareAddStatement(d.getMetadata().keySet()).boundStatementBuilder(); + for (int k = 0; k < primaryKeyValues.size(); ++k) { + SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k); + builder = builder.set(keyColumn.name(), primaryKeyValues.get(k), keyColumn.javaType()); + } + + builder = builder.setString(this.conf.schema.content(), d.getContent()) + .setVector(this.conf.schema.embedding(), CqlVector.newInstance(embedding), Float.class); + + for (var metadataColumn : this.conf.schema.metadataColumns() + .stream() + .filter((mc) -> d.getMetadata().containsKey(mc.name())) + .toList()) { + + builder = builder.set(metadataColumn.name(), d.getMetadata().get(metadataColumn.name()), + metadataColumn.javaType()); + } + futures[i++] = this.conf.session.executeAsync(builder.build()).toCompletableFuture(); + } + CompletableFuture.allOf(futures).join(); + } + + @Override + public Optional delete(List idList) { + CompletableFuture[] futures = new CompletableFuture[idList.size()]; + short i = 0; + for (String id : idList) { + List primaryKeyValues = this.conf.documentIdTranslator.apply(id); + BoundStatement s = this.deleteStmt.bind(primaryKeyValues.toArray()); + futures[i++] = this.conf.session.executeAsync(s).toCompletableFuture(); + } + CompletableFuture.allOf(futures).join(); + return Optional.of(Boolean.TRUE); + } + + @Override + public List similaritySearch(SearchRequest request) { + Preconditions.checkArgument(request.getTopK() <= 1000); + var embedding = this.embeddingClient.embed(request.getQuery()).stream().map(Double::floatValue).toList(); + CqlVector cqlVector = CqlVector.newInstance(embedding); + + String whereClause = ""; + if (request.hasFilterExpression()) { + String expression = this.filterExpressionConverter.convertExpression(request.getFilterExpression()); + if (!expression.isBlank()) { + whereClause = String.format("where %s", expression); + } + } + + String query = String.format(this.similarityStmt, cqlVector, whereClause, cqlVector, request.getTopK()); + List documents = new ArrayList<>(); + logger.trace("Executing {}", query); + + for (Row row : this.conf.session.execute(query)) { + float score = row.getFloat(0); + if (score < request.getSimilarityThreshold()) { + break; + } + Map docFields = new HashMap<>(); + docFields.put(SIMILARITY_FIELD_NAME, score); + for (var metadata : this.conf.schema.metadataColumns()) { + var value = row.get(metadata.name(), metadata.javaType()); + if (null != value) { + docFields.put(metadata.name(), value); + } + } + + documents.add(new Document(getDocumentId(row), row.getString(this.conf.schema.content()), docFields)); + } + return documents; + } + + @Override + public void afterPropertiesSet() { + } + + @Override + public void close() throws Exception { + this.conf.close(); + } + + void checkSchemaValid() { + this.conf.checkSchemaValid(embeddingClient.dimensions()); + } + + private Similarity getIndexSimilarity(TableMetadata metadata) { + + return Similarity.valueOf(metadata.getIndex(this.conf.schema.index()) + .get() + .getOptions() + .getOrDefault("similarity_function", "COSINE") + .toUpperCase()); + } + + private PreparedStatement prepareDeleteStatement() { + Delete stmt = null; + DeleteSelection stmtStart = QueryBuilder.deleteFrom(conf.schema.keyspace(), conf.schema.table()); + + for (var c : this.conf.schema.partitionKeys()) { + stmt = (null != stmt ? stmt : stmtStart).whereColumn(c.name()).isEqualTo(QueryBuilder.bindMarker(c.name())); + } + for (var c : this.conf.schema.clusteringKeys()) { + stmt = stmt.whereColumn(c.name()).isEqualTo(QueryBuilder.bindMarker(c.name())); + } + + return this.conf.session.prepare(stmt.build()); + } + + private PreparedStatement prepareAddStatement(Set metadataFields) { + if (!this.addStmts.containsKey(metadataFields)) { + RegularInsert stmt = null; + InsertInto stmtStart = QueryBuilder.insertInto(this.conf.schema.keyspace(), this.conf.schema.table()); + + for (var c : this.conf.schema.partitionKeys()) { + stmt = (null != stmt ? stmt : stmtStart).value(c.name(), QueryBuilder.bindMarker(c.name())); + } + for (var c : this.conf.schema.clusteringKeys()) { + stmt = stmt.value(c.name(), QueryBuilder.bindMarker(c.name())); + } + + stmt = stmt.value(this.conf.schema.content(), QueryBuilder.bindMarker(this.conf.schema.content())) + .value(this.conf.schema.embedding(), QueryBuilder.bindMarker(this.conf.schema.embedding())); + + for (String metadataField : this.conf.schema.metadataColumns() + .stream() + .map((mc) -> mc.name()) + .filter((mc) -> metadataFields.contains(mc)) + .toList()) { + + stmt = stmt.value(metadataField, QueryBuilder.bindMarker(metadataField)); + } + this.addStmts.putIfAbsent(metadataFields, this.conf.session.prepare(stmt.build())); + } + return this.addStmts.get(metadataFields); + } + + private String similaritySearchStatement() { + StringBuilder ids = new StringBuilder(); + for (var m : this.conf.schema.partitionKeys()) { + ids.append(m.name()).append(','); + } + for (var m : this.conf.schema.clusteringKeys()) { + ids.append(m.name()).append(','); + } + ids.deleteCharAt(ids.length() - 1); + + String similarityFunction = new StringBuilder("similarity_").append(this.similarity.toString().toLowerCase()) + .append('(') + .append(conf.schema.embedding()) + .append(",?)") + .toString(); + + StringBuilder extraSelectFields = new StringBuilder(); + for (var m : this.conf.schema.metadataColumns()) { + extraSelectFields.append(',').append(m.name()); + } + + // java-driver-query-builder doesn't support orderByAnnOf yet + String query = String.format(QUERY_FORMAT, similarityFunction, ids.toString(), this.conf.schema.content(), + extraSelectFields.toString(), this.conf.schema.keyspace(), this.conf.schema.table(), + this.conf.schema.embedding()); + + query = query.replace("?", "%s"); + logger.debug("preparing {}", query); + return query; + } + + private String getDocumentId(Row row) { + List primaryKeyValues = new ArrayList<>(); + for (var m : this.conf.schema.partitionKeys()) { + primaryKeyValues.add(row.get(m.name(), m.javaType())); + } + for (var m : this.conf.schema.clusteringKeys()) { + primaryKeyValues.add(row.get(m.name(), m.javaType())); + } + return this.conf.primaryKeyTranslator.apply(primaryKeyValues); + } + +} diff --git a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java new file mode 100644 index 000000000..d85d1ee4b --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java @@ -0,0 +1,536 @@ +/* + * Copyright 2024 - 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 java.net.InetSocketAddress; +import java.time.Duration; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Stream; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.CqlSessionBuilder; +import com.datastax.oss.driver.api.core.cql.SimpleStatement; +import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata; +import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata; +import com.datastax.oss.driver.api.core.type.DataType; +import com.datastax.oss.driver.api.core.type.DataTypes; +import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry; +import com.datastax.oss.driver.api.core.type.reflect.GenericType; +import com.datastax.oss.driver.api.querybuilder.BuildableQuery; +import com.datastax.oss.driver.api.querybuilder.SchemaBuilder; +import com.datastax.oss.driver.api.querybuilder.schema.AlterTableAddColumn; +import com.datastax.oss.driver.api.querybuilder.schema.AlterTableAddColumnEnd; +import com.datastax.oss.driver.api.querybuilder.schema.CreateTable; +import com.datastax.oss.driver.api.querybuilder.schema.CreateTableStart; +import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting; +import com.datastax.oss.driver.shaded.guava.common.base.Preconditions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Configuration for the Cassandra vector store. + * + * All metadata fields configured to the store will be fetched and added to all queried + * documents. + * + * If you wish to metadata search against a field its 'searchable' argument must be true. + * + * The Cassandra Java Driver is configured via the application.conf resource found in the + * classpath. See + * https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration + * + * @since 1.0.0 + */ +public final class CassandraVectorStoreConfig implements AutoCloseable { + + public static final String DEFAULT_KEYSPACE_NAME = "springframework"; + + public static final String DEFAULT_TABLE_NAME = "ai_vector_store"; + + public static final String DEFAULT_ID_NAME = "id"; + + public static final String DEFAULT_INDEX_NAME = "embedding_index"; + + public static final String DEFAULT_CONTENT_COLUMN_NAME = "content"; + + public static final String DEFAULT_EMBEDDING_COLUMN_NAME = "embedding"; + + private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStore.class); + + record Schema(String keyspace, String table, List partitionKeys, List clusteringKeys, + String content, String embedding, String index, Set metadataColumns) { + + } + + public record SchemaColumn(String name, DataType type, SchemaColumnTags... tags) { + public SchemaColumn(String name, DataType type) { + this(name, type, new SchemaColumnTags[0]); + } + + public GenericType javaType() { + return CodecRegistry.DEFAULT.codecFor(type).getJavaType(); + } + + public boolean indexed() { + for (SchemaColumnTags t : tags) { + if (SchemaColumnTags.INDEXED == t) { + return true; + } + } + return false; + } + } + + public enum SchemaColumnTags { + + INDEXED + + } + + /** + * It is a requirement that an empty {@code List} returns an example formatted + * id + */ + public interface DocumentIdTranslator extends Function> { + + } + + public interface PrimaryKeyTranslator extends Function, String> { + + } + + final CqlSession session; + + final Schema schema; + + final boolean disallowSchemaChanges; + + final DocumentIdTranslator documentIdTranslator; + + final PrimaryKeyTranslator primaryKeyTranslator; + + private final boolean closeSessionOnClose; + + private CassandraVectorStoreConfig(Builder builder) { + this.session = null != builder.session ? builder.session : builder.sessionBuilder.build(); + this.closeSessionOnClose = null == builder.session; + + this.schema = new Schema(builder.keyspace, builder.table, builder.partitionKeys, builder.clusteringKeys, + builder.contentColumnName, builder.embeddingColumnName, builder.indexName, builder.metadataColumns); + + this.disallowSchemaChanges = builder.disallowSchemaCreation; + this.documentIdTranslator = builder.documentIdTranslator; + this.primaryKeyTranslator = builder.primaryKeyTranslator; + } + + public static Builder builder() { + return new Builder(); + } + + @Override + public void close() throws Exception { + if (this.closeSessionOnClose) { + this.session.close(); + } + } + + SchemaColumn getPrimaryKeyColumn(int index) { + return index < this.schema.partitionKeys().size() ? this.schema.partitionKeys().get(index) + : this.schema.clusteringKeys().get(index - this.schema.partitionKeys().size()); + } + + @VisibleForTesting + void dropKeyspace() { + Preconditions.checkState(this.schema.keyspace.startsWith("test_"), "Only test keyspaces can be dropped"); + this.session.execute(SchemaBuilder.dropKeyspace(this.schema.keyspace).ifExists().build()); + } + + public static class Builder { + + private CqlSession session = null; + + private CqlSessionBuilder sessionBuilder = null; + + private String keyspace = DEFAULT_KEYSPACE_NAME; + + private String table = DEFAULT_TABLE_NAME; + + private List partitionKeys = List.of(new SchemaColumn(DEFAULT_ID_NAME, DataTypes.TEXT)); + + private List clusteringKeys = List.of(); + + private String indexName = DEFAULT_INDEX_NAME; + + private String contentColumnName = DEFAULT_CONTENT_COLUMN_NAME; + + private String embeddingColumnName = DEFAULT_EMBEDDING_COLUMN_NAME; + + private Set metadataColumns = new HashSet<>(); + + private boolean disallowSchemaCreation = false; + + private DocumentIdTranslator documentIdTranslator = (String id) -> List.of(id); + + private PrimaryKeyTranslator primaryKeyTranslator = (List primaryKeyColumns) -> { + if (primaryKeyColumns.isEmpty()) { + return "test"; + } + Preconditions.checkArgument(1 == primaryKeyColumns.size()); + return (String) primaryKeyColumns.get(0); + }; + + private Builder() { + } + + public Builder withCqlSession(CqlSession session) { + Preconditions.checkState(null == this.sessionBuilder, + "Cannot call withContactPoint(..) or withLocalDatacenter(..) and this method"); + + this.session = session; + return this; + } + + public Builder addContactPoint(InetSocketAddress contactPoint) { + Preconditions.checkState(null == this.session, "Cannot call withCqlSession(..) and this method"); + if (null == this.sessionBuilder) { + this.sessionBuilder = new CqlSessionBuilder(); + } + this.sessionBuilder.addContactPoint(contactPoint); + return this; + } + + public Builder withLocalDatacenter(String localDC) { + Preconditions.checkState(null == this.session, "Cannot call withCqlSession(..) and this method"); + if (null == this.sessionBuilder) { + this.sessionBuilder = new CqlSessionBuilder(); + } + this.sessionBuilder.withLocalDatacenter(localDC); + return this; + } + + public Builder withKeyspaceName(String keyspace) { + this.keyspace = keyspace; + return this; + } + + public Builder withTableName(String table) { + this.table = table; + return this; + } + + public Builder withPartitionKeys(List partitionKeys) { + this.partitionKeys = partitionKeys; + return this; + } + + public Builder withClusteringKeys(List clusteringKeys) { + this.clusteringKeys = clusteringKeys; + return this; + } + + public Builder withIndexName(String name) { + this.indexName = name; + return this; + } + + public Builder withContentColumnName(String name) { + this.contentColumnName = name; + return this; + } + + public Builder withEmbeddingColumnName(String name) { + this.embeddingColumnName = name; + return this; + } + + public Builder addMetadataColumn(SchemaColumn... fields) { + Builder builder = this; + for (SchemaColumn f : fields) { + builder = builder.addMetadataColumn(f); + } + return builder; + } + + public Builder addMetadataColumn(SchemaColumn field) { + + Preconditions.checkArgument(this.metadataColumns.stream().noneMatch((sc) -> sc.name().equals(field.name())), + "A metadata field with name %s has already been added", field.name()); + + this.metadataColumns.add(field); + return this; + } + + public Builder disallowSchemaChanges() { + this.disallowSchemaCreation = true; + return this; + } + + public Builder withDocumentIdTranslator(DocumentIdTranslator documentIdTranslator) { + this.documentIdTranslator = documentIdTranslator; + return this; + } + + public Builder withPrimaryKeyTranslator(PrimaryKeyTranslator primaryKeyTranslator) { + this.primaryKeyTranslator = primaryKeyTranslator; + return this; + } + + public CassandraVectorStoreConfig build() { + for (SchemaColumn metadata : this.metadataColumns) { + + Preconditions.checkArgument( + !this.partitionKeys.stream().anyMatch((c) -> c.name().equals(metadata.name())), + "metadataColumn %s cannot have same name as a partition key", metadata.name()); + + Preconditions.checkArgument( + !this.clusteringKeys.stream().anyMatch((c) -> c.name().equals(metadata.name())), + "metadataColumn %s cannot have same name as a clustering key", metadata.name()); + + Preconditions.checkArgument(!metadata.name().equals(this.contentColumnName), + "metadataColumn %s cannot have same name as content column name", this.contentColumnName); + + Preconditions.checkArgument(!metadata.name().equals(this.embeddingColumnName), + "metadataColumn %s cannot have same name as embedding column name", this.embeddingColumnName); + + } + { + int primaryKeyColumnsCount = this.partitionKeys.size() + this.clusteringKeys.size(); + String exampleId = this.primaryKeyTranslator.apply(Collections.emptyList()); + List testIdTranslation = this.documentIdTranslator.apply(exampleId); + + Preconditions.checkArgument(testIdTranslation.size() == primaryKeyColumnsCount, + "documentIdTranslator results length %s doesn't match number of primary key columns %s", + String.valueOf(testIdTranslation.size()), String.valueOf(primaryKeyColumnsCount)); + + Preconditions.checkArgument( + exampleId.equals(this.primaryKeyTranslator.apply(this.documentIdTranslator.apply(exampleId))), + "primaryKeyTranslator is not an inverse function to documentIdTranslator"); + } + return new CassandraVectorStoreConfig(this); + } + + } + + void ensureSchemaExists(int vectorDimension) { + if (!this.disallowSchemaChanges) { + ensureKeyspaceExists(); + ensureTableExists(vectorDimension); + ensureTableColumnsExist(vectorDimension); + ensureIndexesExists(); + checkSchemaAgreement(); + } + else { + checkSchemaValid(vectorDimension); + } + } + + private void checkSchemaAgreement() throws IllegalStateException { + if (!this.session.checkSchemaAgreement()) { + logger.warn("Waiting for cluster schema agreement, sleeping 10s…"); + try { + Thread.sleep(Duration.ofSeconds(10).toMillis()); + } + catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(ex); + } + if (!this.session.checkSchemaAgreement()) { + logger.error("no cluster schema agreement still, continuing, let's hope this works…"); + } + } + } + + void checkSchemaValid(int vectorDimension) { + + Preconditions.checkState(this.session.getMetadata().getKeyspace(this.schema.keyspace).isPresent(), + "keyspace %s does not exist", this.schema.keyspace); + + Preconditions.checkState(this.session.getMetadata() + .getKeyspace(this.schema.keyspace) + .get() + .getTable(this.schema.table) + .isPresent(), "table %s does not exist"); + + TableMetadata tableMetadata = this.session.getMetadata() + .getKeyspace(this.schema.keyspace) + .get() + .getTable(this.schema.table) + .get(); + + Preconditions.checkState(tableMetadata.getColumn(this.schema.content).isPresent(), "column %s does not exist", + this.schema.content); + + Preconditions.checkState(tableMetadata.getColumn(this.schema.embedding).isPresent(), "column %s does not exist", + this.schema.embedding); + + for (SchemaColumn m : this.schema.metadataColumns) { + Optional column = tableMetadata.getColumn(m.name()); + Preconditions.checkState(column.isPresent(), "column %s does not exist", m.name()); + + Preconditions.checkArgument(column.get().getType().equals(m.type()), + "Mismatching type on metadata column %s of %s vs %s", m.name(), column.get().getType(), m.type()); + + if (m.indexed()) { + Preconditions.checkState( + tableMetadata.getIndexes().values().stream().anyMatch((i) -> i.getTarget().equals(m.name())), + "index %s does not exist", m.name()); + } + } + + } + + private void ensureIndexesExists() { + { + SimpleStatement indexStmt = SchemaBuilder.createIndex(this.schema.index) + .ifNotExists() + .custom("SAI") + .onTable(this.schema.keyspace, this.schema.table) + .andColumn(this.schema.embedding) + .build(); + + logger.debug("Executing {}", indexStmt.getQuery()); + this.session.execute(indexStmt); + } + Stream + .concat(this.schema.partitionKeys.stream(), + Stream.concat(this.schema.clusteringKeys.stream(), this.schema.metadataColumns.stream())) + .filter((cs) -> cs.indexed()) + .forEach((metadata) -> { + + SimpleStatement indexStmt = SchemaBuilder.createIndex(String.format("%s_idx", metadata.name())) + .ifNotExists() + .custom("SAI") + .onTable(this.schema.keyspace, this.schema.table) + .andColumn(metadata.name()) + .build(); + + logger.debug("Executing {}", indexStmt.getQuery()); + this.session.execute(indexStmt); + }); + } + + private void ensureTableExists(int vectorDimension) { + + CreateTable createTable = null; + + CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table) + .ifNotExists(); + + for (SchemaColumn partitionKey : this.schema.partitionKeys) { + createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name, + partitionKey.type); + } + for (SchemaColumn clusteringKey : this.schema.clusteringKeys) { + createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type); + } + + createTable = createTable.withColumn(this.schema.content, DataTypes.TEXT); + + for (SchemaColumn metadata : this.schema.metadataColumns) { + createTable = createTable.withColumn(metadata.name(), metadata.type()); + } + + // https://datastax-oss.atlassian.net/browse/JAVA-3118 + // .withColumn(config.embedding, new DefaultVectorType(DataTypes.FLOAT, + // vectorDimension)); + + StringBuilder tableStmt = new StringBuilder(createTable.asCql()); + tableStmt.setLength(tableStmt.length() - 1); + tableStmt.append(',') + .append(this.schema.embedding) + .append(" vector)"); + logger.debug("Executing {}", tableStmt.toString()); + this.session.execute(tableStmt.toString()); + } + + private void ensureTableColumnsExist(int vectorDimension) { + + TableMetadata tableMetadata = this.session.getMetadata() + .getKeyspace(this.schema.keyspace) + .get() + .getTable(this.schema.table) + .get(); + + Set newColumns = new HashSet<>(); + boolean addContent = tableMetadata.getColumn(this.schema.content).isEmpty(); + boolean addEmbedding = tableMetadata.getColumn(this.schema.embedding).isEmpty(); + + for (SchemaColumn metadata : this.schema.metadataColumns) { + Optional column = tableMetadata.getColumn(metadata.name()); + if (column.isPresent()) { + + Preconditions.checkArgument(column.get().getType().equals(metadata.type()), + "Cannot change type on metadata field %s from %s to %s", metadata.name(), + column.get().getType(), metadata.type()); + } + else { + newColumns.add(metadata); + } + } + + if (!newColumns.isEmpty() || addContent || addEmbedding) { + AlterTableAddColumn alterTable = SchemaBuilder.alterTable(this.schema.keyspace, this.schema.table); + for (SchemaColumn metadata : newColumns) { + alterTable = alterTable.addColumn(metadata.name(), metadata.type()); + } + if (addContent) { + alterTable = alterTable.addColumn(this.schema.content, DataTypes.TEXT); + } + if (addEmbedding) { + // special case for embedding column, bc JAVA-3118, as above + StringBuilder alterTableStmt = new StringBuilder(((BuildableQuery) alterTable).asCql()); + if (newColumns.isEmpty() && !addContent) { + alterTableStmt.append(" ADD "); + } + else { + alterTableStmt.setLength(alterTableStmt.length() - 1); + alterTableStmt.append(','); + } + alterTableStmt.append(this.schema.embedding) + .append(" vector"); + + logger.debug("Executing {}", alterTableStmt.toString()); + this.session.execute(alterTableStmt.toString()); + } + else { + SimpleStatement stmt = ((AlterTableAddColumnEnd) alterTable).build(); + logger.debug("Executing {}", stmt.getQuery()); + this.session.execute(stmt); + } + } + } + + private void ensureKeyspaceExists() { + + SimpleStatement keyspaceStmt = SchemaBuilder.createKeyspace(this.schema.keyspace) + .ifNotExists() + .withSimpleStrategy(1) + .build(); + + logger.debug("Executing {}", keyspaceStmt.getQuery()); + this.session.execute(keyspaceStmt); + } + +} diff --git a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverterTests.java b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverterTests.java new file mode 100644 index 000000000..9cf501c20 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverterTests.java @@ -0,0 +1,212 @@ +/* + * Copyright 2024 - 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 java.util.Collection; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata; +import com.datastax.oss.driver.api.core.CqlIdentifier; +import com.datastax.oss.driver.api.core.type.DataTypes; +import com.datastax.oss.driver.internal.core.metadata.schema.DefaultColumnMetadata; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +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.Filter.Value; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.GTE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.IN; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.LTE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR; + +/** + * @author Mick Semb Wever + * @since 1.0.0 + */ +class CassandraFilterExpressionConverterTests { + + @Test + void testEQOnPartition() { + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(COLUMNS); + + String vectorExpr = filter.convertExpression(new Expression(EQ, new Key("id"), new Value("BG"))); + + assertThat(vectorExpr).isEqualTo("\"id\" = 'BG'"); + } + + @Test + void testEQ() { + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(COLUMNS); + + String vectorExpr = filter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG"))); + + assertThat(vectorExpr).isEqualTo("\"country\" = 'BG'"); + } + + @Test + void testNoSuchColumn() { + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(COLUMNS); + + Assertions.assertThrows(IllegalArgumentException.class, () -> { + filter.convertExpression(new Expression(EQ, new Key("unknown_column"), new Value("BG"))); + }); + } + + @Test + void tesEqAndGte() { + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(COLUMNS); + + // genre == "drama" AND year >= 2020 + String vectorExpr = filter + .convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")), + new Expression(GTE, new Key("year"), new Value(2020)))); + + assertThat(vectorExpr).isEqualTo("\"genre\" = 'drama' and \"year\" >= 2020"); + } + + @Test + void tesOr() { + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(COLUMNS); + + // genre == "drama" OR year = 2020 + String vectorExpr = filter + .convertExpression(new Expression(OR, new Expression(EQ, new Key("genre"), new Value("drama")), + new Expression(EQ, new Key("year"), new Value(2020)))); + + assertThat(vectorExpr).isEqualTo("\"genre\" = 'drama' or \"year\" = 2020"); + } + + @Test + void tesIn() { + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(COLUMNS); + + // genre in ["comedy", "documentary", "drama"] + String vectorExpr = filter.convertExpression( + new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama")))); + + assertThat(vectorExpr).isEqualTo("\"genre\" IN ('comedy','documentary','drama')"); + } + + @Test + void testNe() { + Set columns = new HashSet(COLUMNS); + + columns.add(new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("city"), DataTypes.TEXT, false)); + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(columns); + + // year >= 2020 OR country == "BG" AND city != "Sofia" + String vectorExpr = filter + .convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)), + new Group(new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")), + new Expression(NE, new Key("city"), new Value("Sofia")))))); + + assertThat(vectorExpr).isEqualTo("\"year\" >= 2020 or \"country\" = 'BG' and \"city\" != 'Sofia'"); + } + + @Test + void testGroup() { + Set columns = new HashSet(COLUMNS); + + columns.add(new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("city"), DataTypes.TEXT, false)); + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(columns); + + // (year >= 2020 OR country == "BG") AND city IN ["Sofia", "Plovdiv"] + String vectorExpr = filter.convertExpression(new Expression(AND, + new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)), + new Expression(EQ, new Key("country"), new Value("BG")))), + new Expression(IN, new Key("city"), new Value(List.of("Sofia", "Plovdiv"))))); + + assertThat(vectorExpr).isEqualTo("\"year\" >= 2020 or \"country\" = 'BG' and \"city\" IN ('Sofia','Plovdiv')"); + } + + @Test + void tesBoolean() { + Set columns = new HashSet(COLUMNS); + + columns.add(new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("isOpen"), DataTypes.BOOLEAN, false)); + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(columns); + + // isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"] + String vectorExpr = filter.convertExpression(new Expression(AND, + new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)), + new Expression(GTE, new Key("year"), new Value(2020))), + new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US"))))); + + assertThat(vectorExpr).isEqualTo("\"isOpen\" = true and \"year\" >= 2020 and \"country\" IN ('BG','NL','US')"); + } + + @Test + void testDecimal() { + Set columns = new HashSet(COLUMNS); + + columns + .add(new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("temperature"), DataTypes.DOUBLE, false)); + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(columns); + + // temperature >= -15.6 && temperature <= +20.13 + String vectorExpr = filter + .convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)), + new Expression(LTE, new Key("temperature"), new Value(20.13)))); + + assertThat(vectorExpr).isEqualTo("\"temperature\" >= -15.6 and \"temperature\" <= 20.13"); + } + + @Test + void testComplexIdentifiers() { + Set columns = new HashSet(COLUMNS); + + columns.add(new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("\"country 1 2 3\""), DataTypes.TEXT, + false)); + + columns + .add(new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("'country 1 2 3'"), DataTypes.TEXT, false)); + + CassandraFilterExpressionConverter filter = new CassandraFilterExpressionConverter(columns); + + String vectorExpr = filter.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG"))); + assertThat(vectorExpr).isEqualTo("\"\"\"country 1 2 3\"\"\" = 'BG'"); + + vectorExpr = filter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG"))); + assertThat(vectorExpr).isEqualTo("\"'country 1 2 3'\" = 'BG'"); + } + + private static final CqlIdentifier T = CqlIdentifier.fromInternal("test"); + + private static final Collection COLUMNS = Set.of( + new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("id"), DataTypes.TEXT, false), + new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("content"), DataTypes.TEXT, false), + new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("country"), DataTypes.TEXT, false), + new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("genre"), DataTypes.TEXT, false), + new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("drama"), DataTypes.TEXT, false), + new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("year"), DataTypes.SMALLINT, false)); + +} diff --git a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java new file mode 100644 index 000000000..91ec53e50 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java @@ -0,0 +1,549 @@ +/* + * Copyright 2024 - 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 java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import org.junit.jupiter.api.Assertions; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.CqlSessionBuilder; +import com.datastax.oss.driver.api.core.servererrors.InvalidQueryException; +import com.datastax.oss.driver.api.core.servererrors.SyntaxError; +import com.datastax.oss.driver.api.core.type.DataTypes; +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testcontainers.containers.CassandraContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.transformers.TransformersEmbeddingClient; +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn; +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.ApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.core.io.DefaultResourceLoader; + +import static java.lang.String.format; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Use `mvn failsafe:integration-test -Dit.test=CassandraRichSchemaVectorStoreIT` + * + * @author Mick Semb Wever + * @since 1.0.0 + */ +@Testcontainers +class CassandraRichSchemaVectorStoreIT { + + static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("cassandra"); + + private static final Logger logger = LoggerFactory.getLogger(CassandraRichSchemaVectorStoreIT.class); + + private static final List documents = List.of( + + new Document("Neptune§¶0", + "Neptune\\n\\nThis article contains special characters. Without proper rendering support, you may see question marks, boxes, or other symbols. Neptune is the eighth and farthest planet from the Sun in the Solar System. It is an ice giant. It is the fourth-largest planet in the system. Neptunes mass is 17 times Earths mass and a little bit more than Uranus mass. Neptune is denser and smaller than Uranus. Because of its greater mass, Neptunes gravity makes its atmosphere smaller and denser. It was named after the Roman god of the sea, Neptune. Neptunes astronomical symbol is ♆, the trident of the god Neptune. Neptunes atmosphere is mostly hydrogen and helium. It also contains small amounts of methane which makes the planet appear blue. Neptunes blue color is similar, but slightly darker, than the color of Uranus. Neptune also has the strongest winds of any planet in the Solar System, as high as 2,100\\xa0km/h or 1,300\\xa0mph. Urbain Le Verrier and John Couch Adams were the astronomers who discovered Neptune. Neptune was not", + Map.of("revision", 9385813, "id", 558)), + + new Document("Neptune§¶1", + "Neptune\\n\\nbut slightly darker, than the color of Uranus. Neptune also has the strongest winds of any planet in the Solar System, as high as 2,100\\xa0km/h or 1,300\\xa0mph. Urbain Le Verrier and John Couch Adams were the astronomers who discovered Neptune. Neptune was not discovered using a telescope. It was the first planet to be discovered using mathematics. In 1821, astronomers saw that Uranus orbit was different from what they expected. Another nearby planets mass was changing Uranus orbit. They found Neptune was the cause. Voyager 2 visited Neptune on 25 August 1989. It was the only spacecraft to visit the planet. Neptune used to have a huge storm known as the \"Great Dark Spot\". Voyager 2 discovered the spot in 1989. The dark spot was not seen in 1994, but new spots were found since then. It is not known why the dark spot disappeared. Visits by other space probes have been planned. Neptune has five rings surrounding it, however, it is hard too see from Earth due to the distance from Neptune. Galileo Galilei was the first", + Map.of("revision", 9385813, "id", 558)), + + new Document("Neptune§¶2", + "Neptune\\n\\nfound since then. It is not known why the dark spot disappeared. Visits by other space probes have been planned. Neptune has five rings surrounding it, however, it is hard too see from Earth due to the distance from Neptune. Galileo Galilei was the first person who saw Neptune. He saw it on 28 December 1612 and 27 January 1613. His drawings showed points near Jupiter where Neptune is placed. But Galileo was not credited for the discovery. He thought Neptune was a \"fixed star\" instead of a planet. Because Neptune slowly moved across the sky, Galileos small telescope was not strong enough to see that Neptune was a planet. In 1821, Alexis Bouvard published the astronomical tables of the orbit of Uranus. Later observations showed that Uranus was moving in an irregular way in its orbit. Some astronomers thought this was caused by another large body. In 1843, John Couch Adams calculated the orbit of an eighth planet that could possibly affect the orbit of Uranus. He sent his calculations to Sir George Airy, the", + Map.of("revision", 9385813, "id", 558))); + + private static final String URANUS_ORBIT_QUERY = "It was the first planet to be discovered using mathematics. In 1821, astronomers saw that Uranus orbit was different from what they expected. Another nearby planets mass was changing Uranus orbit."; + + @Container + static CassandraContainer cassandraContainer = new CassandraContainer(DEFAULT_IMAGE_NAME.withTag("5.0")); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class); + + @Test + void ensureSchemaCreation() { + this.contextRunner.run(context -> { + try (CassandraVectorStore store = createStore(context, false).store()) { + Assertions.assertNotNull(store); + store.checkSchemaValid(); + store.similaritySearch(SearchRequest.query("1843").withTopK(1)); + } + }); + } + + @Test + void ensureSchemaNoCreation() { + this.contextRunner.run(context -> { + executeCqlFile(context, "test_wiki_full_schema.cql"); + var wrapper = createStore(context, List.of(), true, false); + try { + Assertions.assertNotNull(wrapper.store()); + wrapper.store().checkSchemaValid(); + + wrapper.store().similaritySearch(SearchRequest.query("1843").withTopK(1)); + + wrapper.conf().dropKeyspace(); + executeCqlFile(context, "test_wiki_partial_3_schema.cql"); + + // IllegalStateException: column all_minilm_l6_v2_embedding does not exist + IllegalStateException ise = Assertions.assertThrows(IllegalStateException.class, () -> { + createStore(context, List.of(), true, false); + }); + + Assertions.assertEquals("column all_minilm_l6_v2_embedding does not exist", ise.getMessage()); + } + finally { + wrapper.conf().dropKeyspace(); + wrapper.store().close(); + } + }); + } + + @Test + void ensureSchemaPartialCreation() { + this.contextRunner.run(context -> { + for (int i = 0; i < 4; ++i) { + executeCqlFile(context, format("test_wiki_partial_%d_schema.cql", i)); + var wrapper = createStore(context, List.of(), false, false); + try { + Assertions.assertNotNull(wrapper.store()); + wrapper.store().checkSchemaValid(); + + wrapper.store().similaritySearch(SearchRequest.query("1843").withTopK(1)); + wrapper.conf().dropKeyspace(); + } + finally { + wrapper.store().close(); + } + } + }); + } + + @Test + void addAndSearch() { + contextRunner.run(context -> { + try (CassandraVectorStore store = createStore(context, false).store()) { + store.add(documents); + + List results = store + .similaritySearch(SearchRequest.query("Neptunes gravity makes its atmosphere").withTopK(1)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId()); + + assertThat(resultDoc.getContent()).contains("Neptunes gravity makes its atmosphere"); + + assertThat(resultDoc.getMetadata()).hasSize(3); + + assertThat(resultDoc.getMetadata()).containsKeys("id", "revision", + CassandraVectorStore.SIMILARITY_FIELD_NAME); + + // Remove all documents from the createStore + store.delete(documents.stream().map(doc -> doc.getId()).toList()); + + results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + assertThat(results).isEmpty(); + } + }); + } + + @Test + void searchWithPartitionFilter() throws InterruptedException { + contextRunner.run(context -> { + try (CassandraVectorStore store = createStore(context, false).store()) { + store.add(documents); + + List results = store.similaritySearch(SearchRequest.query("Great Dark Spot").withTopK(5)); + assertThat(results).hasSize(3); + + results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY) + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("wiki == 'simplewiki' && language == 'en' && title == 'Neptune'")); + + assertThat(results).hasSize(3); + assertThat(results.get(0).getId()).isEqualTo(documents.get(1).getId()); + + // BUG CASSANDRA-19544 + // should be able to restrict on clustering keys (when filtering isn't + // required) + // + // results = store.similaritySearch(SearchRequest.query("Great Dark Spot") + // .withTopK(5) + // .withSimilarityThresholdAll() + // .withFilterExpression( + // "wiki == 'simplewiki' && language == 'en' && title == 'Neptune' && + // \"chunk_no\" == 0")); + // + // assertThat(results).hasSize(1); + // assertThat(results.get(0).getId()).isEqualTo(documents.get(0).getId()); + + results = store.similaritySearch(SearchRequest.query("Great Dark Spot") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression( + "wiki == 'simplewiki' && language == 'en' && title == 'Neptune' && id == 558")); + + assertThat(results).hasSize(3); + + // cassandra server will throw an error + Assertions.assertThrows(SyntaxError.class, () -> { + store.similaritySearch(SearchRequest.query("Great Dark Spot") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression( + "NOT(wiki == 'simplewiki' && language == 'en' && title == 'Neptune' && id == 1)")); + }); + } + }); + } + + @Test + void unsearchableFilters() throws InterruptedException { + contextRunner.run(context -> { + try (CassandraVectorStore store = createStore(context, false).store()) { + store.add(documents); + + List results = store.similaritySearch(SearchRequest.query("Great Dark Spot").withTopK(5)); + assertThat(results).hasSize(3); + + Assertions.assertThrows(InvalidQueryException.class, () -> { + store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("revision == 9385813")); + }); + } + }); + } + + @Test + void searchWithFilters() throws InterruptedException { + contextRunner.run(context -> { + try (CassandraVectorStore store = createStore(context, false).store()) { + store.add(documents); + + List results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY).withTopK(5)); + assertThat(results).hasSize(3); + + results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY) + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("id == 558")); + + assertThat(results).hasSize(3); + assertThat(results.get(0).getId()).isEqualTo(documents.get(1).getId()); + + results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY) + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("id > 557")); + + assertThat(results).hasSize(3); + assertThat(results.get(0).getId()).isEqualTo(documents.get(1).getId()); + + results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY) + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("id >= 558")); + + assertThat(results).hasSize(3); + assertThat(results.get(0).getId()).isEqualTo(documents.get(1).getId()); + + // cassandra java-driver will throw an error, + // as chunk_no is not searchable (i.e. no SAI index on it) + // note, it is possible to have SAI indexes on primary key columns to + // achieve + // e.g. searchWithFilterOnPrimaryKeys() + Assertions.assertThrows(InvalidQueryException.class, () -> { + store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY) + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("id > 557 && \"chunk_no\" == 1")); + }); + + // cassandra server will throw an error, + // as revision is not searchable (i.e. no SAI index on it) + Assertions.assertThrows(SyntaxError.class, () -> { + store.similaritySearch(SearchRequest.query("Great Dark Spot") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("id == 558 || revision == 2020")); + }); + + // cassandra java-driver will throw an error + Assertions.assertThrows(InvalidQueryException.class, () -> { + store.similaritySearch(SearchRequest.query("Great Dark Spot") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("NOT(id == 557 || revision == 2020)")); + }); + } + }); + } + + @Test + void searchWithFilterOnPrimaryKeys() throws InterruptedException { + contextRunner.run(context -> { + + List overrides = List.of( + new SchemaColumn("title", DataTypes.TEXT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED), + new SchemaColumn("chunk_no", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED)); + + try (CassandraVectorStore store = createStore(context, overrides, false, true).store()) { + + store.add(documents); + + List results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY).withTopK(5)); + assertThat(results).hasSize(3); + + store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY) + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("id > 557 && \"chunk_no\" == 1")); + + assertThat(results).hasSize(3); + assertThat(results.get(0).getId()).isEqualTo(documents.get(1).getId()); + + // Cassandra java-driver bug, not detecting index on title exists + // + // store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY) + // .withTopK(5) + // .withSimilarityThresholdAll() + // .withFilterExpression("id > 557 && title == 'Neptune'")); + // + // assertThat(results).hasSize(3); + // assertThat(results.get(0).getId()).isEqualTo(documents.get(1).getId()); + } + }); + } + + @Test + void documentUpdate() { + contextRunner.run(context -> { + try (CassandraVectorStore store = createStore(context, false).store()) { + store.add(documents); + + List results = store.similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY).withTopK(1)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getContent()).contains(URANUS_ORBIT_QUERY); + assertThat(resultDoc.getMetadata()).containsKey("revision"); + + String newContent = "The World is Big and Salvation Lurks Around the Corner"; + + Document sameIdDocument = new Document(documents.get(1).getId(), newContent, Collections.emptyMap()); + + // BUG in Cassandra 5.0-beta1 + // uncomment when 5.0-beta2 is release and cassandraContainer pulls it + // + // store.add(List.of(sameIdDocument)); + // + // results = + // store.similaritySearch(SearchRequest.query(newContent).withTopK(1)); + // + // assertThat(results).hasSize(1); + // resultDoc = results.get(0); + // assertThat(resultDoc.getId()).isEqualTo(sameIdDocument.getId()); + // assertThat(resultDoc.getContent()).contains(newContent); + // + // // the empty metadata map will not overwrite the row's existing "id" + // and + // // "revision" values + // assertThat(resultDoc.getMetadata()).containsKeys("id", "revision", + // CassandraVectorStore.SIMILARITY_FIELD_NAME); + + store.delete(List.of(sameIdDocument.getId())); + + results = store.similaritySearch(SearchRequest.query(newContent).withTopK(1)); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isNotEqualTo(sameIdDocument.getId()); + assertThat(resultDoc.getContent()).doesNotContain(newContent); + + assertThat(resultDoc.getMetadata()).containsKeys("id", "revision", + CassandraVectorStore.SIMILARITY_FIELD_NAME); + } + }); + } + + @Test + void searchWithThreshold() { + contextRunner.run(context -> { + try (CassandraVectorStore store = createStore(context, false).store()) { + store.add(documents); + + List fullResult = store + .similaritySearch(SearchRequest.query(URANUS_ORBIT_QUERY).withTopK(5).withSimilarityThresholdAll()); + + List distances = fullResult.stream() + .map(doc -> (Float) doc.getMetadata().get(CassandraVectorStore.SIMILARITY_FIELD_NAME)) + .toList(); + + assertThat(distances).hasSize(3); + + float threshold = (distances.get(0) + distances.get(1)) / 2; + + List results = store.similaritySearch( + SearchRequest.query(URANUS_ORBIT_QUERY).withTopK(5).withSimilarityThreshold(threshold)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(1).getId()); + + assertThat(resultDoc.getContent()).contains(URANUS_ORBIT_QUERY); + + assertThat(resultDoc.getMetadata()).containsKeys("id", "revision", + CassandraVectorStore.SIMILARITY_FIELD_NAME); + } + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Bean + public EmbeddingClient embeddingClient() { + // default is ONNX all-MiniLM-L6-v2 + return new TransformersEmbeddingClient(); + } + + @Bean + public CqlSession cqlSession() { + return new CqlSessionBuilder() + // comment next two lines out to connect to a local C* cluster + .addContactPoint(cassandraContainer.getContactPoint()) + .withLocalDatacenter(cassandraContainer.getLocalDatacenter()) + .build(); + } + + } + + private StoreWrapper createStore(ApplicationContext context, + boolean disallowSchemaCreation) throws IOException { + + return createStore(context, List.of(), disallowSchemaCreation, true); + } + + private StoreWrapper createStore(ApplicationContext context, + List extraMetadataFields, boolean disallowSchemaCreation, boolean dropKeyspaceFirst) + throws IOException { + + Optional wikiOverride = extraMetadataFields.stream() + .filter((f) -> "wiki".equals(f.name())) + .findFirst(); + + Optional langOverride = extraMetadataFields.stream() + .filter((f) -> "language".equals(f.name())) + .findFirst(); + + Optional titleOverride = extraMetadataFields.stream() + .filter((f) -> "title".equals(f.name())) + .findFirst(); + + Optional chunkNoOverride = extraMetadataFields.stream() + .filter((f) -> "chunk_no".equals(f.name())) + .findFirst(); + + SchemaColumn wikiSC = wikiOverride.orElse(new SchemaColumn("wiki", DataTypes.TEXT)); + SchemaColumn langSC = langOverride.orElse(new SchemaColumn("language", DataTypes.TEXT)); + SchemaColumn titleSC = titleOverride.orElse(new SchemaColumn("title", DataTypes.TEXT)); + SchemaColumn chunkNoSC = chunkNoOverride.orElse(new SchemaColumn("chunk_no", DataTypes.INT)); + + List partitionKeys = List.of(wikiSC, langSC, titleSC); + List clusteringKeys = List.of(chunkNoSC); + + CassandraVectorStoreConfig.Builder builder = CassandraVectorStoreConfig.builder() + .withCqlSession(context.getBean(CqlSession.class)) + .withKeyspaceName("test_wikidata") + .withTableName("articles") + .withPartitionKeys(partitionKeys) + .withClusteringKeys(clusteringKeys) + .withContentColumnName("body") + .withEmbeddingColumnName("all_minilm_l6_v2_embedding") + .withIndexName("all_minilm_l6_v2_ann") + + .addMetadataColumn(new SchemaColumn("revision", DataTypes.INT), + new SchemaColumn("id", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED)) + + // this store uses '§¶' as a deliminator in the document id between db columns + // 'title' and 'chunk_no' + .withPrimaryKeyTranslator((List primaryKeys) -> { + if (primaryKeys.isEmpty()) { + return "test§¶0"; + } + return format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3)); + }) + .withDocumentIdTranslator((id) -> { + String[] parts = id.split("§¶"); + String title = parts[0]; + int chunk_no = 0 < parts.length ? Integer.parseInt(parts[1]) : 0; + return List.of("simplewiki", "en", title, chunk_no); + }); + + for (SchemaColumn cf : extraMetadataFields) { + if (!partitionKeys.contains(cf) && !clusteringKeys.contains(cf)) { + builder = builder.addMetadataColumn(cf); + } + } + + if (disallowSchemaCreation) { + builder = builder.disallowSchemaChanges(); + } + + CassandraVectorStoreConfig conf = builder.build(); + if (dropKeyspaceFirst) { + conf.dropKeyspace(); + } + return new StoreWrapper(new CassandraVectorStore(conf, context.getBean(EmbeddingClient.class)), conf); + } + + private void executeCqlFile(ApplicationContext context, String filename) throws IOException { + logger.info("executing {}", filename); + + CqlSession session = context.getBean(CqlSession.class); + + String[] cql = new DefaultResourceLoader().getResource(filename) + .getContentAsString(StandardCharsets.UTF_8) + .trim() + .split(";"); + + for (var c : cql) { + session.execute(c.trim()); + } + } + + public record StoreWrapper(K store, V conf) { + } + +} diff --git a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java new file mode 100644 index 000000000..0984a067f --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java @@ -0,0 +1,388 @@ +/* + * Copyright 2024 - 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 java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.CqlSessionBuilder; +import com.datastax.oss.driver.api.core.servererrors.InvalidQueryException; +import com.datastax.oss.driver.api.core.servererrors.SyntaxError; +import com.datastax.oss.driver.api.core.type.DataTypes; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.CassandraContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.transformers.TransformersEmbeddingClient; +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn; +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumnTags; +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.ApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.core.io.DefaultResourceLoader; + +import static java.lang.String.format; +import static java.util.Collections.emptyMap; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Use `mvn failsafe:integration-test -Dit.test=CassandraVectorStoreIT` + * + * @author Mick Semb Wever + * @since 1.0.0 + */ +@Testcontainers +class CassandraVectorStoreIT { + + static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("cassandra"); + + @Container + static CassandraContainer cassandraContainer = new CassandraContainer(DEFAULT_IMAGE_NAME.withTag("5.0")); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class); + + List documents = List.of( + new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")), + new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()), + new Document("3", getText("classpath:/test/data/great.depression.txt"), + Map.of("meta2", "meta2", "something_extra", "blue"))); + + public static String getText(String uri) { + var resource = new DefaultResourceLoader().getResource(uri); + try { + return resource.getContentAsString(StandardCharsets.UTF_8); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Test + void ensureBeanGetsCreated() { + this.contextRunner.run(context -> { + try (CassandraVectorStore store = context.getBean(CassandraVectorStore.class)) { + Assertions.assertNotNull(store); + store.checkSchemaValid(); + } + }); + } + + @Test + void addAndSearch() { + contextRunner.run(context -> { + try (CassandraVectorStore store = createTestStore(context, new SchemaColumn("meta1", DataTypes.TEXT), + new SchemaColumn("meta2", DataTypes.TEXT))) { + store.add(documents); + + List results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId()); + + assertThat(resultDoc.getContent()).contains( + "Spring AI provides abstractions that serve as the foundation for developing AI applications."); + + assertThat(resultDoc.getMetadata()).hasSize(2); + assertThat(resultDoc.getMetadata()).containsKeys("meta1", CassandraVectorStore.SIMILARITY_FIELD_NAME); + + // Remove all documents from the store + store.delete(documents.stream().map(doc -> doc.getId()).toList()); + + results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + assertThat(results).isEmpty(); + } + }); + } + + @Test + void searchWithPartitionFilter() throws InterruptedException { + contextRunner.run(context -> { + + try (CassandraVectorStore store = createTestStore(context, + new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) { + + var bgDocument = new Document("BG", "The World is Big and Salvation Lurks Around the Corner", + Map.of("year", (short) 2020)); + var nlDocument = new Document("NL", "The World is Big and Salvation Lurks Around the Corner", + emptyMap()); + var bgDocument2 = new Document("BG2", "The World is Big and Salvation Lurks Around the Corner", + Map.of("year", (short) 2023)); + + store.add(List.of(bgDocument, nlDocument, bgDocument2)); + + List results = store.similaritySearch(SearchRequest.query("The World").withTopK(5)); + assertThat(results).hasSize(3); + + results = store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression(format("%s == 'NL'", CassandraVectorStoreConfig.DEFAULT_ID_NAME))); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression(format("%s == 'BG2'", CassandraVectorStoreConfig.DEFAULT_ID_NAME))); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument2.getId()); + + results = store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression( + format("%s == 'BG' && year == 2020", CassandraVectorStoreConfig.DEFAULT_ID_NAME))); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId()); + + // cassandra server will throw an error + Assertions.assertThrows(SyntaxError.class, () -> { + store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression( + format("NOT(%s == 'BG' && year == 2020)", CassandraVectorStoreConfig.DEFAULT_ID_NAME))); + }); + } + }); + } + + @Test + void unsearchableFilters() throws InterruptedException { + contextRunner.run(context -> { + try (CassandraVectorStore store = context.getBean(CassandraVectorStore.class)) { + + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", (short) 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", (short) 2023)); + + store.add(List.of(bgDocument, nlDocument, bgDocument2)); + + List results = store.similaritySearch(SearchRequest.query("The World").withTopK(5)); + assertThat(results).hasSize(3); + + Assertions.assertThrows(InvalidQueryException.class, () -> { + store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'NL'")); + }); + } + }); + } + + @Test + void searchWithFilters() throws InterruptedException { + contextRunner.run(context -> { + + try (CassandraVectorStore store = createTestStore(context, + new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED), + new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) { + + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", (short) 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", (short) 2023)); + + store.add(List.of(bgDocument, nlDocument, bgDocument2)); + + List results = store.similaritySearch(SearchRequest.query("The World").withTopK(5)); + assertThat(results).hasSize(3); + + results = store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'NL'")); + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG'")); + + 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 = store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG' && year == 2020")); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId()); + + // cassandra server will throw an error + Assertions.assertThrows(SyntaxError.class, () -> { + store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG' || year == 2020")); + }); + + // cassandra server will throw an error + Assertions.assertThrows(SyntaxError.class, () -> { + store.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("NOT(country == 'BG' && year == 2020)")); + }); + } + }); + } + + @Test + void documentUpdate() { + contextRunner.run(context -> { + try (CassandraVectorStore store = context.getBean(CassandraVectorStore.class)) { + + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")); + + store.add(List.of(document)); + + List results = store.similaritySearch(SearchRequest.query("Spring").withTopK(5)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).containsKey("meta1"); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", + Collections.singletonMap("meta2", "meta2")); + + store.add(List.of(sameIdDocument)); + + results = store.similaritySearch(SearchRequest.query("FooBar").withTopK(5)); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).containsKeys("meta2", CassandraVectorStore.SIMILARITY_FIELD_NAME); + + store.delete(List.of(document.getId())); + } + }); + } + + @Test + void searchWithThreshold() { + contextRunner.run(context -> { + try (CassandraVectorStore store = context.getBean(CassandraVectorStore.class)) { + store.add(documents); + + List fullResult = store + .similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThresholdAll()); + + List distances = fullResult.stream() + .map(doc -> (Float) doc.getMetadata().get(CassandraVectorStore.SIMILARITY_FIELD_NAME)) + .toList(); + + assertThat(distances).hasSize(3); + + float threshold = (distances.get(0) + distances.get(1)) / 2; + + List results = store + .similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThreshold(threshold)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId()); + + assertThat(resultDoc.getContent()).contains( + "Spring AI provides abstractions that serve as the foundation for developing AI applications."); + + assertThat(resultDoc.getMetadata()).containsKeys("meta1", CassandraVectorStore.SIMILARITY_FIELD_NAME); + } + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Bean + public CassandraVectorStore store(CqlSession cqlSession, EmbeddingClient embeddingClient) { + + CassandraVectorStoreConfig conf = storeBuilder(cqlSession) + .addMetadataColumn(new SchemaColumn("meta1", DataTypes.TEXT), new SchemaColumn("meta2", DataTypes.TEXT), + new SchemaColumn("country", DataTypes.TEXT), new SchemaColumn("year", DataTypes.SMALLINT)) + .build(); + + conf.dropKeyspace(); + return new CassandraVectorStore(conf, embeddingClient); + } + + @Bean + public EmbeddingClient embeddingClient() { + return new TransformersEmbeddingClient(); + } + + @Bean + public CqlSession cqlSession() { + return new CqlSessionBuilder() + // comment next two lines out to connect to a local C* cluster + .addContactPoint(cassandraContainer.getContactPoint()) + .withLocalDatacenter(cassandraContainer.getLocalDatacenter()) + .build(); + } + + } + + static CassandraVectorStoreConfig.Builder storeBuilder(CqlSession cqlSession) { + return CassandraVectorStoreConfig.builder() + .withCqlSession(cqlSession) + .withKeyspaceName("test_" + CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME); + } + + private CassandraVectorStore createTestStore(ApplicationContext context, SchemaColumn... metadataFields) { + + CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class)) + .addMetadataColumn(metadataFields); + + CassandraVectorStoreConfig conf = builder.build(); + conf.dropKeyspace(); + return new CassandraVectorStore(conf, context.getBean(EmbeddingClient.class)); + } + +} diff --git a/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java new file mode 100644 index 000000000..76028e650 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java @@ -0,0 +1,135 @@ +/* + * Copyright 2024 - 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 java.util.List; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.CqlSessionBuilder; +import com.datastax.oss.driver.api.core.type.DataTypes; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.transformers.TransformersEmbeddingClient; +import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn; +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 static java.lang.String.format; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Example integration-test to use against the schema and full wiki datasets in sstable + * format available from https://github.com/datastax-labs/colbert-wikipedia-data + * + * Use `mvn failsafe:integration-test -Dit.test=WikiVectorStoreExample` + * + * @author Mick Semb Wever + * @since 1.0.0 + */ +@Testcontainers +class WikiVectorStoreExample { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class); + + @Test + void ensureBeanGetsCreated() { + this.contextRunner.run(context -> { + CassandraVectorStore store = context.getBean(CassandraVectorStore.class); + Assertions.assertNotNull(store); + store.checkSchemaValid(); + + store.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + }); + } + + @Test + void search() { + this.contextRunner.run(context -> { + CassandraVectorStore store = context.getBean(CassandraVectorStore.class); + Assertions.assertNotNull(store); + store.checkSchemaValid(); + + var results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + assertThat(results).hasSize(1); + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Bean + public CassandraVectorStore store(CqlSession cqlSession, EmbeddingClient embeddingClient) { + + CassandraVectorStoreConfig conf = CassandraVectorStoreConfig.builder() + .withCqlSession(cqlSession) + .withKeyspaceName("wikidata") + .withTableName("articles") + + .withPartitionKeys(List.of(new SchemaColumn("wiki", DataTypes.TEXT), + new SchemaColumn("language", DataTypes.TEXT), new SchemaColumn("title", DataTypes.TEXT))) + + .withClusteringKeys(List.of(new SchemaColumn("chunk_no", DataTypes.INT), + new SchemaColumn("bert_embedding_no", DataTypes.INT))) + + .withContentColumnName("body") + .withEmbeddingColumnName("all_minilm_l6_v2_embedding") + .withIndexName("all_minilm_l6_v2_ann") + .disallowSchemaChanges() + + .addMetadataColumn(new SchemaColumn("revision", DataTypes.INT), new SchemaColumn("id", DataTypes.INT)) + + .withPrimaryKeyTranslator((List primaryKeys) -> { + if (primaryKeys.isEmpty()) { + return "test§¶0"; + } + return format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3)); + }) + + .withDocumentIdTranslator((id) -> { + String[] parts = id.split("§¶"); + String title = parts[0]; + int chunk_no = 0 < parts.length ? Integer.parseInt(parts[1]) : 0; + return List.of("simplewiki", "en", title, chunk_no, 0); + }) + .build(); + + return new CassandraVectorStore(conf, embeddingClient()); + } + + @Bean + public EmbeddingClient embeddingClient() { + // default is ONNX all-MiniLM-L6-v2 which is what we want + return new TransformersEmbeddingClient(); + } + + @Bean + public CqlSession cqlSession() { + return new CqlSessionBuilder() + // presumes a local C* cluster is running + .build(); + } + + } + +} diff --git a/vector-stores/spring-ai-cassandra/src/test/resources/application.conf b/vector-stores/spring-ai-cassandra/src/test/resources/application.conf new file mode 100644 index 000000000..cf73fb15e --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/resources/application.conf @@ -0,0 +1,6 @@ +# Reference configuration for the DataStax Java driver for Apache Cassandra®. +# see https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration +datastax-java-driver { + # drop statements in tests can be slow + basic.request.timeout = 20 seconds +} diff --git a/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_full_schema.cql b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_full_schema.cql new file mode 100644 index 000000000..c6f6cf17a --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_full_schema.cql @@ -0,0 +1,19 @@ +CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}; + +CREATE TABLE IF NOT EXISTS test_wikidata.articles ( + wiki text, + language text, + title text, + chunk_no int, + id int, + revision int, + body text, + all_minilm_l6_v2_embedding vector, + PRIMARY KEY ((wiki, language, title), chunk_no) +); + +CREATE CUSTOM INDEX IF NOT EXISTS all_minilm_l6_v2_ann ON test_wikidata.articles(all_minilm_l6_v2_embedding) USING 'SAI' + WITH OPTIONS = { 'similarity_function': 'COSINE' }; + + +CREATE CUSTOM INDEX IF NOT EXISTS id_idx ON test_wikidata.articles(id) USING 'SAI'; \ No newline at end of file diff --git a/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_0_schema.cql b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_0_schema.cql new file mode 100644 index 000000000..d2f3fcd62 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_0_schema.cql @@ -0,0 +1 @@ +CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}; diff --git a/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_1_schema.cql b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_1_schema.cql new file mode 100644 index 000000000..cb1a53582 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_1_schema.cql @@ -0,0 +1,13 @@ +CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}; + +CREATE TABLE IF NOT EXISTS test_wikidata.articles ( + wiki text, + language text, + title text, + chunk_no int, + id int, + revision int, + body text, + all_minilm_l6_v2_embedding vector, + PRIMARY KEY ((wiki, language, title), chunk_no) +); diff --git a/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_2_schema.cql b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_2_schema.cql new file mode 100644 index 000000000..5853b2274 --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_2_schema.cql @@ -0,0 +1,15 @@ +CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}; + +CREATE TABLE IF NOT EXISTS test_wikidata.articles ( + wiki text, + language text, + title text, + chunk_no int, + id int, + all_minilm_l6_v2_embedding vector, + PRIMARY KEY ((wiki, language, title), chunk_no) +); + +CREATE CUSTOM INDEX IF NOT EXISTS all_minilm_l6_v2_ann ON test_wikidata.articles(all_minilm_l6_v2_embedding) USING 'SAI' + WITH OPTIONS = { 'similarity_function': 'COSINE' }; + diff --git a/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_3_schema.cql b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_3_schema.cql new file mode 100644 index 000000000..a605116ca --- /dev/null +++ b/vector-stores/spring-ai-cassandra/src/test/resources/test_wiki_partial_3_schema.cql @@ -0,0 +1,12 @@ +CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}; + +CREATE TABLE IF NOT EXISTS test_wikidata.articles ( + wiki text, + language text, + title text, + chunk_no int, + id int, + revision int, + body text, + PRIMARY KEY ((wiki, language, title), chunk_no) +); \ No newline at end of file