ChromaDB + Autenticação e Coleções

Neste tutorial vamos aprender a trabalhar com o ChromaDB, um banco de dados para gerenciar vetores de embeddings, com autenticação e manipulação de coleções.

Configuração de um servidor ChromaDB

Neste caso, usaremos o Docker para implantar o servidor. Para colocar em funcionamento um servidor ChromaDB com autenticação, é imprescindível definir corretamente as variáveis de autenticação no arquivo `docker-compose.yml`. É fundamental não guardar o token ou a chave secreta em texto puro, por motivos de segurança; esta demonstração serve apenas para ilustrar como fazer a configuração.

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 o servidor ChromaDB, execute o seguinte comando:

docker-compose up

Parâmetros da coleção e da consulta

Definimos alguns parâmetros que usaremos para a coleção e a 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/'

Carregamento e divisão de documentos

Criamos uma função para carregar documentos de um diretório e outra para dividi-los em fragmentos:

Função para carregar documentos de um diretório

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

Função para dividir documentos em 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

Carregamento e divisão dos documentos

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

Geração de embeddings

Geramos embeddings para os fragmentos dos documentos usando um modelo pré-treinado:

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

Cliente do ChromaDB e operações

Para conectar ao servidor ChromaDB e nos autenticar, primeiro precisamos instanciar o cliente. Neste exemplo, configuramos a porta e fornecemos as credenciais de autenticação necessárias.

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

Criação ou recuperação de uma coleção

Aqui veremos como verificar se uma coleção existe e, caso não exista, como criá-la.

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

Geração de embeddings para os documentos

Geramos os embeddings dos documentos, que depois nos permitirão fazer buscas dentro da coleção.

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

Adição dos documentos e seus metadados à coleção

Com os embeddings gerados, adicionamos os documentos e seus metadados à coleção recém-criada ou 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]
)

Execução de uma consulta de exemplo

Fazemos uma consulta de exemplo: primeiro criamos o embedding da consulta e depois usamos esse embedding para obter resultados do banco de dados.

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'}]},
)

Por fim, exibimos os resultados da consulta.

print(f"Results: {results}")

Código completo

A seguir, o código completo usado neste 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}")