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

cognee/tests/unit/modules/retrieval/test_node_edge_vector_search.py

@@ -0,0 +1,273 @@
+import pytest
+from unittest.mock import AsyncMock
+
+from cognee.modules.retrieval.utils.node_edge_vector_search import NodeEdgeVectorSearch
+from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
+
+
+class MockScoredResult:
+    """Mock class for vector search results."""
+
+    def __init__(self, id, score, payload=None):
+        self.id = id
+        self.score = score
+        self.payload = payload or {}
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_single_query_shape():
+    """Test that single query mode produces flat lists (not list-of-lists)."""
+    mock_vector_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = AsyncMock()
+    mock_vector_engine.embedding_engine.embed_text = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    node_results = [MockScoredResult("node1", 0.95), MockScoredResult("node2", 0.87)]
+    edge_results = [MockScoredResult("edge1", 0.92)]
+
+    def search_side_effect(*args, **kwargs):
+        collection_name = kwargs.get("collection_name")
+        if collection_name == "EdgeType_relationship_name":
+            return edge_results
+        return node_results
+
+    mock_vector_engine.search = AsyncMock(side_effect=search_side_effect)
+
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    collections = ["Entity_name", "EdgeType_relationship_name"]
+
+    await vector_search.embed_and_retrieve_distances(
+        query="test query", query_batch=None, collections=collections, wide_search_limit=10
+    )
+
+    assert vector_search.query_list_length is None
+    assert vector_search.edge_distances == edge_results
+    assert vector_search.node_distances["Entity_name"] == node_results
+    mock_vector_engine.embedding_engine.embed_text.assert_called_once_with(["test query"])
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_batch_query_shape_and_empties():
+    """Test that batch query mode produces list-of-lists with correct length and handles empty collections."""
+    mock_vector_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = AsyncMock()
+
+    query_batch = ["query a", "query b"]
+    node_results_query_a = [MockScoredResult("node1", 0.95)]
+    node_results_query_b = [MockScoredResult("node2", 0.87)]
+    edge_results_query_a = [MockScoredResult("edge1", 0.92)]
+    edge_results_query_b = []
+
+    def batch_search_side_effect(*args, **kwargs):
+        collection_name = kwargs.get("collection_name")
+        if collection_name == "EdgeType_relationship_name":
+            return [edge_results_query_a, edge_results_query_b]
+        elif collection_name == "Entity_name":
+            return [node_results_query_a, node_results_query_b]
+        elif collection_name == "MissingCollection":
+            raise CollectionNotFoundError("Collection not found")
+        return [[], []]
+
+    mock_vector_engine.batch_search = AsyncMock(side_effect=batch_search_side_effect)
+
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    collections = [
+        "Entity_name",
+        "EdgeType_relationship_name",
+        "MissingCollection",
+        "EmptyCollection",
+    ]
+
+    await vector_search.embed_and_retrieve_distances(
+        query=None, query_batch=query_batch, collections=collections, wide_search_limit=None
+    )
+
+    assert vector_search.query_list_length == 2
+    assert len(vector_search.edge_distances) == 2
+    assert vector_search.edge_distances[0] == edge_results_query_a
+    assert vector_search.edge_distances[1] == edge_results_query_b
+    assert len(vector_search.node_distances["Entity_name"]) == 2
+    assert vector_search.node_distances["Entity_name"][0] == node_results_query_a
+    assert vector_search.node_distances["Entity_name"][1] == node_results_query_b
+    assert len(vector_search.node_distances["MissingCollection"]) == 2
+    assert vector_search.node_distances["MissingCollection"] == [[], []]
+    assert len(vector_search.node_distances["EmptyCollection"]) == 2
+    assert vector_search.node_distances["EmptyCollection"] == [[], []]
+    mock_vector_engine.embedding_engine.embed_text.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_input_validation_both_provided():
+    """Test that providing both query and query_batch raises ValueError."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    collections = ["Entity_name"]
+
+    with pytest.raises(ValueError, match="Cannot provide both 'query' and 'query_batch'"):
+        await vector_search.embed_and_retrieve_distances(
+            query="test", query_batch=["test1", "test2"], collections=collections
+        )
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_input_validation_neither_provided():
+    """Test that providing neither query nor query_batch raises ValueError."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    collections = ["Entity_name"]
+
+    with pytest.raises(ValueError, match="Must provide either 'query' or 'query_batch'"):
+        await vector_search.embed_and_retrieve_distances(
+            query=None, query_batch=None, collections=collections
+        )
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_extract_relevant_node_ids_single_query():
+    """Test that extract_relevant_node_ids returns IDs for single query mode."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    vector_search.query_list_length = None
+    vector_search.node_distances = {
+        "Entity_name": [MockScoredResult("node1", 0.95), MockScoredResult("node2", 0.87)],
+        "TextSummary_text": [MockScoredResult("node1", 0.90), MockScoredResult("node3", 0.92)],
+    }
+
+    node_ids = vector_search.extract_relevant_node_ids()
+    assert set(node_ids) == {"node1", "node2", "node3"}
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_extract_relevant_node_ids_batch():
+    """Test that extract_relevant_node_ids returns empty list for batch mode."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    vector_search.query_list_length = 2
+    vector_search.node_distances = {
+        "Entity_name": [
+            [MockScoredResult("node1", 0.95)],
+            [MockScoredResult("node2", 0.87)],
+        ],
+    }
+
+    node_ids = vector_search.extract_relevant_node_ids()
+    assert node_ids == []
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_has_results_single_query():
+    """Test has_results returns True when results exist and False when only empties."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+
+    vector_search.edge_distances = [MockScoredResult("edge1", 0.92)]
+    vector_search.node_distances = {}
+    assert vector_search.has_results() is True
+
+    vector_search.edge_distances = []
+    vector_search.node_distances = {"Entity_name": [MockScoredResult("node1", 0.95)]}
+    assert vector_search.has_results() is True
+
+    vector_search.edge_distances = []
+    vector_search.node_distances = {}
+    assert vector_search.has_results() is False
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_has_results_batch():
+    """Test has_results works correctly for batch mode with list-of-lists."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    vector_search.query_list_length = 2
+
+    vector_search.edge_distances = [[MockScoredResult("edge1", 0.92)], []]
+    vector_search.node_distances = {}
+    assert vector_search.has_results() is True
+
+    vector_search.edge_distances = [[], []]
+    vector_search.node_distances = {
+        "Entity_name": [[MockScoredResult("node1", 0.95)], []],
+    }
+    assert vector_search.has_results() is True
+
+    vector_search.edge_distances = [[], []]
+    vector_search.node_distances = {"Entity_name": [[], []]}
+    assert vector_search.has_results() is False
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_single_query_collection_not_found():
+    """Test that CollectionNotFoundError in single query mode returns empty list."""
+    mock_vector_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = AsyncMock()
+    mock_vector_engine.embedding_engine.embed_text = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+    mock_vector_engine.search = AsyncMock(
+        side_effect=CollectionNotFoundError("Collection not found")
+    )
+
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    collections = ["MissingCollection"]
+
+    await vector_search.embed_and_retrieve_distances(
+        query="test query", query_batch=None, collections=collections, wide_search_limit=10
+    )
+
+    assert vector_search.node_distances["MissingCollection"] == []
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_missing_collections_single_query():
+    """Test that missing collections in single-query mode are handled gracefully with empty lists."""
+    mock_vector_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = AsyncMock()
+    mock_vector_engine.embedding_engine.embed_text = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    node_result = MockScoredResult("node1", 0.95)
+
+    def search_side_effect(*args, **kwargs):
+        collection_name = kwargs.get("collection_name")
+        if collection_name == "Entity_name":
+            return [node_result]
+        elif collection_name == "MissingCollection":
+            raise CollectionNotFoundError("Collection not found")
+        return []
+
+    mock_vector_engine.search = AsyncMock(side_effect=search_side_effect)
+
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    collections = ["Entity_name", "MissingCollection", "EmptyCollection"]
+
+    await vector_search.embed_and_retrieve_distances(
+        query="test query", query_batch=None, collections=collections, wide_search_limit=10
+    )
+
+    assert len(vector_search.node_distances["Entity_name"]) == 1
+    assert vector_search.node_distances["Entity_name"][0].id == "node1"
+    assert vector_search.node_distances["Entity_name"][0].score == 0.95
+    assert vector_search.node_distances["MissingCollection"] == []
+    assert vector_search.node_distances["EmptyCollection"] == []
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_has_results_batch_nodes_only():
+    """Test has_results returns True when only node distances are populated in batch mode."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    vector_search.query_list_length = 2
+    vector_search.edge_distances = [[], []]
+    vector_search.node_distances = {
+        "Entity_name": [[MockScoredResult("node1", 0.95)], []],
+    }
+
+    assert vector_search.has_results() is True
+
+
+@pytest.mark.asyncio
+async def test_node_edge_vector_search_has_results_batch_edges_only():
+    """Test has_results returns True when only edge distances are populated in batch mode."""
+    mock_vector_engine = AsyncMock()
+    vector_search = NodeEdgeVectorSearch(vector_engine=mock_vector_engine)
+    vector_search.query_list_length = 2
+    vector_search.edge_distances = [[MockScoredResult("edge1", 0.92)], []]
+    vector_search.node_distances = {}
+
+    assert vector_search.has_results() is True

cognee/tests/unit/modules/retrieval/triplet_retriever_test.py

@@ -31,14 +31,18 @@ async def test_get_context_success(mock_vector_engine):
         "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
+        await retriever.get_completion_from_context("test query", objects, context)
 
     assert context == "Alice knows Bob\nBob works at Tech Corp"
-    mock_vector_engine.search.assert_awaited_once_with("Triplet_text", "test query", limit=5)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "Triplet_text", "test query", limit=5, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
-async def test_get_context_no_collection(mock_vector_engine):
+async def test_get_objects_no_collection(mock_vector_engine):
     """Test that NoDataError is raised when Triplet_text collection doesn't exist."""
     mock_vector_engine.has_collection.return_value = False
 
