ChromaDB + Autenticazione e Collezioni

In questo tutorial impareremo a lavorare con ChromaDB, un database per gestire vettori di embedding, con l’autenticazione e la gestione delle collezioni.

Configurazione di un server ChromaDB

In questo caso useremo Docker per distribuire il server. Per avviare un server ChromaDB con autenticazione è indispensabile definire correttamente le variabili di autenticazione nel file `docker-compose.yml`. È fondamentale non salvare il token o la chiave segreta in chiaro, per motivi di sicurezza; questa dimostrazione serve solo a illustrare come procedere con la configurazione.

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

Per avviare il server ChromaDB, esegui il seguente comando:

docker-compose up

Parametri per la collezione e la query

Definiamo alcuni parametri che useremo per la collezione e la query:

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

Caricamento e divisione dei documenti

Creiamo una funzione per caricare i documenti da una cartella e un’altra per dividerli in frammenti:

Funzione per caricare documenti da una cartella

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

Funzione per dividere i documenti in frammenti

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

Caricamento e divisione dei documenti

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

Generazione degli embedding

Generiamo gli embedding per i frammenti dei documenti usando un modello pre-addestrato:

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

Client di ChromaDB e operazioni

Per collegarci al server ChromaDB e autenticarci, dobbiamo prima istanziare il client. In questo esempio configuriamo la porta e forniamo le credenziali di autenticazione necessarie.

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

Creazione o recupero di una collezione

Qui vedremo come verificare se una collezione esiste e, in caso contrario, come crearla.

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

Generazione degli embedding per i documenti

Generiamo gli embedding dei documenti, che ci permetteranno poi di fare ricerche all’interno della collezione.

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

Aggiunta dei documenti e dei loro metadati alla collezione

Una volta generati gli embedding, aggiungiamo i documenti e i relativi metadati alla collezione appena creata o recuperata.

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

Esecuzione di una query di esempio

Eseguiamo una query di esempio: prima creiamo l’embedding della query, poi lo usiamo per ottenere i risultati dal database.

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

Infine mostriamo i risultati della query.

print(f"Results: {results}")

Codice completo

Ecco il codice completo usato in questo 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}")