ChromaDB + Autenticación y Colecciones

En este tutorial aprenderemos a trabajar con ChromaDB, una base de datos para manejar vectores de embeddings, con autenticación y la manipulación de colecciones.

Configuración de un servidor ChromaDB

En este caso, usaremos Docker para implementar el servidor. Para poner en marcha un servidor de ChromaDB que cuente con sistemas de autenticación, resulta imprescindible definir correctamente las variables de autenticación en el archivo `docker-compose.yml`. Es fundamental no guardar el token o clave secreta en texto plano por motivos de seguridad; la presente demostración tiene únicamente fines ilustrativos sobre cómo proceder con la configuración.

version: "3.9"
services:
  chroma:
    image: ghcr.io/chroma-core/chroma:latest
    environment:
      CHROMA_SERVER_AUTHN_CREDENTIALS: "test-token"
      CHROMA_SERVER_AUTHN_PROVIDER: "chromadb.auth.token_authn.TokenAuthServerProvider"
    volumes:
      - index_data:/chroma/.chroma/index
    ports:
      - 8001:8000
    networks:
      - net
volumes:
  index_data:
    driver: local
  backups:
    driver: local

networks:
  net:
    driver: bridge

Para iniciar el servidor ChromaDB, ejecuta el siguiente comando:

docker-compose up

Parámetros para la colección y la consulta

Definimos algunos parámetros que usaremos para la colección y la consulta:

EMBEDDINGS_MAX_RESULTS = 2
CHROMA_SERVER_AUTHN_CREDENTIALS = os.getenv('CHROMA_SERVER_AUTHN_CREDENTIALS')
CHROMA_CLIENT_AUTHN_PROVIDER = 'chromadb.auth.token_authn.TokenAuthClientProvider'
VECTOR_EMBEDDING_HOST = 'localhost'
directory = 'Zero/'

Carga y división de documentos

Creamos una función para cargar documentos de un directorio y otra para dividirlos en fragmentos:

Función para cargar documentos desde un directorio

def load_docs(directory):
    loader = DirectoryLoader(directory, glob="**/*.txt", loader_cls=TextLoader)
    documents = loader.load()
    return documents

Función para dividir documentos en fragmentos

def split_docs(documents, chunk_size=1000, chunk_overlap=20):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs = text_splitter.split_documents(documents)
    return docs

Carga y división de los documentos

documents = load_docs(directory)
docs = split_docs(documents)

Generación de Embeddings

Generamos embeddings para los fragmentos de los documentos usando un modelo preentrenado:

embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

Cliente de ChromaDB y operaciones

Para conectar con el servidor de ChromaDB y autenticarnos, primero debemos instanciar el cliente. En este ejemplo, configuramos el puerto y proveemos las credenciales de autenticación necesarias.

chroma_client = chromadb.HttpClient(
    port=8001,
    settings=Settings(
    chroma_client_auth_provider=CHROMA_CLIENT_AUTHN_PROVIDER,
    chroma_client_auth_credentials=CHROMA_SERVER_AUTHN_CREDENTIALS)
)

Creación o Recuperación de una Colección

Aquí veremos cómo verificar si una colección existe y, si no es así, cómo crearla.

collection = chroma_client.get_or_create_collection(name='myCollection')

Generación de Embeddings para los Documentos

Generamos los embeddings para los documentos, que nos permitirán luego hacer búsquedas dentro de la colección.

doc_embeddings = embeddings.embed_documents([doc.page_content for doc in docs])

Agregado de Documentos y sus Metadatos a la Colección

Una vez generados los embeddings, agregamos los documentos y sus metadatos correspondientes a la colección recién creada o recuperada.

collection.add(
    ids = [str(uuid.uuid4()) for _ in docs],
    embeddings = doc_embeddings,
    documents = [doc.page_content for doc in docs],
    metadatas = [{'timestamp': timestamp, 'chapter': 'A', 'region': 'AMER', 'book': 'REGISTRY'} for _ in docs]
)

Ejecución de una Consulta de Ejemplo

Realizamos una consulta de ejemplo donde primero creamos el embedding de la consulta y luego usamos este embedding para obtener resultados de la base de datos.

embed_query = embeddings.embed_documents('The search query')

results = collection.query(
    query_embeddings = embed_query,
    n_results= EMBEDDINGS_MAX_RESULTS,
    where= {'$and': [{'chapter': 'A'}, {'region': 'AMER'}]},
)

Finalmente, mostramos los resultados de la consulta.

print(f"Results: {results}")

Código Completo

A continuación, se presenta el código completo utilizado en este tutorial:

from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from dotenv import load_dotenv
import chromadb, uuid, os, datetime
from chromadb.config import Settings
load_dotenv()

EMBEDDINGS_MAX_RESULTS = 2
CHROMA_SERVER_AUTHN_CREDENTIALS = os.getenv('CHROMA_SERVER_AUTHN_CREDENTIALS')
CHROMA_CLIENT_AUTHN_PROVIDER = 'chromadb.auth.token_authn.TokenAuthClientProvider'
VECTOR_EMBEDDING_HOST = 'localhost'
directory = 'Zero/'

def load_docs(directory):
    loader = DirectoryLoader(directory, glob="**/*.txt", loader_cls=TextLoader)
    documents = loader.load()
    return documents

def split_docs(documents, chunk_size=1000, chunk_overlap=20):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs = text_splitter.split_documents(documents)
    return docs

documents = load_docs(directory)
docs = split_docs(documents)
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

try:
    chroma_client = chromadb.HttpClient(
        port=8001,
        settings=Settings(
        chroma_client_auth_provider=CHROMA_CLIENT_AUTHN_PROVIDER,
        chroma_client_auth_credentials=CHROMA_SERVER_AUTHN_CREDENTIALS)
    )

    collection = chroma_client.get_or_create_collection(name='myCollection')

    timestamp = datetime.datetime.now().isoformat()
    doc_embeddings = embeddings.embed_documents([doc.page_content for doc in docs])

    collection.add(
        ids = [str(uuid.uuid4()) for _ in docs],
        embeddings = doc_embeddings,
        documents = [doc.page_content for doc in docs],
        metadatas = [{'timestamp': timestamp, 'chapter': 'A', 'region': 'AMER', 'book': 'REGISTRY'} for _ in docs]
    )

    embed_query = embeddings.embed_documents('The search query')
    results = collection.query(
        query_embeddings = embed_query,
        n_results= EMBEDDINGS_MAX_RESULTS,
        where= {'$and': [{'chapter': 'A'}, {'region': 'AMER'}]},
    )

    print(f"Results: {results}")
except Exception as error:
    print(f"Error: {error}")