Updated CLEAR_SESSION_QUERY

Author: alexthomas93

src/neo4j_graphrag/message_history.py

@@ -30,10 +30,12 @@
 CREATE_SESSION_NODE_QUERY = "MERGE (s:`{node_label}` {{id:$session_id}})"
 
 CLEAR_SESSION_QUERY = (
-    "MATCH (s:`{node_label}`)-[:LAST_MESSAGE]->(last_message) "
-    "WHERE s.id = $session_id MATCH p=(last_message)<-[:NEXT]-() "
-    "WITH p, length(p) AS length ORDER BY length DESC LIMIT 1 "
-    "UNWIND nodes(p) as node DETACH DELETE node;"
+    "MATCH (s:`{node_label}`) "
+    "WHERE s.id = $session_id "
+    "OPTIONAL MATCH p=(s)-[:LAST_MESSAGE]->(:Message)<-[:NEXT*0..]-(:Message) "
+    "WITH CASE WHEN p IS NULL THEN [s] ELSE nodes(p) + [s] END AS nodes "
+    "UNWIND nodes AS node "
+    "DETACH DELETE node;"
 )
 
 GET_MESSAGES_QUERY = (

tests/e2e/test_graphrag_e2e.py

@@ -105,7 +105,6 @@ def test_graphrag_happy_path_with_neo4j_message_history(
         driver=driver,
         session_id="123",
     )
-    message_history.clear()
     message_history.add_messages(
         messages=[
             LLMMessage(role="user", content="initial question"),

tests/e2e/test_message_history_e2e.py

@@ -76,6 +76,16 @@ def test_neo4j_message_history_clear(driver: neo4j.Driver) -> None:
     assert len(message_history.messages) == 0
 
 
+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()
+    results = driver.execute_query(
+        query_="MATCH (s:`Session`) WHERE s.id = '123' RETURN s"
+    )
+    assert results.records == []
+
+
 def test_neo4j_message_window_size(driver: neo4j.Driver) -> None:
     driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
     message_history = Neo4jMessageHistory(session_id="123", driver=driver, window=1)