cognee 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

cognee/tests/unit/modules/graph/cognee_graph_test.py

@@ -1,4 +1,5 @@
 import pytest
+from unittest.mock import AsyncMock
 
 from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError
 from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
@@ -11,6 +12,30 @@ def setup_graph():
     return CogneeGraph()
 
 
+@pytest.fixture
+def mock_adapter():
+    """Fixture to create a mock adapter for database operations."""
+    adapter = AsyncMock()
+    return adapter
+
+
+@pytest.fixture
+def mock_vector_engine():
+    """Fixture to create a mock vector engine."""
+    engine = AsyncMock()
+    engine.search = AsyncMock()
+    return engine
+
+
+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 {}
+
+
 def test_add_node_success(setup_graph):
     """Test successful addition of a node."""
     graph = setup_graph
@@ -73,3 +98,384 @@ def test_get_edges_nonexistent_node(setup_graph):
     graph = setup_graph
     with pytest.raises(EntityNotFoundError, match="Node with id nonexistent does not exist."):
         graph.get_edges_from_node("nonexistent")
+
+
+@pytest.mark.asyncio
+async def test_project_graph_from_db_full_graph(setup_graph, mock_adapter):
+    """Test projecting a full graph from database."""
+    graph = setup_graph
+
+    nodes_data = [
+        ("1", {"name": "Node1", "description": "First node"}),
+        ("2", {"name": "Node2", "description": "Second node"}),
+    ]
+    edges_data = [
+        ("1", "2", "CONNECTS_TO", {"relationship_name": "connects"}),
+    ]
+
+    mock_adapter.get_graph_data = AsyncMock(return_value=(nodes_data, edges_data))
+
+    await graph.project_graph_from_db(
+        adapter=mock_adapter,
+        node_properties_to_project=["name", "description"],
+        edge_properties_to_project=["relationship_name"],
+    )
+
+    assert len(graph.nodes) == 2
+    assert len(graph.edges) == 1
+    assert graph.get_node("1") is not None
+    assert graph.get_node("2") is not None
+    assert graph.edges[0].node1.id == "1"
+    assert graph.edges[0].node2.id == "2"
+
+
+@pytest.mark.asyncio
+async def test_project_graph_from_db_id_filtered(setup_graph, mock_adapter):
+    """Test projecting an ID-filtered graph from database."""
+    graph = setup_graph
+
+    nodes_data = [
+        ("1", {"name": "Node1"}),
+        ("2", {"name": "Node2"}),
+    ]
+    edges_data = [
+        ("1", "2", "CONNECTS_TO", {"relationship_name": "connects"}),
+    ]
+
+    mock_adapter.get_id_filtered_graph_data = AsyncMock(return_value=(nodes_data, edges_data))
+
+    await graph.project_graph_from_db(
+        adapter=mock_adapter,
+        node_properties_to_project=["name"],
+        edge_properties_to_project=["relationship_name"],
+        relevant_ids_to_filter=["1", "2"],
+    )
+
+    assert len(graph.nodes) == 2
+    assert len(graph.edges) == 1
+    mock_adapter.get_id_filtered_graph_data.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_project_graph_from_db_nodeset_subgraph(setup_graph, mock_adapter):
+    """Test projecting a nodeset subgraph filtered by node type and name."""
+    graph = setup_graph
+
+    nodes_data = [
+        ("1", {"name": "Alice", "type": "Person"}),
+        ("2", {"name": "Bob", "type": "Person"}),
+    ]
+    edges_data = [
+        ("1", "2", "KNOWS", {"relationship_name": "knows"}),
+    ]
+
+    mock_adapter.get_nodeset_subgraph = AsyncMock(return_value=(nodes_data, edges_data))
+
+    await graph.project_graph_from_db(
+        adapter=mock_adapter,
+        node_properties_to_project=["name", "type"],
+        edge_properties_to_project=["relationship_name"],
+        node_type="Person",
+        node_name=["Alice"],
+    )
+
+    assert len(graph.nodes) == 2
+    assert graph.get_node("1") is not None
+    assert len(graph.edges) == 1
+    mock_adapter.get_nodeset_subgraph.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_project_graph_from_db_empty_graph(setup_graph, mock_adapter):
+    """Test projecting empty graph raises EntityNotFoundError."""
+    graph = setup_graph
+
+    mock_adapter.get_graph_data = AsyncMock(return_value=([], []))
+
+    with pytest.raises(EntityNotFoundError, match="Empty graph projected from the database."):
+        await graph.project_graph_from_db(
+            adapter=mock_adapter,
+            node_properties_to_project=["name"],
+            edge_properties_to_project=[],
+        )
+
+
+@pytest.mark.asyncio
+async def test_project_graph_from_db_missing_nodes(setup_graph, mock_adapter):
+    """Test that edges referencing missing nodes raise error."""
+    graph = setup_graph
+
+    nodes_data = [
+        ("1", {"name": "Node1"}),
+    ]
+    edges_data = [
+        ("1", "999", "CONNECTS_TO", {"relationship_name": "connects"}),
+    ]
+
+    mock_adapter.get_graph_data = AsyncMock(return_value=(nodes_data, edges_data))
+
+    with pytest.raises(EntityNotFoundError, match="Edge references nonexistent nodes"):
+        await graph.project_graph_from_db(
+            adapter=mock_adapter,
+            node_properties_to_project=["name"],
+            edge_properties_to_project=["relationship_name"],
+        )
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_to_graph_nodes(setup_graph):
+    """Test mapping vector distances to graph nodes."""
+    graph = setup_graph
+
+    node1 = Node("1", {"name": "Node1"})
+    node2 = Node("2", {"name": "Node2"})
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    node_distances = {
+        "Entity_name": [
+            MockScoredResult("1", 0.95),
+            MockScoredResult("2", 0.87),
+        ]
+    }
+
+    await graph.map_vector_distances_to_graph_nodes(node_distances)
+
+    assert graph.get_node("1").attributes.get("vector_distance") == 0.95
+    assert graph.get_node("2").attributes.get("vector_distance") == 0.87
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_partial_node_coverage(setup_graph):
+    """Test mapping vector distances when only some nodes have results."""
+    graph = setup_graph
+
+    node1 = Node("1", {"name": "Node1"})
+    node2 = Node("2", {"name": "Node2"})
+    node3 = Node("3", {"name": "Node3"})
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+
+    node_distances = {
+        "Entity_name": [
+            MockScoredResult("1", 0.95),
+            MockScoredResult("2", 0.87),
+        ]
+    }
+
+    await graph.map_vector_distances_to_graph_nodes(node_distances)
+
+    assert graph.get_node("1").attributes.get("vector_distance") == 0.95
+    assert graph.get_node("2").attributes.get("vector_distance") == 0.87
+    assert graph.get_node("3").attributes.get("vector_distance") == 3.5
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_multiple_categories(setup_graph):
+    """Test mapping vector distances from multiple collection categories."""
+    graph = setup_graph
+
+    # Create nodes
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    node4 = Node("4")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+    graph.add_node(node4)
+
+    node_distances = {
+        "Entity_name": [
+            MockScoredResult("1", 0.95),
+            MockScoredResult("2", 0.87),
+        ],
+        "TextSummary_text": [
+            MockScoredResult("3", 0.92),
+        ],
+    }
+
+    await graph.map_vector_distances_to_graph_nodes(node_distances)
+
+    assert graph.get_node("1").attributes.get("vector_distance") == 0.95
+    assert graph.get_node("2").attributes.get("vector_distance") == 0.87
+    assert graph.get_node("3").attributes.get("vector_distance") == 0.92
+    assert graph.get_node("4").attributes.get("vector_distance") == 3.5
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_to_graph_edges_with_payload(setup_graph):
+    """Test mapping vector distances to edges when edge_distances provided."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    edge = Edge(
+        node1,
+        node2,
+        attributes={"edge_text": "CONNECTS_TO", "relationship_type": "connects"},
+    )
+    graph.add_edge(edge)
+
+    edge_distances = [
+        MockScoredResult("e1", 0.92, payload={"text": "CONNECTS_TO"}),
+    ]
+
+    await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
+
+    assert graph.edges[0].attributes.get("vector_distance") == 0.92
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_partial_edge_coverage(setup_graph):
+    """Test mapping edge distances when only some edges have results."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+
+    edge1 = Edge(node1, node2, attributes={"edge_text": "CONNECTS_TO"})
+    edge2 = Edge(node2, node3, attributes={"edge_text": "DEPENDS_ON"})
+    graph.add_edge(edge1)
+    graph.add_edge(edge2)
+
+    edge_distances = [
+        MockScoredResult("e1", 0.92, payload={"text": "CONNECTS_TO"}),
+    ]
+
+    await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
+
+    assert graph.edges[0].attributes.get("vector_distance") == 0.92
+    assert graph.edges[1].attributes.get("vector_distance") == 3.5
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_edges_fallback_to_relationship_type(setup_graph):
+    """Test that edge mapping falls back to relationship_type when edge_text is missing."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    edge = Edge(
+        node1,
+        node2,
+        attributes={"relationship_type": "KNOWS"},
+    )
+    graph.add_edge(edge)
+
+    edge_distances = [
+        MockScoredResult("e1", 0.85, payload={"text": "KNOWS"}),
+    ]
+
+    await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
+
+    assert graph.edges[0].attributes.get("vector_distance") == 0.85
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_no_edge_matches(setup_graph):
+    """Test edge mapping when no edges match the distance results."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    edge = Edge(
+        node1,
+        node2,
+        attributes={"edge_text": "CONNECTS_TO", "relationship_type": "connects"},
+    )
+    graph.add_edge(edge)
+
+    edge_distances = [
+        MockScoredResult("e1", 0.92, payload={"text": "SOME_OTHER_EDGE"}),
+    ]
+
+    await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
+
+    assert graph.edges[0].attributes.get("vector_distance") == 3.5
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_none_returns_early(setup_graph):
+    """Test that edge_distances=None returns early without error."""
+    graph = setup_graph
+    graph.add_node(Node("1"))
+    graph.add_node(Node("2"))
+    graph.add_edge(Edge(graph.get_node("1"), graph.get_node("2")))
+
+    await graph.map_vector_distances_to_graph_edges(edge_distances=None)
+
+    assert graph.edges[0].attributes.get("vector_distance") == 3.5
+
+
+@pytest.mark.asyncio
+async def test_calculate_top_triplet_importances(setup_graph):
+    """Test calculating top triplet importances by score."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    node4 = Node("4")
+
+    node1.add_attribute("vector_distance", 0.9)
+    node2.add_attribute("vector_distance", 0.8)
+    node3.add_attribute("vector_distance", 0.7)
+    node4.add_attribute("vector_distance", 0.6)
+
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+    graph.add_node(node4)
+
+    edge1 = Edge(node1, node2)
+    edge2 = Edge(node2, node3)
+    edge3 = Edge(node3, node4)
+
+    edge1.add_attribute("vector_distance", 0.85)
+    edge2.add_attribute("vector_distance", 0.75)
+    edge3.add_attribute("vector_distance", 0.65)
+
+    graph.add_edge(edge1)
+    graph.add_edge(edge2)
+    graph.add_edge(edge3)
+
+    top_triplets = await graph.calculate_top_triplet_importances(k=2)
+
+    assert len(top_triplets) == 2
+
+    assert top_triplets[0] == edge3
+    assert top_triplets[1] == edge2
+
+
+@pytest.mark.asyncio
+async def test_calculate_top_triplet_importances_default_distances(setup_graph):
+    """Test calculating importances when nodes/edges have no vector distances."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    edge = Edge(node1, node2)
+    graph.add_edge(edge)
+
+    top_triplets = await graph.calculate_top_triplet_importances(k=1)
+
+    assert len(top_triplets) == 1
+    assert top_triplets[0] == edge

cognee/tests/unit/modules/memify_tasks/test_cognify_session.py

@@ -0,0 +1,111 @@
+import pytest
+from unittest.mock import AsyncMock, patch
+
+from cognee.tasks.memify.cognify_session import cognify_session
+from cognee.exceptions import CogneeValidationError, CogneeSystemError
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_success():
+    """Test successful cognification of session data."""
+    session_data = (
+        "Session ID: test_session\n\nQuestion: What is AI?\n\nAnswer: AI is artificial intelligence"
+    )
+
+    with (
+        patch("cognee.add", new_callable=AsyncMock) as mock_add,
+        patch("cognee.cognify", new_callable=AsyncMock) as mock_cognify,
+    ):
+        await cognify_session(session_data, dataset_id="123")
+
+        mock_add.assert_called_once_with(
+            session_data, dataset_id="123", node_set=["user_sessions_from_cache"]
+        )
+        mock_cognify.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_empty_string():
+    """Test cognification fails with empty string."""
+    with pytest.raises(CogneeValidationError) as exc_info:
+        await cognify_session("")
+
+    assert "Session data cannot be empty" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_whitespace_string():
+    """Test cognification fails with whitespace-only string."""
+    with pytest.raises(CogneeValidationError) as exc_info:
+        await cognify_session("   \n\t  ")
+
+    assert "Session data cannot be empty" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_none_data():
+    """Test cognification fails with None data."""
+    with pytest.raises(CogneeValidationError) as exc_info:
+        await cognify_session(None)
+
+    assert "Session data cannot be empty" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_add_failure():
+    """Test cognification handles cognee.add failure."""
+    session_data = "Session ID: test\n\nQuestion: test?"
+
+    with (
+        patch("cognee.add", new_callable=AsyncMock) as mock_add,
+        patch("cognee.cognify", new_callable=AsyncMock),
+    ):
+        mock_add.side_effect = Exception("Add operation failed")
+
+        with pytest.raises(CogneeSystemError) as exc_info:
+            await cognify_session(session_data)
+
+        assert "Failed to cognify session data" in str(exc_info.value)
+        assert "Add operation failed" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_cognify_failure():
+    """Test cognification handles cognify failure."""
+    session_data = "Session ID: test\n\nQuestion: test?"
+
+    with (
+        patch("cognee.add", new_callable=AsyncMock),
+        patch("cognee.cognify", new_callable=AsyncMock) as mock_cognify,
+    ):
+        mock_cognify.side_effect = Exception("Cognify operation failed")
+
+        with pytest.raises(CogneeSystemError) as exc_info:
+            await cognify_session(session_data)
+
+        assert "Failed to cognify session data" in str(exc_info.value)
+        assert "Cognify operation failed" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_re_raises_validation_error():
+    """Test that CogneeValidationError is re-raised as-is."""
+    with pytest.raises(CogneeValidationError):
+        await cognify_session("")
+
+
+@pytest.mark.asyncio
+async def test_cognify_session_with_special_characters():
+    """Test cognification with special characters."""
+    session_data = "Session: test™ © Question: What's special? Answer: Cognee is special!"
+
+    with (
+        patch("cognee.add", new_callable=AsyncMock) as mock_add,
+        patch("cognee.cognify", new_callable=AsyncMock) as mock_cognify,
+    ):
+        await cognify_session(session_data, dataset_id="123")
+
+        mock_add.assert_called_once_with(
+            session_data, dataset_id="123", node_set=["user_sessions_from_cache"]
+        )
+        mock_cognify.assert_called_once()

