llama-index

1
from typing import Dict, Type
2

3
from llama_index.legacy.vector_stores.chroma import ChromaVectorStore
4
from llama_index.legacy.vector_stores.lantern import LanternVectorStore
5
from llama_index.legacy.vector_stores.pinecone import PineconeVectorStore
6
from llama_index.legacy.vector_stores.postgres import PGVectorStore
7
from llama_index.legacy.vector_stores.qdrant import QdrantVectorStore
8
from llama_index.legacy.vector_stores.types import BasePydanticVectorStore
9
from llama_index.legacy.vector_stores.weaviate import WeaviateVectorStore
10

11
LOADABLE_VECTOR_STORES: Dict[str, Type[BasePydanticVectorStore]] = {
12
    ChromaVectorStore.class_name(): ChromaVectorStore,
13
    QdrantVectorStore.class_name(): QdrantVectorStore,
14
    PineconeVectorStore.class_name(): PineconeVectorStore,
15
    PGVectorStore.class_name(): PGVectorStore,
16
    WeaviateVectorStore.class_name(): WeaviateVectorStore,
17
    LanternVectorStore.class_name(): LanternVectorStore,
18
}
19

20

21
def load_vector_store(data: dict) -> BasePydanticVectorStore:
22
    if isinstance(data, BasePydanticVectorStore):
23
        return data
24
    class_name = data.pop("class_name", None)
25
    if class_name is None:
26
        raise ValueError("class_name is required to load a vector store")
27

28
    if class_name not in LOADABLE_VECTOR_STORES:
29
        raise ValueError(f"Unable to load vector store of type {class_name}")
30

31
    # pop unused keys
32
    data.pop("flat_metadata", None)
33
    data.pop("stores_text", None)
34
    data.pop("is_embedding_query", None)
35

36
    if class_name == WeaviateVectorStore.class_name():
37
        import weaviate
38

39
        auth_config_dict = data.pop("auth_config", None)
40
        if auth_config_dict is not None:
41
            auth_config = None
42
            if "api_key" in auth_config_dict:
43
                auth_config = weaviate.AuthApiKey(**auth_config_dict)
44
            elif "username" in auth_config_dict:
45
                auth_config = weaviate.AuthClientPassword(**auth_config_dict)
46
            else:
47
                raise ValueError(
48
                    "Unable to load weaviate auth config, please use an auth "
49
                    "config with an api_key or username/password."
50
                )
51

52
            data["auth_config"] = auth_config
53

54
    return LOADABLE_VECTOR_STORES[class_name](**data)  # type: ignore
55