restoring graphrag e2e tests for v1

Author: serengil

tests/e2e/test_graphrag_e2e.py

@@ -60,14 +60,7 @@ def test_graphrag_happy_path(
     )
 
     llm.invoke.assert_called_once_with(
-        [
-            {
-                "role": "system",
-                "content": "Answer the user question using the provided context.",
-            },
-            {
-                "role": "user",
-                "content": """Context:
+        """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'}>
 
@@ -79,8 +72,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"
@@ -155,21 +148,13 @@ def test_graphrag_happy_path_with_neo4j_message_history(
     llm.invoke.assert_has_calls(
         [
             call(
-                [
-                    {"role": "system", "content": first_invocation_system_instruction},
-                    {"role": "user", "content": first_invocation_input},
-                ]
+                input=first_invocation_input,
+                system_instruction=first_invocation_system_instruction,
             ),
             call(
-                [
-                    {
-                        "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},
-                ]
+                second_invocation,
+                message_history.messages,
+                system_instruction="Answer the user question using the provided context.",
             ),
         ]
     )
@@ -205,14 +190,7 @@ def test_graphrag_happy_path_return_context(
     )
 
     llm.invoke.assert_called_once_with(
-        [
-            {
-                "role": "system",
-                "content": "Answer the user question using the provided context.",
-            },
-            {
-                "role": "user",
-                "content": """Context:
+        """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'}>
 
@@ -224,8 +202,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"
@@ -258,14 +236,7 @@ def test_graphrag_happy_path_examples(
     )
 
     llm.invoke.assert_called_once_with(
-        [
-            {
-                "role": "system",
-                "content": "Answer the user question using the provided context.",
-            },
-            {
-                "role": "user",
-                "content": """Context:
+        """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'}>
 
@@ -277,8 +248,8 @@ def test_graphrag_happy_path_examples(
 
 Answer:
 """,
-            },
-        ]
+        None,
+        system_instruction="Answer the user question using the provided context.",
     )
     assert result.answer == "some text"