[Graph RAG] Init Commit with GraphRag interfaces

autogen/agentchat/contrib/graph_rag/document.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Optional
+
+
+class DocumentType(Enum):
+    """
+    Enum for supporting document type.
+    """
+
+    TEXT = auto()
+    HTML = auto()
+    PDF = auto()
+
+
+@dataclass
+class Document:
+    """
+    A wrapper of graph store query results.
+    """
+
+    doctype: DocumentType
+    data: Optional[object] = None
+    path_or_url: Optional[str] = ""

autogen/agentchat/contrib/graph_rag/graph_rag_agent.py

@@ -0,0 +1,44 @@
+from typing import List, Protocol
+
+from autogen.agentchat import ConversableAgent
+
+from .document import Document
+from .graph_store import GraphStore
+
+
+class GraphRagAgent(ConversableAgent, Protocol):
+    """
+    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.
+    """
+
+    def _init_db(self, input_doc: List[Document] | 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.
+
+        Args:
+        input_doc: a list of input documents that are used to build the graph in database.
+
+        Returns: GraphStore
+        """
+        pass
+
+    def retrieve(self, question: str, **kwargs):
+        """
+        Retrieve answers with human readable questions.
+        """
+        pass
+
+    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,28 @@
+from dataclasses import dataclass
+from typing import List, Optional, Protocol
+
+
+@dataclass
+class GraphStoreQueryResult:
+    """
+    A wrapper of graph store query results.
+
+    answer: human readable answer to question/query.
+    results: intermediate results to question/query, e.g. node entities.
+    """
+
+    answer: Optional[str] = None
+    results: Optional[List] = []
+
+
+class GraphStore(Protocol):
+    """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.
+    """
+
+    def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryResult:
+        """
+        This method transform a string format question into database query and return the result.
+        """
+        pass

notebook/agentchat_nested_sequential_chats.ipynb

@@ -811,7 +811,9 @@
   "front_matter": {
    "description": "Solve complex tasks with one or more sequence chats nested as inner monologue.",
    "tags": [
-    "nested chat", "sequential chats", "orchestration"
+    "nested chat",
+    "sequential chats",
+    "orchestration"
    ]
   },
   "kernelspec": {