chroma token-based auth

langgenius · Jun 9, 2024 · 9aa2986 · 9aa2986
1 parent f1a75ff
commit 9aa2986
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/api/core/rag/datasource/vdb/chroma/chroma_vector.py b/api/core/rag/datasource/vdb/chroma/chroma_vector.py
@@ -1,9 +1,9 @@
 import json
 import logging
-from typing import Any
+from typing import Any, Optional
 
 import chromadb
-from chromadb import QueryResult
+from chromadb import QueryResult, Settings
 from pydantic import BaseModel
 
 from core.rag.datasource.entity.embedding import Embeddings
@@ -20,14 +20,24 @@ class ChromaConfig(BaseModel):
     port: int
     tenant: str
     database: str
+    auth_provider: Optional[str] = None
+    auth_credentials: Optional[str] = None
 
     def to_chroma_params(self):
+        # settings=Settings(chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider",
+        #                   chroma_client_auth_credentials="test-token")
+        settings = Settings()
+        if self.auth_provider:
+            settings.chroma_client_auth_provider = self.auth_provider
+            settings.chroma_client_auth_credentials = self.auth_credentials
+
         return {
             'host': self.host,
             'port': self.port,
             'ssl': False,
             'tenant': self.tenant,
             'database': self.database,
+            'settings': settings,
         }
 
 
@@ -141,4 +151,4 @@ def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings
                 tenant=config.get('CHROMA_TENANT', chromadb.DEFAULT_TENANT),
                 database=config.get('CHROMA_DATABASE', chromadb.DEFAULT_DATABASE),
             ),
-        )
+        )
diff --git a/api/tests/integration_tests/vdb/chroma/test_chroma.py b/api/tests/integration_tests/vdb/chroma/test_chroma.py
@@ -18,6 +18,8 @@ def __init__(self):
                 port=8000,
                 tenant=chromadb.DEFAULT_TENANT,
                 database=chromadb.DEFAULT_DATABASE,
+                auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider",
+                auth_credentials="difyai123456",
             )
         )
 

diff --git a/docker/docker-compose.chroma.yaml b/docker/docker-compose.chroma.yaml
@@ -7,6 +7,8 @@ services:
     volumes:
       - ./volumes/chroma:/chroma/chroma
     environment:
-      - IS_PERSISTENT=TRUE
+      IS_PERSISTENT: TRUE
+      CHROMA_SERVER_AUTHN_PROVIDER: chromadb.auth.token_authn.TokenAuthenticationServerProvider
+      CHROMA_SERVER_AUTHN_CREDENTIALS: difyai123456
     ports:
-      - 8000:8000
+      - "8000:8000"