第五章：向量数据库

<dependencies>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
    </dependency>

    <dependency>
        <groupId>org.springframework.ai</groupId>
        <artifactId>spring-ai-autoconfigure-model-openai</artifactId>
    </dependency>

    <dependency>
        <groupId>org.springframework.ai</groupId>
        <artifactId>spring-ai-vector-store</artifactId>
    </dependency>

</dependencies>

server:
  port: 8080

spring:
  application:
    name: vector-simple

  ai:
    openai:
      api-key: ${DASHSCOPEAPIKEY}
      base-url: https://dashscope.aliyuncs.com/compatible-mode
      embedding:
        options:
          model: text-embedding-v1

@RestController
@RequestMapping("/vector/simple")
public class VectorSimpleController {
    private static final Logger logger = LoggerFactory.getLogger(VectorSimpleController.class);
    private final SimpleVectorStore simpleVectorStore;
    private final String SAVEPATH = System.getProperty("user.dir") + "/vector/vector-simple/src/main/resources/save.json";

    public VectorSimpleController(EmbeddingModel embeddingModel) {
        this.simpleVectorStore = SimpleVectorStore
                .builder(embeddingModel).build();
    }

    @GetMapping("/add")
    public void add() {
        logger.info("start add data");

        HashMap<String, Object> map = new HashMap<>();
        map.put("year", 2025);
        map.put("name", "yingzi");
        List<Document> documents = List.of(
                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("year", 2024)),
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", map),
                new Document("1", "test content", map));
        simpleVectorStore.add(documents);
    }

    @GetMapping("/delete")
    public void delete() {
        logger.info("start delete data");
        simpleVectorStore.delete(List.of("1"));
    }

    @GetMapping("/save")
    public void save() {
        logger.info("start save data: {}", SAVEPATH);
        File file = new File(SAVEPATH);
        if (file.exists()) {
            file.delete();
        }
        simpleVectorStore.save(file);
    }

    @GetMapping("/load")
    public void load() {
        logger.info("start load data: {}", SAVEPATH);
        File file = new File(SAVEPATH);
        simpleVectorStore.load(file);
    }

    @GetMapping("/search")
    public List<Document> search() {
        logger.info("start search data");
        return simpleVectorStore.similaritySearch(SearchRequest
                .builder()
                .query("Spring")
                .topK(2)
                .build());
    }

    @GetMapping("/search-filter")
    public List<Document> searchFilter() {
        logger.info("start search  filter data");
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        Filter.Expression expression = b.and(
                b.in("year", 2025, 2024),
                b.eq("name", "yingzi")
        ).build();

        return simpleVectorStore.similaritySearch(SearchRequest
                .builder()
                .query("Spring")
                .topK(2)
                .filterExpression(expression).build());
    }
}

<dependency>
        <groupId>org.springframework.ai</groupId>
        <artifactId>spring-ai-starter-vector-store-redis</artifactId>
    </dependency>

Redis
server:
  port: 8080

spring:
  application:
    name: vector-redis

  ai:
    openai:
      api-key: ${DASHSCOPEAPIKEY}
      base-url: https://dashscope.aliyuncs.com/compatible-mode
      embedding:
        options:
          model: text-embedding-v1

    vectorstore:
      redis:
        initialize-schema: true
        prefix: prefixyingzi
        index-name: indexyingzi

  data:
    redis:
      host: localhost
      port: 6379 # Redis服务器连接端口

package com.spring.ai.tutorial.vector.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import redis.clients.jedis.JedisPooled;

@Configuration
public class RedisConfig {

    private static final Logger logger = LoggerFactory.getLogger(RedisConfig.class);

    @Value("${spring.data.redis.host}")
    private String host;
    @Value("${spring.data.redis.port}")
    private int port;
    @Value("${spring.ai.vectorstore.redis.prefix}")
    private String prefix;
    @Value("${spring.ai.vectorstore.redis.index-name}")
    private String indexName;

    @Bean
    public JedisPooled jedisPooled() {
        logger.info("Redis host: {}, port: {}", host, port);
        return new JedisPooled(host, port);
    }

    @Bean
    @Qualifier("redisVectorStoreCustom")
    public RedisVectorStore vectorStore(JedisPooled jedisPooled, EmbeddingModel embeddingModel) {
        logger.info("create redis vector store");
        return RedisVectorStore.builder(jedisPooled, embeddingModel)
                .indexName(indexName)                // Optional: defaults to "spring-ai-index"
                .prefix(prefix)                  // Optional: defaults to "embedding:"
                .metadataFields(                         // Optional: define metadata fields for filtering
                        RedisVectorStore.MetadataField.tag("name"),
                        RedisVectorStore.MetadataField.numeric("year"))
                .initializeSchema(true)                   // Optional: defaults to false
                .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy
                .build();
    }

}

