Add Azure AI Search Vector Store

- Use Azure AI Search end point to implement the VectorStore interface.
 - Add ITs and README.
 - Create/update vector index on after properties set.
 - Add boot auto-configuration and ITs.
 - Add boot starter for the vector store.

 Resolves: #82
This commit is contained in:
meyerg
2023-11-06 12:07:08 -06:00
committed by Christian Tzolov
parent 621aa7d53e
commit 0aa8eb651e
14 changed files with 1152 additions and 15 deletions

View File

@@ -33,6 +33,7 @@
<module>embedding-clients/transformers-embedding</module>
<module>vector-stores/spring-ai-pinecone</module>
<module>vector-stores/spring-ai-chroma</module>
<module>vector-stores/spring-ai-azure</module>
</modules>
@@ -92,6 +93,8 @@
<milvus.version>2.3.3</milvus.version>
<pinecone.version>0.6.0</pinecone.version>
<protobuf-java-util.version>3.24.4</protobuf-java-util.version>
<fastjson.version>2.0.42</fastjson.version>
<azure-search.version>11.6.0</azure-search.version>
<!-- testing dependecies -->
<testcontainers.version>1.19.0</testcontainers.version>

View File

@@ -102,6 +102,14 @@
<optional>true</optional>
</dependency>
<!-- Azure Vector Store -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-azure-vector-store</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
@@ -157,7 +165,6 @@
<dependency>
<groupId>org.awaitility</groupId>
<artifactId>awaitility</artifactId>
<version>3.0.0</version>
<scope>test</scope>
</dependency>

View File

@@ -0,0 +1,70 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.vectorstore.azure;
import com.azure.core.credential.AzureKeyCredential;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.AzureVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
/**
* @author Christian Tzolov
*/
@AutoConfiguration
@ConditionalOnClass({ EmbeddingClient.class, SearchIndexClient.class })
@EnableConfigurationProperties({ AzureVectorStoreProperties.class })
@ConditionalOnProperty(prefix = "spring.ai.vectorstore.azure", value = { "url", "api-key", "index-name" })
public class AzureVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public SearchIndexClient searchIndexClient(AzureVectorStoreProperties properties) {
return new SearchIndexClientBuilder().endpoint(properties.getUrl())
.credential(new AzureKeyCredential(properties.getApiKey()))
.buildClient();
}
@Bean
@ConditionalOnMissingBean
public VectorStore vectorStore(SearchIndexClient searchIndexClient, EmbeddingClient embeddingClient,
AzureVectorStoreProperties properties) {
var vectorStore = new AzureVectorStore(searchIndexClient, embeddingClient);
vectorStore.setIndexName(properties.getIndexName());
if (properties.getDefaultTopK() >= 0) {
vectorStore.setDefaultTopK(properties.getDefaultTopK());
}
if (properties.getDefaultSimilarityThreshold() >= 0.0) {
vectorStore.setDefaultSimilarityThreshold(properties.getDefaultSimilarityThreshold());
}
return vectorStore;
}
}

View File

@@ -0,0 +1,80 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.vectorstore.azure;
import org.springframework.ai.vectorstore.AzureVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* @author Christian Tzolov
*/
@ConfigurationProperties(AzureVectorStoreProperties.CONFIG_PREFIX)
public class AzureVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.azure";
private String url;
private String apiKey;
private String indexName = AzureVectorStore.DEFAULT_INDEX_NAME;
private int defaultTopK = -1;
private double defaultSimilarityThreshold = -1;
public String getUrl() {
return url;
}
public void setUrl(String endpointUrl) {
this.url = endpointUrl;
}
public String getApiKey() {
return apiKey;
}
public void setApiKey(String apiKey) {
this.apiKey = apiKey;
}
public String getIndexName() {
return indexName;
}
public void setIndexName(String indexName) {
this.indexName = indexName;
}
public int getDefaultTopK() {
return defaultTopK;
}
public void setDefaultTopK(int defaultTopK) {
this.defaultTopK = defaultTopK;
}
public double getDefaultSimilarityThreshold() {
return defaultSimilarityThreshold;
}
public void setDefaultSimilarityThreshold(double defaultSimilarityThreshold) {
this.defaultSimilarityThreshold = defaultSimilarityThreshold;
}
}

