ChromaDB + Authentification et Collections

Dans ce tutoriel, nous allons apprendre à travailler avec ChromaDB, une base de données pour gérer des vecteurs d’embeddings, avec authentification et manipulation de collections.

Configuration d’un serveur ChromaDB

Ici, nous utiliserons Docker pour déployer le serveur. Pour lancer un serveur ChromaDB avec authentification, il est indispensable de définir correctement les variables d’authentification dans le fichier `docker-compose.yml`. Il est essentiel de ne pas stocker le token ou la clé secrète en clair, pour des raisons de sécurité ; cette démonstration sert uniquement à illustrer la marche à suivre pour la configuration.

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

Pour démarrer le serveur ChromaDB, exécutez la commande suivante :

docker-compose up

Paramètres de la collection et de la requête

Nous définissons quelques paramètres que nous utiliserons pour la collection et la requête :

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/'

Chargement et découpage des documents

Nous créons une fonction pour charger les documents d’un répertoire et une autre pour les découper en fragments :

Fonction pour charger les documents d’un répertoire

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

Fonction pour découper les documents en fragments

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

Chargement et découpage des documents

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

Génération des embeddings

Nous générons des embeddings pour les fragments des documents à l’aide d’un modèle pré-entraîné :

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

Client ChromaDB et opérations

Pour se connecter au serveur ChromaDB et s’authentifier, il faut d’abord instancier le client. Dans cet exemple, nous configurons le port et fournissons les identifiants d’authentification nécessaires.

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

Création ou récupération d’une collection

Nous allons voir comment vérifier si une collection existe et, si ce n’est pas le cas, comment la créer.

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

Génération des embeddings des documents

Nous générons les embeddings des documents, qui nous permettront ensuite d’effectuer des recherches dans la collection.

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

Ajout des documents et de leurs métadonnées à la collection

Une fois les embeddings générés, nous ajoutons les documents et leurs métadonnées à la collection que nous venons de créer ou de récupérer.

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]
)

Exécution d’une requête d’exemple

Nous exécutons une requête d’exemple : nous créons d’abord l’embedding de la requête, puis nous l’utilisons pour obtenir des résultats de la base de données.

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

Enfin, nous affichons les résultats de la requête.

print(f"Results: {results}")

Code complet

Voici le code complet utilisé dans ce tutoriel :

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}")