Spring AI Embeddings & Vector Stores

float[] a = embeddingModel.embed("How do I reset my password?");
float[] b = embeddingModel.embed("I forgot my login credentials");
float[] c = embeddingModel.embed("What is the refund policy?");
 
// cosineSimilarity(a, b) ≈ 0.85  — different words, same intent
// cosineSimilarity(a, c) ≈ 0.20  — unrelated

services:
  postgres:
    image: pgvector/pgvector:pg17
    environment:
      POSTGRES_DB: javaagentic
      POSTGRES_USER: app
      POSTGRES_PASSWORD: secret
    ports:
      - '5432:5432'
    volumes:
      - pgdata:/var/lib/postgresql/data
    healthcheck:
      test: ['CMD-SHELL', 'pg_isready -U app -d javaagentic']
      interval: 5s
      retries: 10
 
volumes:
  pgdata:

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-starter-vector-store-pgvector</artifactId>
</dependency>
<dependency>
  <groupId>org.springframework.boot</groupId>
  <artifactId>spring-boot-starter-jdbc</artifactId>
</dependency>
<dependency>
  <groupId>org.postgresql</groupId>
  <artifactId>postgresql</artifactId>
  <scope>runtime</scope>
</dependency>

spring:
  datasource:
    url: jdbc:postgresql://localhost:5432/javaagentic
    username: app
    password: secret
  ai:
    openai:
      api-key: ${OPENAI_API_KEY}
      embedding:
        options:
          model: text-embedding-3-small   # 1536 dimensions
    vectorstore:
      pgvector:
        # Convenient in development; manage the schema with Flyway or
        # Liquibase in production so migrations are reviewable.
        initialize-schema: true
        # Must equal the embedding model's output dimension.
        dimensions: 1536
        # HNSW: fast, accurate, higher build cost. The right default.
        index-type: HNSW
        # COSINE_DISTANCE pairs with normalised embeddings, which is what
        # the common hosted models produce.
        distance-type: COSINE_DISTANCE
        max-document-batch-size: 1000

CREATE EXTENSION IF NOT EXISTS vector;
 
CREATE TABLE IF NOT EXISTS vector_store (
    id          uuid PRIMARY KEY DEFAULT uuid_generate_v4(),
    content     text,
    metadata    jsonb,
    embedding   vector(1536)
);
 
CREATE INDEX ON vector_store USING hnsw (embedding vector_cosine_ops);

package com.javaagentic.demo.rag;
 
import java.util.List;
import java.util.Map;
 
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
 
@Service
public class DocumentIngestionService {
 
    private final VectorStore vectorStore;
    private final TokenTextSplitter splitter;
 
    public DocumentIngestionService(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
        // Chunk sizing is the single highest-impact tuning knob in retrieval.
        // ~400 tokens with ~80 overlap suits prose documentation; code and
        // tables usually want larger chunks that keep structures intact.
        this.splitter = TokenTextSplitter.builder()
                .withChunkSize(400)
                .withMinChunkSizeChars(200)
                .build();
    }
 
    @Transactional
    public int ingest(String sourceId, String tenantId, String text) {
        Document document = new Document(text, Map.of(
                // Metadata written at ingest time is the only thing you can
                // filter on later. Err on the side of recording more.
                "sourceId", sourceId,
                "tenantId", tenantId,
                "ingestedAt", java.time.Instant.now().toString()));
 
        List<Document> chunks = splitter.apply(List.of(document));
        vectorStore.add(chunks);   // embeds and persists in one call
        return chunks.size();
    }
}

package com.javaagentic.demo.rag;
 
import java.util.List;
 
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
 
@Service
public class RetrievalService {
 
    private final VectorStore vectorStore;
 
    public RetrievalService(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }
 
    public List<Document> retrieve(String query, String tenantId) {
        return vectorStore.similaritySearch(SearchRequest.builder()
                .query(query)
                .topK(5)
                // Without a threshold you always get five results, however
                // irrelevant — and irrelevant context makes answers worse,
                // not merely unhelpful.
                .similarityThreshold(0.7)
                // Tenant isolation belongs in the filter, not in a prompt
                // instruction. Similarity is not an access control mechanism.
                .filterExpression("tenantId == '%s'".formatted(tenantId))
                .build());
    }
}