Implementing Semantic Search Using Spring AI and PGVector

1. Overview

Searching is a fundamental concept in software, aimed at finding relevant information in a large dataset. It involves finding a specific item in a collection of items.

In this tutorial, we’ll explore how to implement semantic search using Spring AI, PGVector, and Ollama.

2. Background

Semantic search is an advanced search technique that uses the meaning of words to find the most relevant results. To build a semantic search application, we need to understand some key concepts:

• Word Embeddings: Word embeddings are a type of word representation that allows words with similar meanings to have similar representations. Word embeddings convert words into numerical vectors that can be used in machine-learning models.
• Semantic Similarity: Semantic similarity is a measure of how similar two pieces of text are in terms of meaning. It’s used to compare the meaning of words, sentences, or documents.
• Vector Space Model: The vector space model is a mathematical model used to represent text documents as vectors in a high-dimensional space. In this model, each word is represented as a vector, and the similarity between two words is calculated based on the distance between their vectors.
• Cosine Similarity: Cosine similarity is a similarity measure between two non-zero vectors of an inner product space that measures the cosine of the angle between them. It calculates the similarity between two vectors in the vector space model.

Now let’s get to building an application that demonstrates this.

3. Prerequisites

First, we should have Docker installed on our machine to run PGVector and Ollama.

Then, we need the Spring AI Ollama and PGVector dependencies in our Spring application:

<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>

We’ll also add Spring Boot’s Docker Compose support to manage the Ollama and PGVector Docker containers:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-docker-compose</artifactId>
    <version>3.1.1</version>
</dependency>

In addition to the dependency, we’ll put these together by describing the two services in a docker-compose.yml file:

services:
  postgres:
    image: pgvector/pgvector:pg17
    environment:
      POSTGRES_DB: vectordb
      POSTGRES_USER: postgres
      POSTGRES_PASSWORD: postgres
    ports:
      - "5434:5432"
    healthcheck:
      test: [ "CMD-SHELL", "pg_isready -U postgres" ]
      interval: 10s
      timeout: 5s
      retries: 5

  ollama:
    image: ollama/ollama:latest
    ports:
      - "11435:11434"
    volumes:
      - ollama_data:/root/.ollama
    healthcheck:
      test: [ "CMD", "curl", "-f", "http://localhost:11435/api/health" ]
      interval: 10s
      timeout: 5s
      retries: 10

volumes:
  ollama_data:

4. Configuring the Application

Next, we need to configure our Spring Boot application to use the Ollama and PGVector services. In the application.yml file, we define several properties. Let’s note specifically what’s chosen for the ollama and vectorstore properties:

spring:
  ai:
    ollama:
      init:
        pull-model-strategy: when_missing
        chat:
          include: true
      embedding:
        options:
          model: nomic-embed-text
    vectorstore:
      pgvector:
        initialize-schema: true
        dimensions: 768
        index-type: hnsw
  docker:
    compose:
      file: docker-compose.yml
      enabled: true
  datasource:
    url: jdbc:postgresql://localhost:5434/vectordb
    username: postgres
    password: postgres
    driver-class-name: org.postgresql.Driver
  jpa:
    database-platform: org.hibernate.dialect.PostgreSQLDialect

We chose nomic-embed-text for our Ollama model. Spring AI pulls it for us if we don’t have that one downloaded.

The PGVector settings ensure proper vector storage setup by initializing the database schema (initialize-schema: true), aligning the vector dimensions with common embedding sizes (dimensions: 768), and optimizing search efficiency using the Hierarchical Navigable Small World (HNSW) index (index-type: hnsw) for fast approximate nearest neighbor searches.

5. Performing Semantic Search

Now that our infrastructure is ready, we can implement a simple semantic search application. Our use case will be a smart book search engine that allows users to search for books based on their content.

Initially, we’ll build a simple search functionality using PGVector and later, we’ll enhance it with Ollama to provide more context-aware responses.

Let’s define a Book class that represents a book entity:

public record Book(String title, String author, String description) {
}

Before we can search for books, we need to ingest book data into the PGVector store. The following method adds some sample book data:

