cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/tests/integration/retrieval/test_graph_completion_retriever_context_extension.py

@@ -164,13 +164,18 @@ async def setup_test_environment_empty():
 async def test_graph_completion_extension_context_simple(setup_test_environment_simple):
     """Integration test: verify GraphCompletionContextExtensionRetriever can retrieve context (simple)."""
     retriever = GraphCompletionContextExtensionRetriever()
+    query = "Who works at Canva?"
 
-    context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     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"
 
-    answer = await retriever.get_completion("Who works at Canva?")
+    answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
 
     assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
     assert all(isinstance(item, str) and item.strip() for item in answer), (
@@ -182,16 +187,19 @@ async def test_graph_completion_extension_context_simple(setup_test_environment_
 async def test_graph_completion_extension_context_complex(setup_test_environment_complex):
     """Integration test: verify GraphCompletionContextExtensionRetriever can retrieve context (complex)."""
     retriever = GraphCompletionContextExtensionRetriever(top_k=20)
+    query = "Who works at Figma and drives Tesla?"
 
-    context = await resolve_edges_to_text(
-        await retriever.get_context("Who works at Figma and drives Tesla?")
-    )
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert "Mike Rodger --[works_for]--> Figma" in context, "Failed to get Mike Rodger"
     assert "Ike Loma --[works_for]--> Figma" in context, "Failed to get Ike Loma"
     assert "Jason Statham --[works_for]--> Figma" in context, "Failed to get Jason Statham"
 
-    answer = await retriever.get_completion("Who works at Figma?")
+    answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
 
     assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
     assert all(isinstance(item, str) and item.strip() for item in answer), (
@@ -203,24 +211,18 @@ async def test_graph_completion_extension_context_complex(setup_test_environment
 async def test_get_graph_completion_extension_context_on_empty_graph(setup_test_environment_empty):
     """Integration test: verify GraphCompletionContextExtensionRetriever handles empty graph correctly."""
     retriever = GraphCompletionContextExtensionRetriever()
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("Who works at Figma?")
-    assert context == [], "Context should be empty on an empty graph"
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
+    assert context == "", "Context should be empty on an empty graph"
 
-    answer = await retriever.get_completion("Who works at Figma?")
+    answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
 
     assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
     assert all(isinstance(item, str) and item.strip() for item in answer), (
         "Answer must contain only non-empty strings"
     )
-
-
-@pytest.mark.asyncio
-async def test_graph_completion_extension_get_triplets_empty(setup_test_environment_empty):
-    """Integration test: verify GraphCompletionContextExtensionRetriever get_triplets handles empty graph."""
-    retriever = GraphCompletionContextExtensionRetriever()
-
-    triplets = await retriever.get_triplets("Who works at Figma?")
-
-    assert isinstance(triplets, list), "Triplets should be a list"
-    assert len(triplets) == 0, "Should return empty list on empty graph"

cognee/tests/integration/retrieval/test_graph_completion_retriever_cot.py

@@ -1,12 +1,11 @@
-import os
 import pytest
 import pathlib
 import pytest_asyncio
 from typing import Optional, Union
 import cognee
+from cognee.exceptions import CogneeValidationError
 
 from cognee.low_level import setup, DataPoint
-from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
 from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
 
@@ -158,13 +157,18 @@ async def setup_test_environment_empty():
 async def test_graph_completion_cot_context_simple(setup_test_environment_simple):
     """Integration test: verify GraphCompletionCotRetriever can retrieve context (simple)."""
     retriever = GraphCompletionCotRetriever()
+    query = "Who works at Canva?"
 
-    context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     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"
 
-    answer = await retriever.get_completion("Who works at Canva?")
+    answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
 
     assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
     assert all(isinstance(item, str) and item.strip() for item in answer), (
@@ -176,14 +180,19 @@ async def test_graph_completion_cot_context_simple(setup_test_environment_simple
 async def test_graph_completion_cot_context_complex(setup_test_environment_complex):
     """Integration test: verify GraphCompletionCotRetriever can retrieve context (complex)."""
     retriever = GraphCompletionCotRetriever(top_k=20)
+    query = "Who works at Figma?"
 
-    context = await resolve_edges_to_text(await retriever.get_context("Who works at Figma?"))
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert "Mike Rodger --[works_for]--> Figma" in context, "Failed to get Mike Rodger"
     assert "Ike Loma --[works_for]--> Figma" in context, "Failed to get Ike Loma"
     assert "Jason Statham --[works_for]--> Figma" in context, "Failed to get Jason Statham"
 
-    answer = await retriever.get_completion("Who works at Figma?")
+    answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
 
     assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
     assert all(isinstance(item, str) and item.strip() for item in answer), (
@@ -195,24 +204,14 @@ async def test_graph_completion_cot_context_complex(setup_test_environment_compl
 async def test_get_graph_completion_cot_context_on_empty_graph(setup_test_environment_empty):
     """Integration test: verify GraphCompletionCotRetriever handles empty graph correctly."""
     retriever = GraphCompletionCotRetriever()
+    query = "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?")
-
-    assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-    assert all(isinstance(item, str) and item.strip() for item in answer), (
-        "Answer must contain only non-empty strings"
-    )
-
-
-@pytest.mark.asyncio
-async def test_graph_completion_cot_get_triplets_empty(setup_test_environment_empty):
-    """Integration test: verify GraphCompletionCotRetriever get_triplets handles empty graph."""
-    retriever = GraphCompletionCotRetriever()
+    triplets = await retriever.get_retrieved_objects(query)
 
-    triplets = await retriever.get_triplets("Who works at Figma?")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
+    assert context == "", "Context should be empty on an empty graph"
 
-    assert isinstance(triplets, list), "Triplets should be a list"
-    assert len(triplets) == 0, "Should return empty list on empty graph"
+    with pytest.raises(CogneeValidationError):
+        await retriever.get_completion_from_context(
+            query=query, retrieved_objects=triplets, context=context
+        )

cognee/tests/integration/retrieval/test_rag_completion_retriever.py

@@ -83,6 +83,18 @@ async def setup_test_environment_with_chunks_simple():
     try:
         await cognee.prune.prune_data()
         await cognee.prune.prune_system(metadata=True)
+
+        from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+        from cognee.infrastructure.databases.vector.create_vector_engine import (
+            _create_vector_engine,
+        )
+        from cognee.infrastructure.databases.relational.create_relational_engine import (
+            create_relational_engine,
+        )
+
+        _create_graph_engine.cache_clear()
+        _create_vector_engine.cache_clear()
+        create_relational_engine.cache_clear()
     except Exception:
         pass
 
@@ -174,6 +186,18 @@ async def setup_test_environment_with_chunks_complex():
     try:
         await cognee.prune.prune_data()
         await cognee.prune.prune_system(metadata=True)
+
+        from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+        from cognee.infrastructure.databases.vector.create_vector_engine import (
+            _create_vector_engine,
+        )
+        from cognee.infrastructure.databases.relational.create_relational_engine import (
+            create_relational_engine,
+        )
+
+        _create_graph_engine.cache_clear()
+        _create_vector_engine.cache_clear()
+        create_relational_engine.cache_clear()
     except Exception:
         pass
 
@@ -195,11 +219,33 @@ async def setup_test_environment_empty():
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
 
+    from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+    from cognee.infrastructure.databases.vector.create_vector_engine import _create_vector_engine
+    from cognee.infrastructure.databases.relational.create_relational_engine import (
+        create_relational_engine,
+    )
+
+    _create_graph_engine.cache_clear()
+    _create_vector_engine.cache_clear()
+    create_relational_engine.cache_clear()
+
     yield
 
     try:
         await cognee.prune.prune_data()
         await cognee.prune.prune_system(metadata=True)
+
+        from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+        from cognee.infrastructure.databases.vector.create_vector_engine import (
+            _create_vector_engine,
+        )
+        from cognee.infrastructure.databases.relational.create_relational_engine import (
+            create_relational_engine,
+        )
+
+        _create_graph_engine.cache_clear()
+        _create_vector_engine.cache_clear()
+        create_relational_engine.cache_clear()
     except Exception:
         pass
 
@@ -208,8 +254,13 @@ async def setup_test_environment_empty():
 async def test_rag_completion_context_simple(setup_test_environment_with_chunks_simple):
     """Integration test: verify CompletionRetriever can retrieve context (simple)."""
     retriever = CompletionRetriever()
+    query = "Mike"
 
-    context = await retriever.get_context("Mike")
+    retrieved_objects = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(
+        query=query, retrieved_objects=retrieved_objects
+    )
 
     assert isinstance(context, str), "Context should be a string"
     assert "Mike Broski" in context, "Failed to get Mike Broski"
@@ -219,8 +270,13 @@ async def test_rag_completion_context_simple(setup_test_environment_with_chunks_
 async def test_rag_completion_context_multiple_chunks(setup_test_environment_with_chunks_simple):
     """Integration test: verify CompletionRetriever can retrieve context from multiple chunks."""
     retriever = CompletionRetriever()
+    query = "Steve"
+
+    retrieved_objects = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("Steve")
+    context = await retriever.get_context_from_objects(
+        query=query, retrieved_objects=retrieved_objects
+    )
 
     assert isinstance(context, str), "Context should be a string"
     assert "Steve Rodger" in context, "Failed to get Steve Rodger"
@@ -231,8 +287,13 @@ async def test_rag_completion_context_complex(setup_test_environment_with_chunks
     """Integration test: verify CompletionRetriever can retrieve context (complex)."""
     # TODO: top_k doesn't affect the output, it should be fixed.
     retriever = CompletionRetriever(top_k=20)
+    query = "Christina"
 
-    context = await retriever.get_context("Christina")
+    retrieved_objects = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(
+        query=query, retrieved_objects=retrieved_objects
+    )
 
     assert context[0:15] == "Christina Mayer", "Failed to get Christina Mayer"
 
@@ -241,14 +302,18 @@ async def test_rag_completion_context_complex(setup_test_environment_with_chunks
 async def test_get_rag_completion_context_on_empty_graph(setup_test_environment_empty):
     """Integration test: verify CompletionRetriever handles empty graph correctly."""
     retriever = CompletionRetriever()
+    query = "Christina Mayer"
 
     with pytest.raises(NoDataError):
-        await retriever.get_context("Christina Mayer")
+        await retriever.get_retrieved_objects(query)
 
     vector_engine = get_vector_engine()
     await vector_engine.create_collection(
         "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
     )
 
-    context = await retriever.get_context("Christina Mayer")
+    retrieved_objects = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(
+        query=query, retrieved_objects=retrieved_objects
+    )
     assert context == "", "Returned context should be empty on an empty graph"

cognee/tests/integration/retrieval/test_structured_output.py

@@ -43,84 +43,128 @@ def _assert_structured_answer(answer: list[TestAnswer]):
 
 async def _test_get_structured_graph_completion_cot():
     retriever = GraphCompletionCotRetriever()
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     # Test with string response model (default)
-    string_answer = await retriever.get_completion("Who works at Figma?")
+    string_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
     _assert_string_answer(string_answer)
 
+    retriever.response_model = TestAnswer
     # Test with structured response model
-    structured_answer = await retriever.get_completion(
-        "Who works at Figma?", response_model=TestAnswer
+    triplets = await retriever.get_retrieved_objects(query)
+    structured_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
     )
     _assert_structured_answer(structured_answer)
 
 
 async def _test_get_structured_graph_completion():
     retriever = GraphCompletionRetriever()
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     # Test with string response model (default)
-    string_answer = await retriever.get_completion("Who works at Figma?")
+    string_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
     _assert_string_answer(string_answer)
 
+    retriever.response_model = TestAnswer
     # Test with structured response model
-    structured_answer = await retriever.get_completion(
-        "Who works at Figma?", response_model=TestAnswer
+    structured_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
     )
     _assert_structured_answer(structured_answer)
 
 
 async def _test_get_structured_graph_completion_temporal():
     retriever = TemporalRetriever()
+    query = "When did Steve start working at Figma?"
 
     # Test with string response model (default)
-    string_answer = await retriever.get_completion("When did Steve start working at Figma?")
+    triplets = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
+
+    # Test with string response model (default)
+    string_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
     _assert_string_answer(string_answer)
 
+    retriever.response_model = TestAnswer
     # Test with structured response model
-    structured_answer = await retriever.get_completion(
-        "When did Steve start working at Figma??", response_model=TestAnswer
+    structured_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
     )
     _assert_structured_answer(structured_answer)
 
 
 async def _test_get_structured_graph_completion_rag():
     retriever = CompletionRetriever()
+    query = "Where does Steve work?"
+
+    triplets = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     # Test with string response model (default)
-    string_answer = await retriever.get_completion("Where does Steve work?")
+    string_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
     _assert_string_answer(string_answer)
 
+    retriever.response_model = TestAnswer
     # Test with structured response model
-    structured_answer = await retriever.get_completion(
-        "Where does Steve work?", response_model=TestAnswer
+    structured_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
     )
     _assert_structured_answer(structured_answer)
 
 
 async def _test_get_structured_graph_completion_context_extension():
     retriever = GraphCompletionContextExtensionRetriever()
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     # Test with string response model (default)
-    string_answer = await retriever.get_completion("Who works at Figma?")
+    string_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
     _assert_string_answer(string_answer)
 
+    retriever.response_model = TestAnswer
     # Test with structured response model
-    structured_answer = await retriever.get_completion(
-        "Who works at Figma?", response_model=TestAnswer
+    structured_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
     )
     _assert_structured_answer(structured_answer)
 
 
 async def _test_get_structured_entity_completion():
     retriever = EntityCompletionRetriever(DummyEntityExtractor(), DummyContextProvider())
+    query = "Who is Albert Einstein?"
+
+    entities = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=entities)
 
     # Test with string response model (default)
-    string_answer = await retriever.get_completion("Who is Albert Einstein?")
+    string_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=entities, context=context
+    )
     _assert_string_answer(string_answer)
 
+    retriever.response_model = TestAnswer
     # Test with structured response model
-    structured_answer = await retriever.get_completion(
-        "Who is Albert Einstein?", response_model=TestAnswer
+    structured_answer = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=entities, context=context
     )
     _assert_structured_answer(structured_answer)
 

cognee/tests/integration/retrieval/test_summaries_retriever.py

@@ -158,27 +158,38 @@ async def setup_test_environment_empty():
 
 
 @pytest.mark.asyncio
-async def test_summaries_retriever_context(setup_test_environment_with_summaries):
+async def test_summaries_retriever(setup_test_environment_with_summaries):
     """Integration test: verify SummariesRetriever can retrieve summary context."""
     retriever = SummariesRetriever(top_k=20)
+    query = "Christina"
+    summaries = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=summaries)
 
-    context = await retriever.get_context("Christina")
+    completion = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=summaries, context=context
+    )
 
-    assert isinstance(context, list), "Context should be a list"
-    assert len(context) > 0, "Context should not be empty"
-    assert context[0]["text"] == "C.M.", "Failed to get Christina Mayer"
+    assert isinstance(completion, list), "Context should be a list"
+    assert len(completion) > 0, "Context should not be empty"
+    assert completion[0]["text"] == "C.M.", "Failed to get Christina Mayer"
 
 
 @pytest.mark.asyncio
-async def test_summaries_retriever_context_on_empty_graph(setup_test_environment_empty):
+async def test_summaries_retriever_on_empty_graph(setup_test_environment_empty):
     """Integration test: verify SummariesRetriever handles empty graph correctly."""
     retriever = SummariesRetriever()
+    query = "Christina Mayer"
 
     with pytest.raises(NoDataError):
-        await retriever.get_context("Christina Mayer")
+        await retriever.get_retrieved_objects(query)
 
     vector_engine = get_vector_engine()
     await vector_engine.create_collection("TextSummary_text", payload_schema=TextSummary)
 
-    context = await retriever.get_context("Christina Mayer")
-    assert context == [], "Returned context should be empty on an empty graph"
+    summaries = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=summaries)
+    completion = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=summaries, context=context
+    )
+
+    assert completion == [], "Returned context should be empty on an empty graph"

cognee/tests/integration/retrieval/test_temporal_retriever.py

@@ -197,8 +197,11 @@ async def setup_test_environment_empty():
 async def test_temporal_retriever_context_with_time_range(setup_test_environment_with_events):
     """Integration test: verify TemporalRetriever can retrieve events within time range."""
     retriever = TemporalRetriever(top_k=5)
+    query = "What happened in January 2021?"
 
-    context = await retriever.get_context("What happened in January 2021?")
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert isinstance(context, str), "Context should be a string"
     assert len(context) > 0, "Context should not be empty"
@@ -211,8 +214,11 @@ async def test_temporal_retriever_context_with_time_range(setup_test_environment
 async def test_temporal_retriever_context_with_single_time(setup_test_environment_with_events):
     """Integration test: verify TemporalRetriever can retrieve events at specific time."""
     retriever = TemporalRetriever(top_k=5)
+    query = "What happened in July 2021?"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("What happened in July 2021?")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert isinstance(context, str), "Context should be a string"
     assert len(context) > 0, "Context should not be empty"
@@ -227,8 +233,11 @@ async def test_temporal_retriever_context_fallback_to_triplets(
 ):
     """Integration test: verify TemporalRetriever falls back to triplets when no time extracted."""
     retriever = TemporalRetriever(top_k=5)
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("Who works at Figma?")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert isinstance(context, str), "Context should be a string"
     assert len(context) > 0, "Context should not be empty"
@@ -241,8 +250,11 @@ async def test_temporal_retriever_context_fallback_to_triplets(
 async def test_temporal_retriever_context_empty_graph(setup_test_environment_empty):
     """Integration test: verify TemporalRetriever handles empty graph correctly."""
     retriever = TemporalRetriever()
+    query = "What happened?"
 
-    context = await retriever.get_context("What happened?")
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert isinstance(context, str), "Context should be a string"
     assert len(context) >= 0, "Context should be a string (possibly empty)"
@@ -252,8 +264,14 @@ async def test_temporal_retriever_context_empty_graph(setup_test_environment_emp
 async def test_temporal_retriever_get_completion(setup_test_environment_with_events):
     """Integration test: verify TemporalRetriever can generate completions."""
     retriever = TemporalRetriever()
+    query = "What happened in January 2021?"
+
+    triplets = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
-    completion = await retriever.get_completion("What happened in January 2021?")
+    completion = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
 
     assert isinstance(completion, list), "Completion should be a list"
     assert len(completion) > 0, "Completion should not be empty"
@@ -266,8 +284,14 @@ async def test_temporal_retriever_get_completion(setup_test_environment_with_eve
 async def test_temporal_retriever_get_completion_fallback(setup_test_environment_with_graph_data):
     """Integration test: verify TemporalRetriever get_completion works with triplet fallback."""
     retriever = TemporalRetriever()
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
-    completion = await retriever.get_completion("Who works at Figma?")
+    completion = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=triplets, context=context
+    )
 
     assert isinstance(completion, list), "Completion should be a list"
     assert len(completion) > 0, "Completion should not be empty"
@@ -280,8 +304,11 @@ async def test_temporal_retriever_get_completion_fallback(setup_test_environment
 async def test_temporal_retriever_top_k_limit(setup_test_environment_with_events):
     """Integration test: verify TemporalRetriever respects top_k parameter."""
     retriever = TemporalRetriever(top_k=2)
+    query = "What happened in 2021?"
 
-    context = await retriever.get_context("What happened in 2021?")
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert isinstance(context, str), "Context should be a string"
     separator_count = context.count("#####################")
@@ -292,8 +319,11 @@ async def test_temporal_retriever_top_k_limit(setup_test_environment_with_events
 async def test_temporal_retriever_multiple_events(setup_test_environment_with_events):
     """Integration test: verify TemporalRetriever can retrieve multiple events."""
     retriever = TemporalRetriever(top_k=10)
+    query = "What events occurred in 2021?"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("What events occurred in 2021?")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert isinstance(context, str), "Context should be a string"
     assert len(context) > 0, "Context should not be empty"

cognee/tests/integration/retrieval/test_triplet_retriever.py

@@ -78,8 +78,11 @@ async def setup_test_environment_empty():
 async def test_triplet_retriever_context_simple(setup_test_environment_with_triplets):
     """Integration test: verify TripletRetriever can retrieve triplet context."""
     retriever = TripletRetriever(top_k=5)
+    query = "Alice"
 
-    context = await retriever.get_context("Alice")
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert "Alice knows Bob" in context, "Failed to get Alice triplet"
     assert isinstance(context, str), "Context should be a string"
@@ -90,8 +93,11 @@ async def test_triplet_retriever_context_simple(setup_test_environment_with_trip
 async def test_triplet_retriever_context_multiple_triplets(setup_test_environment_with_triplets):
     """Integration test: verify TripletRetriever can retrieve multiple triplets."""
     retriever = TripletRetriever(top_k=5)
+    query = "Bob"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("Bob")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert "Alice knows Bob" in context or "Bob works at Tech Corp" in context, (
         "Failed to get Bob-related triplets"
@@ -102,8 +108,11 @@ async def test_triplet_retriever_context_multiple_triplets(setup_test_environmen
 async def test_triplet_retriever_top_k_limit(setup_test_environment_with_triplets):
     """Integration test: verify TripletRetriever respects top_k parameter."""
     retriever = TripletRetriever(top_k=1)
+    query = "Alice"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("Alice")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert isinstance(context, str), "Context should be a string"
 
@@ -116,4 +125,4 @@ async def test_triplet_retriever_context_empty(setup_test_environment_empty):
     retriever = TripletRetriever()
 
     with pytest.raises(NoDataError):
-        await retriever.get_context("Alice")
+        await retriever.get_retrieved_objects("Alice")