gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update test_search_db.py

Browse source
This commit is contained in:
hajdul88 2025-12-11 18:36:21 +01:00
parent e766a2d78c
commit 84058d4525
1 changed files with 4 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
cognee/tests/test_search_db.py
View file

@ -62,7 +62,7 @@ async def setup_search_db_environment():
if has_collection: if has_collection:
collection = await vector_engine.get_collection("Triplet_text") collection = await vector_engine.get_collection("Triplet_text")
count = await collection.count_rows() if hasattr(collection, 'count_rows') else "unknown" count = await collection.count_rows() if hasattr(collection, "count_rows") else "unknown"
logger.info(f"Triplet_text collection row count: {count}") logger.info(f"Triplet_text collection row count: {count}")
graph_engine = await get_graph_engine() graph_engine = await get_graph_engine()
@ -71,7 +71,9 @@ async def setup_search_db_environment():
query = "Next to which country is Germany located?" query = "Next to which country is Germany located?"
gk_retriever_context = await GraphCompletionCotRetriever().get_context(query=query) gk_retriever_context = await GraphCompletionCotRetriever().get_context(query=query)
gk_cot_retriever_context = await GraphCompletionCotRetriever().get_context(query=query) gk_cot_retriever_context = await GraphCompletionCotRetriever().get_context(query=query)
gk_ext_retriever_context = await GraphCompletionContextExtensionRetriever().get_context(query=query) gk_ext_retriever_context = await GraphCompletionContextExtensionRetriever().get_context(
query=query
)
gk_sum_retriever_context = await GraphSummaryCompletionRetriever().get_context(query=query) gk_sum_retriever_context = await GraphSummaryCompletionRetriever().get_context(query=query)
triplet_retriever_context = await TripletRetriever().get_context(query=query) triplet_retriever_context = await TripletRetriever().get_context(query=query)
@ -148,7 +150,6 @@ async def test_graph_vector_engine_consistency(setup_search_db_environment):
async def test_retriever_contexts(setup_search_db_environment): async def test_retriever_contexts(setup_search_db_environment):
"""Test that all retrievers return valid contexts with expected content.""" """Test that all retrievers return valid contexts with expected content."""
context_gk = setup_search_db_environment["gk_retriever_context"] context_gk = setup_search_db_environment["gk_retriever_context"]
context_gk_cot = setup_search_db_environment["gk_cot_retriever_context"] context_gk_cot = setup_search_db_environment["gk_cot_retriever_context"]
context_gk_ext = setup_search_db_environment["gk_ext_retriever_context"] context_gk_ext = setup_search_db_environment["gk_ext_retriever_context"]
@ -178,4 +179,3 @@ async def test_retriever_contexts(setup_search_db_environment):
assert "germany" in lower_triplet or "netherlands" in lower_triplet, ( assert "germany" in lower_triplet or "netherlands" in lower_triplet, (
f"TripletRetriever: Context did not contain 'germany' or 'netherlands'; got: {context_triplet!r}" f"TripletRetriever: Context did not contain 'germany' or 'netherlands'; got: {context_triplet!r}"
) )