Embeddings y RAG en aplicaciones web

Quiero esta vez escribir un poco sobre el proceso de construcción de una pequeña idea que use para poder comprender el uso de embeddings y RAG en aplicaciones web que requieran este tipo de tecnologías.

Buscando entender un poco estos conceptos encontré la idea de poder ayudar a mejorar la experiencia de usuarios que necesiten llevar al siguiente nivel sus procesos de capacitación o uso de procesos sistemáticos que se encuentren alojados en archivos PDF o word a través de una plataforma que ponga a disposición un chatbot que contenga toda la información de estos archivos y el usuario pueda realizar preguntas sobre ellos de forma natural y logre recibir una respuesta.

Segun esta necesidad encontre que los embeddings y el uso de RAG en chatbots con IA logran no solo cubrir esta necesidad sino abrir un abanico de opciones para poder llevar esta idea a niveles superiores.

En un primer momento el flujo fue muy sencillo:

En la anterior imágen encontramos

• 3 nodos que representan a openai
• 2 nodos que representan endpoints de nuestra aplicación
• 2 nodos que representa nuestra base de datos vectorial
• 1 nodo que representa una base de datos tipo postgreSQL

Ahora bien, lo primero que debemos resaltar son los nodos de openai ya que estos nos permiten generar los embeddings y además de esto la respuesta final que le entregaremos al usuario.

Los nodos que representan la base de datos vectorial nos permiten almacenar estos datos vectorizados, en este caso estamos usando postgreSQL, para poder consultarlos y de esta forma alimentar a openai con esta informacion y lograr dar con una respuesta concreta al usuario.

El nodo de postgreSQL nos permite almacenar datos del documento almacenado para poder mostrar al usuario en una UI el listado de documentos en el sistema.

Por último los 2 nodos que representan los endpoints son las puertas de entrada y salida para poder consumir todo estos recursos desde nuestra aplicación.

Paremos un momento e intentemos revisar detalladamente el papel de los 3 nodos de openai. Estos, como ya lo nombramos, cumplen dos funciones: Realizar el proceso de creación del embedding y generar la respuesta final del usuario a través del chat. Pero ¿Qué es un embedding y porque debemos de utilizar esta técnica para lograr nuestro objetivo?

Un embedding en pocas palabras es la representación numérica de cada palabra de nuestro elemento a almacenar. Esta representación no es arbitraria ya que, usando matematica de vectores, logra llegar a una representación numerica de cada palabra.

Para lograr esto, empresas como Openai, Anthropic o Google nos proveen de modelos especializados en embeddings. Muy similar a los modelos que siempre escuchamos como GPT-4o, Sonnet o Gemini existen modelos para embeddings. En este caso estaremos utilizando el modelo text-embedding-3-small de openai sin embargo existen muchos otros y su uso dependera del objetivo de nuestra aplicación y el presupuesto que tengamos a disposición.

Por ejemplo text-embedding-3-small es un modelo de uso general y economico con 1536 valores por vector pero si vamos a modelos como text-embedding-3-large encontramos un modelo mas preciso con 3072 valores por vector pero en ese mismo sentido el procesamiento es mas costoso.

Ademas de esto debemos tener en cuenta factores como el dominio de entrenamiento del modelo, soporte en varios idiomas o capacidad semantica.

Una vez acalarado esto debemos almacenar nuestros vectores en una base de datos que nos permita almacenar este tipo de datos y para ello usamos PostgreSQL a través de su extensión pgvector

create extension if not exists vector;

Sin embargo, tenemos bases de datos que soportan nativamente vectores como es el caso de Pinecone o Chroma. Para este caso decidimos usar PostgreSQL porque estamos utilizando Supabase como BaaS y así proveernos de sus diferentes features para un desarrollo más rápido.

Dentro de esta tabla de Postgres, utilizando nuestra extensión de pgvector, ya podemos generar columnas que nos almacenen este nuevo tipo de dato. Por ejemplo, para nuestro caso tenemos nuestra tabla definida así:

create table if not exists document_sections (
  id               bigserial primary key,
  document_id      bigint not null references documents(id) on delete cascade,
  section_order    int not null,                      
  section_content  text not null,                     
  embedding        vector(1536) not null,            
  meta             jsonb not null default '{}'::jsonb,
  created_at       timestamptz not null default now()
);

Como puedes observar nuestro embedding se almacenará aquí:

embedding vector(1536) not null

En donde el número 1536 representa la cantidad de elementos que existirán dentro de cada vector.

Vale, ahora debemos de hablar de algo muy importante en este punto y es algo llamado chunk. Tanto nuestro modelo de embedding como nuestra tabla de vectores no almacena todo el documento de golpe ya que si el usuario llega a cargar un documento muy grande la exigencia de computo sería enorme, la precision a la hora de realizar busquedas se veria afectada y, además de esto, debemos de tener en cuenta que estos modelos de embeddings tambien tienen un máximo de tokens por lo tanto no podemos procesar tokens de manera infinita.

Para ello nos vemos en la necesidad de procesar los datos por chunks. Para lograr esto tenemos múltiples caminos. Implementar nuestro propio "chunker" o utilizar uno de alguna libreria. Para este caso decidí crear mi propio "chunker" ya que no es un producto con pretensiones comerciales pero facilmente puedes utilizar sentence-splitter.

La tarea de realizar chunks es muy importante ya que de esta depende el nivel de precision que existira en los vectores. Debes de tener en cuenta que realizar cortes muy grandes genera distorsion y sesgo a la hora de realizar busquedas vectoriales pero si realizamos cortes muy pequeños perdemos contexto global de nuestro texto, responder a preguntas generales se vuelve complejo y debemos de procesar más embeddings lo cual nos genera mas costos.

Vale, entonces hagamos un paso a paso del funcionamiento de nuestra aplicación. En esta primera parte describiremos el proceso de uploading:

1. El usuario hace un upload del documento desde la UI
2. Se envían estos datos al servidor usando el endpoint send-file
3. Una vez dentro de send-file ejecutamos el chunk de nuestro archivo y por cada chunk generamos un embedding.
4. Insertamos nuestro embedding en la base de datos vectorial.
5. Insertamos en nuestra tabla convencional de postgres los datos del archivo para mostrar en la UI.

En este punto ya con nuestro archivo en nuestra base de datos podemos consultarlo y para ello recurrimos a los siguientes pasos

1. El usuario realiza una pregunta en el chat
2. Esta pregunta la transformamos en un embedding
3. Usamos este embedding generado y realizamos un busqueda vectorial en nuestra base de datos vectorial con el fin de encontrar similitudes entre el embedding generado por la pregunta y los embeddings almacenados en nuestra base de datos
4. Usamos un LLM para procesar tanto la respuesta de la busqueda vectorial utilizada aqui como contexto y la pregunta a responder enviada inicialmente por el usuario.
5. La respuesta generada la devolvemos al cliente.

De esta forma cerramos el ciclo de nuestra aplicación y como puedes ver estamos utilizando un modelo de LLM alimentando con los datos de nuestra base de datos vectorial para generar una respuesta que se limite solamente a estos datos. En nuestra aplicación esta última sección luce algo así:

    const aiResponse = await openai.chat.completions.create({
      model: "gpt-4o-mini",
      temperature: 1,
      max_completion_tokens: 150,
      messages: [
        {
          role: "system",
          content:
            "Use the context if it exists. If the context contains 'NO_COINCIDENCE', thank the user and kindly explain that you cannot respond with the information available.",
        },
        {
          role: "user",
          content: `
                Using the following information please answer the question: 
                Context: ${context}
                Question: ${message}
          `,
        },
      ],
    });

Habras notado que utilizamos un termino busqueda vectorial al cual no nos hemos referido hasta el momento. Existen 3 tipos de busquedas vectoriales que podemos llegar a aplicar:

1. Similitud del coseno (Cosine similarity)
2. Producto interno (Dot product)
3. Distancia euclidiana (Euclidian distance)

Cada una de estas tres opciones tiene su objetivo, para nuestro caso aplicamos la busqueda por similitud del coseno. Este algoritmo nos permite realizar busquedas por angulo entre dos vectores e ignora la longitud de los mismos. En otras palabras dentro del plano vectorial si el angulo entre dos vectores es pequeño representa similitud, por el contrario si es mayor o totalmente opuesto representan menos similitud o directamente significados totalmente opuestos.

Por lo general en embeddings se tiende a utilizar la similitud por coseno ya que nos interesa mas la direccion semantica (angulo) más que la longitud del vector ya que este algoritmo normaliza esta longitud por defecto haciendo que este valor pierda importancia.

Se debe de tener en cuenta que esta normalización tiene un costo y, si bien es el algoritmo más utilizado, el producto interno tiene un mayor rendimiento pero la longitud al no estar normalizada puede llegar a introducir sesgos o alucinaciones.

En nuestro caso hemos diseñado esta consulta a traves de un procedimiento almacenado en supabase utilizando la similitud del coseno.

Una vez repasado de manera somera el flujo de nuestra aplicación podemos empezar a ver que factores extra nos podemos llegar a ecnontrar y como los podemos resolver.

Uno de los principales inconvenientes que podemos llegar a enfrentar es no bloquear al usuario durante todo este poceso ya que, como puedes llegar a imaginar, el tiempo que trascurre entre subir el archivo y generar los embeddings puede tomar un tiempo dependiendo de la velocidad de red que haya en ese momento y la disponibilidad del servicio de openai.

En ese sentido, para no bloquear al usuario podemos generar un sistema de colas que nos permitan desacoplar este proceso y, a través de un estado de carga, avisar al usuario cuando el archivo ya se encuentra listo para su consulta.

Para lograr esto podemos utilizar servicios como SQS de AWS o algun broker de mensajeria como RabbitMQ. Ya depende de tus necesidades y presupuesto.

Para nuestro caso hemos decidido mantenerlo simple y usamos un servicio de mensajeria de Upstash llamado Qstash lo cual nos permitio desacoplar nuestro sistema sin necesidad de crear una infraestrcutura manteniendo simple nuestra implementación.

Otro posible problema que puedes llegar a enfrentar es que si trabajas con servicios serverless como Lambdas o endpoints de Next estos tienen un limite de datos que se pueden llegar a transmitir en cada petición, por lo tanto enviar archivos PDF o docx puede generarte problemas de tipo HTTP 413 "Content Too Large". Para solucionar esto puedes usar servicios como S3 de AWS o Google Storage en GCP. Lo que implementamos nosotros fue: Desde el cliente subir el archivo a un bucket en supabase con un storagePath y, en el punto en el que necesitamos generar el embedding ya del lado del server, consultamos este bucket, descargamos el archivo y generamos el embedding y asi no tenemos que recurrir a envios de red muy grandes que pueden afectar el rendimiento de nuestra aplicación.

Por último te dejo el diagrama final de la aplicacion en donde se detalla la interacción de cada componente: