undo formating

deepsense-ai · Oct 18, 2024 · 21ea0f4 · 21ea0f4
1 parent b46edba
commit 21ea0f4
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/packages/ragbits-core/src/ragbits/core/vector_store/chromadb_store.py b/packages/ragbits-core/src/ragbits/core/vector_store/chromadb_store.py
@@ -20,8 +20,8 @@ class ChromaDBStore(VectorStore):
     def __init__(
         self,
         index_name: str,
-        chroma_client: chromadb.ClientAPI,
-        embedding_function: Union[Embeddings, chromadb.EmbeddingFunction],
+        chroma_client: "chromadb.ClientAPI",
+        embedding_function: Union[Embeddings, "chromadb.EmbeddingFunction"],
         max_distance: Optional[float] = None,
         distance_method: Literal["l2", "ip", "cosine"] = "l2",
     ):
@@ -72,7 +72,7 @@ def from_config(cls, config: dict) -> "ChromaDBStore":
             distance_method=config.get("distance_method", "l2"),
         )
 
-    def _get_chroma_collection(self) -> chromadb.Collection:
+    def _get_chroma_collection(self) -> "chromadb.Collection":
         """
         Based on the selected embedding_function, chooses how to retrieve the ChromaDB collection.
         If the collection doesn't exist, it creates one.
@@ -116,7 +116,7 @@ def _process_db_entry(self, entry: VectorDBEntry) -> tuple[str, list[float], dic
         return doc_id, embedding, metadata
 
     @property
-    def embedding_function(self) -> Union[Embeddings, chromadb.EmbeddingFunction]:
+    def embedding_function(self) -> Union[Embeddings, "chromadb.EmbeddingFunction"]:
         """
         Returns the embedding function.