Skip to content

Commit

Permalink
feat: chroma DB token
Browse files Browse the repository at this point in the history
  • Loading branch information
david20571015 committed Jun 11, 2024
1 parent ef7ee81 commit 9cb2a48
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion sync_crawler/writer/chromadb_writer.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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,
Expand All @@ -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)
Expand All @@ -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)

Expand Down

0 comments on commit 9cb2a48

Please sign in to comment.