package com.spring.ai.tutorial.vector.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/redis")
public class RedisController {

    private static final Logger logger = LoggerFactory.getLogger(RedisController.class);
    private final RedisVectorStore redisVectorStore;

    @Autowired
    public RedisController(@Qualifier("redisVectorStoreCustom") RedisVectorStore redisVectorStore) {
        this.redisVectorStore = redisVectorStore;
    }

    @GetMapping("/add")
    public void add() {
        logger.info("start add data");

        HashMap<String, Object> map = new HashMap<>();
        map.put("id", "12345");
        map.put("year", 2025);
        map.put("name", "yingzi");
        List<Document> documents = List.of(
                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("year", 2024)),
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", map));
        redisVectorStore.add(documents);
    }

    @GetMapping("/search")
    public List<Document> search() {
        logger.info("start search data");
        return redisVectorStore.similaritySearch(SearchRequest
                .builder()
                .query("Spring")
                .topK(2)
                .build());
    }

    @GetMapping("delete-filter")
    public void searchFilter() {
        logger.info("start delete data with filter");
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        Filter.Expression expression = b.eq("name", "yingzi").build();

        redisVectorStore.delete(expression);
    }
}

<properties>
        <es.version>8.17.4</es.version>
    </properties>

    // 额外添加的es依赖
    <dependency>
        <groupId>org.springframework.ai</groupId>
        <artifactId>spring-ai-starter-vector-store-elasticsearch</artifactId>
    </dependency>

    <dependency>
        <groupId>co.elastic.clients</groupId>
        <artifactId>elasticsearch-java</artifactId>
        <version>${es.version}</version>
        <scope>compile</scope>
    </dependency>

server:
  port: 8080

spring:
  application:
    name: vector-elasticsearch

  ai:
    openai:
      api-key: ${DASHSCOPEAPIKEY}
      base-url: https://dashscope.aliyuncs.com/compatible-mode
      embedding:
        options:
          model: text-embedding-v1

    vectorstore:
      elasticsearch:
        initialize-schema: true
        index-name: yingzi
        similarity: cosine
        dimensions: 1536

  elasticsearch:
    uris: http://127.0.0.1:9200
    username: elastic
    password: yingzi

package com.spring.ai.tutorial.vector.config;

import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.client.RestClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions;
import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class ElasticsearchConfig {

    private static final Logger logger = LoggerFactory.getLogger(ElasticsearchConfig.class);


    @Value("${spring.elasticsearch.uris}")
    private String url;
    @Value("${spring.elasticsearch.username}")
    private String username;
    @Value("${spring.elasticsearch.password}")
    private String password;

    @Value("${spring.ai.vectorstore.elasticsearch.index-name}")
    private String indexName;
    @Value("${spring.ai.vectorstore.elasticsearch.similarity}")
    private SimilarityFunction similarityFunction;
    @Value("${spring.ai.vectorstore.elasticsearch.dimensions}")
    private int dimensions;


    @Bean
    public RestClient restClient() {
        // 解析URL
        String[] urlParts = url.split("://");
        String protocol = urlParts[0];
        String hostAndPort = urlParts[1];
        String[] hostPortParts = hostAndPort.split(":");
        String host = hostPortParts[0];
        int port = Integer.parseInt(hostPortParts[1]);

        // 创建凭证提供者
        CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(AuthScope.ANY,
            new UsernamePasswordCredentials(username, password));

        logger.info("create elasticsearch rest client");
        // 构建RestClient
        return RestClient.builder(new HttpHost(host, port, protocol))
                .setHttpClientConfigCallback(httpClientBuilder -> {
                    httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                    return httpClientBuilder;
                })
                .build();
    }

    @Bean
    @Qualifier("elasticsearchVectorStore")
    public ElasticsearchVectorStore vectorStore(RestClient restClient, EmbeddingModel embeddingModel) {
        logger.info("create elasticsearch vector store");

        ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
        options.setIndexName(indexName);    // Optional: defaults to "spring-ai-document-index"
        options.setSimilarity(similarityFunction);           // Optional: defaults to COSINE
        options.setDimensions(dimensions);             // Optional: defaults to model dimensions or 1536

        return ElasticsearchVectorStore.builder(restClient, embeddingModel)
                .options(options)                     // Optional: use custom options
                .initializeSchema(true)               // Optional: defaults to false
                .batchingStrategy(new TokenCountBatchingStrategy())// Optional: defaults to TokenCountBatchingStrategy
                .build();
    }

}

