diff --git a/sync_crawler/writer/chromadb_writer.py b/sync_crawler/writer/chromadb_writer.py index 7402208..94149f5 100644 --- a/sync_crawler/writer/chromadb_writer.py +++ b/sync_crawler/writer/chromadb_writer.py @@ -1,6 +1,7 @@ from typing import override import chromadb +from chromadb.config import Settings from llama_index.core import Document, VectorStoreIndex from llama_index.embeddings.huggingface import HuggingFaceEmbedding from llama_index.vector_stores.chroma import ChromaVectorStore @@ -14,6 +15,7 @@ def __init__( self, host: str = "localhost", port: str = "8000", + token: str = "", collection: str = "news", embedding_model: str = "sentence-transformers/distiluse-base-multilingual-cased-v1", in_memory: bool = False, @@ -23,6 +25,7 @@ def __init__( Args: host: Host of ChromaDB server. port: Port of ChromaDB server. + token: Token for authentication. collection: Name of collection. embedding_model: Name of embedding model. All available models can be found [here](https://huggingface.co/models?language=zh) @@ -32,7 +35,14 @@ def __init__( if in_memory: self._client = chromadb.EphemeralClient() else: - self._client = chromadb.HttpClient(host=host, port=port) + self._client = chromadb.HttpClient( + host=host, + port=port, + settings=Settings( + chroma_client_auth_provider="chromadb.auth.token.TokenAuthClientProvider", + chroma_client_auth_credentials=token, + ), + ) self._collection = self._client.get_or_create_collection(collection)