Added test_graphrag_happy_path_with_neo4j_message_history

Author: alexthomas93

src/neo4j_graphrag/generation/graphrag.py

@@ -103,7 +103,8 @@ def search(
 
         Args:
             query_text (str): The user question.
-            message_history (Optional[list]): A collection previous messages, with each message having a specific role assigned.
+            message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
+                with each message having a specific role assigned.
             examples (str): Examples added to the LLM prompt.
             retriever_config (Optional[dict]): Parameters passed to the retriever.
                 search method; e.g.: top_k
@@ -175,7 +176,7 @@ def chat_summary_prompt(
         if isinstance(message_history, list):
             message_history = InMemoryMessageHistory(messages=message_history)
         message_list = [
-            ": ".join([f"{value}" for _, value in message.items()])
+            f"{message['role']}: {message['content']}"
             for message in message_history.messages
         ]
         history = "\n".join(message_list)

tests/e2e/test_graphrag_e2e.py

@@ -13,16 +13,18 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call
 
 import neo4j
 import pytest
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.generation.graphrag import GraphRAG
 from neo4j_graphrag.generation.types import RagResultModel
 from neo4j_graphrag.llm import LLMResponse
+from neo4j_graphrag.llm.types import LLMMessage
+from neo4j_graphrag.message_history import Neo4jMessageHistory
 from neo4j_graphrag.retrievers import VectorCypherRetriever
-from neo4j_graphrag.types import RetrieverResult
+from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem
 
 from tests.e2e.conftest import BiologyEmbedder
 from tests.e2e.utils import build_data_objects, populate_neo4j
@@ -79,6 +81,93 @@ def test_graphrag_happy_path(
     assert result.retriever_result is None
 
 
+@pytest.mark.usefixtures("populate_neo4j_db")
+def test_graphrag_happy_path_with_neo4j_message_history(
+    retriever_mock: MagicMock,
+    llm: MagicMock,
+    driver: neo4j.Driver,
+) -> None:
+    rag = GraphRAG(
+        retriever=retriever_mock,
+        llm=llm,
+    )
+    retriever_mock.search.return_value = RetrieverResult(
+        items=[
+            RetrieverResultItem(content="item content 1"),
+            RetrieverResultItem(content="item content 2"),
+        ]
+    )
+    llm.invoke.side_effect = [
+        LLMResponse(content="llm generated summary"),
+        LLMResponse(content="llm generated text"),
+    ]
+    message_history = Neo4jMessageHistory(
+        driver=driver,
+        session_id="123",
+        node_label="Message",
+    )
+    message_history.clear()
+    message_history.add_messages(
+        messages=[
+            LLMMessage(role="user", content="initial question"),
+            LLMMessage(role="assistant", content="answer to initial question"),
+        ]
+    )
+    res = rag.search(
+        query_text="question",
+        message_history=message_history,
+    )
+    expected_retriever_query_text = """
+Message Summary:
+llm generated summary
+
+Current Query:
+question
+"""
+
+    first_invocation_input = """
+Summarize the message history:
+
+user: initial question
+assistant: answer to initial question
+"""
+    first_invocation_system_instruction = "You are a summarization assistant. Summarize the given text in no more than 300 words."
+    second_invocation = """Context:
+item content 1
+item content 2
+
+Examples:
+
+
+Question:
+question
+
+Answer:
+"""
+    retriever_mock.search.assert_called_once_with(
+        query_text=expected_retriever_query_text
+    )
+    assert llm.invoke.call_count == 2
+    llm.invoke.assert_has_calls(
+        [
+            call(
+                input=first_invocation_input,
+                system_instruction=first_invocation_system_instruction,
+            ),
+            call(
+                second_invocation,
+                message_history.messages,
+                system_instruction="Answer the user question using the provided context.",
+            ),
+        ]
+    )
+
+    assert isinstance(res, RagResultModel)
+    assert res.answer == "llm generated text"
+    assert res.retriever_result is None
+    message_history.clear()
+
+
 @pytest.mark.usefixtures("populate_neo4j_db")
 def test_graphrag_happy_path_return_context(
     driver: MagicMock, llm: MagicMock, biology_embedder: BiologyEmbedder
@@ -127,7 +216,7 @@ def test_graphrag_happy_path_return_context(
 
 @pytest.mark.usefixtures("populate_neo4j_db")
 def test_graphrag_happy_path_examples(
-    driver: MagicMock, llm: MagicMock, biology_embedder: BiologyEmbedder
+    driver: MagicMock, llm: MagicMock, biology_embedder: MagicMock
 ) -> None:
     retriever = VectorCypherRetriever(
         driver,