Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Graph RAG] Init Commit with GraphRag interfaces #3388

Merged
merged 10 commits into from
Sep 4, 2024
Empty file.
44 changes: 44 additions & 0 deletions autogen/agentchat/contrib/graph_rag/graph_rag_agent.py
Original file line number Diff line number Diff line change
@@ -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):
randombet marked this conversation as resolved.
Show resolved Hide resolved
"""
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:
randombet marked this conversation as resolved.
Show resolved Hide resolved
"""
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):
randombet marked this conversation as resolved.
Show resolved Hide resolved
"""
Retrieve answers with human readable questions.
"""
pass

@abstractmethod
def add_records(self, new_records: List) -> bool:
randombet marked this conversation as resolved.
Show resolved Hide resolved
"""
Add new records to the underlying database and add to the graph if required.
"""
pass
25 changes: 25 additions & 0 deletions autogen/agentchat/contrib/graph_rag/graph_store.py
Original file line number Diff line number Diff line change
@@ -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:
randombet marked this conversation as resolved.
Show resolved Hide resolved
"""
This method transform a string format question into database query and return the result.
"""
pass
Loading