ChromadbRM: Provide optional client param

Author: XiaoConstantine

dspy/retrieve/chromadb_rm.py

@@ -26,9 +26,6 @@
     from chromadb.config import Settings
     from chromadb.utils import embedding_functions
 except ImportError:
-    chromadb = None
-
-if chromadb is None:
     raise ImportError(
         "The chromadb library is required to use ChromadbRM. Install it with `pip install dspy-ai[chromadb]`",
     )
@@ -73,9 +70,10 @@ def __init__(
         embedding_function: Optional[
             EmbeddingFunction[Embeddable]
         ] = ef.DefaultEmbeddingFunction(),
+        client: Optional[chromadb.Client] = None,
         k: int = 7,
     ):
-        self._init_chromadb(collection_name, persist_directory)
+        self._init_chromadb(collection_name, persist_directory, client=client)
         self.ef = embedding_function
 
         super().__init__(k=k)
@@ -84,22 +82,26 @@ def _init_chromadb(
         self,
         collection_name: str,
         persist_directory: str,
+        client: Optional[chromadb.Client] = None
     ) -> chromadb.Collection:
         """Initialize chromadb and return the loaded index.
 
         Args:
             collection_name (str): chromadb collection name
             persist_directory (str): chromadb persist directory
+            client (chromadb.Client): A chromadb client provided by user
 
-
-        Returns:
+        Returns: collection per collection_name
         """
 
-        self._chromadb_client = chromadb.Client(
-            Settings(
-                persist_directory=persist_directory,
-                is_persistent=True,
-            ),
+        if client:
+            self._chromadb_client = client
+        else:
+            self._chromadb_client = chromadb.Client(
+                Settings(
+                    persist_directory=persist_directory,
+                    is_persistent=True,
+                ),
         )
         self._chromadb_collection = self._chromadb_client.get_or_create_collection(
             name=collection_name,