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 index 72d26dfb4..6bee093e3 100644 --- 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.