package com.spring.ai.tutorial.vector.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/vector/es")
public class VectorEsController {

    private static final Logger logger = LoggerFactory.getLogger(VectorEsController.class);
    private final ElasticsearchVectorStore elasticsearchVectorStore;

    @Autowired
    public VectorEsController(@Qualifier("elasticsearchVectorStore") ElasticsearchVectorStore elasticsearchVectorStore) {
        this.elasticsearchVectorStore = elasticsearchVectorStore;
    }

    @GetMapping("/add")
    public void add() {
        logger.info("start import data");

        HashMap<String, Object> map = new HashMap<>();
        map.put("id", "12345");
        map.put("year", 2025);
        map.put("name", "yingzi");
        List<Document> documents = List.of(
                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("year", 2024)),
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", map));
        elasticsearchVectorStore.add(documents);
    }

    @GetMapping("/search")
    public List<Document> search() {
        logger.info("start search data");
        return elasticsearchVectorStore.similaritySearch(SearchRequest
                .builder()
                .query("Spring")
                .topK(2)
                .build());
    }

    @GetMapping("delete-filter")
    public void searchFilter() {
        logger.info("start search data with filter");
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        Filter.Expression expression = b.and(
                b.in("year", 2025, 2024),
                b.eq("name", "yingzi")
        ).build();

        elasticsearchVectorStore.delete(expression);
    }
}

package org.springframework.ai.document;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonCreator.Mode;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import org.springframework.ai.content.Media;
import org.springframework.ai.document.id.IdGenerator;
import org.springframework.ai.document.id.RandomIdGenerator;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

@JsonIgnoreProperties({"contentFormatter", "embedding"})
public class Document {
    public static final ContentFormatter DEFAULTCONTENTFORMATTER = DefaultContentFormatter.defaultConfig();
    private final String id;

    private final String text;
    private final Media media;
    private final Map<String, Object> metadata;
    @Nullable
    private final Double score;
    @JsonIgnore
    private ContentFormatter contentFormatter;

    @JsonCreator(
        mode = Mode.PROPERTIES
    )
    public Document(@JsonProperty("content") String content) {
        this((String)content, new HashMap());
    }

    public Document(String text, Map<String, Object> metadata) {
        this((new RandomIdGenerator()).generateId(new Object[0]), text, (Media)null, metadata, (Double)null);
    }

    public Document(String id, String text, Map<String, Object> metadata) {
        this(id, text, (Media)null, metadata, (Double)null);
    }

    public Document(Media media, Map<String, Object> metadata) {
        this((new RandomIdGenerator()).generateId(new Object[0]), (String)null, media, metadata, (Double)null);
    }

    public Document(String id, Media media, Map<String, Object> metadata) {
        this(id, (String)null, media, metadata, (Double)null);
    }

    private Document(String id, String text, Media media, Map<String, Object> metadata, @Nullable Double score) {
        this.contentFormatter = DEFAULTCONTENTFORMATTER;
        Assert.hasText(id, "id cannot be null or empty");
        Assert.notNull(metadata, "metadata cannot be null");
        Assert.noNullElements(metadata.keySet(), "metadata cannot have null keys");
        Assert.noNullElements(metadata.values(), "metadata cannot have null values");
        Assert.isTrue(text != null ^ media != null, "exactly one of text or media must be specified");
        this.id = id;
        this.text = text;
        this.media = media;
        this.metadata = new HashMap(metadata);
        this.score = score;
    }

    public static Builder builder() {
        return new Builder();
    }

    public String getId() {
        return this.id;
    }

    @Nullable
    public String getText() {
        return this.text;
    }

    public boolean isText() {
        return this.text != null;
    }

    @Nullable
    public Media getMedia() {
        return this.media;
    }

    @JsonIgnore
    public String getFormattedContent() {
        return this.getFormattedContent(MetadataMode.ALL);
    }

    public String getFormattedContent(MetadataMode metadataMode) {
        Assert.notNull(metadataMode, "Metadata mode must not be null");
        return this.contentFormatter.format(this, metadataMode);
    }