cognee/tests/unit/modules/memify_tasks/test_extract_user_sessions.py

@@ -0,0 +1,175 @@
+import sys
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from cognee.tasks.memify.extract_user_sessions import extract_user_sessions
+from cognee.exceptions import CogneeSystemError
+from cognee.modules.users.models import User
+
+# Get the actual module object (not the function) for patching
+extract_user_sessions_module = sys.modules["cognee.tasks.memify.extract_user_sessions"]
+
+
+@pytest.fixture
+def mock_user():
+    """Create a mock user."""
+    user = MagicMock(spec=User)
+    user.id = "test-user-123"
+    return user
+
+
+@pytest.fixture
+def mock_qa_data():
+    """Create mock Q&A data."""
+    return [
+        {
+            "question": "What is cognee?",
+            "context": "context about cognee",
+            "answer": "Cognee is a knowledge graph solution",
+            "time": "2025-01-01T12:00:00",
+        },
+        {
+            "question": "How does it work?",
+            "context": "how it works context",
+            "answer": "It processes data and creates graphs",
+            "time": "2025-01-01T12:05:00",
+        },
+    ]
+
+
+@pytest.mark.asyncio
+async def test_extract_user_sessions_success(mock_user, mock_qa_data):
+    """Test successful extraction of sessions."""
+    mock_cache_engine = AsyncMock()
+    mock_cache_engine.get_all_qas.return_value = mock_qa_data
+
+    with (
+        patch.object(extract_user_sessions_module, "session_user") as mock_session_user,
+        patch.object(
+            extract_user_sessions_module, "get_cache_engine", return_value=mock_cache_engine
+        ),
+    ):
+        mock_session_user.get.return_value = mock_user
+
+        sessions = []
+        async for session in extract_user_sessions([{}], session_ids=["test_session"]):
+            sessions.append(session)
+
+        assert len(sessions) == 1
+        assert "Session ID: test_session" in sessions[0]
+        assert "Question: What is cognee?" in sessions[0]
+        assert "Answer: Cognee is a knowledge graph solution" in sessions[0]
+        assert "Question: How does it work?" in sessions[0]
+        assert "Answer: It processes data and creates graphs" in sessions[0]
+
+
+@pytest.mark.asyncio
+async def test_extract_user_sessions_multiple_sessions(mock_user, mock_qa_data):
+    """Test extraction of multiple sessions."""
+    mock_cache_engine = AsyncMock()
+    mock_cache_engine.get_all_qas.return_value = mock_qa_data
+
+    with (
+        patch.object(extract_user_sessions_module, "session_user") as mock_session_user,
+        patch.object(
+            extract_user_sessions_module, "get_cache_engine", return_value=mock_cache_engine
+        ),
+    ):
+        mock_session_user.get.return_value = mock_user
+
+        sessions = []
+        async for session in extract_user_sessions([{}], session_ids=["session1", "session2"]):
+            sessions.append(session)
+
+        assert len(sessions) == 2
+        assert mock_cache_engine.get_all_qas.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_extract_user_sessions_no_data(mock_user, mock_qa_data):
+    """Test extraction handles empty data parameter."""
+    mock_cache_engine = AsyncMock()
+    mock_cache_engine.get_all_qas.return_value = mock_qa_data
+
+    with (
+        patch.object(extract_user_sessions_module, "session_user") as mock_session_user,
+        patch.object(
+            extract_user_sessions_module, "get_cache_engine", return_value=mock_cache_engine
+        ),
+    ):
+        mock_session_user.get.return_value = mock_user
+
+        sessions = []
+        async for session in extract_user_sessions(None, session_ids=["test_session"]):
+            sessions.append(session)
+
+        assert len(sessions) == 1
+
+
+@pytest.mark.asyncio
+async def test_extract_user_sessions_no_session_ids(mock_user):
+    """Test extraction handles no session IDs provided."""
+    mock_cache_engine = AsyncMock()
+
+    with (
+        patch.object(extract_user_sessions_module, "session_user") as mock_session_user,
+        patch.object(
+            extract_user_sessions_module, "get_cache_engine", return_value=mock_cache_engine
+        ),
+    ):
+        mock_session_user.get.return_value = mock_user
+
+        sessions = []
+        async for session in extract_user_sessions([{}], session_ids=None):
+            sessions.append(session)
+
+        assert len(sessions) == 0
+        mock_cache_engine.get_all_qas.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_extract_user_sessions_empty_qa_data(mock_user):
+    """Test extraction handles empty Q&A data."""
+    mock_cache_engine = AsyncMock()
+    mock_cache_engine.get_all_qas.return_value = []
+
+    with (
+        patch.object(extract_user_sessions_module, "session_user") as mock_session_user,
+        patch.object(
+            extract_user_sessions_module, "get_cache_engine", return_value=mock_cache_engine
+        ),
+    ):
+        mock_session_user.get.return_value = mock_user
+
+        sessions = []
+        async for session in extract_user_sessions([{}], session_ids=["empty_session"]):
+            sessions.append(session)
+
+        assert len(sessions) == 0
+
+
+@pytest.mark.asyncio
+async def test_extract_user_sessions_cache_error_handling(mock_user, mock_qa_data):
+    """Test extraction continues on cache error for specific session."""
+    mock_cache_engine = AsyncMock()
+    mock_cache_engine.get_all_qas.side_effect = [
+        mock_qa_data,
+        Exception("Cache error"),
+        mock_qa_data,
+    ]
+
+    with (
+        patch.object(extract_user_sessions_module, "session_user") as mock_session_user,
+        patch.object(
+            extract_user_sessions_module, "get_cache_engine", return_value=mock_cache_engine
+        ),
+    ):
+        mock_session_user.get.return_value = mock_user
+
+        sessions = []
+        async for session in extract_user_sessions(
+            [{}], session_ids=["session1", "session2", "session3"]
+        ):
+            sessions.append(session)
+
+        assert len(sessions) == 2