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.embeddings import Embeddings
|
||||||
from langchain_core.retrievers import BaseRetriever
|
from langchain_core.retrievers import BaseRetriever
|
||||||
from langchain_core.vectorstores import VectorStore
|
from langchain_core.vectorstores import VectorStore
|
||||||
|
from langflow.base.vectorstores.utils import chroma_collection_to_records
|
||||||
from langflow.custom import CustomComponent
|
from langflow.custom import CustomComponent
|
||||||
from langflow.schema import Record
|
from langflow.schema import Record
|
||||||
|
|
||||||
|
|
@ -121,4 +121,7 @@ class ChromaComponent(CustomComponent):
|
||||||
client=client,
|
client=client,
|
||||||
embedding_function=embedding,
|
embedding_function=embedding,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
store = chroma.get()
|
||||||
|
self.status = chroma_collection_to_records(store)
|
||||||
return chroma
|
return chroma
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue