cognee 0.5.1__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/test_user_qa_feedback.py

@@ -0,0 +1,312 @@
+import pytest
+from unittest.mock import AsyncMock, patch, MagicMock
+from uuid import UUID, NAMESPACE_OID, uuid5
+
+from cognee.modules.retrieval.user_qa_feedback import UserQAFeedback
+from cognee.modules.retrieval.utils.models import UserFeedbackEvaluation, UserFeedbackSentiment
+from cognee.modules.engine.models import NodeSet
+
+
+@pytest.fixture
+def mock_feedback_evaluation():
+    """Create a mock feedback evaluation."""
+    evaluation = MagicMock(spec=UserFeedbackEvaluation)
+    evaluation.evaluation = MagicMock()
+    evaluation.evaluation.value = "positive"
+    evaluation.score = 4.5
+    return evaluation
+
+
+@pytest.fixture
+def mock_graph_engine():
+    """Create a mock graph engine."""
+    engine = AsyncMock()
+    engine.get_last_user_interaction_ids = AsyncMock(return_value=[])
+    engine.add_edges = AsyncMock()
+    engine.apply_feedback_weight = AsyncMock()
+    return engine
+
+
+class TestUserQAFeedback:
+    @pytest.mark.asyncio
+    async def test_init_default(self):
+        """Test UserQAFeedback initialization with default last_k."""
+        retriever = UserQAFeedback()
+        assert retriever.last_k == 1
+
+    @pytest.mark.asyncio
+    async def test_init_custom_last_k(self):
+        """Test UserQAFeedback initialization with custom last_k."""
+        retriever = UserQAFeedback(last_k=5)
+        assert retriever.last_k == 5
+
+    @pytest.mark.asyncio
+    async def test_add_feedback_success_with_relationships(
+        self, mock_feedback_evaluation, mock_graph_engine
+    ):
+        """Test add_feedback successfully creates feedback with relationships."""
+        interaction_id_1 = str(UUID("550e8400-e29b-41d4-a716-446655440000"))
+        interaction_id_2 = str(UUID("550e8400-e29b-41d4-a716-446655440001"))
+        mock_graph_engine.get_last_user_interaction_ids = AsyncMock(
+            return_value=[interaction_id_1, interaction_id_2]
+        )
+
+        feedback_text = "This answer was helpful"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_feedback_evaluation,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.get_graph_engine",
+                return_value=mock_graph_engine,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.add_data_points",
+                new_callable=AsyncMock,
+            ) as mock_add_data,
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.index_graph_edges",
+                new_callable=AsyncMock,
+            ) as mock_index_edges,
+        ):
+            retriever = UserQAFeedback(last_k=2)
+            result = await retriever.add_feedback(feedback_text)
+
+            assert result == [feedback_text]
+            mock_add_data.assert_awaited_once()
+            mock_graph_engine.add_edges.assert_awaited_once()
+            mock_index_edges.assert_awaited_once()
+            mock_graph_engine.apply_feedback_weight.assert_awaited_once()
+
+            # Verify add_edges was called with correct relationships
+            call_args = mock_graph_engine.add_edges.call_args[0][0]
+            assert len(call_args) == 2
+            assert call_args[0][0] == uuid5(NAMESPACE_OID, name=feedback_text)
+            assert call_args[0][1] == UUID(interaction_id_1)
+            assert call_args[0][2] == "gives_feedback_to"
+            assert call_args[0][3]["relationship_name"] == "gives_feedback_to"
+            assert call_args[0][3]["ontology_valid"] is False
+
+            # Verify apply_feedback_weight was called with correct node IDs
+            weight_call_args = mock_graph_engine.apply_feedback_weight.call_args[1]["node_ids"]
+            assert len(weight_call_args) == 2
+            assert interaction_id_1 in weight_call_args
+            assert interaction_id_2 in weight_call_args
+
+    @pytest.mark.asyncio
+    async def test_add_feedback_success_no_relationships(
+        self, mock_feedback_evaluation, mock_graph_engine
+    ):
+        """Test add_feedback successfully creates feedback without relationships."""
+        mock_graph_engine.get_last_user_interaction_ids = AsyncMock(return_value=[])
+
+        feedback_text = "This answer was helpful"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_feedback_evaluation,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.get_graph_engine",
+                return_value=mock_graph_engine,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.add_data_points",
+                new_callable=AsyncMock,
+            ) as mock_add_data,
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.index_graph_edges",
+                new_callable=AsyncMock,
+            ) as mock_index_edges,
+        ):
+            retriever = UserQAFeedback(last_k=1)
+            result = await retriever.add_feedback(feedback_text)
+
+            assert result == [feedback_text]
+            mock_add_data.assert_awaited_once()
+            # Should not call add_edges or index_graph_edges when no relationships
+            mock_graph_engine.add_edges.assert_not_awaited()
+            mock_index_edges.assert_not_awaited()
+            mock_graph_engine.apply_feedback_weight.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_add_feedback_creates_correct_feedback_node(
+        self, mock_feedback_evaluation, mock_graph_engine
+    ):
+        """Test add_feedback creates CogneeUserFeedback with correct attributes."""
+        mock_graph_engine.get_last_user_interaction_ids = AsyncMock(return_value=[])
+
+        feedback_text = "This was a negative experience"
+        mock_feedback_evaluation.evaluation.value = "negative"
+        mock_feedback_evaluation.score = -3.0
+
+        with (
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_feedback_evaluation,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.get_graph_engine",
+                return_value=mock_graph_engine,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.add_data_points",
+                new_callable=AsyncMock,
+            ) as mock_add_data,
+        ):
+            retriever = UserQAFeedback()
+            await retriever.add_feedback(feedback_text)
+
+            # Verify add_data_points was called with correct CogneeUserFeedback
+            call_args = mock_add_data.call_args[1]["data_points"]
+            assert len(call_args) == 1
+            feedback_node = call_args[0]
+            assert feedback_node.id == uuid5(NAMESPACE_OID, name=feedback_text)
+            assert feedback_node.feedback == feedback_text
+            assert feedback_node.sentiment == "negative"
+            assert feedback_node.score == -3.0
+            assert isinstance(feedback_node.belongs_to_set, NodeSet)
+            assert feedback_node.belongs_to_set.name == "UserQAFeedbacks"
+
+    @pytest.mark.asyncio
+    async def test_add_feedback_calls_llm_with_correct_prompt(
+        self, mock_feedback_evaluation, mock_graph_engine
+    ):
+        """Test add_feedback calls LLM with correct sentiment analysis prompt."""
+        mock_graph_engine.get_last_user_interaction_ids = AsyncMock(return_value=[])
+
+        feedback_text = "Great answer!"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_feedback_evaluation,
+            ) as mock_llm,
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.get_graph_engine",
+                return_value=mock_graph_engine,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.add_data_points",
+                new_callable=AsyncMock,
+            ),
+        ):
+            retriever = UserQAFeedback()
+            await retriever.add_feedback(feedback_text)
+
+            mock_llm.assert_awaited_once()
+            call_kwargs = mock_llm.call_args[1]
+            assert call_kwargs["text_input"] == feedback_text
+            assert "sentiment analysis assistant" in call_kwargs["system_prompt"]
+            assert call_kwargs["response_model"] == UserFeedbackEvaluation
+
+    @pytest.mark.asyncio
+    async def test_add_feedback_uses_last_k_parameter(
+        self, mock_feedback_evaluation, mock_graph_engine
+    ):
+        """Test add_feedback uses last_k parameter when getting interaction IDs."""
+        mock_graph_engine.get_last_user_interaction_ids = AsyncMock(return_value=[])
+
+        feedback_text = "Test feedback"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_feedback_evaluation,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.get_graph_engine",
+                return_value=mock_graph_engine,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.add_data_points",
+                new_callable=AsyncMock,
+            ),
+        ):
+            retriever = UserQAFeedback(last_k=5)
+            await retriever.add_feedback(feedback_text)
+
+            mock_graph_engine.get_last_user_interaction_ids.assert_awaited_once_with(limit=5)
+
+    @pytest.mark.asyncio
+    async def test_add_feedback_with_single_interaction(
+        self, mock_feedback_evaluation, mock_graph_engine
+    ):
+        """Test add_feedback with single interaction ID."""
+        interaction_id = str(UUID("550e8400-e29b-41d4-a716-446655440000"))
+        mock_graph_engine.get_last_user_interaction_ids = AsyncMock(return_value=[interaction_id])
+
+        feedback_text = "Test feedback"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_feedback_evaluation,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.get_graph_engine",
+                return_value=mock_graph_engine,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.add_data_points",
+                new_callable=AsyncMock,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.index_graph_edges",
+                new_callable=AsyncMock,
+            ),
+        ):
+            retriever = UserQAFeedback()
+            result = await retriever.add_feedback(feedback_text)
+
+            assert result == [feedback_text]
+            # Should create relationship for the interaction
+            call_args = mock_graph_engine.add_edges.call_args[0][0]
+            assert len(call_args) == 1
+            assert call_args[0][1] == UUID(interaction_id)
+
+    @pytest.mark.asyncio
+    async def test_add_feedback_applies_weight_correctly(
+        self, mock_feedback_evaluation, mock_graph_engine
+    ):
+        """Test add_feedback applies feedback weight with correct score."""
+        interaction_id = str(UUID("550e8400-e29b-41d4-a716-446655440000"))
+        mock_graph_engine.get_last_user_interaction_ids = AsyncMock(return_value=[interaction_id])
+        mock_feedback_evaluation.score = 4.5
+
+        feedback_text = "Positive feedback"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_feedback_evaluation,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.get_graph_engine",
+                return_value=mock_graph_engine,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.add_data_points",
+                new_callable=AsyncMock,
+            ),
+            patch(
+                "cognee.modules.retrieval.user_qa_feedback.index_graph_edges",
+                new_callable=AsyncMock,
+            ),
+        ):
+            retriever = UserQAFeedback()
+            await retriever.add_feedback(feedback_text)
+
+            mock_graph_engine.apply_feedback_weight.assert_awaited_once_with(
+                node_ids=[interaction_id], weight=4.5
+            )