Artículo

Construcción de una aplicación RAG con MongoDB y Spring AI

Ricardo Romero Benitez
Ricardo Romero
Construcción de una aplicación RAG con MongoDB y Spring AI

1. Visión general

El uso de tecnologías de IA se está convirtiendo en una habilidad clave en el desarrollo moderno. En este artículo construiremos una aplicación Wiki RAG que pueda responder preguntas basadas en documentos almacenados.

Utilizaremos Spring AI para integrar nuestra aplicación con la base de datos vectorial MongoDB y el LLM.

2. Aplicaciones RAG

Usamos aplicaciones de Generación Aumentada por Recuperación (RAG) cuando la generación de lenguaje natural necesita basarse en datos contextuales. Un componente clave de las aplicaciones RAG es la base de datos vectorial, que juega un papel crucial en la gestión y recuperación efectiva de estos datos:

Flujos de población de documentos y de prompting en aplicaciones RAG

Usamos un modelo de incrustación para procesar los documentos fuente. El modelo de incrustación convierte el texto de nuestros documentos en vectores de alta dimensión. Estos vectores capturan el significado semántico del contenido, lo que nos permite comparar y recuperar contenido similar basado en el contexto en lugar de solo coincidencia de palabras clave. Luego almacenamos los documentos en el almacén vectorial.

Una vez que hemos guardado los documentos, podemos enviar prompts basados en ellos de la siguiente manera:

  • Primero, usamos el modelo de incrustación para procesar la pregunta, convirtiéndola en un vector que captura su significado semántico.
  • A continuación, realizamos una búsqueda de similitud, comparando el vector de la pregunta con los vectores de los documentos almacenados en el almacén vectorial.
  • A partir de los documentos más relevantes, construimos un contexto para la pregunta.
  • Finalmente, enviamos tanto la pregunta como su contexto al LLM, que construye una respuesta relevante a la consulta y enriquecida por el contexto proporcionado.

3. MongoDB Atlas Vector Search

En este tutorial utilizaremos MongoDB Atlas Search como nuestro almacén vectorial. Proporciona capacidades de búsqueda vectorial que cubren nuestras necesidades en este proyecto. Para configurar la instancia local de MongoDB Atlas Search con fines de prueba, usaremos el contenedor mongodb-atlas-local docker. Creamos un archivo docker-compose.yml:

version: '3.1'

services:
  my-mongodb:
    image: mongodb/mongodb-atlas-local:7.0.9
    container_name: my-mongodb
    environment:
      - MONGODB_INITDB_ROOT_USERNAME=wikiuser
      - MONGODB_INITDB_ROOT_PASSWORD=password
    ports:
      - 27017:27017

4. Dependencias y Configuración

Empecemos añadiendo las dependencias necesarias. Como nuestra aplicación ofrecerá una API HTTP, incluiremos la dependencia spring-boot-starter-web:

<dependency>
    <groupId>org.springframework.boot</groupId> 
    <artifactId>spring-boot-starter-web</artifactId>
    <version>LATEST_VERSION</version>
</dependency>

Además, utilizaremos el cliente de la API de OpenAI para conectarnos al LLM, así que añadimos también su dependencia:

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

Finalmente, añadimos la dependencia del MongoDB Atlas Store:

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-mongodb-atlas-store-spring-boot-starter</artifactId>
    <version>LATEST_VERSION</version>
</dependency>

Ahora añadamos las propiedades de configuración para nuestra aplicación:

spring:
  data:
    mongodb:
      uri: mongodb://wikiuser:password@localhost:27017/admin
      database: wiki
  ai:
    vectorstore:
      mongodb:
        collection-name: vector_store
        initialize-schema: true
        path-name: embedding
        indexName: vector_index
    openai:
      api-key: ${OPENAI_API_KEY}
      chat:
        options:
          model: gpt-3.5-turbo

