Updated Neo4jMessageHistory to allow for optional session node deletion

Author: alexthomas93

src/neo4j_graphrag/message_history.py

@@ -27,7 +27,7 @@
 
 CREATE_SESSION_NODE_QUERY = "MERGE (s:`{node_label}` {{id:$session_id}})"
 
-CLEAR_SESSION_QUERY = (
+DELETE_SESSION_AND_MESSAGES_QUERY = (
     "MATCH (s:`{node_label}`) "
     "WHERE s.id = $session_id "
     "OPTIONAL MATCH p=(s)-[:LAST_MESSAGE]->(:Message)<-[:NEXT*0..]-(:Message) "
@@ -36,6 +36,14 @@
     "DETACH DELETE node;"
 )
 
+DELETE_MESSAGES_QUERY = (
+    "MATCH (s:`{node_label}`)-[:LAST_MESSAGE]->(last_message:Message) "
+    "WHERE s.id = $session_id "
+    "MATCH p=(last_message)<-[:NEXT*0..]-(:Message) "
+    "UNWIND nodes(p) as node "
+    "DETACH DELETE node;"
+)
+
 GET_MESSAGES_QUERY = (
     "MATCH (s:`{node_label}`)-[:LAST_MESSAGE]->(last_message) "
     "WHERE s.id = $session_id MATCH p=(last_message)<-[:NEXT*0.."
@@ -202,9 +210,19 @@ def add_message(self, message: LLMMessage) -> None:
             },
         )
 
-    def clear(self) -> None:
-        """Clear the message history."""
-        self._driver.execute_query(
-            query_=CLEAR_SESSION_QUERY.format(node_label="Session"),
-            parameters_={"session_id": self._session_id},
-        )
+    def clear(self, delete_session_node: bool = False) -> None:
+        """Clear the message history.
+
+        Args:
+            delete_session_node (bool): Whether to delete the session node. Defaults to False.
+        """
+        if delete_session_node:
+            self._driver.execute_query(
+                query_=DELETE_SESSION_AND_MESSAGES_QUERY.format(node_label="Session"),
+                parameters_={"session_id": self._session_id},
+            )
+        else:
+            self._driver.execute_query(
+                query_=DELETE_MESSAGES_QUERY.format(node_label="Session"),
+                parameters_={"session_id": self._session_id},
+            )

tests/e2e/test_message_history_e2e.py

@@ -62,7 +62,7 @@ def test_neo4j_message_history_add_messages(driver: neo4j.Driver) -> None:
     )
 
 
-def test_neo4j_message_history_clear(driver: neo4j.Driver) -> None:
+def test_neo4j_message_history_clear_messages(driver: neo4j.Driver) -> None:
     driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
     message_history = Neo4jMessageHistory(session_id="123", driver=driver)
     message_history.add_messages(
@@ -74,12 +74,38 @@ def test_neo4j_message_history_clear(driver: neo4j.Driver) -> None:
     assert len(message_history.messages) == 2
     message_history.clear()
     assert len(message_history.messages) == 0
+    # Test that the session node is not deleted
+    results = driver.execute_query(
+        query_="MATCH (s:`Session`) WHERE s.id = '123' RETURN s"
+    )
+    assert len(results.records) == 1
+    assert results.records[0]["s"]["id"] == "123"
+    assert list(results.records[0]["s"].labels) == ["Session"]
+
+
+def test_neo4j_message_history_clear_session_and_messages(driver: neo4j.Driver) -> None:
+    driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
+    message_history = Neo4jMessageHistory(session_id="123", driver=driver)
+    message_history.add_messages(
+        [
+            LLMMessage(role="system", content="You are a helpful assistant."),
+            LLMMessage(role="user", content="Hello"),
+        ]
+    )
+    assert len(message_history.messages) == 2
+    message_history.clear(delete_session_node=True)
+    assert len(message_history.messages) == 0
+    # Test that the session node is deleted
+    results = driver.execute_query(
+        query_="MATCH (s:`Session`) WHERE s.id = '123' RETURN s"
+    )
+    assert results.records == []
 
 
 def test_neo4j_message_history_clear_no_messages(driver: neo4j.Driver) -> None:
     driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
     message_history = Neo4jMessageHistory(session_id="123", driver=driver)
-    message_history.clear()
+    message_history.clear(delete_session_node=True)
     results = driver.execute_query(
         query_="MATCH (s:`Session`) WHERE s.id = '123' RETURN s"
     )