jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor ChromaSearchComponent to return records instead of documents

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-07 22:33:19 -03:00
commit b1ef28da7b
1 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/backend/langflow/components/vectorstores/ChromaSearch.py
View file

@ -4,6 +4,7 @@ import chromadb # type: ignore
from langchain_community.vectorstores.chroma import Chroma from langchain_community.vectorstores.chroma import Chroma
from langflow import CustomComponent from langflow import CustomComponent
from langflow.field_typing import Document, Embeddings, Text from langflow.field_typing import Document, Embeddings, Text
from langflow.schema import Record, docs_to_records
class ChromaSearchComponent(CustomComponent): class ChromaSearchComponent(CustomComponent):
@ -61,7 +62,7 @@ class ChromaSearchComponent(CustomComponent):
chroma_server_host: Optional[str] = None, chroma_server_host: Optional[str] = None,
chroma_server_port: Optional[int] = None, chroma_server_port: Optional[int] = None,
chroma_server_grpc_port: Optional[int] = None, chroma_server_grpc_port: Optional[int] = None,
) -> List[Document]: ) -> List[Record]:
""" """
Builds the Vector Store or BaseRetriever object. Builds the Vector Store or BaseRetriever object.
@ -113,4 +114,4 @@ class ChromaSearchComponent(CustomComponent):
docs = chroma.search(query=inputs, search_type=search_type.lower()) docs = chroma.search(query=inputs, search_type=search_type.lower())
else: else:
raise ValueError("Invalid inputs provided.") raise ValueError("Invalid inputs provided.")
return docs return docs_to_records(docs)