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

Refactor ChromaComponent to accept inputs as Records

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-05 19:17:04 -03:00
commit 7c91e69004
1 changed files with 9 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

11
src/backend/langflow/components/vectorstores/Chroma.py
View file

@ -2,11 +2,12 @@ from typing import List, Optional, Union
import chromadb # type: ignore import chromadb # type: ignore
from langchain.embeddings.base import Embeddings from langchain.embeddings.base import Embeddings
from langchain.schema import BaseRetriever, Document from langchain.schema import BaseRetriever
from langchain_community.vectorstores import VectorStore from langchain_community.vectorstores import VectorStore
from langchain_community.vectorstores.chroma import Chroma from langchain_community.vectorstores.chroma import Chroma
from langflow import CustomComponent from langflow import CustomComponent
from langflow.schema.schema import Record
class ChromaComponent(CustomComponent): class ChromaComponent(CustomComponent):
@ -55,7 +56,7 @@ class ChromaComponent(CustomComponent):
embedding: Embeddings, embedding: Embeddings,
chroma_server_ssl_enabled: bool, chroma_server_ssl_enabled: bool,
index_directory: Optional[str] = None, index_directory: Optional[str] = None,
documents: Optional[List[Document]] = None, inputs: Optional[List[Record]] = None,
chroma_server_cors_allow_origins: Optional[str] = None, chroma_server_cors_allow_origins: Optional[str] = None,
chroma_server_host: Optional[str] = None, chroma_server_host: Optional[str] = None,
chroma_server_port: Optional[int] = None, chroma_server_port: Optional[int] = None,
@ -98,6 +99,12 @@ class ChromaComponent(CustomComponent):
if index_directory is not None: if index_directory is not None:
index_directory = self.resolve_path(index_directory) index_directory = self.resolve_path(index_directory)
documents = []
for _input in inputs:
if isinstance(_input, Record):
documents.append(_input.to_lc_document())
else:
documents.append(_input)
if documents is not None and embedding is not None: if documents is not None and embedding is not None:
if len(documents) == 0: if len(documents) == 0:
raise ValueError( raise ValueError(