Implement FalkorGraphQueryEngine

.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

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):

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