diff --git a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py
index 3e7f6626f..26ae2f883 100644
--- a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py
+++ b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py
@@ -51,7 +51,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context, _ = await retriever.get_context("Who works at Canva?")
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
@@ -129,7 +129,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
 
         print(context)
 
@@ -167,7 +167,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         await setup()
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
         assert context == "", "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
diff --git a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py
index ff92dfd8f..be25299aa 100644
--- a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py
+++ b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py
@@ -47,7 +47,7 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context, _ = await retriever.get_context("Who works at Canva?")
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
@@ -124,7 +124,7 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
 
         print(context)
 
@@ -162,7 +162,7 @@ class TestGraphCompletionCoTRetriever:
 
         await setup()
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
         assert context == "", "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
diff --git a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py
index 976b69e69..50784f94a 100644
--- a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py
+++ b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py
@@ -67,7 +67,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context, _ = await retriever.get_context("Who works at Canva?")
 
         # Ensure the top-level sections are present
         assert "Nodes:" in context, "Missing 'Nodes:' section in context"
@@ -191,7 +191,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
 
         print(context)
 
@@ -222,5 +222,5 @@ class TestGraphCompletionRetriever:
 
         await setup()
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
         assert context == "", "Context should be empty on an empty graph"