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 d0670de102..b2940c99a6 100644
--- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
+++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
@@ -1,3 +1,7 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
 import os
 from dataclasses import dataclass, field
 from typing import List
diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py
index 1860047500..2e99595e93 100644
--- a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py
+++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py
@@ -1,3 +1,7 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
 from typing import Dict, Union
 
 from autogen import UserProxyAgent
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 2f56507be9..6f54126e3e 100644
--- a/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
+++ b/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
@@ -9,6 +9,7 @@
 import pytest
 from conftest import reason, skip_openai  # noqa: E402
 from graphrag_sdk.schema import Schema
+
 try:
     from autogen.agentchat.contrib.graph_rag.document import (
         Document,
@@ -38,7 +39,7 @@ def test_falkor_db_query_engine():
     3. Query it with a question and verify the result contains the critical information.
     """
     # Arrange
-    test_schema = Ontology()
+    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)