Getting Started With Spring AI and PostgreSQL PGVector

The Spring AI is a new project of the Spring ecosystem that streamlines the creation of AI applications in Java. By using Spring AI together with PostgreSQL pgvector, you can build generative AI applications that draw insights from your data.

First, this article introduces you to the Spring AI ChatClient that uses the OpenAI GPT-4 model to generate recommendations based on user prompts. Next, the article shows how to deploy PostgreSQL with the PGVector extension and perform vector similarity searches using the Spring AI EmbeddingClient and Spring JdbcClient.

Adding Spring AI Dependency

Spring AI supports many large language model (LLM) providers, with each LLM having its own Spring AI dependency.

Let's assume that you prefer working with OpenAI models and APIs. Then, you need to add the following dependency to a project:

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
    <version>{latest.version}</version>
</dependency>

Also, at the time of writing, Spring AI was in active development, with the framework artifacts being released in the Spring Milestone and/or Snapshot repositories. Thus, if you still can't find Spring AI on https://start.spring.io/, then add the repositories to the pom.xml file:

<repositories>
    <repository>
      <id>spring-milestones</id>
      <name>Spring Milestones</name>
      <url>https://repo.spring.io/milestone</url>
      <snapshots>
        <enabled>false</enabled>
      </snapshots>
    </repository>
    <repository>
      <id>spring-snapshots</id>
      <name>Spring Snapshots</name>
      <url>https://repo.spring.io/snapshot</url>
      <releases>
        <enabled>false</enabled>
      </releases>
    </repository>
</repositories>

Setting Up OpenAI Module

The OpenAI module comes with several configuration properties, allowing the management of connectivity-related settings and fine-tuning the behavior of OpenAI models.

At a minimum, you need to provide your OpenAI API key, which will be used by Spring AI to access GPT and embedding models. Once the key is created, add it to the application.properties file:

spring.ai.openai.api-key=sk-...

Then, if necessary, you can select particular GPT and embedding models:

spring.ai.openai.chat.model=gpt-4
spring.ai.openai.embedding.model=text-embedding-ada-002

In the end, you can test that the OpenAI module is configured properly by implementing a simple assistant with Spring AI's ChatClient:

// Inject the ChatClient bean
@Autowired
private ChatClient aiClient;

// Create a system message for ChatGPT explaining the task
private static final SystemMessage SYSTEM_MESSAGE = new SystemMessage(
    """
    You're an assistant who helps to find lodging in San Francisco.
    Suggest three options. Send back a JSON object in the format below.
    [{\"name\": \"<hotel name>\", \"description\": \"<hotel description>\", \"price\": <hotel price>}]
    Don't add any other text to the response. Don't add the new line or any other symbols to the response. Send back the raw JSON.
    """);

public void searchPlaces(String prompt) {
    // Create a Spring AI prompt with the system message and the user message
    Prompt chatPrompt = new Prompt(List.of(SYSTEM_MESSAGE, new UserMessage(prompt)));

    // Send the prompt to ChatGPT and get the response
    ChatResponse response = aiClient.generate(chatPrompt);

    // Get the raw JSON from the response and print it
    String rawJson = response.getGenerations().get(0).getContent();

    System.out.println(rawJson);
}

For the sake of the experiment, if you pass the "I'd like to stay near the Golden Gate Bridge" prompt, then the searchPlaces the method might provide lodging recommendations as follows:

[
 {"name": "Cavallo Point", "description": "Historic hotel offering refined rooms, some with views of the Golden Gate Bridge, plus a spa & dining.", "price": 450}, 
 {"name": "Argonaut Hotel", "description": "Upscale, nautical-themed hotel offering Golden Gate Bridge views, plus a seafood restaurant.", "price": 300}, 
 {"name": "Hotel Del Sol", "description": "Colorful, retro hotel with a pool, offering complimentary breakfast & an afternoon cookies reception.", "price": 200}
]

Starting Postgres With PGVector

If you run the previous code snippet with the ChatClient, you'll notice that it usually takes over 10 seconds for the OpenAI GPT model to generate a response. The model has a broad and deep knowledge base, and it takes time to produce a relevant response.

Apart from the high latency, the GPT model might not have been trained on data that is relevant to your application workload. Thus, it might generate responses that are far from being satisfactory for the user.

However, you can always expedite the search and provide users with accurate responses if you generate embeddings on a subset of your data and then let Postgres work with those embeddings.

The pgvector extension allows storing and querying vector embeddings in Postgres. The easiest way to start with PGVector is by starting a Postgres instance with the extension in Docker:

mkdir ~/postgres-volume/

docker run --name postgres \
    -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=password \
    -p 5432:5432 \
    -v ~/postgres-volume/:/var/lib/postgresql/data -d ankane/pgvector:latest

Once started, you can connect to the container and enable the extension by executing the CREATE EXTENSION vector statement: 

docker exec -it postgres psql -U postgres -c 'CREATE EXTENSION vector'

 Lastly, add the Postgres JDBC driver dependency to the pom.xml file:

<dependency>
  <groupId>org.postgresql</groupId>
  <artifactId>postgresql</artifactId>
  <version>{latest.version}</version>
</dependency>

Configure the Spring DataSource by adding the following settings to the application.properties file:

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

Performing Vector Similarity Search With Spring AI

At a minimum, the vector similarity search is a two-step process.

First, you need to use an embedding model to generate a vector/embedding for a provided user prompt or other text. Spring AI supports the EmbeddingClient that connects to OpenAI's or other providers' embedding models and generates a vectorized representation for the text input:

// Inject the Spring AI Embedding client
@Autowired
private EmbeddingClient aiClient;

public List<Place> searchPlaces(String prompt) {
    // Use the Embedding client to generate a vector for the user prompt
    List<Double> promptEmbedding = aiClient.embed(prompt);
    ...
}

Second, you use the generated embedding to perform a similarity search across vectors stored in the Postgres database. For instance, you can use the Spring JdbcClient for this task:

@Autowired
private JdbcClient jdbcClient;

// Inject the Spring AI Embedding client
@Autowired
private EmbeddingClient aiClient;


public List<Place> searchPlaces(String prompt) {
    // Use the Embedding client to generate a vector for the user prompt
    List<Double> promptEmbedding = aiClient.embed(prompt);

    // Perform the vector similarity search
    StatementSpec query = jdbcClient.sql(
        "SELECT name, description, price " +
        "FROM airbnb_listing WHERE 1 - (description_embedding <=> :user_promt::vector) > 0.7 " +
        "ORDER BY description_embedding <=> :user_promt::vector LIMIT 3")
        .param("user_promt", promptEmbedding.toString());

    // Return the recommended places
    return query.query(Place.class).list();
}

• The description_embedding column stores embeddings that were pre-generated for Airbnb listing overviews from the description column. The Airbnb embeddings were produced by the same model that is used by Spring AI's EmbeddingClient for the user prompts.
• Postgres uses PGVector to calculate the cosine distance (<=>) between the Airbnb and user prompt embeddings (description_embedding <=> :user_prompt::vector) and then returns only those Airbnb listings whose description is > 0.7 similar to the provided user prompt. The similarity is measured as a value in the range from 0 to 1. The closer the similarity to 1, the more related the vectors are.

What's Next

Spring AI and PostgreSQL PGVector provide all the essential capabilities needed for building generative AI applications in Java. If you're curious to learn more, watch this hands-on tutorial. It guides you through the process of creating a lodging recommendation service in Java from scratch, optimizing similarity searches with specialized indexes, and scaling with distributed Postgres (YugabyteDB):