From 17d0798694951a0dfeca604de8f5f2efbc9479fb Mon Sep 17 00:00:00 2001
From: Aristo <6344553+randombet@users.noreply.github.com>
Date: Fri, 6 Sep 2024 22:16:31 +0000
Subject: [PATCH] Implement FalkorGraphQueryEngine

---
 .github/workflows/contrib-graph-rag-tests.yml | 11 ++++
 .../graph_rag/falkor_graph_query_engine.py    |  2 +-
 .../graph_rag/test_falkor_graph_rag.py        | 52 +++++++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/contrib-graph-rag-tests.yml b/.github/workflows/contrib-graph-rag-tests.yml
index d5cb89ed5e6b..647f7fd39ad7 100644
--- a/.github/workflows/contrib-graph-rag-tests.yml
+++ b/.github/workflows/contrib-graph-rag-tests.yml
@@ -22,6 +22,7 @@ permissions:
   # contents: read
   # deployments: read
 jobs:
+<<<<<<< HEAD
   GraphRagUnitTest:
     runs-on: ${{ matrix.os }}
     strategy:
@@ -61,6 +62,8 @@ jobs:
           file: ./coverage.xml
           flags: unittests
 
+=======
+>>>>>>> 53214b0f72 (Implement FalkorGraphQueryEngine)
   GraphRagIntegrationTest-FalkorDB-Ubuntu:
     runs-on: ubuntu-latest
     strategy:
@@ -90,6 +93,14 @@ jobs:
         run: |
           echo "AUTOGEN_USE_DOCKER=False" >> $GITHUB_ENV
       - name: Coverage
+<<<<<<< HEAD
+=======
+        env:
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY }}
+          AZURE_OPENAI_API_BASE: ${{ secrets.AZURE_OPENAI_API_BASE }}
+          OAI_CONFIG_LIST: ${{ secrets.OAI_CONFIG_LIST }}
+>>>>>>> 53214b0f72 (Implement FalkorGraphQueryEngine)
         run: |
           pip install pytest-cov>=5
           pytest test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py --skip-openai
diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
index c10a6b40a566..ffc300585393 100644
--- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
+++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
@@ -6,7 +6,7 @@
 from graphrag_sdk.schema import Schema
 
 from .document import Document
-from .graph_query_engine import GraphQueryEngine, GraphStoreQueryResult
+from .graph_query_engine import GraphStoreQueryResult
 
 
 class FalkorGraphQueryResult(GraphStoreQueryResult):
diff --git a/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py b/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
index e69de29bb2d1..36ac51a3eae7 100644
--- a/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
+++ b/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
@@ -0,0 +1,52 @@
+import os
+import sys
+
+import pytest
+from graphrag_sdk import KnowledgeGraph, Source
+from graphrag_sdk.schema import Schema
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "../../.."))
+from conftest import reason, skip_openai  # noqa: E402
+
+try:
+    from autogen.agentchat.contrib.graph_rag.document import (
+        Document,
+        DocumentType,
+    )
+    from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import (
+        FalkorGraphQueryEngine,
+        GraphStoreQueryResult,
+    )
+except ImportError:
+    skip = True
+else:
+    skip = False
+
+reason = "do not run on MacOS or windows OR dependency is not installed OR " + reason
+
+
+@pytest.mark.skipif(
+    sys.platform in ["darwin", "win32"] or skip or skip_openai,
+    reason=reason,
+)
+def test_falkor_db_query_engine():
+    # Arrange
+    test_schema = Schema()
+    actor = test_schema.add_entity("Actor").add_attribute("name", str, unique=True)
+    movie = test_schema.add_entity("Movie").add_attribute("title", str, unique=True)
+    test_schema.add_relation("ACTED", actor, movie)
+
+    query_engine = FalkorGraphQueryEngine(schema=test_schema)
+
+    source_file = "test/agentchat/contrib/graph_rag/the_matrix.txt"
+    input_docs = [Document(doctype=DocumentType.TEXT, path_or_url=source_file)]
+
+    question = "Name a few actors who've played in 'The Matrix'"
+
+    # Act
+    query_engine.init_db(input_doc=input_docs)
+
+    query_result: GraphStoreQueryResult = query_engine.query(question=question)
+
+    # Assert
+    assert query_result.answer.find("Keanu Reeves") >= 0