void run() {
    var books = List.of(
            new Book("The Great Gatsby", "F. Scott Fitzgerald", "The Great Gatsby is a 1925 novel by American writer F. Scott Fitzgerald. Set in the Jazz Age on Long Island, near New York City, the novel depicts first-person narrator Nick Carraway's interactions with mysterious millionaire Jay Gatsby and Gatsby's obsession to reunite with his former lover, Daisy Buchanan."),
            new Book("To Kill a Mockingbird", "Harper Lee", "To Kill a Mockingbird is a novel by the American author Harper Lee. It was published in 1960 and was instantly successful. In the United States, it is widely read in high schools and middle schools."),
            new Book("1984", "George Orwell", "Nineteen Eighty-Four: A Novel, often referred to as 1984, is a dystopian social science fiction novel by the English novelist George Orwell. It was published on 8 June 1949 by Secker & Warburg as Orwell's ninth and final book completed in his lifetime."),
            new Book("The Catcher in the Rye", "J. D. Salinger", "The Catcher in the Rye is a novel by J. D. Salinger, partially published in serial form in 1945–1946 and as a novel in 1951. It was originally intended for adults but is often read by adolescents for its themes of angst, alienation, and as a critique on superficiality in society."),
            new Book("Lord of the Flies", "William Golding", "Lord of the Flies is a 1954 novel by Nobel Prize-winning British author William Golding. The book focuses on a group of British")
    );

    List<Document> documents = books.stream()
            .map(book -> new Document(book.toString()))
            .toList();

    vectorStore.add(documents);
}

Now that we have added the sample book data to the PGVector store, we can implement the semantic search functionality.

5.1. Semantic Search

Our goal is to implement a semantic search API that allows users to find books based on their content.
Let’s define a controller that interacts with PGVector to perform similarity searches:

@RequestMapping("/books")
class BookSearchController {
    final VectorStore vectorStore;
    final ChatClient chatClient;

    BookSearchController(VectorStore vectorStore, ChatClient.Builder chatClientBuilder) {
        this.vectorStore = vectorStore;
        this.chatClient = chatClientBuilder.build();
    }
...

Next, we’ll create a POST /search endpoint that accepts search criteria from the user and returns a list of matching books:

@PostMapping("/search")
List<String> semanticSearch(@RequestBody String query) {
    return vectorStore.similaritySearch(SearchRequest.builder()
        .query(query)
        .topK(3)
        .build())
       .stream()
      .map(Document::getText)
      .toList();
}

Notice that we used VectorStore#similaritySearch. This performs a semantic search across the books we ingested earlier.

After starting the application, we’re ready to perform a search. Let’s use cURL to search for instances of 1984:

curl -X POST --data "1984" http://localhost:8080/books/search

The response contains three books: one with an exact match and two with partial matches:

[
  "Book[title=1984, author=George Orwell, description=Nineteen Eighty-Four: A Novel, often referred to as 1984, is a dystopian social science fiction novel by the English novelist George Orwell.]",
  "Book[title=The Catcher in the Rye, author=J. D. Salinger, description=The Catcher in the Rye is a novel by J. D. Salinger, partially published in serial form in 1945–1946 and as a novel in 1951.]",
  "Book[title=To Kill a Mockingbird, author=Harper Lee, description=To Kill a Mockingbird is a novel by the American author Harper Lee.]"
]

5.2. Enhancing Semantic Search with Ollama

We can integrate Ollama to generate paraphrased responses that provide additional context to improve the semantic search results, following these three steps:

1. Retrieve the top three matching book descriptions from the search query.
2. Feed these descriptions into Ollama to generate a more natural, context-aware response.
3. Deliver a response that includes summarized and paraphrased information, providing clearer and more relevant insights.

Let’s create a new method in the BookSearchController that uses Ollama to generate paraphrases of the query:

@PostMapping("/enhanced-search")
String enhancedSearch(@RequestBody String query) {
    String context = vectorStore.similaritySearch(SearchRequest.builder()
        .query(query)
        .topK(3)
        .build())
      .stream()
      .map(Document::getText)
      .reduce("", (a, b) -> a + b + "\n");

    return chatClient.prompt()
      .system(context)
      .user(query)
      .call()
      .content();
}

And now let’s test the enhanced semantic search functionality by sending a POST request to the /books/enhanced-search endpoint:

curl -X POST --data "1984" http://localhost:8080/books/enhanced-search
1984 is a classic dystopian novel written by George Orwell. Here's an excerpt from the book:

"He loved Big Brother. He even admired him. After all, who wouldn't? Big Brother was all-powerful, all-knowing, and infinitely charming. And now that he had given up all his money in bank accounts with his names on them, and his credit cards, and his deposit slips, he felt free."

This excerpt sets the tone for the novel, which depicts a totalitarian society where the government exercises total control over its citizens. The protagonist, Winston Smith, is a low-ranking member of the ruling Party who begins to question the morality of their regime.

Would you like to know more about the book or its themes?

Instead of returning three separate book descriptions as the simple semantic search, Ollama synthesizes the most relevant information from the search results. In this case, 1984 is the most relevant match, so Ollama focuses on providing a detailed summary rather than listing unrelated books. This mimics human-like search assistance, making the results more engaging and insightful.

6. Conclusion

In this article, we explored how to implement semantic search using Spring AI, PGVector, and Ollama. We compared two endpoints; one that performed a semantic search of our book catalog and another that fed and enhanced that search result with an Ollama LLM.