[Graph RAG] Init Commit with GraphRag interfaces

autogen/agentchat/contrib/graph_rag/graph_rag_agent.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+from autogen.agentchat import ConversableAgent
+
+from .graph_store import GraphStore
+
+
+class GraphRagAgent(ConversableAgent, ABC):
+    """
+    A graph rag agent is a conversable agent which could query graph database for answers.
+
+    An implementing agent class would
+    1. create a graph in the underlying database with input documents
+    2. use the retrieve() method to retrieve information.
+    3. use the retrieved information to generate and send back messages.
+    """
+
+    @abstractmethod
+    def _init_db(self, input_doc: List | None = None) -> GraphStore:
+        """
+        This method initializes graph database with the input documents or records.
+        Usually, it takes the following steps,
+        1. connecting to a graph database.
+        2. extract graph nodes, edges based on input data, graph schema and etc.
+        3. build indexes etc.
+
+        return: GraphStore
+        """
+        pass
+
+    @abstractmethod
+    def retrieve(self, question: str, **kwargs):
+        """
+        Retrieve answers with human readable questions.
+        """
+        pass
+
+    @abstractmethod
+    def add_records(self, new_records: List) -> bool:
+        """
+        Add new records to the underlying database and add to the graph if required.
+        """
+        pass

autogen/agentchat/contrib/graph_rag/graph_store.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass
+class GraphStoreQueryResult:
+    """
+    A wrapper of graph store query results.
+    """
+
+    answer: str
+
+
+class GraphStore(ABC):
+    """An abstract base class that represents a underlying graph database.
+
+    This interface defines the basic methods which are required by implementing graph rag from graph database.
+    """
+
+    @abstractmethod
+    def query(self, question: str, **kwargs) -> GraphStoreQueryResult:
+        """
+        This method transform a string format question into database query and return the result.
+        """
+        pass