Fix ChromaComponent to use embedding correctly

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-26 16:07:08 -03:00
commit 66b0a9046f

View file

@ -5,6 +5,7 @@ from langchain.embeddings.base import Embeddings
from langchain.schema import BaseRetriever, Document
from langchain_community.vectorstores import VectorStore
from langchain_community.vectorstores.chroma import Chroma
from langflow import CustomComponent
@ -83,7 +84,8 @@ class ChromaComponent(CustomComponent):
if chroma_server_host is not None:
chroma_settings = chromadb.config.Settings(
chroma_server_cors_allow_origins=chroma_server_cors_allow_origins or None,
chroma_server_cors_allow_origins=chroma_server_cors_allow_origins
or None,
chroma_server_host=chroma_server_host,
chroma_server_port=chroma_server_port or None,
chroma_server_grpc_port=chroma_server_grpc_port or None,
@ -98,14 +100,20 @@ class ChromaComponent(CustomComponent):
if documents is not None and embedding is not None:
if len(documents) == 0:
raise ValueError("If documents are provided, there must be at least one document.")
raise ValueError(
"If documents are provided, there must be at least one document."
)
chroma = Chroma.from_documents(
documents=documents, # type: ignore
persist_directory=index_directory,
collection_name=collection_name,
embedding=embedding,
embedding_function=embedding,
client_settings=chroma_settings,
)
else:
chroma = Chroma(persist_directory=index_directory, client_settings=chroma_settings)
chroma = Chroma(
persist_directory=index_directory,
client_settings=chroma_settings,
embedding_function=embedding,
)
return chroma