feat: add document structure into GraphRAG

dbgpt/app/knowledge/service.py

@@ -650,12 +650,12 @@ def query_graph(self, space_name, limit):
                 {
                     "id": node.vid,
                     "communityId": node.get_prop("_community_id"),
-                    "name": node.vid,
-                    "type": "",
+                    "name": node.name,
+                    "type": node.get_prop("type") or ""
                 }
             )
         for edge in graph.edges():
             res["edges"].append(
-                {"source": edge.sid, "target": edge.tid, "name": edge.name, "type": ""}
+                {"source": edge.sid, "target": edge.tid, "name": edge.name, "type": edge.get_prop("type") or ""}
             )
         return res

dbgpt/rag/transformer/graph_extractor.py

@@ -65,7 +65,7 @@ def _parse_response(self, text: str, limit: Optional[int] = None) -> List[Graph]
                     match = re.match(r"\((.*?)#(.*?)\)", line)
                     if match:
                         name, summary = [part.strip() for part in match.groups()]
-                        graph.upsert_vertex(Vertex(name, description=summary))
+                        graph.upsert_vertex(Vertex(name, description=summary, vertex_type='entity'))
                 elif current_section == "Relationships":
                     match = re.match(r"\((.*?)#(.*?)#(.*?)#(.*?)\)", line)
                     if match:
@@ -74,7 +74,7 @@ def _parse_response(self, text: str, limit: Optional[int] = None) -> List[Graph]
                         ]
                         edge_count += 1
                         graph.append_edge(
-                            Edge(source, target, name, description=summary)
+                            Edge(source, target, name, description=summary, edge_type='relation')
                         )
 
             if limit and edge_count >= limit:

dbgpt/storage/graph_store/base.py

@@ -5,7 +5,7 @@
 
 from dbgpt._private.pydantic import BaseModel, ConfigDict, Field
 from dbgpt.core import Embeddings
-from dbgpt.storage.graph_store.graph import Direction, Graph
+from dbgpt.storage.graph_store.graph import Direction, Graph, Vertex
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +44,10 @@ def get_vertex_type(self) -> str:
     def get_edge_type(self) -> str:
         """Get the edge type."""
 
+    @abstractmethod
+    def get_document_vertex(self, doc_name:str) -> Vertex:
+        """Add graph."""
+
     @abstractmethod
     def insert_triplet(self, sub: str, rel: str, obj: str):
         """Add triplet."""
@@ -60,6 +64,10 @@ def get_triplets(self, sub: str) -> List[Tuple[str, str]]:
     def delete_triplet(self, sub: str, rel: str, obj: str):
         """Delete triplet."""
 
+    @abstractmethod
+    def delete_document(self, doc_name: str):
+        """Delete document."""
+
     @abstractmethod
     def truncate(self):
         """Truncate Graph."""
@@ -87,6 +95,14 @@ def explore(
     ) -> Graph:
         """Explore on graph."""
 
+    def explore_text_link(
+        self,
+        subs: List[str],
+        depth: Optional[int] = None,
+        limit: Optional[int] = None,
+    ) -> Graph:
+        """Explore  text link on graph."""
+
     @abstractmethod
     def query(self, query: str, **args) -> Graph:
         """Execute a query."""

dbgpt/storage/graph_store/graph.py

@@ -188,11 +188,11 @@ def get_neighbor_edges(
         """Get neighbor edges."""
 
     @abstractmethod
-    def vertices(self) -> Iterator[Vertex]:
+    def vertices(self, vertex_prop:Optional[str] = None) -> Iterator[Vertex]:
         """Get vertex iterator."""
 
     @abstractmethod
-    def edges(self) -> Iterator[Edge]:
+    def edges(self, edge_prop:Optional[str] = None) -> Iterator[Edge]:
         """Get edge iterator."""
 
     @abstractmethod
@@ -335,13 +335,21 @@ def unique_elements(elements):
 
         return itertools.islice(es, limit) if limit else es
 
-    def vertices(self) -> Iterator[Vertex]:
+    def vertices(self, vertex_type: Optional[str] = None) -> Iterator[Vertex]:
         """Return vertices."""
-        return iter(self._vs.values())
+        return (
+            item for item in self._vs.values()
+            if vertex_type is None or item.get_prop('vertex_type') == vertex_type
+        )
 
-    def edges(self) -> Iterator[Edge]:
+    def edges(self, edge_type: Optional[str] = None) -> Iterator[Edge]:
         """Return edges."""
-        return iter(e for nbs in self._oes.values() for es in nbs.values() for e in es)
+        return (
+            e for nbs in self._oes.values()
+            for es in nbs.values()
+            for e in es
+            if edge_type is None or e.get_prop('edge_type') == edge_type
+        )
 
     def del_vertices(self, *vids: str):
         """Delete specified vertices."""