Add comments and finalize the tests

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -14,6 +14,9 @@ class FalkorGraphQueryResult(GraphStoreQueryResult):
 
 
 class FalkorGraphQueryEngine:
+    """
+    This is a wrapper for Falkor DB KnowledgeGraph.
+    """
 
     def __init__(
         self,
@@ -29,6 +32,9 @@ def __init__(
         self.knowledge_graph = KnowledgeGraph(name, host, port, username, password, model, schema)
 
     def init_db(self, input_doc: List[Document] | None):
+        """
+        Build the knowledge graph with input documents.
+        """
         sources = []
         for doc in input_doc:
             if os.path.exists(doc.path_or_url):
@@ -41,6 +47,17 @@ def add_records(self, new_records: List) -> bool:
         raise NotImplementedError("This method is not supported by Falkor DB SDK yet.")
 
     def query(self, question: str, n_results: int = 1, **kwargs) -> FalkorGraphQueryResult:
+        """
+        Query the knowledage graph with a question and optional message history.
+
+        Args:
+        question: a human input question.
+        n_results: number of returned results.
+        kwargs:
+            messages: a list of message history.
+
+        Returns: FalkorGraphQueryResult
+        """
         messages = kwargs.pop("messages", [])
         answer, messages = self.knowledge_graph.ask(question, messages)
         return FalkorGraphQueryResult(answer=answer, results=[], messages=messages)

test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py

@@ -30,6 +30,12 @@
     reason=reason,
 )
 def test_falkor_db_query_engine():
+    """
+    Test Falkor DB Query Engine.
+    1. create a test Falkor DB Query Engine with a schema.
+    2. Initialize it with an input txt file.
+    3. Query it with a question and verify the result contains the critical information.
+    """
     # Arrange
     test_schema = Schema()
     actor = test_schema.add_entity("Actor").add_attribute("name", str, unique=True)
@@ -50,3 +56,9 @@ def test_falkor_db_query_engine():
 
     # Assert
     assert query_result.answer.find("Keanu Reeves") >= 0
+    for message in query_result.messages:
+        if isinstance(message, dict) and "role" in message and message["role"] == "user":
+            assert "content" in message
+            assert message["content"] is question
+            return
+    pytest.fail("Question not found in message history.")