Crear una aplicación RAG (Generación Aumentada por Recuperación) con Redis y Spring AI

1. Visión general

En este tutorial, construiremos un ChatBot usando el framework Spring AI y la técnica de RAG (Retrieval Augmented Generation). Con la ayuda de Spring AI, nos integraremos con la base de datos vectorial Redis para almacenar y recuperar datos que mejoren el prompt para el LLM (Large Language Model). Una vez que el LLM reciba el prompt con los datos relevantes, generará eficazmente una respuesta con los datos más recientes en lenguaje natural a la consulta del usuario.

2. ¿Qué es RAG?

Los LLM son modelos de Machine Learning pre-entrenados con grandes conjuntos de datos de internet. Para que un LLM funcione dentro de una empresa privada, debemos ajustarlo con la base de conocimientos específica de la organización. Sin embargo, el ajuste fino suele ser un proceso que consume mucho tiempo y requiere recursos informáticos sustanciales. Además, existe una alta probabilidad de que un LLM ajustado genere respuestas irrelevantes o engañosas a las consultas. Este comportamiento se conoce comúnmente como alucinaciones del LLM.

En tales escenarios, RAG es una técnica excelente para restringir o contextualizar las respuestas del LLM. Una DB vectorial desempeña un papel importante en la arquitectura RAG para proporcionar información contextual al LLM. Pero, antes de que una aplicación pueda usarlo en la arquitectura RAG, un proceso ETL (Extract, Transform, Load) debe poblarlo:

El lector recupera los documentos de la base de conocimientos de la organización desde diferentes fuentes. Luego, el transformador divide los documentos recuperados en fragmentos más pequeños y utiliza un modelo de embeddings para vectorizar los contenidos. Finalmente, el escritor carga los vectores o embeddings en la DB vectorial. Las DB vectoriales son bases de datos especializadas que pueden almacenar estos embeddings en un espacio multidimensional.

En RAG, los LLM pueden responder a datos casi en tiempo real si la DB vectorial se actualiza periódicamente a partir de la base de conocimientos de la organización.

Una vez que la DB vectorial está lista con los datos, la aplicación puede usarla para recuperar la información contextual para las consultas de los usuarios:

La aplicación forma el prompt combinando la consulta del usuario con los datos contextuales de la DB vectorial y finalmente lo envía al LLM. El LLM genera la respuesta en lenguaje natural dentro de los límites de los datos contextuales y la envía de vuelta a la aplicación.

3. Implementar RAG con Spring AI y Redis

La pila Redis ofrece servicios de búsqueda vectorial; usaremos el framework Spring AI para integrarnos con ella y construir una aplicación de ChatBot basada en RAG. Además, utilizaremos el modelo LLM GPT-3.5 Turbo de OpenAI para generar la respuesta final.

3.1. Requisitos previos

Para el servicio ChatBot, para autenticar el servicio OpenAI necesitaremos la clave secreta de la API. Crearemos una, después de crear una cuenta OpenAI:

También crearemos una cuenta Redis Cloud para acceder a una Redis Vector DB gratuita:

Para la integración con la Redis Vector DB y el servicio OpenAI, actualizaremos las dependencias Maven con las librerías de Spring AI:

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
    <version>1.0.0-M1</version>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
    <version>1.0.0-M1</version>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-redis-spring-boot-starter</artifactId>
    <version>1.0.0-M1</version>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-pdf-document-reader</artifactId>
    <version>1.0.0-M1</version>
</dependency>

3.2. Clases clave para cargar datos en Redis

En una aplicación Spring Boot, crearemos componentes para cargar y recuperar datos de la Redis Vector DB. Por ejemplo, cargaremos un documento PDF del manual del empleado en la DB Redis.

Ahora, echemos un vistazo a las clases involucradas:

DocumentReader es una interfaz de Spring AI para leer documentos. Usaremos la implementación lista para usar PagePdfDocumentReader de DocumentReader. De manera similar, DocumentWriter y VectorStore son interfaces para escribir datos en sistemas de almacenamiento. RedisVectorStore es una de las muchas implementaciones listas para usar de VectorStore, que utilizaremos para cargar y buscar datos en Redis Vector DB. Escribiremos el DataLoaderService utilizando las clases del framework Spring AI discutidas hasta ahora.

3.3. Implementar el servicio de carga de datos

Entendamos el método load() en la clase DataLoaderService:

@Service
public class DataLoaderService {
    private static final Logger logger = LoggerFactory.getLogger(DataLoaderService.class);

    @Value("classpath:/data/Employee_Handbook.pdf")
    private Resource pdfResource;

    @Autowired
    private VectorStore vectorStore;

    public void load() {
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(this.pdfResource,
            PdfDocumentReaderConfig.builder()
              .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                .withNumberOfBottomTextLinesToDelete(3)
                .withNumberOfTopPagesToSkipBeforeDelete(1)
                .build())
            .withPagesPerDocument(1)
            .build());

        var tokenTextSplitter = new TokenTextSplitter();
        this.vectorStore.accept(tokenTextSplitter.apply(pdfReader.get()));
    }
}

El método load() utiliza la clase PagePdfDocumentReader para leer un archivo PDF y cargarlo en la Redis Vector DB. El framework Spring AI configura automáticamente la interfaz VectorStore usando las propiedades de configuración en el namespace spring.ai.vectorstore:

spring:
  ai:
    vectorstore:
      redis:
        uri: redis://:PQzkkZLOgOXXX@redis-19438.c330.asia-south1-1.gce.redns.redis-cloud.com:19438
        index: faqs
        prefix: "faq:"
        initialize-schema: true

El framework inyecta el objeto RedisVectorStore, una implementación de la interfaz VectorStore, en el DataLoaderService.

La clase TokenTextSplitter divide el documento y, finalmente, la clase VectorStore carga los fragmentos en la Redis Vector DB.

3.4. Clases clave para generar la respuesta final

Una vez que la Redis Vector DB esté lista, podemos recuperar la información contextual relevante para la consulta del usuario. Después, este contexto se utiliza para formar el prompt del LLM y generar la respuesta final. Veamos las clases clave:

El método searchData() en la clase DataRetrievalService recibe la consulta y luego recupera los datos de contexto del VectorStore. El ChatBotService utiliza estos datos para formar el prompt usando la clase PromptTemplate y luego lo envía al servicio OpenAI. El framework Spring Boot lee las propiedades relacionadas con OpenAI del archivo application.yml y luego autoconfigura el objeto OpenAIChatModel. En este artículo, configuramos el perfil activo de Spring como "airag".

Veamos la implementación para entenderla en detalle.

3.5. Implementar el servicio de Chat Bot

Echemos un vistazo a la clase ChatBotService:

@Service
public class ChatBotService {
    @Qualifier("openAiChatModel")
    @Autowired
    private ChatModel chatClient;
    @Autowired
    private DataRetrievalService dataRetrievalService;

    private final String PROMPT_BLUEPRINT = """
      Answer the query strictly referring the provided context:
      {context}
      Query:
      {query}
      In case you don't have any answer from the context provided, just say:
      I'm sorry I don't have the information you are looking for.
    """;

    public String chat(String query) {
        return chatClient.call(createPrompt(query, dataRetrievalService.searchData(query)));
    }

    private String createPrompt(String query, List<Document> context) {
        PromptTemplate promptTemplate = new PromptTemplate(PROMPT_BLUEPRINT);
        promptTemplate.add("query", query);
        promptTemplate.add("context", context);
        return promptTemplate.render();
    }
}

El framework SpringAI crea el bean ChatModel usando las propiedades de configuración de OpenAI en el namespace spring.ai.openai:

spring:
  ai:
    vectorstore:
      redis:
        # Propiedades relacionadas con la base de datos vectorial Redis...
    openai:
      temperature: 0.3
      api-key: ${SPRING_AI_OPENAI_API_KEY}
      model: gpt-3.5-turbo
      #embedding-base-url: https://api.openai.com
      #embedding-api-key: ${SPRING_AI_OPENAI_API_KEY}
      #embedding-model: text-embedding-ada-002

El framework también puede leer la clave API desde la variable de entorno SPRING_AI_OPENAI_API_KEY, que es una opción mucho más segura. Podemos habilitar las claves que comienzan con el texto embedding para crear el bean OpenAiEmbeddingModel, que se usa para crear embeddings vectoriales a partir de los documentos de la base de conocimientos.

El prompt para el servicio OpenAI debe ser inequívoco. Por eso, hemos instruido estrictamente en el blueprint de prompt PROMPT_BLUEPRINT que la respuesta se forme solo a partir de la información contextual.

En el método chat() recuperamos los documentos que coinciden con la consulta desde la Redis Vector DB. Luego, usamos estos documentos y la consulta del usuario para generar el prompt en el método createPrompt(). Finalmente, invocamos el método call() de la clase ChatModel para recibir la respuesta del servicio OpenAI.

Ahora, veamos el servicio de chatbot en acción preguntándole algo del manual del empleado cargado anteriormente en la Redis Vector DB:

@Test
void whenQueryAskedWithinContext_thenAnswerFromTheContext() {
    String response = chatBotService.chat("How are employees supposed to dress?");
    assertNotNull(response);
    logger.info("Response from LLM: {}", response);
}

Entonces, veremos la salida:

Response from LLM: Employees are supposed to dress appropriately for their individual work responsibilities and position.

La salida coincide con el documento PDF del manual del empleado cargado en la Redis Vector DB.

Veamos qué sucede si preguntamos algo que no está en el manual del empleado:

@Test
void whenQueryAskedOutOfContext_thenDontAnswer() {
    String response = chatBotService.chat("What should employees eat?");
    assertEquals("I'm sorry I don't have the information you are looking for.", response);
    logger.info("Response from the LLM: {}", response);
}

Aquí está la salida resultante:

Response from the LLM: I'm sorry I don't have the information you are looking for.

El LLM no encontró nada en el contexto proporcionado y, por lo tanto, no pudo responder a la consulta.

4. Conclusión

En este artículo, discutimos cómo implementar una aplicación basada en la arquitectura RAG usando el framework Spring AI. Formar el prompt con la información contextual es esencial para generar la respuesta correcta del LLM. Por eso, la Redis Vector DB es una excelente solución para almacenar y realizar búsquedas de similitud sobre los vectores de documentos. Además, dividir los documentos es igualmente importante para obtener los registros correctos y restringir el costo de los tokens del prompt.