From 541377a9de8a275a68d5b025dd372921ed732533 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Tue, 23 Sep 2025 12:51:40 +0200 Subject: [PATCH] chore: Delete unused test, fix formatting. --- .../databases/vector/vector_db_interface.py | 6 ++- cognee/tests/test_chromadb.py | 9 +++- cognee/tests/test_lancedb.py | 9 +++- cognee/tests/test_pgvector.py | 8 +++- .../databases/vector/test_vector_engine.py | 46 ------------------- 5 files changed, 28 insertions(+), 50 deletions(-) delete mode 100644 cognee/tests/unit/infrastructure/databases/vector/test_vector_engine.py diff --git a/cognee/infrastructure/databases/vector/vector_db_interface.py b/cognee/infrastructure/databases/vector/vector_db_interface.py index a630f102b..3a3df62eb 100644 --- a/cognee/infrastructure/databases/vector/vector_db_interface.py +++ b/cognee/infrastructure/databases/vector/vector_db_interface.py @@ -106,7 +106,11 @@ class VectorDBInterface(Protocol): @abstractmethod async def batch_search( - self, collection_name: str, query_texts: List[str], limit: Optional[int], with_vectors: bool = False + self, + collection_name: str, + query_texts: List[str], + limit: Optional[int], + with_vectors: bool = False, ): """ Perform a batch search using multiple text queries against a collection. diff --git a/cognee/tests/test_chromadb.py b/cognee/tests/test_chromadb.py index 6fca37547..f785c660a 100644 --- a/cognee/tests/test_chromadb.py +++ b/cognee/tests/test_chromadb.py @@ -66,8 +66,14 @@ async def test_getting_of_documents(dataset_name_1): f"Number of expected documents doesn't match {len(document_ids)} != 2" ) + async def test_vector_engine_search_none_limit(): - file_path = os.path.join(pathlib.Path(__file__).resolve().parent.parent.parent, "examples", "data", "alice_in_wonderland.txt") + file_path = os.path.join( + pathlib.Path(__file__).resolve().parent.parent.parent, + "examples", + "data", + "alice_in_wonderland.txt", + ) await cognee.prune.prune_data() await cognee.prune.prune_system(metadata=True) @@ -93,6 +99,7 @@ async def test_vector_engine_search_none_limit(): # Check that we did not accidentally use any default value for limit in vector search along the way (like 5, 10, or 15) assert len(result) > 15 + async def main(): cognee.config.set_vector_db_config( { diff --git a/cognee/tests/test_lancedb.py b/cognee/tests/test_lancedb.py index a1f14d63c..675c5ee38 100644 --- a/cognee/tests/test_lancedb.py +++ b/cognee/tests/test_lancedb.py @@ -66,8 +66,14 @@ async def test_getting_of_documents(dataset_name_1): f"Number of expected documents doesn't match {len(document_ids)} != 2" ) + async def test_vector_engine_search_none_limit(): - file_path = os.path.join(pathlib.Path(__file__).resolve().parent.parent.parent, "examples", "data", "alice_in_wonderland.txt") + file_path = os.path.join( + pathlib.Path(__file__).resolve().parent.parent.parent, + "examples", + "data", + "alice_in_wonderland.txt", + ) await cognee.prune.prune_data() await cognee.prune.prune_system(metadata=True) @@ -93,6 +99,7 @@ async def test_vector_engine_search_none_limit(): # Check that we did not accidentally use any default value for limit in vector search along the way (like 5, 10, or 15) assert len(result) > 15 + async def main(): cognee.config.set_vector_db_config( { diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index 441e8a2c5..71601522f 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -67,8 +67,14 @@ async def test_getting_of_documents(dataset_name_1): f"Number of expected documents doesn't match {len(document_ids)} != 2" ) + async def test_vector_engine_search_none_limit(): - file_path = os.path.join(pathlib.Path(__file__).resolve().parent.parent.parent, "examples", "data", "alice_in_wonderland.txt") + file_path = os.path.join( + pathlib.Path(__file__).resolve().parent.parent.parent, + "examples", + "data", + "alice_in_wonderland.txt", + ) await cognee.prune.prune_data() await cognee.prune.prune_system(metadata=True) diff --git a/cognee/tests/unit/infrastructure/databases/vector/test_vector_engine.py b/cognee/tests/unit/infrastructure/databases/vector/test_vector_engine.py deleted file mode 100644 index 0afa4cc94..000000000 --- a/cognee/tests/unit/infrastructure/databases/vector/test_vector_engine.py +++ /dev/null @@ -1,46 +0,0 @@ -import os -import pathlib -import pytest - -import cognee - - -class TestVectorEngine: - # Test that vector engine search works well with limit=None. - # Search should return all entities that exist in a collection. Used Alice for a bit larger test. - @pytest.mark.asyncio - async def test_vector_engine_search_none_limit(self): - system_directory_path = os.path.join( - pathlib.Path(__file__).parent, ".cognee_system/test_vector_engine_search_none_limit" - ) - cognee.config.system_root_directory(system_directory_path) - data_directory_path = os.path.join( - pathlib.Path(__file__).parent, ".data_storage/test_vector_engine_search_none_limit" - ) - cognee.config.data_root_directory(data_directory_path) - - file_path = os.path.join(pathlib.Path(__file__).resolve().parent, "data", "alice_in_wonderland.txt") - - await cognee.prune.prune_data() - await cognee.prune.prune_system(metadata=True) - - await cognee.add(file_path) - - await cognee.cognify() - - query_text = "List me all the important characters in Alice in Wonderland." - - from cognee.infrastructure.databases.vector import get_vector_engine - - vector_engine = get_vector_engine() - - collection_name = "Entity_name" - - query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0] - - result = await vector_engine.search( - collection_name=collection_name, query_vector=query_vector, limit=None - ) - - # Check that we did not accidentally use any default value for limit in vector search along the way (like 5, 10, or 15) - assert len(result) > 15 \ No newline at end of file