From f98be96a8c4d7d3a6fc0ab73017f0c6fd95b7fec Mon Sep 17 00:00:00 2001
From: Florian <fanzhidongyzby@163.com>
Date: Fri, 16 Aug 2024 11:43:44 +0800
Subject: [PATCH] add CommunityMetastore

---
 .../community/community_metastore.py          | 44 +++++++++++++++++++
 dbgpt/storage/graph_store/community_store.py  |  2 +
 dbgpt/storage/graph_store/factory.py          |  2 +-
 dbgpt/storage/graph_store/graph.py            |  2 +-
 dbgpt/storage/graph_store/memgraph_store.py   |  2 +-
 dbgpt/storage/graph_store/neo4j_store.py      |  2 +-
 dbgpt/storage/graph_store/tugraph_store.py    |  2 +-
 7 files changed, 51 insertions(+), 5 deletions(-)
 create mode 100644 dbgpt/storage/graph_store/community/community_metastore.py

diff --git a/dbgpt/storage/graph_store/community/community_metastore.py b/dbgpt/storage/graph_store/community/community_metastore.py
new file mode 100644
index 000000000..027c05200
--- /dev/null
+++ b/dbgpt/storage/graph_store/community/community_metastore.py
@@ -0,0 +1,44 @@
+"""Community metastore."""
+from abc import ABC, abstractmethod
+from typing import List
+
+from dbgpt.storage.graph_store.community_store import Community, CommunityStore
+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 PGVectorCommunityMetastore(CommunityMetastore):
+    """Community metastore with vector storage."""
+
+    vectorStore: PGVectorStore
+
+    def get(self, community_id: str) -> Community:
+        pass
+
+    def search(self, query: str) -> List[Community]:
+        pass
+
+    def save(self, community: Community):
+        pass
+
+    def drop(self, community_id: str):
+        pass
diff --git a/dbgpt/storage/graph_store/community_store.py b/dbgpt/storage/graph_store/community_store.py
index fbae04e29..97c861f41 100644
--- a/dbgpt/storage/graph_store/community_store.py
+++ b/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
diff --git a/dbgpt/storage/graph_store/factory.py b/dbgpt/storage/graph_store/factory.py
index 1e3d70804..cf598966a 100644
--- a/dbgpt/storage/graph_store/factory.py
+++ b/dbgpt/storage/graph_store/factory.py
@@ -1,4 +1,4 @@
-"""Connector for vector store."""
+"""Graph store factory."""
 import logging
 from typing import Tuple, Type
 
diff --git a/dbgpt/storage/graph_store/graph.py b/dbgpt/storage/graph_store/graph.py
index dedb3576b..2e9f88a95 100644
--- a/dbgpt/storage/graph_store/graph.py
+++ b/dbgpt/storage/graph_store/graph.py
@@ -1,4 +1,4 @@
-"""Graph store base class."""
+"""Graph definition."""
 import itertools
 import json
 import logging
diff --git a/dbgpt/storage/graph_store/memgraph_store.py b/dbgpt/storage/graph_store/memgraph_store.py
index 0685323d9..f5c1f4f76 100644
--- a/dbgpt/storage/graph_store/memgraph_store.py
+++ b/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
diff --git a/dbgpt/storage/graph_store/neo4j_store.py b/dbgpt/storage/graph_store/neo4j_store.py
index ddcb0baf0..d40869164 100644
--- a/dbgpt/storage/graph_store/neo4j_store.py
+++ b/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
 
diff --git a/dbgpt/storage/graph_store/tugraph_store.py b/dbgpt/storage/graph_store/tugraph_store.py
index 5226dc671..227f80641 100644
--- a/dbgpt/storage/graph_store/tugraph_store.py
+++ b/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