    public String getFormattedContent(ContentFormatter formatter, MetadataMode metadataMode) {
        Assert.notNull(formatter, "formatter must not be null");
        Assert.notNull(metadataMode, "Metadata mode must not be null");
        return formatter.format(this, metadataMode);
    }

    public Map<String, Object> getMetadata() {
        return this.metadata;
    }

    @Nullable
    public Double getScore() {
        return this.score;
    }

    public ContentFormatter getContentFormatter() {
        return this.contentFormatter;
    }

    public void setContentFormatter(ContentFormatter contentFormatter) {
        this.contentFormatter = contentFormatter;
    }

    public Builder mutate() {
        return (new Builder()).id(this.id).text(this.text).media(this.media).metadata(this.metadata).score(this.score);
    }

    public boolean equals(Object o) {
        if (o != null && this.getClass() == o.getClass()) {
            Document document = (Document)o;
            return Objects.equals(this.id, document.id) && Objects.equals(this.text, document.text) && Objects.equals(this.media, document.media) && Objects.equals(this.metadata, document.metadata) && Objects.equals(this.score, document.score);
        } else {
            return false;
        }
    }

    public int hashCode() {
        return Objects.hash(new Object[]{this.id, this.text, this.media, this.metadata, this.score});
    }

    public String toString() {
        String var10000 = this.id;
        return "Document{id='" + var10000 + "', text='" + this.text + "', media='" + String.valueOf(this.media) + "', metadata=" + String.valueOf(this.metadata) + ", score=" + this.score + "}";
    }

    public static class Builder {
        private String id;
        private String text;
        private Media media;
        private Map<String, Object> metadata = new HashMap();
        @Nullable
        private Double score;
        private IdGenerator idGenerator = new RandomIdGenerator();

        public Builder idGenerator(IdGenerator idGenerator) {
            Assert.notNull(idGenerator, "idGenerator cannot be null");
            this.idGenerator = idGenerator;
            return this;
        }

        public Builder id(String id) {
            Assert.hasText(id, "id cannot be null or empty");
            this.id = id;
            return this;
        }

        public Builder text(@Nullable String text) {
            this.text = text;
            return this;
        }

        public Builder media(@Nullable Media media) {
            this.media = media;
            return this;
        }

        public Builder metadata(Map<String, Object> metadata) {
            Assert.notNull(metadata, "metadata cannot be null");
            this.metadata = metadata;
            return this;
        }

        public Builder metadata(String key, Object value) {
            Assert.notNull(key, "metadata key cannot be null");
            Assert.notNull(value, "metadata value cannot be null");
            this.metadata.put(key, value);
            return this;
        }

        public Builder score(@Nullable Double score) {
            this.score = score;
            return this;
        }

        public Document build() {
            if (!StringUtils.hasText(this.id)) {
                this.id = this.idGenerator.generateId(new Object[]{this.text, this.metadata});
            }

            return new Document(this.id, this.text, this.media, this.metadata, this.score);
        }
    }
}

package org.springframework.ai.document;

import java.util.List;
import java.util.function.Consumer;

public interface DocumentWriter extends Consumer<List<Document>> {
    default void write(List<Document> documents) {
        this.accept(documents);
    }
}

public interface BatchingStrategy {
    List<List<Document>> batch(List<Document> documents);
}

package org.springframework.ai.embedding;

import com.knuddels.jtokkit.api.EncodingType;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import org.springframework.ai.document.ContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
import org.springframework.ai.tokenizer.TokenCountEstimator;
import org.springframework.util.Assert;

public class TokenCountBatchingStrategy implements BatchingStrategy {
    private static final int MAXINPUTTOKENCOUNT = 8191;
    private static final double DEFAULTTOKENCOUNTRESERVEPERCENTAGE = 0.1;
    private final TokenCountEstimator tokenCountEstimator;
    private final int maxInputTokenCount;
    private final ContentFormatter contentFormatter;
    private final MetadataMode metadataMode;

    public TokenCountBatchingStrategy() {
        this(EncodingType.CL100KBASE, 8191, 0.1);
    }

