refactor: add utility function for converting chroma collection to records

This commit is contained in:
ogabrielluiz 2024-06-10 12:02:53 -03:00
commit 7557c7979c
3 changed files with 28 additions and 1 deletions

View file

@ -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

View file

@ -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