Especificamos la URL de MongoDB y la base de datos, y también configuramos nuestro almacén vectorial estableciendo el nombre de la colección, el nombre del campo de incrustación y el nombre del índice vectorial. Gracias a la propiedad initialize-schema, todos estos artefactos se crearán automáticamente por el framework Spring AI.

Por último, añadimos la clave de API de OpenAI y la versión del modelo.

5. Guardar Documentos en el Almacén Vectorial

Ahora, añadimos el proceso para guardar datos en nuestro almacén vectorial. Nuestra aplicación será responsable de proporcionar respuestas a preguntas de usuarios basadas en documentos existentes, funcionando esencialmente como una especie de wiki.

Añadimos un modelo que almacenará el contenido de los archivos junto con la ruta del archivo:

public class WikiDocument {
    private String filePath;
    private String content;

    // getters y setters estándar
}

El siguiente paso es añadir el WikiDocumentsRepository. En este repositorio, encapsulamos toda la lógica de persistencia:

import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;

@Component
public class WikiDocumentsRepository {
    private final VectorStore vectorStore;

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

    public void saveWikiDocument(WikiDocument wikiDocument) {

        Map<String, Object> metadata = new HashMap<>();
        metadata.put("filePath", wikiDocument.getFilePath());
        Document document = new Document(wikiDocument.getContent(), metadata);
        List<Document> documents = new TokenTextSplitter().apply(List.of(document));

        vectorStore.add(documents);
    }
}

Aquí, inyectamos el bean de la interfaz VectorStore, que será implementado por MongoDBAtlasVectorStore provisto por el spring-ai-mongodb-atlas-store-spring-boot-starter. En el método saveWikiDocument creamos una instancia de Document y la llenamos con el contenido y los metadatos.

Luego utilizamos TokenTextSplitter para dividir el documento en fragmentos más pequeños y guardarlos en nuestro almacén vectorial. Ahora creamos un WikiDocumentsServiceImpl:

@Service
public class WikiDocumentsServiceImpl {
    private final WikiDocumentsRepository wikiDocumentsRepository;

    // constructores