    public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double reservePercentage) {
        this(encodingType, maxInputTokenCount, reservePercentage, Document.DEFAULTCONTENTFORMATTER, MetadataMode.NONE);
    }

    public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double reservePercentage, ContentFormatter contentFormatter, MetadataMode metadataMode) {
        Assert.notNull(encodingType, "EncodingType must not be null");
        Assert.isTrue(maxInputTokenCount > 0, "MaxInputTokenCount must be greater than 0");
        Assert.isTrue(reservePercentage >= (double)0.0F && reservePercentage < (double)1.0F, "ReservePercentage must be in range [0, 1)");
        Assert.notNull(contentFormatter, "ContentFormatter must not be null");
        Assert.notNull(metadataMode, "MetadataMode must not be null");
        this.tokenCountEstimator = new JTokkitTokenCountEstimator(encodingType);
        this.maxInputTokenCount = (int)Math.round((double)maxInputTokenCount * ((double)1.0F - reservePercentage));
        this.contentFormatter = contentFormatter;
        this.metadataMode = metadataMode;
    }

    public TokenCountBatchingStrategy(TokenCountEstimator tokenCountEstimator, int maxInputTokenCount, double reservePercentage, ContentFormatter contentFormatter, MetadataMode metadataMode) {
        Assert.notNull(tokenCountEstimator, "TokenCountEstimator must not be null");
        Assert.isTrue(maxInputTokenCount > 0, "MaxInputTokenCount must be greater than 0");
        Assert.isTrue(reservePercentage >= (double)0.0F && reservePercentage < (double)1.0F, "ReservePercentage must be in range [0, 1)");
        Assert.notNull(contentFormatter, "ContentFormatter must not be null");
        Assert.notNull(metadataMode, "MetadataMode must not be null");
        this.tokenCountEstimator = tokenCountEstimator;
        this.maxInputTokenCount = (int)Math.round((double)maxInputTokenCount * ((double)1.0F - reservePercentage));
        this.contentFormatter = contentFormatter;
        this.metadataMode = metadataMode;
    }

    public List<List<Document>> batch(List<Document> documents) {
        List<List<Document>> batches = new ArrayList();
        int currentSize = 0;
        List<Document> currentBatch = new ArrayList();
        Map<Document, Integer> documentTokens = new LinkedHashMap();

        for(Document document : documents) {
            int tokenCount = this.tokenCountEstimator.estimate(document.getFormattedContent(this.contentFormatter, this.metadataMode));
            if (tokenCount > this.maxInputTokenCount) {
                throw new IllegalArgumentException("Tokens in a single document exceeds the maximum number of allowed input tokens");
            }

            documentTokens.put(document, tokenCount);
        }

        for(Document document : documentTokens.keySet()) {
            Integer tokenCount = (Integer)documentTokens.get(document);
            if (currentSize + tokenCount > this.maxInputTokenCount) {
                batches.add(currentBatch);
                currentBatch = new ArrayList();
                currentSize = 0;
            }

            currentBatch.add(document);
            currentSize += tokenCount;
        }

        if (!currentBatch.isEmpty()) {
            batches.add(currentBatch);
        }

        return batches;
    }
}

package org.springframework.ai.vectorstore;

import java.util.Objects;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

public class SearchRequest {
    public static final double SIMILARITYTHRESHOLDACCEPTALL = (double)0.0F;
    public static final int DEFAULTTOPK = 4;
    private String query = "";
    private int topK = 4;
    private double similarityThreshold = (double)0.0F;
    @Nullable
    private Filter.Expression filterExpression;

    public static Builder from(SearchRequest originalSearchRequest) {
        return builder().query(originalSearchRequest.getQuery()).topK(originalSearchRequest.getTopK()).similarityThreshold(originalSearchRequest.getSimilarityThreshold()).filterExpression(originalSearchRequest.getFilterExpression());
    }

    public SearchRequest() {
    }

    protected SearchRequest(SearchRequest original) {
        this.query = original.query;
        this.topK = original.topK;
        this.similarityThreshold = original.similarityThreshold;
        this.filterExpression = original.filterExpression;
    }

    public String getQuery() {
        return this.query;
    }

    public int getTopK() {
        return this.topK;
    }

    public double getSimilarityThreshold() {
        return this.similarityThreshold;
    }

    @Nullable
    public Filter.Expression getFilterExpression() {
        return this.filterExpression;
    }

    public boolean hasFilterExpression() {
        return this.filterExpression != null;
    }

    public String toString() {
        String var10000 = this.query;
        return "SearchRequest{query='" + var10000 + "', topK=" + this.topK + ", similarityThreshold=" + this.similarityThreshold + ", filterExpression=" + String.valueOf(this.filterExpression) + "}";
    }

    public boolean equals(Object o) {
        if (this == o) {
            return true;
        } else if (o != null && this.getClass() == o.getClass()) {
            SearchRequest that = (SearchRequest)o;
            return this.topK == that.topK && Double.compare(that.similarityThreshold, this.similarityThreshold) == 0 && Objects.equals(this.query, that.query) && Objects.equals(this.filterExpression, that.filterExpression);
        } else {
            return false;
        }
    }

