add CommunityMetastore

dbgpt/storage/graph_store/community/community_metastore.py

@@ -0,0 +1,40 @@
+"""Community metastore."""
+from abc import ABC, abstractmethod
+from typing import List
+
+from dbgpt.datasource.rdbms.base import RDBMSConnector
+from dbgpt.storage.graph_store.community_store import Community
+from dbgpt.storage.vector_store.base import VectorStoreBase
+from dbgpt.storage.vector_store.pgvector_store import PGVectorStore
+
+
+class CommunityMetastore(ABC):
+    """Community metastore class."""
+
+    @abstractmethod
+    def get(self, community_id: str) -> Community:
+        """Get community."""
+
+    @abstractmethod
+    def search(self, query: str) -> List[Community]:
+        """search communities relevant to query."""
+
+    @abstractmethod
+    def save(self, community: Community):
+        """Upsert community."""
+
+    @abstractmethod
+    def drop(self, community_id: str):
+        """Drop community."""
+
+
+class BuiltinCommunityMetastore(CommunityMetastore):
+    rdb: RDBMSConnector
+
+    vb: VectorStoreBase
+
+
+class PGVectorCommunityMetastore(CommunityMetastore):
+    """Community metastore with vector storage."""
+
+    pg: PGVectorStore

dbgpt/storage/graph_store/community_store.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
 from typing import Dict, Generator, List
 
 from openai import OpenAI
@@ -61,6 +62,7 @@ def __del__(self):
             self.connector.close()
 
 
+@dataclass
 class Community:
     id: str
     level: str

dbgpt/storage/graph_store/factory.py

@@ -1,4 +1,4 @@
-"""Connector for vector store."""
+"""Graph store factory."""
 import logging
 from typing import Tuple, Type
 

dbgpt/storage/graph_store/graph.py

@@ -1,4 +1,4 @@
-"""Graph store base class."""
+"""Graph definition."""
 import itertools
 import json
 import logging

dbgpt/storage/graph_store/memgraph_store.py

@@ -1,4 +1,4 @@
-"""Graph store base class."""
+"""Memory graph store."""
 import json
 import logging
 from typing import List, Optional, Tuple, Generator

dbgpt/storage/graph_store/neo4j_store.py

@@ -1,4 +1,4 @@
-"""Neo4j vector store."""
+"""Neo4j store."""
 import logging
 from typing import List, Optional, Tuple, Generator
 

dbgpt/storage/graph_store/tugraph_store.py

@@ -1,4 +1,4 @@
-"""TuGraph vector store."""
+"""TuGraph store."""
 import logging
 import os
 from typing import List, Optional, Tuple, Generator, Iterator