55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
from typing import Dict, Type
|
|
|
|
from llama_index.vector_stores.chroma import ChromaVectorStore
|
|
from llama_index.vector_stores.lantern import LanternVectorStore
|
|
from llama_index.vector_stores.pinecone import PineconeVectorStore
|
|
from llama_index.vector_stores.postgres import PGVectorStore
|
|
from llama_index.vector_stores.qdrant import QdrantVectorStore
|
|
from llama_index.vector_stores.types import BasePydanticVectorStore
|
|
from llama_index.vector_stores.weaviate import WeaviateVectorStore
|
|
|
|
LOADABLE_VECTOR_STORES: Dict[str, Type[BasePydanticVectorStore]] = {
|
|
ChromaVectorStore.class_name(): ChromaVectorStore,
|
|
QdrantVectorStore.class_name(): QdrantVectorStore,
|
|
PineconeVectorStore.class_name(): PineconeVectorStore,
|
|
PGVectorStore.class_name(): PGVectorStore,
|
|
WeaviateVectorStore.class_name(): WeaviateVectorStore,
|
|
LanternVectorStore.class_name(): LanternVectorStore,
|
|
}
|
|
|
|
|
|
def load_vector_store(data: dict) -> BasePydanticVectorStore:
|
|
if isinstance(data, BasePydanticVectorStore):
|
|
return data
|
|
class_name = data.pop("class_name", None)
|
|
if class_name is None:
|
|
raise ValueError("class_name is required to load a vector store")
|
|
|
|
if class_name not in LOADABLE_VECTOR_STORES:
|
|
raise ValueError(f"Unable to load vector store of type {class_name}")
|
|
|
|
# pop unused keys
|
|
data.pop("flat_metadata", None)
|
|
data.pop("stores_text", None)
|
|
data.pop("is_embedding_query", None)
|
|
|
|
if class_name == WeaviateVectorStore.class_name():
|
|
import weaviate
|
|
|
|
auth_config_dict = data.pop("auth_config", None)
|
|
if auth_config_dict is not None:
|
|
auth_config = None
|
|
if "api_key" in auth_config_dict:
|
|
auth_config = weaviate.AuthApiKey(**auth_config_dict)
|
|
elif "username" in auth_config_dict:
|
|
auth_config = weaviate.AuthClientPassword(**auth_config_dict)
|
|
else:
|
|
raise ValueError(
|
|
"Unable to load weaviate auth config, please use an auth "
|
|
"config with an api_key or username/password."
|
|
)
|
|
|
|
data["auth_config"] = auth_config
|
|
|
|
return LOADABLE_VECTOR_STORES[class_name](**data) # type: ignore
|