    public int hashCode() {
        return Objects.hash(new Object[]{this.query, this.topK, this.similarityThreshold, this.filterExpression});
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder {
        private final SearchRequest searchRequest = new SearchRequest();

        public Builder query(String query) {
            Assert.notNull(query, "Query can not be null.");
            this.searchRequest.query = query;
            return this;
        }

        public Builder topK(int topK) {
            Assert.isTrue(topK >= 0, "TopK should be positive.");
            this.searchRequest.topK = topK;
            return this;
        }

        public Builder similarityThreshold(double threshold) {
            Assert.isTrue(threshold >= (double)0.0F && threshold <= (double)1.0F, "Similarity threshold must be in [0,1] range.");
            this.searchRequest.similarityThreshold = threshold;
            return this;
        }

        public Builder similarityThresholdAll() {
            this.searchRequest.similarityThreshold = (double)0.0F;
            return this;
        }

        public Builder filterExpression(@Nullable Filter.Expression expression) {
            this.searchRequest.filterExpression = expression;
            return this;
        }

        public Builder filterExpression(@Nullable String textExpression) {
            this.searchRequest.filterExpression = textExpression != null ? (new FilterExpressionTextParser()).parse(textExpression) : null;
            return this;
        }

        public SearchRequest build() {
            return this.searchRequest;
        }
    }
}

package org.springframework.ai.vectorstore;

import io.micrometer.observation.ObservationRegistry;
import java.util.List;
import java.util.Optional;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentWriter;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

public interface VectorStore extends DocumentWriter {
    default String getName() {
        return this.getClass().getSimpleName();
    }

    void add(List<Document> documents);

    default void accept(List<Document> documents) {
        this.add(documents);
    }

    void delete(List<String> idList);

    void delete(Filter.Expression filterExpression);

    default void delete(String filterExpression) {
        SearchRequest searchRequest = SearchRequest.builder().filterExpression(filterExpression).build();
        Filter.Expression textExpression = searchRequest.getFilterExpression();
        Assert.notNull(textExpression, "Filter expression must not be null");
        this.delete(textExpression);
    }

    @Nullable
    List<Document> similaritySearch(SearchRequest request);

    @Nullable
    default List<Document> similaritySearch(String query) {
        return this.similaritySearch(SearchRequest.builder().query(query).build());
    }

    default <T> Optional<T> getNativeClient() {
        return Optional.empty();
    }

    public interface Builder<T extends Builder<T>> {
        T observationRegistry(ObservationRegistry observationRegistry);

        T customObservationConvention(VectorStoreObservationConvention convention);

        T batchingStrategy(BatchingStrategy batchingStrategy);

        VectorStore build();
    }
}

package org.springframework.ai.vectorstore.observation;

import io.micrometer.observation.ObservationRegistry;
import java.util.List;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Operation;
import org.springframework.lang.Nullable;

public abstract class AbstractObservationVectorStore implements VectorStore {
    private static final VectorStoreObservationConvention DEFAULTOBSERVATIONCONVENTION = new DefaultVectorStoreObservationConvention();
    private final ObservationRegistry observationRegistry;
    @Nullable
    private final VectorStoreObservationConvention customObservationConvention;
    protected final EmbeddingModel embeddingModel;
    protected final BatchingStrategy batchingStrategy;

    private AbstractObservationVectorStore(EmbeddingModel embeddingModel, ObservationRegistry observationRegistry, @Nullable VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
        this.embeddingModel = embeddingModel;
        this.observationRegistry = observationRegistry;
        this.customObservationConvention = customObservationConvention;
        this.batchingStrategy = batchingStrategy;
    }

    public AbstractObservationVectorStore(AbstractVectorStoreBuilder<?> builder) {
        this(builder.getEmbeddingModel(), builder.getObservationRegistry(), builder.getCustomObservationConvention(), builder.getBatchingStrategy());
    }

    public void add(List<Document> documents) {
        VectorStoreObservationContext observationContext = this.createObservationContextBuilder(Operation.ADD.value()).build();
        VectorStoreObservationDocumentation.AIVECTORSTORE.observation(this.customObservationConvention, DEFAULTOBSERVATIONCONVENTION, () -> observationContext, this.observationRegistry).observe(() -> this.doAdd(documents));
    }

