refactor: add utility function for converting chroma collection to records
This commit is contained in:
parent
91e5aa1dc5
commit
7557c7979c
3 changed files with 28 additions and 1 deletions
0
src/backend/base/langflow/base/vectorstores/__init__.py
Normal file
0
src/backend/base/langflow/base/vectorstores/__init__.py
Normal file
24
src/backend/base/langflow/base/vectorstores/utils.py
Normal file
24
src/backend/base/langflow/base/vectorstores/utils.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue