PostgreSQL con pgvector como base de datos vectorial para RAG

Cómo implementar búsquedas vectoriales y RAG usando PostgreSQL

Hemos explorado la búsqueda vectorial y Retrieval-Augmented Generation (RAG) en artículos anteriores, como RAG desde cero y Crea un chatbot con IA alimentado por tus datos. Esta última incluye una implementación completa de un chatbot con RAG usando Redis como base de datos vectorial.

Pero hay otra base de datos muy conocida y que probablemente hayas usado: PostgreSQL. ¿Y si pudiéramos utilizarla también como base de datos vectorial?

PostgreSQL es una base de datos relacional de código abierto ampliamente utilizada. Es increíblemente versátil, permitiendo además almacenar y manipular datos en formato JSON (similar a las bases de datos NoSQL y documentales) y ofreciendo numerosas extensiones con funcionalidades adicionales, como PostGIS para datos geoespaciales o pgcron para programar tareas.

Gracias a la extensión pgvector, Postgres también puede realizar búsquedas de similitud vectorial de forma eficiente. Esto abre grandes posibilidades para aplicaciones de RAG y de IA, con la ventaja añadida de utilizar una base de datos familiar que probablemente ya tengas en tu stack. También significa que puedes combinar datos relacionales, datos JSON y vectores (embeddings) en un mismo sistema, permitiendo consultas complejas que involucren tanto datos estructurados como búsquedas vectoriales.

Existen bases de datos vectoriales especializadas como Qdrant, Pinecone o Weaviate, que ofrecen un rendimiento optimizado para grandes volúmenes de datos así como funcionalidades más avanzadas. Sin embargo, Postgres con pgvector es una alternativa muy interesante si quieres mantener todos los datos de tu aplicación más integrados y minimizar el número de bases de datos en tu infraestructura, reduciendo costes y complejidad.

En este artículo, exploraremos cómo configurar Postgres como base de datos vectorial y cómo utilizarla en búsquedas vectoriales y aplicaciones de RAG en Python.

Configuración de PostgreSQL y pgvector

Antes de comenzar, necesitamos instalar PostgreSQL, pgvector y las librerías de Python que usaremos:

  1. Descarga e instala PostgreSQL siguiendo las instrucciones oficiales para tu sistema operativo.

  2. Instala la extensión pgvector siguiendo las notas de instalación en el repositorio de pgvector.

  3. Instala las dependencias de Python necesarias. Además de la librería para usar pgvector en Python, utilizaremos SQLAlchemy como ORM y asyncpg como driver para conectarnos a Postgres de forma asíncrona con asyncio:

    pip install sqlalchemy pgvector asyncpg

  4. Crea una nueva base de datos y habilita la extensión pgvector:

    createdb rag_db
psql rag_db
CREATE EXTENSION vector;

Creación de una base de datos vectorial con PostgreSQL

Ahora que la base de datos está configurada, vamos a crear un modelo de SQLAlchemy para representar nuestros datos vectoriales:

from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import Text
from sqlalchemy.dialects.postgresql import JSONB
from pgvector.sqlalchemy import Vector

class Base(DeclarativeBase):
    pass

class Vector(Base):
    __tablename__ = 'vectors'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    text: Mapped[str] = mapped_column(Text)
    vector = mapped_column(Vector(1024))
    metadata_: Mapped[dict | None] = mapped_column('metadata', JSONB)

    def __repr__(self):
        return f'Vector(id={self.id}, text={self.text[:50]}..., metadata={self.metadata_})'

Estamos utilizando la versión 2.0 de SQLAlchemy, que nos permite usar las “anotaciones de tipo” (type hints) de Python con Mapped y mapped_column para inferir los tipos de las columnas de la base de datos.

El modelo Vector define las siguientes columnas en la tabla vectors:

  • id: Un identificador único (y clave primaria) para cada vector.
  • vector: Almacena el vector embedding. Especificamos las dimensiones del embedding (1024 en este caso) para que coincida con nuestro modelo de embedding elegido.
  • text: Contiene el texto original que representa el vector. Esto nos permite recuperar fácilmente el texto correspondiente al realizar búsquedas de similitud vectorial.
  • metadata: Almacena propiedades adicionales para cada vector como JSON. Aquí podemos guardar información como el nombre del documento de origen, el índice del fragmento dentro del documento o cualquier otro metadato relevante para realizar búsquedas vectoriales o filtrados.

