llama-index

1
from enum import Enum
2
from typing import Dict, Type
3

4
from llama_index.legacy.graph_stores.falkordb import FalkorDBGraphStore
5
from llama_index.legacy.graph_stores.kuzu import KuzuGraphStore
6
from llama_index.legacy.graph_stores.nebulagraph import NebulaGraphStore
7
from llama_index.legacy.graph_stores.neo4j import Neo4jGraphStore
8
from llama_index.legacy.graph_stores.simple import SimpleGraphStore
9
from llama_index.legacy.graph_stores.types import GraphStore
10

11

12
class GraphStoreType(str, Enum):
13
    SIMPLE = "simple_kg"
14
    NEBULA = "nebulagraph"
15
    KUZU = "kuzu"
16
    NEO4J = "neo4j"
17
    FALKORDB = "falkordb"
18

19

20
GRAPH_STORE_TYPE_TO_GRAPH_STORE_CLASS: Dict[GraphStoreType, Type[GraphStore]] = {
21
    GraphStoreType.SIMPLE: SimpleGraphStore,
22
    GraphStoreType.NEBULA: NebulaGraphStore,
23
    GraphStoreType.KUZU: KuzuGraphStore,
24
    GraphStoreType.NEO4J: Neo4jGraphStore,
25
    GraphStoreType.FALKORDB: FalkorDBGraphStore,
26
}
27

28
GRAPH_STORE_CLASS_TO_GRAPH_STORE_TYPE: Dict[Type[GraphStore], GraphStoreType] = {
29
    cls_: type_ for type_, cls_ in GRAPH_STORE_TYPE_TO_GRAPH_STORE_CLASS.items()
30
}
31