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

Refactor ChromaSearchComponent configuration

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-15 18:48:29 -03:00
commit dc055c5ee6
1 changed files with 21 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -3,7 +3,7 @@ from typing import List, Optional
import chromadb # type: ignore 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 Embeddings, Text
from langflow.schema import Record, docs_to_records from langflow.schema import Record, docs_to_records
@ -15,6 +15,7 @@ class ChromaSearchComponent(CustomComponent):
display_name: str = "Chroma Search" display_name: str = "Chroma Search"
description: str = "Search a Chroma collection for similar documents." description: str = "Search a Chroma collection for similar documents."
beta: bool = True beta: bool = True
icon = "Chroma"
def build_config(self): def build_config(self):
""" """
@ -25,13 +26,19 @@ class ChromaSearchComponent(CustomComponent):
""" """
return { return {
"inputs": {"display_name": "Input"}, "inputs": {"display_name": "Input"},
"search_type": {"display_name": "Search Type", "options": ["Similarity", "MMR"]}, "search_type": {
"display_name": "Search Type",
"options": ["Similarity", "MMR"],
},
"collection_name": {"display_name": "Collection Name", "value": "langflow"}, "collection_name": {"display_name": "Collection Name", "value": "langflow"},
"persist": {"display_name": "Persist"}, # "persist": {"display_name": "Persist"},
"persist_directory": {"display_name": "Persist Directory"}, "index_directory": {"display_name": "Index Directory"},
"code": {"show": False, "display_name": "Code"}, "code": {"show": False, "display_name": "Code"},
"documents": {"display_name": "Documents", "is_list": True}, "documents": {"display_name": "Documents", "is_list": True},
"embedding": {"display_name": "Embedding"}, "embedding": {
"display_name": "Embedding",
"info": "Embedding model to vectorize inputs (make sure to use same as index)",
},
"chroma_server_cors_allow_origins": { "chroma_server_cors_allow_origins": {
"display_name": "Server CORS Allow Origins", "display_name": "Server CORS Allow Origins",
"advanced": True, "advanced": True,
@ -53,11 +60,9 @@ class ChromaSearchComponent(CustomComponent):
inputs: Text, inputs: Text,
search_type: str, search_type: str,
collection_name: str, collection_name: str,
persist: bool,
embedding: Embeddings, embedding: Embeddings,
chroma_server_ssl_enabled: bool, chroma_server_ssl_enabled: bool,
persist_directory: Optional[str] = None, index_directory: Optional[str] = None,
documents: Optional[List[Document]] = 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,
@ -87,26 +92,20 @@ class ChromaSearchComponent(CustomComponent):
if chroma_server_host is not None: if chroma_server_host is not None:
chroma_settings = chromadb.config.Settings( 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_host=chroma_server_host,
chroma_server_port=chroma_server_port or None, chroma_server_port=chroma_server_port or None,
chroma_server_grpc_port=chroma_server_grpc_port or None, chroma_server_grpc_port=chroma_server_grpc_port or None,
chroma_server_ssl_enabled=chroma_server_ssl_enabled, chroma_server_ssl_enabled=chroma_server_ssl_enabled,
) )
# If documents, then we need to create a Chroma instance using .from_documents chroma = Chroma(
if documents is not None and embedding is not None: embedding_function=embedding,
if len(documents) == 0: collection_name=collection_name,
raise ValueError("If documents are provided, there must be at least one document.") persist_directory=index_directory,
chroma = Chroma.from_documents( client_settings=chroma_settings,
documents=documents, # type: ignore )
persist_directory=persist_directory if persist else None,
collection_name=collection_name,
embedding=embedding,
client_settings=chroma_settings,
)
else:
chroma = Chroma(persist_directory=persist_directory, client_settings=chroma_settings)
# Validate the inputs # Validate the inputs
docs = [] docs = []