View File

@@ -6,3 +6,4 @@ org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreAutoCon
org.springframework.ai.autoconfigure.embedding.transformer.TransformersEmbeddingClientAutoConfiguration
org.springframework.ai.autoconfigure.huggingface.HuggingfaceAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.chroma.ChromaVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.azure.AzureVectorStoreAutoConfiguration

View File

@@ -0,0 +1,137 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.vectorstore.azure;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;
import org.springframework.ai.vectorstore.AzureVectorStore;
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 org.springframework.core.io.DefaultResourceLoader;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.hasSize;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AZURE_AI_SEARCH_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "AZURE_AI_SEARCH_ENDPOINT", matches = ".+")
public class AzureVectorStoreAutoConfigurationIT {
List<Document> documents = List.of(
new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")),
new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()),
new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad")));
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);
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(AzureVectorStoreAutoConfiguration.class))
.withUserConfiguration(Config.class)
.withPropertyValues("spring.ai.vectorstore.azure.apiKey=" + System.getenv("AZURE_AI_SEARCH_API_KEY"),
"spring.ai.vectorstore.azure.url=" + System.getenv("AZURE_AI_SEARCH_ENDPOINT"));
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
Awaitility.setDefaultPollDelay(Duration.ZERO);
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
}
@Test
public void addAndSearchTest() {
contextRunner
.withPropertyValues("spring.ai.vectorstore.azure.indexName=my_test_index",
"spring.ai.vectorstore.azure.defaultTopK=6",
"spring.ai.vectorstore.azure.defaultSimilarityThreshold=0.75")
.run(context -> {
var properties = context.getBean(AzureVectorStoreProperties.class);
assertThat(properties.getUrl()).isEqualTo(System.getenv("AZURE_AI_SEARCH_ENDPOINT"));
assertThat(properties.getApiKey()).isEqualTo(System.getenv("AZURE_AI_SEARCH_API_KEY"));
assertThat(properties.getDefaultTopK()).isEqualTo(6);
assertThat(properties.getDefaultSimilarityThreshold()).isEqualTo(0.75);
assertThat(properties.getIndexName()).isEqualTo("my_test_index");
VectorStore vectorStore = context.getBean(VectorStore.class);
assertThat(vectorStore).isInstanceOf(AzureVectorStore.class);
vectorStore.add(documents);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
}, hasSize(1));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKeys("spring", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
}, hasSize(0));
});
}
@Configuration(proxyBeanMethods = false)
static class Config {
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}

View File

@@ -23,12 +23,11 @@ import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.awaitility.Awaitility;
import org.awaitility.Duration;
import java.time.Duration;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;
@@ -76,7 +75,7 @@ public class PineconeVectorStoreAutoConfigurationIT {
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
Awaitility.setDefaultPollDelay(Duration.ZERO);
Awaitility.setDefaultTimeout(Duration.ONE_MINUTE);
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
}
@Test

View File