    public void delete(List<String> deleteDocIds) {
        VectorStoreObservationContext observationContext = this.createObservationContextBuilder(Operation.DELETE.value()).build();
        VectorStoreObservationDocumentation.AIVECTORSTORE.observation(this.customObservationConvention, DEFAULTOBSERVATIONCONVENTION, () -> observationContext, this.observationRegistry).observe(() -> this.doDelete(deleteDocIds));
    }

    public void delete(Filter.Expression filterExpression) {
        VectorStoreObservationContext observationContext = this.createObservationContextBuilder(Operation.DELETE.value()).build();
        VectorStoreObservationDocumentation.AIVECTORSTORE.observation(this.customObservationConvention, DEFAULTOBSERVATIONCONVENTION, () -> observationContext, this.observationRegistry).observe(() -> this.doDelete(filterExpression));
    }

    @Nullable
    public List<Document> similaritySearch(SearchRequest request) {
        VectorStoreObservationContext searchObservationContext = this.createObservationContextBuilder(Operation.QUERY.value()).queryRequest(request).build();
        return (List)VectorStoreObservationDocumentation.AIVECTORSTORE.observation(this.customObservationConvention, DEFAULTOBSERVATIONCONVENTION, () -> searchObservationContext, this.observationRegistry).observe(() -> {
            List<Document> documents = this.doSimilaritySearch(request);
            searchObservationContext.setQueryResponse(documents);
            return documents;
        });
    }

    public abstract void doAdd(List<Document> documents);

    public abstract void doDelete(List<String> idList);

    protected void doDelete(Filter.Expression filterExpression) {
        throw new UnsupportedOperationException();
    }

    public abstract List<Document> doSimilaritySearch(SearchRequest request);

    public abstract VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName);
}

package org.springframework.ai.vectorstore;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.fasterxml.jackson.databind.json.JsonMapper;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileAlreadyExistsException;
import java.nio.file.Files;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Predicate;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.util.JacksonUtils;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.filter.converter.SimpleVectorStoreFilterExpressionConverter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.core.io.Resource;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;

public class SimpleVectorStore extends AbstractObservationVectorStore {
    private static final Logger logger = LoggerFactory.getLogger(SimpleVectorStore.class);
    private final ObjectMapper objectMapper = ((JsonMapper.Builder)JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules())).build();
    private final ExpressionParser expressionParser = new SpelExpressionParser();
    private final FilterExpressionConverter filterExpressionConverter = new SimpleVectorStoreFilterExpressionConverter();
    protected Map<String, SimpleVectorStoreContent> store = new ConcurrentHashMap();

    protected SimpleVectorStore(SimpleVectorStoreBuilder builder) {
        super(builder);
    }

    public static SimpleVectorStoreBuilder builder(EmbeddingModel embeddingModel) {
        return new SimpleVectorStoreBuilder(embeddingModel);
    }

    public void doAdd(List<Document> documents) {
        Objects.requireNonNull(documents, "Documents list cannot be null");
        if (documents.isEmpty()) {
            throw new IllegalArgumentException("Documents list cannot be empty");
        } else {
            for(Document document : documents) {
                logger.info("Calling EmbeddingModel for document id = {}", document.getId());
                float[] embedding = this.embeddingModel.embed(document);
                SimpleVectorStoreContent storeContent = new SimpleVectorStoreContent(document.getId(), document.getText(), document.getMetadata(), embedding);
                this.store.put(document.getId(), storeContent);
            }

        }
    }

    public void doDelete(List<String> idList) {
        for(String id : idList) {
            this.store.remove(id);
        }

    }

    public List<Document> doSimilaritySearch(SearchRequest request) {
        Predicate<SimpleVectorStoreContent> documentFilterPredicate = this.doFilterPredicate(request);
        float[] userQueryEmbedding = this.getUserQueryEmbedding(request.getQuery());
        return this.store.values().stream().filter(documentFilterPredicate).map((content) -> content.toDocument(SimpleVectorStore.EmbeddingMath.cosineSimilarity(userQueryEmbedding, content.getEmbedding()))).filter((document) -> document.getScore() >= request.getSimilarityThreshold()).sorted(Comparator.comparing(Document::getScore).reversed()).limit((long)request.getTopK()).toList();
    }

    private Predicate<SimpleVectorStoreContent> doFilterPredicate(SearchRequest request) {
        return request.hasFilterExpression() ? (document) -> {
            StandardEvaluationContext context = new StandardEvaluationContext();
            context.setVariable("metadata", document.getMetadata());
            return (Boolean)this.expressionParser.parseExpression(this.filterExpressionConverter.convertExpression(request.getFilterExpression())).getValue(context, Boolean.class);
        } : (document) -> true;
    }

    public void save(File file) {
        String json = this.getVectorDbAsJson();

        try {
            if (!file.exists()) {
                logger.info("Creating new vector store file: {}", file);

                try {
                    Files.createFile(file.toPath());
                } catch (FileAlreadyExistsException e) {
                    throw new RuntimeException("File already exists: " + String.valueOf(file), e);
                } catch (IOException e) {
                    throw new RuntimeException("Failed to create new file: " + String.valueOf(file) + ". Reason: " + e.getMessage(), e);
                }
            } else {
                logger.info("Overwriting existing vector store file: {}", file);
            }

            try (
                OutputStream stream = new FileOutputStream(file);
                Writer writer = new OutputStreamWriter(stream, StandardCharsets.UTF8);
            ) {
                writer.write(json);
                writer.flush();
            }

        } catch (IOException ex) {
            logger.error("IOException occurred while saving vector store file.", ex);
            throw new RuntimeException(ex);
        } catch (SecurityException ex) {
            logger.error("SecurityException occurred while saving vector store file.", ex);
            throw new RuntimeException(ex);
        } catch (NullPointerException ex) {
            logger.error("NullPointerException occurred while saving vector store file.", ex);
            throw new RuntimeException(ex);
        }
    }

    public void load(File file) {
        TypeReference<HashMap<String, SimpleVectorStoreContent>> typeRef = new TypeReference<HashMap<String, SimpleVectorStoreContent>>() {
        };

        try {
            this.store = (Map)this.objectMapper.readValue(file, typeRef);
        } catch (IOException ex) {
            throw new RuntimeException(ex);
        }
    }

    public void load(Resource resource) {
        TypeReference<HashMap<String, SimpleVectorStoreContent>> typeRef = new TypeReference<HashMap<String, SimpleVectorStoreContent>>() {
        };

        try {
            this.store = (Map)this.objectMapper.readValue(resource.getInputStream(), typeRef);
        } catch (IOException ex) {
            throw new RuntimeException(ex);
        }
    }

    private String getVectorDbAsJson() {
        ObjectWriter objectWriter = this.objectMapper.writerWithDefaultPrettyPrinter();

        try {
            return objectWriter.writeValueAsString(this.store);
        } catch (JsonProcessingException e) {
            throw new RuntimeException("Error serializing documentMap to JSON.", e);
        }
    }

    private float[] getUserQueryEmbedding(String query) {
        return this.embeddingModel.embed(query);
    }

    public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
        return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE.value(), operationName).dimensions(this.embeddingModel.dimensions()).collectionName("in-memory-map").similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
    }

    public static final class EmbeddingMath {
        private EmbeddingMath() {
            throw new UnsupportedOperationException("This is a utility class and cannot be instantiated");
        }

        public static double cosineSimilarity(float[] vectorX, float[] vectorY) {
            if (vectorX != null && vectorY != null) {
                if (vectorX.length != vectorY.length) {
                    throw new IllegalArgumentException("Vectors lengths must be equal");
                } else {
                    float dotProduct = dotProduct(vectorX, vectorY);
                    float normX = norm(vectorX);
                    float normY = norm(vectorY);
                    if (normX != 0.0F && normY != 0.0F) {
                        return (double)dotProduct / (Math.sqrt((double)normX) * Math.sqrt((double)normY));
                    } else {
                        throw new IllegalArgumentException("Vectors cannot have zero norm");
                    }
                }
            } else {
                throw new RuntimeException("Vectors must not be null");
            }
        }

        public static float dotProduct(float[] vectorX, float[] vectorY) {
            if (vectorX.length != vectorY.length) {
                throw new IllegalArgumentException("Vectors lengths must be equal");
            } else {
                float result = 0.0F;

                for(int i = 0; i < vectorX.length; ++i) {
                    result += vectorX[i] * vectorY[i];
                }

                return result;
            }
        }

        public static float norm(float[] vector) {
            return dotProduct(vector, vector);
        }
    }

    public static final class SimpleVectorStoreBuilder extends AbstractVectorStoreBuilder<SimpleVectorStoreBuilder> {
        private SimpleVectorStoreBuilder(EmbeddingModel embeddingModel) {
            super(embeddingModel);
        }

        public SimpleVectorStore build() {
            return new SimpleVectorStore(this);
        }
    }
}