Fixed tests

Author: alexthomas93

src/neo4j_graphrag/generation/graphrag.py

@@ -16,7 +16,7 @@
 
 import logging
 import warnings
-from typing import Any, Optional
+from typing import Any, List, Optional, Union
 
 from pydantic import ValidationError
 
@@ -28,6 +28,7 @@
 from neo4j_graphrag.generation.types import RagInitModel, RagResultModel, RagSearchModel
 from neo4j_graphrag.llm import LLMInterface
 from neo4j_graphrag.llm.types import LLMMessage
+from neo4j_graphrag.message_history import InMemoryMessageHistory, MessageHistory
 from neo4j_graphrag.retrievers.base import Retriever
 from neo4j_graphrag.types import RetrieverResult
 
@@ -84,7 +85,7 @@ def __init__(
     def search(
         self,
         query_text: str = "",
-        message_history: Optional[list[LLMMessage]] = None,
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         examples: str = "",
         retriever_config: Optional[dict[str, Any]] = None,
         return_context: bool | None = None,
@@ -127,6 +128,8 @@ def search(
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors())
+        if isinstance(message_history, list):
+            message_history = InMemoryMessageHistory(messages=message_history)
         query = self.build_query(validated_data.query_text, message_history)
         retriever_result: RetrieverResult = self.retriever.search(
             query_text=query, **validated_data.retriever_config
@@ -139,7 +142,7 @@ def search(
         logger.debug(f"RAG: prompt={prompt}")
         answer = self.llm.invoke(
             prompt,
-            message_history,
+            message_history.messages if message_history else None,
             system_instruction=self.prompt_template.system_instructions,
         )
         result: dict[str, Any] = {"answer": answer.content}
@@ -148,10 +151,14 @@ def search(
         return RagResultModel(**result)
 
     def build_query(
-        self, query_text: str, message_history: Optional[list[LLMMessage]] = None
+        self,
+        query_text: str,
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
     ) -> str:
         summary_system_message = "You are a summarization assistant. Summarize the given text in no more than 300 words."
         if message_history:
+            if isinstance(message_history, list):
+                message_history = InMemoryMessageHistory(messages=message_history)
             summarization_prompt = self.chat_summary_prompt(
                 message_history=message_history
             )
@@ -162,10 +169,14 @@ def build_query(
             return self.conversation_prompt(summary=summary, current_query=query_text)
         return query_text
 
-    def chat_summary_prompt(self, message_history: list[LLMMessage]) -> str:
+    def chat_summary_prompt(
+        self, message_history: Union[List[LLMMessage], MessageHistory]
+    ) -> str:
+        if isinstance(message_history, list):
+            message_history = InMemoryMessageHistory(messages=message_history)
         message_list = [
             ": ".join([f"{value}" for _, value in message.items()])
-            for message in message_history
+            for message in message_history.messages
         ]
         history = "\n".join(message_list)
         return f"""

src/neo4j_graphrag/message_history.py

@@ -57,8 +57,8 @@ def clear(self) -> None: ...
 
 
 class InMemoryMessageHistory(MessageHistory):
-    def __init__(self, messages: List[LLMMessage] = []) -> None:
-        self._messages = messages
+    def __init__(self, messages: Optional[List[LLMMessage]] = None) -> None:
+        self._messages = messages or []
 
     @property
     def messages(self) -> List[LLMMessage]:

tests/unit/test_graphrag.py

@@ -21,6 +21,8 @@
 from neo4j_graphrag.generation.prompts import RagTemplate
 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 InMemoryMessageHistory
 from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem
 
 
@@ -114,14 +116,14 @@ def test_graphrag_happy_path_with_message_history(
 question
 """
 
-    first_invokation_input = """
+    first_invocation_input = """
 Summarize the message history:
 
 user: initial question
 assistant: answer to initial question
 """
-    first_invokation_system_instruction = "You are a summarization assistant. Summarize the given text in no more than 300 words."
-    second_invokation = """Context:
+    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
 
@@ -141,11 +143,11 @@ def test_graphrag_happy_path_with_message_history(
     llm.invoke.assert_has_calls(
         [
             call(
-                input=first_invokation_input,
-                system_instruction=first_invokation_system_instruction,
+                input=first_invocation_input,
+                system_instruction=first_invocation_system_instruction,
             ),
             call(
-                second_invokation,
+                second_invocation,
                 message_history,
                 system_instruction="Answer the user question using the provided context.",
             ),
@@ -157,6 +159,82 @@ def test_graphrag_happy_path_with_message_history(
     assert res.retriever_result is None
 
 
+def test_graphrag_happy_path_with_in_memory_message_history(
+    retriever_mock: MagicMock, llm: MagicMock
+) -> 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 = InMemoryMessageHistory(
+        messages=[
+            LLMMessage(role="user", content="initial question"),
+            LLMMessage(role="assistant", content="answer to initial question"),
+        ]
+    )
+    res = rag.search("question", 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
+
+
 def test_graphrag_happy_path_custom_system_instruction(
     retriever_mock: MagicMock, llm: MagicMock
 ) -> None:

tests/unit/test_message_history.py

@@ -62,9 +62,7 @@ def test_neo4j_message_history_invalid_session_id(driver: MagicMock) -> None:
 def test_neo4j_message_history_invalid_driver() -> None:
     with pytest.raises(ValidationError) as exc_info:
         Neo4jMessageHistory(session_id="123", driver=1.5, node_label="123", window=1)  # type: ignore[arg-type]
-    assert "Input should be a valid dictionary or instance of Neo4jDriver" in str(
-        exc_info.value
-    )
+    assert "Input should be an instance of Driver" in str(exc_info.value)
 
 
 def test_neo4j_message_history_invalid_node_label(driver: MagicMock) -> None:
@@ -81,8 +79,8 @@ def test_neo4j_message_history_invalid_window(driver: MagicMock) -> None:
     assert "Input should be greater than 0" in str(exc_info.value)
 
 
-def test_neo4j_message_history_messages_setter(neo4j_driver: MagicMock) -> None:
-    message_history = Neo4jMessageHistory(session_id="123", driver=neo4j_driver)
+def test_neo4j_message_history_messages_setter(driver: MagicMock) -> None:
+    message_history = Neo4jMessageHistory(session_id="123", driver=driver)
     with pytest.raises(NotImplementedError) as exc_info:
         message_history.messages = [
             LLMMessage(role="user", content="may thy knife chip and shatter"),