@@ -49,7 +53,7 @@ async def test_get_context_no_collection(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(NoDataError, match="create_triplet_embeddings"):
-            await retriever.get_context("test query")
+            await retriever.get_retrieved_objects("test query")
 
 
 @pytest.mark.asyncio
@@ -63,13 +67,13 @@ async def test_get_context_empty_results(mock_vector_engine):
         "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        context = await retriever.get_context_from_objects("test query", [])
 
     assert context == ""
 
 
 @pytest.mark.asyncio
-async def test_get_context_collection_not_found_error(mock_vector_engine):
+async def test_get_objects_collection_not_found_error(mock_vector_engine):
     """Test that CollectionNotFoundError is converted to NoDataError."""
     mock_vector_engine.has_collection.side_effect = CollectionNotFoundError("Collection not found")
 
@@ -80,7 +84,7 @@ async def test_get_context_collection_not_found_error(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(NoDataError, match="No data found"):
-            await retriever.get_context("test query")
+            await retriever.get_retrieved_objects("test query")
 
 
 @pytest.mark.asyncio
@@ -98,7 +102,8 @@ async def test_get_context_empty_payload_text(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(KeyError):
-            await retriever.get_context("test query")
+            objects = await retriever.get_retrieved_objects("test query")
+            await retriever.get_context_from_objects("test query", retrieved_objects=objects)
 
 
 @pytest.mark.asyncio
@@ -115,7 +120,8 @@ async def test_get_context_single_triplet(mock_vector_engine):
         "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
 
     assert context == "Single triplet"
 
@@ -172,7 +178,7 @@ async def test_get_completion_without_context(mock_vector_engine):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query")
+        completion = await retriever.get_completion_from_context("test query", None, None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -195,7 +201,9 @@ async def test_get_completion_with_provided_context(mock_vector_engine):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", context="Provided context")
+        completion = await retriever.get_completion_from_context(
+            "test query", None, context="Provided context"
+        )
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -210,7 +218,7 @@ async def test_get_completion_with_session(mock_vector_engine):
     mock_vector_engine.has_collection.return_value = True
     mock_vector_engine.search.return_value = [mock_result]
 
-    retriever = TripletRetriever()
+    retriever = TripletRetriever(session_id="test_session")
 
     mock_user = MagicMock()
     mock_user.id = "test-user-id"
@@ -243,7 +251,9 @@ async def test_get_completion_with_session(mock_vector_engine):
         mock_cache_config.return_value = mock_config
         mock_session_user.get.return_value = mock_user
 
-        completion = await retriever.get_completion("test query", session_id="test_session")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -278,7 +288,9 @@ async def test_get_completion_with_session_no_user_id(mock_vector_engine):
         mock_cache_config.return_value = mock_config
         mock_session_user.get.return_value = None  # No user
 
-        completion = await retriever.get_completion("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -297,7 +309,7 @@ async def test_get_completion_with_response_model(mock_vector_engine):
     mock_vector_engine.has_collection.return_value = True
     mock_vector_engine.search.return_value = [mock_result]
 
-    retriever = TripletRetriever()
+    retriever = TripletRetriever(response_model=TestModel)
 
     with (
         patch(
@@ -314,7 +326,9 @@ async def test_get_completion_with_response_model(mock_vector_engine):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", response_model=TestModel)
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
 
     assert isinstance(completion, list)
     assert len(completion) == 1

cognee/tests/unit/modules/search/test_get_search_type_retriever_instance.py

@@ -0,0 +1,125 @@
+import pytest
+
+from cognee.modules.search.exceptions import UnsupportedSearchTypeError
+from cognee.modules.search.types import SearchType
+
+
+class _DummyCommunityRetriever:
+    def __init__(self, *args, **kwargs):
+        self.kwargs = kwargs
+
+    def get_completion(self, *args, **kwargs):
+        return {"kind": "completion", "init": self.kwargs, "args": args, "kwargs": kwargs}
+
+    def get_context(self, *args, **kwargs):
+        return {"kind": "context", "init": self.kwargs, "args": args, "kwargs": kwargs}
+
+
+@pytest.mark.asyncio
+async def test_feeling_lucky_delegates_to_select_search_type(monkeypatch):
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
+
+    async def _fake_select_search_type(query_text: str):
+        assert query_text == "hello"
+        return SearchType.CHUNKS
+
+    monkeypatch.setattr(mod, "select_search_type", _fake_select_search_type)
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.FEELING_LUCKY, query_text="hello"
+    )
+    assert isinstance(retriever_instance, ChunksRetriever)
+
+
+@pytest.mark.asyncio
+async def test_disallowed_cypher_search_types_raise(monkeypatch):
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+
+    monkeypatch.setenv("ALLOW_CYPHER_QUERY", "false")
+
+    with pytest.raises(UnsupportedSearchTypeError, match="disabled"):
+        await mod.get_search_type_retriever_instance(
+            SearchType.CYPHER, query_text="MATCH (n) RETURN n"
+        )
+
+    with pytest.raises(UnsupportedSearchTypeError, match="disabled"):
+        await mod.get_search_type_retriever_instance(
+            SearchType.NATURAL_LANGUAGE, query_text="Find nodes"
+        )
+
+
+@pytest.mark.asyncio
+async def test_allowed_cypher_search_types_return_tools(monkeypatch):
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.cypher_search_retriever import CypherSearchRetriever
+
+    monkeypatch.setenv("ALLOW_CYPHER_QUERY", "true")
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.CYPHER, query_text="q"
+    )
+    assert isinstance(retriever_instance, CypherSearchRetriever)
+
+
+@pytest.mark.asyncio
+async def test_registered_community_retriever_is_used(monkeypatch):
+    """
+    Integration point: community retrievers are loaded from the registry module and should
+    override the default mapping when present.
+    """
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval import registered_community_retrievers as registry
+
+    monkeypatch.setattr(
+        registry,
+        "registered_community_retrievers",
+        {SearchType.SUMMARIES: _DummyCommunityRetriever},
+    )
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.SUMMARIES, query_text="q", top_k=7
+    )
+
+    assert isinstance(retriever_instance, _DummyCommunityRetriever)
+
+
+@pytest.mark.asyncio
+async def test_unknown_query_type_raises_unsupported():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+
+    with pytest.raises(UnsupportedSearchTypeError, match="UNKNOWN_TYPE"):
+        await mod.get_search_type_retriever_instance("UNKNOWN_TYPE", query_text="q")
+
+
+@pytest.mark.asyncio
+async def test_default_mapping_passes_top_k_to_retrievers():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.SUMMARIES, query_text="q", top_k=4
+    )
+    assert isinstance(retriever_instance, SummariesRetriever)
+
+
+@pytest.mark.asyncio
+async def test_chunks_lexical_returns_jaccard_tools():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.jaccard_retrival import JaccardChunksRetriever
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.CHUNKS_LEXICAL, query_text="q", top_k=3
+    )
+    assert isinstance(retriever_instance, JaccardChunksRetriever)
+
+
+@pytest.mark.asyncio
+async def test_coding_rules_uses_node_name_as_rules_nodeset_name():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.coding_rules_retriever import CodingRulesRetriever
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.CODING_RULES, query_text="q", node_name=[]
+    )
+    assert isinstance(retriever_instance, CodingRulesRetriever)