fix test

ag2ai · Nov 26, 2024 · e9ea6f7 · e9ea6f7
1 parent 63b4802
commit e9ea6f7
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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)