🚀 feat(vector_store.py): add support for Weaviate and FAISS vector stores

This commit adds support for Weaviate and FAISS vector stores to the existing Pinecone, Qdrant, and Chroma vector stores. The `initialize_weaviate` function initializes Weaviate and returns the class object. The `initialize_faiss` function initializes FAISS and returns the class object. These functions are used to initialize the respective vector stores. The `docs_in_params` function is used to check if there are documents in the parameters.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-22 19:59:33 -03:00
commit 9b4e65c121

View file

@ -1,5 +1,6 @@
import json
from typing import Type
from langchain.vectorstores import Pinecone, Qdrant, Chroma
from langchain.vectorstores import Pinecone, Qdrant, Chroma, FAISS, Weaviate
def docs_in_params(params: dict) -> bool:
@ -10,6 +11,45 @@ def docs_in_params(params: dict) -> bool:
)
def initialize_weaviate(class_object: Type[Weaviate], params: dict):
"""Initialize weaviate and return the class object"""
if not docs_in_params(params):
import weaviate
client_kwargs_json = params.get("client_kwargs", "{}")
client_kwargs = json.loads(client_kwargs_json)
client_params = {
"url": params.get("weaviate_url"),
}
client_params.update(client_kwargs)
weaviate_client = weaviate.Client(**client_params)
new_params = {
"client": weaviate_client,
"index_name": params.get("index_name"),
"text_key": params.get("text_key"),
}
weaviate = class_object(**new_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_faiss(class_object: Type[FAISS], params: dict):
"""Initialize faiss and return the class object"""
if not docs_in_params(params):
return class_object.load_local
save_local = params.get("save_local")
faiss_index = class_object(**params)
if save_local:
faiss_index.save_local(folder_path=save_local)
return faiss_index
def initialize_pinecone(class_object: Type[Pinecone], params: dict):
"""Initialize pinecone and return the class object"""