@@ -0,0 +1,247 @@
# Azure AI Search VectorStore
This README will walk you through setting up the `AzureVectorStore`` to store document embeddings and perform similarity searches using the Azure AI Search Service.
[Azure AI Search](https://azure.microsoft.com/en-us/products/ai-services/cognitive-search) is a versatile cloud hosted cloud information retrieval system that is part of Microsoft's larger AI platform.
Among other features, it allows users to query information using vector based storage and retrieval.
## Prerequisites
1. Azure Subscription: You will need an [Azure subscription](https://azure.microsoft.com/en-us/free/) to use any Azure service.
2. Azure AI Search Service: Create an [AI Search service](https://portal.azure.com/#create/Microsoft.Search). Once the service is created,
obtain the admin apiKey from the `Keys` section under `Settings` and retrieve the endpoint from the `Url` field under the `Overview` section.
3. (Optional) Azure OpenAI Service: Create an an Azure [OpenAI service](https://portal.azure.com/#create/Microsoft.AIServicesOpenAI).
**NOTE:** You may have to fill out a separate form to gain access to Azure Open AI services.
Once the service is created, obtain the endpoint and apiKey from the `Keys and Endpoint` section under `Resource Management`
## Configuration
On startup the `AzureVectorStore` will attempt to create a new index within your AI Search service instance.
Alternatively you create the index, manually as explained in [Appendix A](appendix_a).
To set up an AzureVectorStore, you will need the settings retrieved from the prerequisites above along with your index name:
* Azure AI Search Endpoint
* Azure AI Search Key
* (optional) Azure OpenAI API Endpoint
* (optional) Azure OpenAI API Key
You can provide these values as OS environment variables.
```bash
export 'AZURE_AI_SEARCH_API_KEY=<My AI Search API Key>'
export 'AZURE_AI_SEARCH_ENDPOINT=<My AI Search Index>'
export 'OPENAI_API_KEY=<My Azure AI API Key>' (Optional)
```
**NOTE** You can replace Azure Open AI implementation with any valid OpenAI implementation that supports the Embeddings interface. For example, you could use Spring AIs Open AI or TransformersEmbedding implementations for embeddings instead of the Azure implementation.
## Dependencies
Add these dependencies to your project:
1. Select an Embeddings interface implementation.
You can choose between:
* OpenAI Embedding:
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
* Or Azure AI Embedding:
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-azure-openai-spring-boot-starter</artifactId>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
* Or Local Sentence Transformers Embedding:
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-transformers-embedding-spring-boot-starter</artifactId>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
2. Azure (AI Search) Vector Store
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-azure-vector-store</artifactId>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
## Sample Code
To configure an Azure `SearchIndexClient` in your application, you can use the following code:
```java
@Bean
public SearchIndexClient searchIndexClient() {
return new SearchIndexClientBuilder().endpoint(System.getenv("AZURE_AI_SEARCH_ENDPOINT"))
.credential(new AzureKeyCredential(System.getenv("AZURE_AI_SEARCH_API_KEY")))
.buildClient();
}
```
To create a vector store, you can use the following code by injecting the `SearchIndexClient` bean created in the above sample along with and `EmbeddingClient` provided by Spring AI library that's implements the desired Embeddings interface.
```java
@Bean
public VectorStore vectorStore(SearchIndexClient searchIndexClient, EmbeddingClient embeddingClient) {
return new AzureVectorStore(searchIndexClient, embeddingClient);
}
```
In your main code, create some documents
```java
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
new Document("The World is Big and Salvation Lurks Around the Corner"),
new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));
```
Add the documents to your vector store:
```java
vectorStore.add(List.of(document));
```
And finally, retrieve documents similar to a query:
```java
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
```
If all goes well, you should retrieve the document containing the text "Spring AI rocks!!".
## Integration With Azure OpenAI Studio Data Ingestion
Azure Open AI services provides a convenient method to upload documents into an Index as described in this Microsoft
[learning document](https://learn.microsoft.com/en-us/azure/ai-services/openai/use-your-data-quickstart?tabs=command-line&pivots=programming-language-csharp).
The `AzureVectorStore` implementation is compatible with indexes that use this methodology facilitating an *easier* way to integrate with your existing documents for the purpose of searching and integrating with the AI system.
## <a name="appendix_a" /> Appendix A: Create Vector Store Search Index
The easiest way to crate a search index manually, is to create one from a JSON document.
This can be done by clicking on the `Indexes` link under the `Search management` section.
From the Indexes page, click `+ Add index` and select `Add index (JSON)`. In the
`Add index (JSON)` window of the right side of your screen, enter the following JSON replacing `<INDEX NAME>` with the name you would like to give your index and click
`save`.
```json
{
"name": "<INDEX NAME>",
"defaultScoringProfile": null,
"fields": [
{
"name": "id",
"type": "Edm.String",
"searchable": false,
"filterable": false,
"retrievable": true,
"sortable": false,
"facetable": false,
"key": true,
"indexAnalyzer": null,
"searchAnalyzer": null,
"analyzer": null,
"normalizer": null,
"dimensions": null,
"vectorSearchConfiguration": null,
"synonymMaps": []
},
{
"name": "embedding",
"type": "Collection(Edm.Single)",
"searchable": true,
"filterable": false,
"retrievable": true,
"sortable": false,
"facetable": false,
"key": false,
"indexAnalyzer": null,
"searchAnalyzer": null,
"analyzer": null,
"normalizer": null,
"dimensions": 1536, // set the dimensions for the configured Embedding Client. It defaults to to OpenAI's 1536 size.
"vectorSearchConfiguration": "default",
"synonymMaps": []
},
{
"name": "content",
"type": "Edm.String",
"searchable": true,
"filterable": false,
"retrievable": true,
"sortable": false,
"facetable": false,
"key": false,
"indexAnalyzer": null,
"searchAnalyzer": null,
"analyzer": null,
"normalizer": null,
"dimensions": null,
"vectorSearchConfiguration": null,
"synonymMaps": []
},
{
"name": "metadata",
"type": "Edm.String",
"searchable": true,
"filterable": true,
"retrievable": true,
"sortable": true,
"facetable": true,
"key": false,
"indexAnalyzer": null,
"searchAnalyzer": null,
"analyzer": null,
"normalizer": null,
"dimensions": null,
"vectorSearchConfiguration": null,
"synonymMaps": []
}
],
"scoringProfiles": [],
"corsOptions": null,
"suggesters": [],
"analyzers": [],
"normalizers": [],
"tokenizers": [],
"tokenFilters": [],
"charFilters": [],
"encryptionKey": null,
"semantic": null,
"vectorSearch": {
"algorithmConfigurations": [
{
"name": "default",
"kind": "hnsw",
"hnswParameters": {
"metric": "cosine",
"m": 4,
"efConstruction": 400,
"efSearch": 1000
},
"exhaustiveKnnParameters": null
}
]
}
}
```

View File

@@ -0,0 +1,86 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-azure-vector-store</artifactId>
<packaging>jar</packaging>
<name>Spring AI - Azure AI Search Vector Store</name>
<description>Spring AI - Azure AI Search Vector Store</description>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
<scm>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
<connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
</scm>
<properties>
<maven.compiler.target>17</maven.compiler.target>
<maven.compiler.source>17</maven.compiler.source>
</properties>
<dependencies>
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-core</artifactId>
<version>${parent.version}</version>
</dependency>
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-search-documents</artifactId>
<version>${azure-search.version}</version>
<exclusions>
<!-- exclude this to avoid changing the default serializer and the null-value behavior -->
<exclusion>
<groupId>com.azure</groupId>
<artifactId>azure-core-serializer-json-jackson</artifactId>
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>com.alibaba.fastjson2</groupId>
<artifactId>fastjson2</artifactId>
<version>${fastjson.version}</version>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>transformers-embedding</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<!-- Contains smaple test data -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-configuration-processor</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.awaitility</groupId>
<artifactId>awaitility</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@@ -0,0 +1,283 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.TypeReference;
import com.azure.core.util.Context;
import com.azure.search.documents.SearchClient;
import com.azure.search.documents.SearchDocument;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.models.HnswAlgorithmConfiguration;
import com.azure.search.documents.indexes.models.HnswParameters;
import com.azure.search.documents.indexes.models.SearchField;
import com.azure.search.documents.indexes.models.SearchFieldDataType;
import com.azure.search.documents.indexes.models.SearchIndex;
import com.azure.search.documents.indexes.models.VectorSearch;
import com.azure.search.documents.indexes.models.VectorSearchAlgorithmMetric;
import com.azure.search.documents.indexes.models.VectorSearchProfile;
import com.azure.search.documents.models.IndexDocumentsResult;
import com.azure.search.documents.models.IndexingResult;
import com.azure.search.documents.models.SearchOptions;
import com.azure.search.documents.models.VectorSearchOptions;
import com.azure.search.documents.models.VectorizedQuery;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
/**
* Uses Azure Cognitive Search as a backing vector store. Documents can be preloaded into
* a Cognitive Search index and managed via Azure tools or added and managed through this
* VectorStore. The underlying index is configured in the provided Azure
* SearchIndexClient.
*
* @author Greg Meyer
* @author Xiangyang Yu
* @author Christian Tzolov
*/
public class AzureVectorStore implements VectorStore, InitializingBean {
public static final String DEFAULT_INDEX_NAME = "spring_ai_azure_vector_store";
private static final String ID_FIELD_NAME = "id";
private static final String CONTENT_FIELD_NAME = "content";
private static final String EMBEDDING_FIELD_NAME = "embedding";
private static final String METADATA_FIELD_NAME = "metadata";
private static final String DISTANCE_METADATA_FIELD_NAME = "distance";
private static final int DEFAULT_TOP_K = 4;
private static final Double DEFAULT_SIMILARITY_THRESHOLD = 0.0;
private final SearchIndexClient searchIndexClient;
private final EmbeddingClient embeddingClient;
private SearchClient searchClient;
private int defaultTopK = DEFAULT_TOP_K;
private Double defaultSimilarityThreshold = DEFAULT_SIMILARITY_THRESHOLD;
private String indexName = DEFAULT_INDEX_NAME;
/**
* Constructs a new AzureCognitiveSearchVectorStore.
* @param searchIndexClient A pre-configured Azure {@link SearchIndexClient} that CRUD
* for Azure search indexes and factory for {@link SearchClient}.
* @param embeddingClient The client for embedding operations.
*/
public AzureVectorStore(SearchIndexClient searchIndexClient, EmbeddingClient embeddingClient) {
Assert.notNull(embeddingClient, "The embedding client can not be null.");
Assert.notNull(searchIndexClient, "The search index client can not be null.");
this.searchIndexClient = searchIndexClient;
this.embeddingClient = embeddingClient;
}
/**
* Change the Index Name.
* @param indexName The Azure VectorStore index name to use.
*/
public void setIndexName(String indexName) {
Assert.hasText(indexName, "The index name can not be empty.");
this.indexName = indexName;
}
/**
* Sets the a default maximum number of similar documents returned.
* @param topK The default maximum number of similar documents returned.
*/
public void setDefaultTopK(int topK) {
Assert.isTrue(topK >= 0, "The topK should be positive value.");
this.defaultTopK = topK;
}
/**
* Sets the a default similarity threshold for returned documents.
* @param similarityThreshold The a default similarity threshold for returned
* documents.
*/
public void setDefaultSimilarityThreshold(Double similarityThreshold) {
Assert.isTrue(similarityThreshold >= 0.0 && similarityThreshold <= 1.0,
"The similarity threshold must be in range [0.0:1.00].");
this.defaultSimilarityThreshold = similarityThreshold;
}
@Override
public void add(List<Document> documents) {
Assert.notNull(documents, "The document list should not be null.");
if (CollectionUtils.isEmpty(documents)) {
return; // nothing to do;
}
final var searchDocuments = documents.stream().map(document -> {
final var embeddings = this.embeddingClient.embed(document);
SearchDocument searchDocument = new SearchDocument();
searchDocument.put(ID_FIELD_NAME, document.getId());
searchDocument.put(EMBEDDING_FIELD_NAME, embeddings);
searchDocument.put(CONTENT_FIELD_NAME, document.getContent());
// TODO: Consider alternate/native field type for metadata
searchDocument.put(METADATA_FIELD_NAME, new JSONObject(document.getMetadata()).toJSONString());
return searchDocument;
}).toList();
IndexDocumentsResult result = this.searchClient.uploadDocuments(searchDocuments);
for (IndexingResult indexingResult : result.getResults()) {
Assert.isTrue(indexingResult.isSucceeded(),
String.format("Document with key %s upload is not successfully", indexingResult.getKey()));
}
}
@Override
public Optional<Boolean> delete(List<String> documentIds) {
Assert.notNull(documentIds, "The document ID list should not be null.");
if (CollectionUtils.isEmpty(documentIds)) {
return Optional.of(true); // nothing to do;
}
final var searchDocumentIds = documentIds.stream().map(documentId -> {
SearchDocument searchDocument = new SearchDocument();
searchDocument.put(ID_FIELD_NAME, documentId);
return searchDocument;
}).toList();
var results = this.searchClient.deleteDocuments(searchDocumentIds);
boolean resSuccess = true;
for (IndexingResult result : results.getResults()) {
if (!result.isSucceeded()) {
resSuccess = false;
break;
}
}
return Optional.of(resSuccess);
}
@Override
public List<Document> similaritySearch(String query) {
return this.similaritySearch(SearchRequest.query(query)
.withTopK(this.defaultTopK)
.withSimilarityThreshold(this.defaultSimilarityThreshold));
}
@Override
public List<Document> similaritySearch(SearchRequest request) {
Assert.notNull(request, "The search request must not be null.");
if (request.getFilterExpression() != null) {
throw new UnsupportedOperationException(
"The [" + this.getClass() + "] doesn't support metadata filtering!");
}
var searchEmbedding = toFloatList(embeddingClient.embed(request.getQuery()));
final var vectorQuery = new VectorizedQuery(searchEmbedding).setKNearestNeighborsCount(request.getTopK())
// Set the fields to compare the vector against. This is a comma-delimited
// list of field names.
.setFields(EMBEDDING_FIELD_NAME);
final var searchResults = searchClient.search(null,
new SearchOptions().setVectorSearchOptions(new VectorSearchOptions().setQueries(vectorQuery)),
Context.NONE);
return searchResults.stream()
.filter(result -> result.getScore() >= request.getSimilarityThreshold())
.map(result -> {
final AzureSearchDocument entry = result.getDocument(AzureSearchDocument.class);
Map<String, Object> metadata = (StringUtils.hasText(entry.metadata()))
? JSONObject.parseObject(entry.metadata(), new TypeReference<Map<String, Object>>() {
}) : Map.of();
metadata.put(DISTANCE_METADATA_FIELD_NAME, 1 - (float) result.getScore());
final Document doc = new Document(entry.id(), entry.content(), metadata);
doc.setEmbedding(entry.embedding());
return doc;
})
.collect(Collectors.toList());
}
private List<Float> toFloatList(List<Double> doubleList) {
return doubleList.stream().map(Double::floatValue).toList();
}
/**
* Internal data structure for retrieving and and storing documents.
*/
private record AzureSearchDocument(String id, String content, List<Double> embedding, String metadata) {
}
@Override
public void afterPropertiesSet() throws Exception {
int dimensions = this.embeddingClient.dimensions();
SearchIndex searchIndex = new SearchIndex(this.indexName).setFields(
new SearchField(ID_FIELD_NAME, SearchFieldDataType.STRING).setKey(true)
.setFilterable(true)
.setSortable(true),
new SearchField(EMBEDDING_FIELD_NAME, SearchFieldDataType.collection(SearchFieldDataType.SINGLE))
.setSearchable(true)
.setVectorSearchDimensions(dimensions)
// This must match a vector search configuration name.
.setVectorSearchProfileName("my-vector-profile"),
new SearchField(CONTENT_FIELD_NAME, SearchFieldDataType.STRING).setSearchable(true).setFilterable(true),
new SearchField(METADATA_FIELD_NAME, SearchFieldDataType.STRING).setSearchable(true)
.setFilterable(true))
// VectorSearch configuration is required for a vector field. The name used
// for the vector search
// algorithm configuration must match the configuration used by the search
// field used for vector search.
.setVectorSearch(new VectorSearch()
.setProfiles(
Collections.singletonList(new VectorSearchProfile("my-vector-profile", "my-vector-config")))
.setAlgorithms(Collections.singletonList(
new HnswAlgorithmConfiguration("my-vector-config").setParameters(new HnswParameters().setM(4)
.setEfConstruction(400)
.setEfSearch(1000)
.setMetric(VectorSearchAlgorithmMetric.COSINE)))));
var index = this.searchIndexClient.createOrUpdateIndex(searchIndex);
this.searchClient = this.searchIndexClient.getSearchClient(this.indexName);
}
}

View File

@@ -0,0 +1,234 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import com.azure.core.credential.AzureKeyCredential;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasSize;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AZURE_AI_SEARCH_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "AZURE_AI_SEARCH_ENDPOINT", matches = ".+")
public class AzureVectorStoreIT {
List<Document> 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")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
Awaitility.setDefaultPollDelay(Duration.ZERO);
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
}
@Test
public void addAndSearchTest() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
}, hasSize(1));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (19291939) was an economic shock");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(SearchRequest.query("Hello").withTopK(1));
}, hasSize(0));
});
}
@Test
public void documentUpdateTest() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
vectorStore.add(List.of(document));
SearchRequest springSearchRequest = SearchRequest.query("Spring").withTopK(5);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(springSearchRequest);
}, hasSize(1));
List<Document> results = vectorStore.similaritySearch(springSearchRequest);
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");
assertThat(resultDoc.getMetadata()).containsKey("distance");
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
Collections.singletonMap("meta2", "meta2"));
vectorStore.add(List.of(sameIdDocument));
SearchRequest fooBarSearchRequest = SearchRequest.query("FooBar").withTopK(5);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(fooBarSearchRequest).get(0).getContent();
}, equalTo("The World is Big and Salvation Lurks Around the Corner"));
results = vectorStore.similaritySearch(fooBarSearchRequest);
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()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(List.of(document.getId()));
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(fooBarSearchRequest);
}, hasSize(0));
});
}
@Test
public void searchThresholdTest() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
Awaitility.await().until(() -> {
return vectorStore
.similaritySearch(SearchRequest.query("Depression").withTopK(50).withSimilarityThresholdAll());
}, hasSize(3));
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Depression").withTopK(5).withSimilarityThresholdAll());
List<Float> distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
assertThat(distances).hasSize(3);
float threshold = (distances.get(0) + distances.get(1)) / 2;
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("Depression").withTopK(5).withSimilarityThreshold(1 - threshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (19291939) was an economic shock");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(SearchRequest.query("Hello").withTopK(1));
}, hasSize(0));
});
}
@SpringBootConfiguration
@EnableAutoConfiguration
public static class Config {
@Bean
public SearchIndexClient searchIndexClient() {
return new SearchIndexClientBuilder().endpoint(System.getenv("AZURE_AI_SEARCH_ENDPOINT"))
.credential(new AzureKeyCredential(System.getenv("AZURE_AI_SEARCH_API_KEY")))
.buildClient();
}
@Bean
public VectorStore vectorStore(SearchIndexClient searchIndexClient, EmbeddingClient embeddingClient) {
return new AzureVectorStore(searchIndexClient, embeddingClient);
}
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
private 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);
}
}
}

View File

@@ -36,7 +36,6 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.DockerComposeContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -44,7 +43,6 @@ import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
@@ -55,6 +53,7 @@ import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.util.FileSystemUtils;
import static org.assertj.core.api.Assertions.assertThat;

View File

@@ -299,14 +299,6 @@ public class PineconeVectorStore implements VectorStore {
@Override
public List<Document> similaritySearch(SearchRequest request) {
// return this.internalSimilaritySearch(request.getQuery(), request.getTopK(),
// request.getSimilarityThreshold(),
// request.getFilterExpression());
// }
// List<Document> internalSimilaritySearch(String query, int topK, double
// similarityThreshold,
// Filter.Expression filterExpression) {
String nativeExpressionFilters = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";

View File

@@ -30,7 +30,6 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;