diff --git a/src/backend/langflow/interface/initialize/vector_store.py b/src/backend/langflow/interface/initialize/vector_store.py index f3a10d0ba..c1ce5c227 100644 --- a/src/backend/langflow/interface/initialize/vector_store.py +++ b/src/backend/langflow/interface/initialize/vector_store.py @@ -1,15 +1,20 @@ -from typing import Any +from typing import Type +from langchain.vectorstores import Pinecone, Qdrant, Chroma def docs_in_params(params: dict) -> bool: - """Check if params has documents OR texts and one of them is not an empty list""" - return ("documents" in params and params["documents"] != []) or ( - "texts" in params and params["texts"] != [] + """Check if params has documents OR texts and one of them is not an empty list, + If any of them is not an empty list, return True, else return False""" + return ("documents" in params and params["documents"]) or ( + "texts" in params and params["texts"] ) -def initialize_pinecone(class_object: Any, params: dict): +def initialize_pinecone(class_object: Type[Pinecone], params: dict): """Initialize pinecone and return the class object""" + + import pinecone + PINECONE_API_KEY = params.get("pinecone_api_key") PINECONE_ENV = params.get("pinecone_env") @@ -18,23 +23,31 @@ def initialize_pinecone(class_object: Any, params: dict): "Pinecone API key and environment must be provided in the params" ) - if not docs_in_params(params): - import pinecone + # initialize pinecone + pinecone.init( + api_key=PINECONE_API_KEY, # find at app.pinecone.io + environment=PINECONE_ENV, # next to api key in console + ) - # initialize pinecone - pinecone.init( - api_key=PINECONE_API_KEY, # find at app.pinecone.io - environment=PINECONE_ENV, # next to api key in console - ) - # If there are docs in the index, delete them - return class_object.from_existing_index(**params) + # If there are no docs in the params, return an existing index + # but first remove any texts or docs keys from the params + if not docs_in_params(params): + existing_index_params = { + "embedding": params.pop("embedding"), + } + if "index_name" in params: + existing_index_params["index_name"] = params.pop("index_name") + if "namespace" in params: + existing_index_params["namespace"] = params.pop("namespace") + + return class_object.from_existing_index(**existing_index_params) # If there are docs in the params, create a new index if "texts" in params: params["documents"] = params.pop("texts") return class_object.from_documents(**params) -def initialize_chroma(class_object, params): +def initialize_chroma(class_object: Type[Chroma], params: dict): """Initialize a ChromaDB object from the params""" persist = params.pop("persist", False) if not docs_in_params(params): @@ -54,7 +67,7 @@ def initialize_chroma(class_object, params): return chromadb -def initialize_qdrant(class_object, params): +def initialize_qdrant(class_object: Type[Qdrant], params: dict): if not docs_in_params(params): if "location" not in params and "api_key" not in params: raise ValueError("Location and API key must be provided in the params")