DB-GPT/dbgpt/storage/graph_store/factory.py

"""Connector for vector store."""
import logging
from typing import Tuple, Type

from dbgpt.storage import graph_store
from dbgpt.storage.graph_store.base import GraphStoreBase, GraphStoreConfig

logger = logging.getLogger(__name__)


class GraphStoreFactory:
    """Factory for graph store."""

    @staticmethod
    def create(graph_store_type: str, graph_store_configure=None) -> GraphStoreBase:
        """Create a GraphStore instance.

        Args:
            - graph_store_type: graph store type Memory, TuGraph, Neo4j
            - graph_store_config: graph store config
        """
        store_cls, cfg_cls = GraphStoreFactory.__find_type(graph_store_type)

        try:
            config = cfg_cls()
            if graph_store_configure:
                graph_store_configure(config)
            return store_cls(config)
        except Exception as e:
            logger.error("create graph store failed: %s", e)
            raise e

    @staticmethod
    def __find_type(graph_store_type: str) -> Tuple[Type, Type]:
        for t in graph_store.__all__:
            if t.lower() == graph_store_type.lower():
                store_cls, cfg_cls = getattr(graph_store, t)
                if issubclass(store_cls, GraphStoreBase) and issubclass(
                    cfg_cls, GraphStoreConfig
                ):
                    return store_cls, cfg_cls
        raise Exception(f"Graph store {graph_store_type} not supported")