From b1ef28da7b04370557f92628bb3db41725b6300b Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 7 Feb 2024 22:33:19 -0300 Subject: [PATCH] Refactor ChromaSearchComponent to return records instead of documents --- src/backend/langflow/components/vectorstores/ChromaSearch.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/components/vectorstores/ChromaSearch.py b/src/backend/langflow/components/vectorstores/ChromaSearch.py index 4a68d4fa1..ce75a10d3 100644 --- a/src/backend/langflow/components/vectorstores/ChromaSearch.py +++ b/src/backend/langflow/components/vectorstores/ChromaSearch.py @@ -4,6 +4,7 @@ import chromadb # type: ignore from langchain_community.vectorstores.chroma import Chroma from langflow import CustomComponent from langflow.field_typing import Document, Embeddings, Text +from langflow.schema import Record, docs_to_records class ChromaSearchComponent(CustomComponent): @@ -61,7 +62,7 @@ class ChromaSearchComponent(CustomComponent): chroma_server_host: Optional[str] = None, chroma_server_port: Optional[int] = None, chroma_server_grpc_port: Optional[int] = None, - ) -> List[Document]: + ) -> List[Record]: """ Builds the Vector Store or BaseRetriever object. @@ -113,4 +114,4 @@ class ChromaSearchComponent(CustomComponent): docs = chroma.search(query=inputs, search_type=search_type.lower()) else: raise ValueError("Invalid inputs provided.") - return docs + return docs_to_records(docs)