⚠️ Ten en cuenta que usamos “metadata_” como nombre del atributo en el modelo porque “metadata” es una palabra reservada en los modelos de SQLAlchemy, pero el nombre de la columna en la base de datos sí será “metadata”.

Para crear la tabla de la base de datos definida en nuestro modelo, podemos usar el siguiente código de SQLAlchemy:

from sqlalchemy.ext.asyncio import create_async_engine

DB_URL = 'postgresql+asyncpg://admin:postgres@localhost:5432/rag_db'

engine = create_async_engine(DB_URL)

async def db_create():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

El prefijo 'postgresql+asyncpg' en la URL de la base de datos es necesario porque estamos usando el driver asyncpg para habilitar conexiones asíncronas con asyncio.

Una vez creada la tabla vectors, puedes usar pgAdmin para explorar la estructura de la tabla y ejecutar consultas con una interfaz gráfica.

Almacenando documentos como embeddings en PostgreSQL

El siguiente paso es procesar, vectorizar (generando embeddings) y almacenar la información en nuestra base de datos vectorial Postgres. Si estás familiarizado con sistemas RAG, este proceso incluye los siguientes pasos:

  • Extraer el texto de los documentos de origen.
  • Dividir el texto en fragmentos más pequeños.
  • Convertir estos fragmentos en representaciones vectoriales (embeddings) que capturan su significado.
  • Almacenar los embeddings junto con los textos originales y cualquier metadato relevante en nuestra base de datos vectorial.

Veamos una función de ejemplo que realiza estos pasos:

engine = create_async_engine(DB_URL)
Session = async_sessionmaker(engine, expire_on_commit=False)

async def add_document_to_vector_db(doc_path):
    text = extract_text(docpath)
    doc_name = os.path.splitext(os.path.basename(doc_path))[0]
    
    chunks = []
    text_splitter = TextSplitter(chunk_size=512)
    text_chunks = text_splitter.split(text)
    for idx, text_chunk in enumerate(text_chunks):
        chunks.append({
            'text': text_chunk,
            'metadata_': {'doc': doc_name, 'index': idx}
        })

    vectors = await create_embeddings([chunk['text'] for chunk in chunks])

    for chunk, vector in zip(chunks, vectors):
        chunk['vector'] = vector
    
    async with Session() as db:
        for chunk in chunks:
            db.add(Vector(**chunk))
        await db.commit()

En el código anterior:

  • La función extract_text se encarga de extraer texto del documento. Puedes usar librerías como pdfminer o pypdf para extraer texto de archivos PDF, o docx2txt para documentos Word.
  • Usamos un TextSplitter para dividir el documento en fragmentos más pequeños de texto. En este ejemplo, estamos creando fragmentos de 512 tokens. Puedes leer más sobre la funcionalidad de división (chunking) y diferentes estrategias en RAG desde cero.
  • La función create_embeddings convierte nuestros fragmentos de texto en embeddings.
  • Finalmente, creamos una sesión de base de datos y usamos el modelo Vector definido anteriormente para añadir cada embedding, junto con su texto y metadatos, a nuestra base de datos Postgres.

Para generar estos embeddings, puedes utilizar el modelo que prefieras. Pero asegúrate de que las dimensiones del modelo de embedding coincidan con las dimensiones de la columna vector del modelo (1024 en nuestro ejemplo).

Aquí tienes un ejemplo de implementación usando el modelo text-embedding-3-large de OpenAI:

from openai import AsyncOpenAI

client = AsyncOpenAI(api_key=os.environ['OPENAI_API_KEY'])

async def get_embeddings(input):
    res = await client.embeddings.create(input=input, model='text-embedding-3-large', dimensions=1024)
    return [item.embedding for item in res.data]

Búsqueda vectorial con PostgreSQL

Ahora que tenemos los documentos vectorizados y almacenados en PostgreSQL, podemos realizar búsquedas vectoriales para extraer la información más relevante para nuestras consultas. Este es un paso clave en la construcción de un sistema de Retrieval-Augmented Generation (RAG).

La siguiente función muestra cómo podemos implementar una búsqueda de similitud vectorial en PostgreSQL con pgvector:

async def vector_search(query_vector, top_k=3):
    async with Session() as db:
        query = (
            select(Vector.text, Vector.metadata_, Vector.vector.cosine_distance(query_vector).label('distance'))
            .order_by('distance')
            .limit(top_k)
        )
        res = await db.execute(query)
        return [{
            'text': text,
            'metadata': metadata,
            'score': 1 - distance
        } for text, metadata, distance in res]

Analicemos esta función y el proceso de búsqueda:

  • El parámetro query_vector es el vector embedding de nuestra consulta de búsqueda. Se genera usando el mismo modelo de embedding que utilizamos anteriormente para los fragmentos de los documentos.
  • Usamos la función cosine_distance proporcionada por pgvector para calcular la distancia entre nuestro vector de consulta y cada vector en nuestra base de datos. La distancia del coseno es una medida de disimilitud: cuanto menor sea la distancia entre vectores, mayor será la similitud con la consulta del usuario.
  • Ordenamos los resultados por distancia (ascendente) y recuperamos los top_k fragmentos más similares a nuestra consulta.
  • Para cada uno de los top_k fragmentos más similares, devolvemos el texto del fragmento, los metadatos y una puntuación de similitud que se calcula como 1 - distance.

Es importante señalar que, por defecto, pgvector realiza búsquedas exactas de vecinos más próximos (nearest neighbor search), lo que garantiza un recall perfecto (encuentra todos los vecinos más cercanos), pero puede resultar más lento cuando el volumen de datos es elevado.

En esos casos, también podemos crear un índice para acelerar las búsquedas y sacrificar un poco de exactitud a cambio de búsquedas más rápidas. Los índices de este tipo disponibles en pgvector son IVFFlat (Inverted File Flat) y HNSW (Hierarchical Navigable Small World). Puedes leer más sobre estos índices aquí.

RAG en acción

Con la funcionalidad de búsqueda vectorial lista, podemos ya integrarlo todo para crear un ejemplo básico de RAG, usando el modelo GPT-4o de OpenAI, que responda preguntas sobre los documentos que hemos procesado y almacenado en la base de datos.

Primero, veamos los prompts que vamos a utilizar:

SYSTEM_PROMPT = """
You are an AI assistant that answers questions about documents in your knowledge base.
"""

RAG_PROMPT = """
Use the following pieces of context to answer the user question.
You must only use the facts from the context to answer.
If the answer cannot be found in the context, say that you don't have enough information to answer the question and provide any relevant facts found in the context.

Context:
{context}

User Question:
{question}
"""

Y así es como podemos implementar un sistema RAG básico:

from openai import AsyncOpenAI

client = AsyncOpenAI(api_key=os.environ['OPENAI_API_KEY'])

async def answer_question_with_rag(question):
    query_vector = await get_embedding(question)
    top_chunks = await vector_search(query_vector, top_k=3)
    context = '\n\n---\n\n'.join([chunk['text'] for chunk in top_chunks]) + '\n\n---'
    user_message = RAG_PROMPT.format(context=context, question=question)
    messages=[
        {'role': 'system', 'content': SYSTEM_PROMPT},
        {'role': 'user', 'content': user_message}
    ]
    response = await client.chat.completions.create(model='gpt-4o', messages=messages)
    return response.choices[0].message.content

Esta función muestra cómo funciona la técnica RAG: convierte la pregunta del usuario en un vector embedding, realiza una búsqueda de similitud vectorial en Postgres y añade la información extraída como contexto para que GPT-4o genere una respuesta fundamentada.

Puedes usarla así:

question = "What are the main challenges in renewable energy adoption?"
answer = await answer_question_with_rag(question)
print(answer)

Ahora puedes adaptar todo el código que hemos visto a tus propias aplicaciones. Puedes procesar y almacenar tus propios documentos y usar la combinación de PostgreSQL con pgvector y GPT-4o (o cualquier otro LLM de tu elección) para responder preguntas basadas en esos documentos.

Y también puedes aprovechar estas ideas para construir aplicaciones más avanzadas, como chatbots o asistentes de IA, con una arquitectura simple que se beneficia de la potencia y versatilidad de PostgreSQL, manteniendo todos tus datos integrados en un único lugar.