    public void saveWikiDocument(String filePath) {
        try {
            String content = Files.readString(Path.of(filePath));
            WikiDocument wikiDocument = new WikiDocument();
            wikiDocument.setFilePath(filePath);
            wikiDocument.setContent(content);

            wikiDocumentsRepository.saveWikiDocument(wikiDocument);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
}

En la capa de servicio, recuperamos el contenido del archivo, creamos instancias de WikiDocument y las enviamos al repositorio para su persistencia.

En el controlador simplemente pasamos la ruta del archivo a la capa de servicio y devolvemos un código de estado 201 si el documento se guarda correctamente:

@RestController
@RequestMapping("wiki")
public class WikiDocumentsController {
    private final WikiDocumentsServiceImpl wikiDocumentsService;

    // constructores

    @PostMapping
    public ResponseEntity<Void> saveDocument(@RequestParam String filePath) {
        wikiDocumentsService.saveWikiDocument(filePath);

        return ResponseEntity.status(201).build();
    }
}

Debemos prestar atención a los aspectos de seguridad de este endpoint. Existe una vulnerabilidad potencial en la que los usuarios podrían cargar archivos inesperados, como archivos de configuración o del sistema, usando este endpoint. Como solución, podemos limitar los directorios desde los que se pueden cargar archivos. Ahora, iniciemos nuestra aplicación y veamos cómo funciona nuestro flujo. Añadimos la dependencia de prueba de Spring Boot, que nos permitirá configurar un contexto web de prueba:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-test</artifactId>
    <version>LATEST_VERSION</version>
</dependency>

A continuación, arrancamos la instancia de prueba de la aplicación y llamamos al endpoint POST para dos documentos:

@AutoConfigureMockMvc
@ExtendWith(SpringExtension.class)
@SpringBootTest
class RAGMongoDBApplicationManualTest {

    @Autowired
    private MockMvc mockMvc;

    @Test
    void givenMongoDBVectorStore_whenCallingPostDocumentEndpoint_thenExpectedResponseCodeShouldBeReturned() throws Exception {
        mockMvc.perform(post("/wiki?filePath={filePath}",
          "src/test/resources/documentation/owl-documentation.md"))
          .andExpect(status().isCreated());

        mockMvc.perform(post("/wiki?filePath={filePath}",
          "src/test/resources/documentation/rag-documentation.md"))
          .andExpect(status().isCreated());
    }
}

Ambas llamadas deberían devolver un código de estado 201, lo que indica que los documentos fueron añadidos. Podemos usar MongoDB Compass para confirmar que los documentos se guardaron correctamente en el almacén vectorial:

Vista del documento almacenado en MongoDB Compass

Como podemos ver, ambos documentos se guardaron. Podemos ver el contenido original así como un arreglo de incrustaciones.

6. Búsqueda de Similitud

Añadamos la funcionalidad de búsqueda de similitud. Incluiremos un método findSimilarDocuments en nuestro repositorio:

@Component
public class WikiDocumentsRepository {
    private final VectorStore vectorStore;

    public List<WikiDocument> findSimilarDocuments(String searchText) {

        return vectorStore
          .similaritySearch(SearchRequest
            .query(searchText)
            .withSimilarityThreshold(0.87)
            .withTopK(10))
          .stream()
          .map(document -> {
              WikiDocument wikiDocument = new WikiDocument();
              wikiDocument.setFilePath((String) document.getMetadata().get("filePath"));
              wikiDocument.setContent(document.getContent());

              return wikiDocument;
          })
          .toList();
    }
}

Llamamos al método similaritySearch del VectorStore. Además del texto de búsqueda, especificamos un límite de resultados y un umbral de similitud. El parámetro umbral de similitud nos permite controlar cuán cerca debe coincidir el contenido del documento con nuestro texto de búsqueda.

En la capa de servicio, proximizamos la llamada al repositorio:

public List<WikiDocument> findSimilarDocuments(String searchText) {
    return wikiDocumentsRepository.findSimilarDocuments(searchText);
}

En el controlador, añadimos un endpoint GET que recibe el texto de búsqueda como parámetro y lo pasa al servicio:

@RestController
@RequestMapping("/wiki")
public class WikiDocumentsController {
    @GetMapping
    public List<WikiDocument> get(@RequestParam("searchText") String searchText) {
        return wikiDocumentsService.findSimilarDocuments(searchText);
    }
}

Ahora llamemos a nuestro nuevo endpoint y veamos cómo funciona la búsqueda de similitud:

@Test
void givenMongoDBVectorStoreWithDocuments_whenMakingSimilaritySearch_thenExpectedDocumentShouldBePresent() throws Exception {
    String responseContent = mockMvc.perform(get("/wiki?searchText={searchText}", "RAG Application"))
      .andExpect(status().isOk())
      .andReturn()
      .getResponse()
      .getContentAsString();

    assertThat(responseContent)
      .contains("RAG AI Application is responsible for storing the documentation");
}

Llamamos al endpoint con un texto de búsqueda que no era una coincidencia exacta en el documento. Sin embargo, todavía recuperamos el documento con contenido similar y confirmamos que contenía el texto que almacenamos en nuestro archivo rag-documentation.md.

7. Endpoint de Prompt

Comencemos a construir el flujo de prompt, que es la funcionalidad central de nuestra aplicación. Empezaremos con la AdvisorConfiguration:

@Configuration
public class AdvisorConfiguration {

    @Bean
    public QuestionAnswerAdvisor questionAnswerAdvisor(VectorStore vectorStore) {
        return new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults());
    }
}

Creamos un bean QuestionAnswerAdvisor, responsable de construir la solicitud de prompt, incluyendo la pregunta inicial. Además, adjuntará la respuesta de búsqueda de similitud del almacén vectorial como contexto para la pregunta. Ahora, añadamos el endpoint de búsqueda a nuestra API:

@RestController
@RequestMapping("/wiki")
public class WikiDocumentsController {
    private final WikiDocumentsServiceImpl wikiDocumentsService;
    private final ChatClient chatClient;
    private final QuestionAnswerAdvisor questionAnswerAdvisor;

    public WikiDocumentsController(WikiDocumentsServiceImpl wikiDocumentsService,
                                   @Qualifier("openAiChatModel") ChatModel chatModel,
                                   QuestionAnswerAdvisor questionAnswerAdvisor) {
        this.wikiDocumentsService = wikiDocumentsService;
        this.questionAnswerAdvisor = questionAnswerAdvisor;
        this.chatClient = ChatClient.builder(chatModel).build();
    }

    @GetMapping("/search")
    public String getWikiAnswer(@RequestParam("question") String question) {
        return chatClient.prompt()
          .user(question)
          .advisors(questionAnswerAdvisor)
          .call()
          .content();
    }
}

Aquí, construimos una solicitud de prompt añadiendo la entrada del usuario al prompt y adjuntando nuestro QuestionAnswerAdvisor.

Finalmente, llamemos a nuestro endpoint y veamos qué nos dice sobre las aplicaciones RAG:

@Test
void givenMongoDBVectorStoreWithDocumentsAndLLMClient_whenAskQuestionAboutRAG_thenExpectedResponseShouldBeReturned() throws Exception {
    String responseContent = mockMvc.perform(get("/wiki/search?question={question}", "Explain the RAG Applications"))
      .andExpect(status().isOk())
      .andReturn()
      .getResponse()
      .getContentAsString();

    logger.atInfo().log(responseContent);

    assertThat(responseContent).isNotEmpty();
}

Enviamos la pregunta "Explain the RAG applications" a nuestro endpoint y registramos la respuesta de la API:

b.s.r.m.RAGMongoDBApplicationManualTest : Based on the context provided, the RAG AI Application is a tool 
used for storing documentation and enabling users to search for specific information efficiently...

Como vemos, el endpoint devolvió información sobre las aplicaciones RAG basada en los archivos de documentación que guardamos previamente en la base de datos vectorial.

Ahora intentemos preguntar algo que seguramente no tengamos en nuestra base de conocimientos:

@Test
void givenMongoDBVectorStoreWithDocumentsAndLLMClient_whenAskUnknownQuestion_thenExpectedResponseShouldBeReturned() throws Exception {
    String responseContent = mockMvc.perform(get("/wiki/search?question={question}", "Explain the Economic theory"))
      .andExpect(status().isOk())
      .andReturn()
      .getResponse()
      .getContentAsString();

    logger.atInfo().log(responseContent);

    assertThat(responseContent).isNotEmpty();
}

Ahora hemos preguntado sobre teoría económica, y esta es la respuesta:

b.s.r.m.RAGMongoDBApplicationManualTest : I'm sorry, but the economic theory is not directly related to the information provided about owls and the RAG AI Application.
If you have a specific question about economic theory, please feel free to ask.

Esta vez, nuestra aplicación no encontró documentos relacionados y no utilizó otras fuentes para proporcionar una respuesta.

8. Conclusión

En este artículo implementamos con éxito una aplicación RAG utilizando el framework Spring AI, que es una herramienta excelente para integrar diversas tecnologías de IA. Además, MongoDB demuestra ser una opción sólida para manejar el almacenamiento vectorial.

Con esta poderosa combinación, podemos construir aplicaciones modernas basadas en IA para diversos propósitos, incluidos chatbots, sistemas wiki automatizados y motores de búsqueda.

Como siempre, el código está disponible en GitHub.

Newsletter Semanal de Java

Cada viernes recibe lo más nuevo del ecosistema Java: frameworks, herramientas y mejores prácticas.

Sin spam. Cancela cuando quieras.

Escrito por
Ricardo Romero
Compartir artículo