Update falkor db test to use 0.1.3b0 version

ag2ai · Nov 26, 2024 · 7425700 · 7425700
1 parent d3f8352
commit 7425700
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py b/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
@@ -2,8 +2,7 @@
 
 import pytest
 from conftest import reason, skip_openai  # noqa: E402
-from graphrag_sdk import KnowledgeGraph, Source
-from graphrag_sdk.ontology import Ontology
+from graphrag_sdk.schema import Schema
 
 try:
     from autogen.agentchat.contrib.graph_rag.document import (
@@ -34,7 +33,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)