From 3d0c32f366be93d146f8a15639d3a27a948e869b Mon Sep 17 00:00:00 2001
From: Diwakar Singh Maurya <diwakar.s.maurya@gmail.com>
Date: Sun, 27 Oct 2024 14:01:00 +0700
Subject: [PATCH] feature to support authentication in chromadb

---
 .../apps/retrieval/vector/dbs/chroma.py           | 15 +++++++++++++--
 backend/open_webui/config.py                      |  2 ++
 backend/requirements.txt                          |  2 +-
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/backend/open_webui/apps/retrieval/vector/dbs/chroma.py b/backend/open_webui/apps/retrieval/vector/dbs/chroma.py
index 7782671a2..86169c07e 100644
--- a/backend/open_webui/apps/retrieval/vector/dbs/chroma.py
+++ b/backend/open_webui/apps/retrieval/vector/dbs/chroma.py
@@ -13,11 +13,22 @@ from open_webui.config import (
     CHROMA_HTTP_SSL,
     CHROMA_TENANT,
     CHROMA_DATABASE,
+    CHROMA_CLIENT_AUTH_PROVIDER,
+    CHROMA_CLIENT_AUTH_CREDENTIALS,
 )
 
 
 class ChromaClient:
     def __init__(self):
+        settings_dict = {
+            "allow_reset": True,
+            "anonymized_telemetry": False,
+        }
+        if CHROMA_CLIENT_AUTH_PROVIDER is not None:
+            settings_dict["chroma_client_auth_provider"] = CHROMA_CLIENT_AUTH_PROVIDER
+        if CHROMA_CLIENT_AUTH_CREDENTIALS is not None:
+            settings_dict["chroma_client_auth_credentials"] = CHROMA_CLIENT_AUTH_CREDENTIALS
+
         if CHROMA_HTTP_HOST != "":
             self.client = chromadb.HttpClient(
                 host=CHROMA_HTTP_HOST,
@@ -26,12 +37,12 @@ class ChromaClient:
                 ssl=CHROMA_HTTP_SSL,
                 tenant=CHROMA_TENANT,
                 database=CHROMA_DATABASE,
-                settings=Settings(allow_reset=True, anonymized_telemetry=False),
+                settings=Settings(**settings_dict),
             )
         else:
             self.client = chromadb.PersistentClient(
                 path=CHROMA_DATA_PATH,
-                settings=Settings(allow_reset=True, anonymized_telemetry=False),
+                settings=Settings(**settings_dict),
                 tenant=CHROMA_TENANT,
                 database=CHROMA_DATABASE,
             )
diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py
index 9d1bd72d8..4b08a292b 100644
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@@ -937,6 +937,8 @@ CHROMA_TENANT = os.environ.get("CHROMA_TENANT", chromadb.DEFAULT_TENANT)
 CHROMA_DATABASE = os.environ.get("CHROMA_DATABASE", chromadb.DEFAULT_DATABASE)
 CHROMA_HTTP_HOST = os.environ.get("CHROMA_HTTP_HOST", "")
 CHROMA_HTTP_PORT = int(os.environ.get("CHROMA_HTTP_PORT", "8000"))
+CHROMA_CLIENT_AUTH_PROVIDER = os.environ.get("CHROMA_CLIENT_AUTH_PROVIDER", "")
+CHROMA_CLIENT_AUTH_CREDENTIALS = os.environ.get("CHROMA_CLIENT_AUTH_CREDENTIALS", "")
 # Comma-separated list of header=value pairs
 CHROMA_HTTP_HEADERS = os.environ.get("CHROMA_HTTP_HEADERS", "")
 if CHROMA_HTTP_HEADERS:
diff --git a/backend/requirements.txt b/backend/requirements.txt
index 561f291cb..565c75075 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -40,7 +40,7 @@ langchain-community==0.2.12
 langchain-chroma==0.1.4
 
 fake-useragent==1.5.1
-chromadb==0.5.9
+chromadb==0.5.15
 pymilvus==2.4.7
 qdrant-client~=1.12.0