Spring AI and PgVectorStore Configuration Examples

Learn to configure Postgres PgVectorStore to store the vectors generated with OpenAI and Ollama embedding models in a Spring AI project. These embeddings are used to make a RAG-style application such as ChatBots.

1. Installing Postgres PgVector Database

The pgvector is an open-source extension for PostgreSQL designed to efficiently store vector data within the database and perform similarity searches. It is designed to work seamlessly with other PostgreSQL features, including indexing and querying.

1.1. Using Docker

The easiest way to install the pgvector on a machine is by using Docker.

docker run -it --rm --name postgres -p 5432:5432 \
  -e POSTGRES_USER=postgres \
  -e POSTGRES_PASSWORD=postgres \
  -e VECTOR_STORE_TYPE=pgVector \
  pgvector/pgvector:pg16

When using Docker Compose, add the following service with additional custom configurations:

services:
  postgres:
    image: pgvector/pgvector:pg16
    container_name: postgres
    ports:
      - "5432:5432"
    environment:
      POSTGRES_USER: postgres
      POSTGRES_PASSWORD: postgres
      VECTOR_STORE_TYPE: pgVector

If you have an existing Postgres installation, you must enable the vector, hstore and uuid-ossp extensions.

1.2. Configuring PgVector with Spring AI

Spring AI helps here by enabling the required extensions automatically during startup. Spring AI runs the following schema file when the property spring.ai.vectorstore.pgvector.initialize-schema is set to true.

CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS hstore;
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";

CREATE TABLE IF NOT EXISTS vector_store (
  id uuid DEFAULT uuid_generate_v4() PRIMARY KEY,
  content text,
  metadata json,
  embedding vector(1536) // 1536 is the default embedding dimension
);

CREATE INDEX ON vector_store USING HNSW (embedding vector_cosine_ops);

To customize the index-type, distance-tye and dimensions, we can change the respective properties in the application.properties file.

spring.datasource.url=jdbc:postgresql://localhost:5432/postgres
spring.datasource.username=postgres
spring.datasource.password=postgres

spring.ai.vectorstore.pgvector.initialize-schema=true
spring.ai.vectorstore.pgvector.index-type=HNSW
spring.ai.vectorstore.pgvector.distance-type=COSINE_DISTANCE
spring.ai.vectorstore.pgvector.dimensions=1536
spring.ai.vectorstore.pgvector.schema-validation=true
spring.ai.vectorstore.pgvector.remove-existing-vector-store-table=false

2. Configuring PgVectorStore with OpenAI Embedding Models

Let us build a simple demo application that uses OpenAI models: text-embedding-ada-002 for generating embeddings and gpt-4 for chat capability.

2.1. Maven

Start with adding the following dependencies:

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
</dependency>
<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-pgvector-store-spring-boot-starter</artifactId>
</dependency>

2.2. Properties Configuration

Next, we must specify the models for generating vectors and chat conversations. Additionally, we will specify the PgVector connection details with other optional configurations.

spring.ai.openai.api-key=${OPENAI_API_KEY}
spring.ai.openai.chat.model=gpt-4
spring.ai.openai.embedding.model=text-embedding-ada-002

spring.datasource.url=jdbc:postgresql://localhost:5432/postgres
spring.datasource.username=postgres
spring.datasource.password=postgres

spring.ai.vectorstore.pgvector.initialize-schema=true
spring.ai.vectorstore.pgvector.index-type=HNSW
spring.ai.vectorstore.pgvector.distance-type=COSINE_DISTANCE
spring.ai.vectorstore.pgvector.dimensions=1536
spring.ai.vectorstore.pgvector.schema-validation=true
spring.ai.vectorstore.pgvector.remove-existing-vector-store-table=false

2.3. Demo App

On application startup, Spring boot auto-configuration creates the ChatClient, EmbeddingModel, and VectorStore beans with configured properties. The VectorStore bean is of type PgVectorStore class.

We can use the configured VectorStore bean to store the embeddings generated using the OpenAI model and make similarity searches in the Postgres database.

import java.util.List;
import java.util.Map;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

@SpringBootApplication
public class App implements CommandLineRunner {

  public static void main(String[] args) {
    SpringApplication.run(App.class);
  }

  private final VectorStore vectorStore;

  public App(VectorStore vectorStore) {
    this.vectorStore = vectorStore;
  }

  @Override
  public void run(String... args) {
    List<Document> documents = List.of(
      new Document("Java is a high-level, object-oriented programming language known for its platform independence."),
      new Document("It is widely used for developing enterprise applications, Android apps, and big data processing systems."),
      new Document("Java's strong typing, automatic memory management, and extensive libraries contribute to its popularity.", Map.of("reason", "popularity")));

    // Add the documents to PGVector
    vectorStore.add(documents);

    // Retrieve documents similar to a query
    List<Document> results = vectorStore
      .similaritySearch(SearchRequest.query("programming language").withTopK(1));

    results.stream()
      .map(Document::getContent)
      .forEach(System.out::println);
  }
}

The program output:

Java is a high-level, object-oriented programming language known for its platform independence.

3. Configuring PgVectorStore with Ollama Embedding Models

When using PgVectorStore to interact with locally deployed LLM models, we need to use Ollama, as it makes downloading and running the models super easy.

You can find detailed instructions for installing Ollama in Ollama Local Setup and Spring AI Integration Example. In our demo, we used Docker to run Ollama locally.

docker run -d --gpus all -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama

Once Ollama is running, we can download and run a model with a simple run command:

ollama run mistral

3.1. Maven

Start with adding the following dependencies. These starters will import all the necessary dependencies.

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-ollama-spring-boot-starter</artifactId>
</dependency>
<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-pgvector-store-spring-boot-starter</artifactId>
</dependency>

3.2. Properties Configuration

After running a model and setting up the dependencies, we can configure the application properties for models to generate embeddings and connect to the PgVector database.

spring.ai.ollama.embedding.options.model=mistral
spring.ai.ollama.chat.options.model=mistral

spring.datasource.url=jdbc:postgresql://localhost:5432/postgres
spring.datasource.username=postgres
spring.datasource.password=postgres

spring.ai.vectorstore.pgvector.initialize-schema=true
spring.ai.vectorstore.pgvector.index-type=HNSW
spring.ai.vectorstore.pgvector.distance-type=COSINE_DISTANCE
spring.ai.vectorstore.pgvector.dimensions=4096
#spring.ai.vectorstore.pgvector.schema-validation=true
#spring.ai.vectorstore.pgvector.remove-existing-vector-store-table=false

3.3. Demo App

Spring boot will initialize and autoconfigure the EmbeddingModel and VectorStore beans on application startup. We can use these beans to generate embeddings and make similarity searches similar to the previous demo application code.

@SpringBootApplication
public class PgVectorOllamaApp implements CommandLineRunner {

  public static void main(String[] args) {
    SpringApplication.run(PgVectorOllamaApp.class);
  }

  private final VectorStore vectorStore;

  public PgVectorOllamaApp(VectorStore vectorStore) {
    this.vectorStore = vectorStore;
  }

  @Override
  public void run(String... args) {
    List<Document> documents = List.of(
      new Document("Java is a high-level, object-oriented programming language known for its platform independence."),
      new Document("It is widely used for developing enterprise applications, Android apps, and big data processing systems."),
      new Document("Java's strong typing, automatic memory management, and extensive libraries contribute to its popularity.", Map.of("reason", "popularity")));

    // Add the documents to PGVector
    vectorStore.add(documents);

    // Retrieve documents similar to a query
    List<Document> results = vectorStore
      .similaritySearch(SearchRequest.query("programming language").withTopK(1));

    results.stream()
      .map(Document::getContent)
      .forEach(System.out::println);
  }
}

The program output:

Java is a high-level, object-oriented programming language known for its platform independence.

4. ERROR: column cannot have more than 2000 dimensions

It is worth noting the PgVector’s embedding column cannot have more than 2000 dimensions for the index. Most LLM models generate vectors with dimensions greater than 2000. So, if you use the index types HNSW or ivfflat, you will get the runtime error while storing the generated embeddings.

spring.ai.vectorstore.pgvector.index-type=ivfflat
spring.ai.vectorstore.pgvector.dimensions=4096

org.springframework.beans.factory.UnsatisfiedDependencyException: Error creating bean with name 
'pgVectorOllamaApp': Unsatisfied dependency expressed through constructor parameter 0: Error creating bean with
 name 'vectorStore' defined in class path resource [org/springframework/ai/autoconfigure/vectorstore/pgvector/
 PgVectorStoreAutoConfiguration.class]: StatementCallback; uncategorized SQLException for SQL [CREATE INDEX IF NOT 
 EXISTS spring_ai_vector_index ON public.vector_store USING IVFFLAT (embedding vector_cosine_ops) 
 ]; SQL state [XX000]; error code [0]; ERROR: column cannot have more than 2000 dimensions for ivfflat index

There are two solutions to fix this error:

• Use any other vector store, such as ChromDB.
• Use index-type as none. However, this impacts performance. It can be useful in a development environment for testing purposes, but do not use it in production.

spring.ai.vectorstore.pgvector.index-type=none # Do not use in production. Suitable only for demo purposes.
spring.ai.vectorstore.pgvector.dimensions=4096

After making this change, regenerate the schema again, and you will face the “column cannot have more than 2000 dimensions” error.

spring.ai.vectorstore.pgvector.initialize-schema=true
spring.ai.vectorstore.pgvector.remove-existing-vector-store-table=true

5. Summary

This Spring AI PgVector example discussed the following:

• Installing Postgres vector database using Docker.
• Enabling PgVector extension in Postgres database.
• Spring AI autoconfiguration options to connect with pgvector and custom properties.
• How to configure PgVector and OpenAi LLM in RAG-style applications.
• How to configure PgVector and Ollama LLM.
• Fixing dimension mismatch error when using the index types as HNSW or ivfflat.

Happy Learning !!

Source Code on Github