From bfef9495f19183bb38c441f609070660d9c099fe Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 22 Jun 2023 19:03:53 -0300 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(vector=5Fstore.py):=20add=20su?= =?UTF-8?q?pport=20for=20initializing=20Pinecone=20and=20ChromaDB=20object?= =?UTF-8?q?s=20from=20existing=20indexes=20The=20code=20has=20been=20refac?= =?UTF-8?q?tored=20to=20improve=20readability=20by=20adding=20type=20hints?= =?UTF-8?q?=20and=20reformatting=20the=20code.=20The=20initialize=5Fpineco?= =?UTF-8?q?ne=20and=20initialize=5Fchroma=20functions=20have=20been=20upda?= =?UTF-8?q?ted=20to=20support=20initializing=20objects=20from=20existing?= =?UTF-8?q?=20indexes.=20If=20there=20are=20no=20docs=20in=20the=20params,?= =?UTF-8?q?=20the=20functions=20will=20return=20an=20existing=20index.=20I?= =?UTF-8?q?f=20there=20are=20docs=20in=20the=20params,=20the=20functions?= =?UTF-8?q?=20will=20create=20a=20new=20index.=20The=20initialize=5Fqdrant?= =?UTF-8?q?=20function=20has=20not=20been=20changed.=20=F0=9F=94=A8=20refa?= =?UTF-8?q?ctor(vector=5Fstore.py):=20improve=20code=20readability=20by=20?= =?UTF-8?q?adding=20type=20hints=20and=20reformatting=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../interface/initialize/vector_store.py | 45 ++++++++++++------- 1 file changed, 29 insertions(+), 16 deletions(-) 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")