diff --git a/src/backend/base/langflow/base/vectorstores/__init__.py b/src/backend/base/langflow/base/vectorstores/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/backend/base/langflow/base/vectorstores/utils.py b/src/backend/base/langflow/base/vectorstores/utils.py new file mode 100644 index 000000000..e8d93ff37 --- /dev/null +++ b/src/backend/base/langflow/base/vectorstores/utils.py @@ -0,0 +1,24 @@ +from langflow.schema import Record + + +def chroma_collection_to_records(collection_dict: dict): + """ + Converts a collection of chroma vectors into a list of records. + + Args: + collection_dict (dict): A dictionary containing the collection of chroma vectors. + + Returns: + list: A list of records, where each record represents a document in the collection. + """ + records = [] + for i, doc in enumerate(collection_dict["documents"]): + record_dict = { + "id": collection_dict["ids"][i], + "document": doc, + } + if "metadatas" in collection_dict: + for key, value in collection_dict["metadatas"][i].items(): + record_dict[key] = value + records.append(Record(**record_dict)) + return records diff --git a/src/backend/base/langflow/components/vectorstores/Chroma.py b/src/backend/base/langflow/components/vectorstores/Chroma.py index 5742aad7b..b7cfad7d4 100644 --- a/src/backend/base/langflow/components/vectorstores/Chroma.py +++ b/src/backend/base/langflow/components/vectorstores/Chroma.py @@ -6,7 +6,7 @@ from langchain_chroma import Chroma from langchain_core.embeddings import Embeddings from langchain_core.retrievers import BaseRetriever from langchain_core.vectorstores import VectorStore - +from langflow.base.vectorstores.utils import chroma_collection_to_records from langflow.custom import CustomComponent from langflow.schema import Record @@ -121,4 +121,7 @@ class ChromaComponent(CustomComponent): client=client, embedding_function=embedding, ) + + store = chroma.get() + self.status = chroma_collection_to_records(store) return chroma