more tests for llminterfacev2

Author: serengil

src/neo4j_graphrag/generation/graphrag.py

@@ -206,10 +206,24 @@ def _build_query(
             summarization_prompt = self._chat_summary_prompt(
                 message_history=message_history
             )
-            summary = self.llm.invoke(
-                input=summarization_prompt,
+            messages = legacy_inputs_to_messages(
+                summarization_prompt,
                 system_instruction=summary_system_message,
-            ).content
+            )
+            if isinstance(self.llm, LLMInterfaceV2) or self.llm.__module__.startswith(
+                "langchain"
+            ):
+                summary = self.llm.invoke(
+                    input=messages,
+                ).content
+            elif isinstance(self.llm, LLMInterface):
+                summary = self.llm.invoke(
+                    input=summarization_prompt,
+                    system_instruction=summary_system_message,
+                ).content
+            else:
+                raise ValueError(f"Type {type(self.llm)} of LLM is not supported.")
+
             return self.conversation_prompt(summary=summary, current_query=query_text)
         return query_text
 

tests/e2e/test_graphrag_e2e.py

@@ -60,7 +60,14 @@ def test_graphrag_happy_path(
     )
 
     llm.invoke.assert_called_once_with(
-        """Context:
+        [
+            {
+                "role": "system",
+                "content": "Answer the user question using the provided context.",
+            },
+            {
+                "role": "user",
+                "content": """Context:
 <Record node={'question': 'In 1953 Watson & Crick built a model of the molecular structure of this, the gene-carrying substance'}>
 <Record node={'question': 'This organ removes excess glucose from the blood & stores it as glycogen'}>
 
@@ -72,8 +79,8 @@ def test_graphrag_happy_path(
 
 Answer:
 """,
-        None,
-        system_instruction="Answer the user question using the provided context.",
+            },
+        ]
     )
     assert isinstance(result, RagResultModel)
     assert result.answer == "some text"
@@ -148,13 +155,21 @@ def test_graphrag_happy_path_with_neo4j_message_history(
     llm.invoke.assert_has_calls(
         [
             call(
-                input=first_invocation_input,
-                system_instruction=first_invocation_system_instruction,
+                [
+                    {"role": "system", "content": first_invocation_system_instruction},
+                    {"role": "user", "content": first_invocation_input},
+                ]
             ),
             call(
-                second_invocation,
-                message_history.messages,
-                system_instruction="Answer the user question using the provided context.",
+                [
+                    {
+                        "role": "system",
+                        "content": "Answer the user question using the provided context.",
+                    },
+                    {"role": "user", "content": "initial question"},
+                    {"role": "assistant", "content": "answer to initial question"},
+                    {"role": "user", "content": second_invocation},
+                ]
             ),
         ]
     )
@@ -190,7 +205,14 @@ def test_graphrag_happy_path_return_context(
     )
 
     llm.invoke.assert_called_once_with(
-        """Context:
+        [
+            {
+                "role": "system",
+                "content": "Answer the user question using the provided context.",
+            },
+            {
+                "role": "user",
+                "content": """Context:
 <Record node={'question': 'In 1953 Watson & Crick built a model of the molecular structure of this, the gene-carrying substance'}>
 <Record node={'question': 'This organ removes excess glucose from the blood & stores it as glycogen'}>
 
@@ -202,8 +224,8 @@ def test_graphrag_happy_path_return_context(
 
 Answer:
 """,
-        None,
-        system_instruction="Answer the user question using the provided context.",
+            },
+        ],
     )
     assert isinstance(result, RagResultModel)
     assert result.answer == "some text"
@@ -236,7 +258,14 @@ def test_graphrag_happy_path_examples(
     )
 
     llm.invoke.assert_called_once_with(
-        """Context:
+        [
+            {
+                "role": "system",
+                "content": "Answer the user question using the provided context.",
+            },
+            {
+                "role": "user",
+                "content": """Context:
 <Record node={'question': 'In 1953 Watson & Crick built a model of the molecular structure of this, the gene-carrying substance'}>
 <Record node={'question': 'This organ removes excess glucose from the blood & stores it as glycogen'}>
 
@@ -248,8 +277,8 @@ def test_graphrag_happy_path_examples(
 
 Answer:
 """,
-        None,
-        system_instruction="Answer the user question using the provided context.",
+            },
+        ]
     )
     assert result.answer == "some text"