Implementación de búsqueda semántica con Spring AI y PGVector

1. Visión general

La búsqueda es un concepto fundamental en el software, destinado a encontrar información relevante en un conjunto de datos amplio. Implica localizar un elemento específico dentro de una colección de elementos.

En este tutorial, exploraremos cómo implementar la búsqueda semántica utilizando Spring AI, PGVector y Ollama.

2. Antecedentes

La búsqueda semántica es una técnica de búsqueda avanzada que utiliza el significado de las palabras para encontrar los resultados más relevantes. Para construir una aplicación de búsqueda semántica, necesitamos comprender algunos conceptos clave:

Embeddings de palabras: Los embeddings de palabras son un tipo de representación de palabras que permite que las palabras con significados similares tengan representaciones similares. Los embeddings de palabras convierten las palabras en vectores numéricos que pueden utilizarse en modelos de aprendizaje automático.
Similaridad semántica: La similaridad semántica es una medida de cuán similares son dos piezas de texto en términos de significado. Se utiliza para comparar el significado de palabras, oraciones o documentos.
Modelo de espacio vectorial: El modelo de espacio vectorial es un modelo matemático utilizado para representar documentos de texto como vectores en un espacio de alta dimensión. En este modelo, cada palabra se representa como un vector, y la similaridad entre dos palabras se calcula en función de la distancia entre sus vectores.
Similitud de coseno: La similitud de coseno es una medida de similitud entre dos vectores no nulos de un espacio de producto interno que mide el coseno del ángulo entre ellos. Calcula la similitud entre dos vectores en el modelo de espacio vectorial.

Ahora pasemos a construir una aplicación que demuestre esto.

3. Prerrequisitos

Primero, deberíamos tener Docker instalado en nuestra máquina para ejecutar PGVector y Ollama.

Luego, necesitamos las dependencias de Spring AI Ollama y PGVector en nuestra aplicación Spring:

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

También agregaremos el soporte de Docker Compose de Spring Boot para gestionar los contenedores Docker de Ollama y PGVector:

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

Además de la dependencia, las juntaremos describiendo los dos servicios en un archivo docker-compose.yml:

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. Configuración de la aplicación

A continuación, necesitamos configurar nuestra aplicación Spring Boot para usar los servicios de Ollama y PGVector. En el archivo application.yml, definimos varias propiedades. Observemos específicamente lo seleccionado para las propiedades ollama y vectorstore:

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

Elegimos nomic-embed-text como nuestro modelo Ollama. Spring AI lo descarga por nosotros si no lo tenemos descargado.

La configuración de PGVector garantiza una correcta configuración del almacenamiento vectorial al inicializar el esquema de la base de datos (initialize-schema: true), alinear las dimensiones vectoriales con tamaños de embedding comunes (dimensions: 768) y optimizar la eficiencia de búsqueda utilizando el índice Hierarchical Navigable Small World (HNSW) (index-type: hnsw) para búsquedas de vecino más cercano aproximado rápidas.

5. Realizando búsqueda semántica

Ahora que nuestra infraestructura está lista, podemos implementar una aplicación de búsqueda semántica simple. Nuestro caso de uso será un motor de búsqueda inteligente de libros que permite a los usuarios buscar libros en función de su contenido.

Inicialmente, construiremos una funcionalidad de búsqueda simple usando PGVector y, posteriormente, la mejoraremos con Ollama para proporcionar respuestas más conscientes del contexto.

Definamos una clase Book que representa una entidad de libro:

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

Antes de poder buscar libros, necesitamos ingerir datos de libros en la tienda PGVector. El siguiente método agrega algunos datos de libros de muestra:

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);
}

Ahora que hemos agregado los datos de libros de muestra al almacén PGVector, podemos implementar la funcionalidad de búsqueda semántica.

5.1. Búsqueda semántica

Nuestro objetivo es implementar una API de búsqueda semántica que permita a los usuarios encontrar libros en función de su contenido.

Definamos un controlador que interactúe con PGVector para realizar búsquedas de similitud:

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

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

A continuación, crearemos un endpoint POST /search que acepte criterios de búsqueda del usuario y devuelva una lista de libros coincidentes:

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

Nótese que utilizamos VectorStore# similaritySearch. Esto realiza una búsqueda semántica entre los libros que ingerimos anteriormente.

Después de iniciar la aplicación, estamos listos para realizar una búsqueda. Utilicemos cURL para buscar instancias de 1984:

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

La respuesta contiene tres libros: uno con una coincidencia exacta y dos con coincidencias parciales:

[
  "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. Mejorando la búsqueda semántica con Ollama

Podemos integrar Ollama para generar respuestas parafraseadas que proporcionen contexto adicional para mejorar los resultados de la búsqueda semántica, siguiendo estos tres pasos:

Recuperar las tres descripciones de libros coincidentes superiores de la consulta de búsqueda.
Alimentar estas descripciones a Ollama para generar una respuesta más natural y consciente del contexto.
Entregar una respuesta que incluya información resumida y parafraseada, proporcionando ideas más claras y relevantes.

Creemos un nuevo método en BookSearchController que utilice Ollama para generar parafraseos de la consulta:

@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();
}

Y ahora probemos la funcionalidad de búsqueda semántica mejorada enviando una solicitud POST al endpoint /books/enhanced-search:

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?

En lugar de devolver tres descripciones de libros separadas como la búsqueda semántica simple, Ollama sintetiza la información más relevante de los resultados de búsqueda. En este caso, 1984 es la coincidencia más relevante, por lo que Ollama se centra en proporcionar un resumen detallado en lugar de listar libros no relacionados. Esto imita la asistencia de búsqueda humana, haciendo que los resultados sean más atractivos y perspicaces.

6. Conclusión

En este artículo, exploramos cómo implementar búsqueda semántica usando Spring AI, PGVector y Ollama. Comparamos dos endpoints; uno que realiza una búsqueda semántica de nuestro catálogo de libros y otro que alimenta y mejora ese resultado de búsqueda con un LLM de Ollama.
El código que respalda este artículo está disponible en GitHub. Una vez que estés conectado como un Miembro Baeldung Pro, comienza a aprender y programar en el proyecto.