Repository: neo4j/neo4j-graphrag-python

Issue:
No handling of JSON Array LLM output
If the LLM generates a JSON Array instead of a JSON Object (I think this is what gives rise to the Exception at least), the entity_relation_extrator fails:

```console
Traceback (most recent call last):
  File "/<...>/.venv/bin/grag", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/<...>/src/grag/main.py", line 43, in main
    asyncio.run(process(args, args.data.rglob("*.pdf")))
  File "/usr/lib/python3.12/asyncio/runners.py", line 194, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/asyncio/base_events.py", line 687, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/<...>/src/grag/main.py", line 36, in process
    await kg_builder.run_async(file_path=file)
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/kg_builder.py", line 126, in run_async
    return await self.runner.run({"file_path": file_path, "text": text})
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/config/runner.py", line 130, in run
    result = await self.pipeline.run(data=run_param)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 640, in run
    await orchestrator.run(data)
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 327, in run
    await asyncio.gather(*tasks)
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 153, in run_task
    await self.on_task_complete(data=data, task=task, result=res)
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 192, in on_task_complete
    await asyncio.gather(*[self.run_task(n, data) async for n in self.next(task)])
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 153, in run_task
    await self.on_task_complete(data=data, task=task, result=res)
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 192, in on_task_complete
    await asyncio.gather(*[self.run_task(n, data) async for n in self.next(task)])
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 153, in run_task
    await self.on_task_complete(data=data, task=task, result=res)
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 192, in on_task_complete
    await asyncio.gather(*[self.run_task(n, data) async for n in self.next(task)])
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 150, in run_task
    res = await task.run(inputs)
          ^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 105, in run
    res = await self.execute(**inputs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/pipeline/pipeline.py", line 95, in execute
    component_result = await self.component.run(**kwargs)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/pydantic/_internal/_validate_call.py", line 33, in wrapper_function
    return await wrapper(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/components/entity_relation_extractor.py", line 336, in run
    chunk_graphs: list[Neo4jGraph] = list(await asyncio.gather(*tasks))
                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/components/entity_relation_extractor.py", line 280, in run_for_chunk
    chunk_graph = await self.extract_for_chunk(schema, examples, chunk)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<...>/.venv/lib/python3.12/site-packages/neo4j_graphrag/experimental/components/entity_relation_extractor.py", line 228, in extract_for_chunk
    chunk_graph = Neo4jGraph(**result)
                  ^^^^^^^^^^^^^^^^^^^^
TypeError: neo4j_graphrag.experimental.components.types.Neo4jGraph() argument after ** must be a mapping, not list
```

Not sure what the proper way to handle this would be, but what I do know is that it isn't fun too lose all progress on larger documents when this happens :(

Discussion:
@stellasia: Hey Mathias,

Thank you for reporting the issue, and sorry for the inconvenience. A fix is on its way to the main branch, and we are likely to release a new version that will include the fix next week, I'll keep you posted.
@stellasia: The fix for this issue is included in the release `1.4.3`. Note that in these cases where the LLM produces an array, the chunk won't have any entities (but it won't raise error and stop the process anymore).

Solution patch:
diff --git a/docs/source/user_guide_rag.rst b/docs/source/user_guide_rag.rst
index f569505..1eb43bb 100644
--- a/docs/source/user_guide_rag.rst
+++ b/docs/source/user_guide_rag.rst
@@ -561,12 +561,9 @@ certain movie properties, the retrieval query can be structured as follows:
 .. code:: python
 
     retrieval_query = """
-        MATCH
-        (actor:Actor)-[:ACTED_IN]->(node)
-        RETURN
-        node.title AS movie_title,
-        node.plot AS movie_plot,
-        collect(actor.name) AS actors;
+    RETURN  node.title as movieTitle,
+            node.plot as moviePlot,
+            collect { MATCH (actor:Actor)-[:ACTED_IN]->(node) RETURN actor.name } AS actors
     """
     retriever = VectorCypherRetriever(
         driver,
@@ -593,8 +590,8 @@ The `content` field is a formatted string containing the key information intende
 .. code:: python
 
     def result_formatter(record: neo4j.Record) -> RetrieverResultItem:
+        content=f"Movie title: {record.get('movieTitle')}, description: {record.get('movieDescription')}, actors: {record.get('actors')}",
         return RetrieverResultItem(
-            content=f"Movie title: {record.get('movieTitle')}, description: {record.get('movieDescription')}, actors: {record.get('actors')}",
             metadata={
                 "title": record.get('movieTitle'),
                 "score": record.get("score"),
@@ -604,7 +601,7 @@ The `content` field is a formatted string containing the key information intende
     retriever = VectorCypherRetriever(
         driver,
         index_name=INDEX_NAME,
-        retrieval_query="MATCH (node)<-[:ACTED_IN]-(p:Person) RETURN node.title as movieTitle, node.plot as movieDescription, collect(p.name) as actors, score",
+        retrieval_query="OPTIONAL MATCH (node)<-[:ACTED_IN]-(p:Person) RETURN node.title as movieTitle, node.plot as movieDescription, collect(p.name) as actors, score",
         result_formatter=result_formatter,
     )
 
@@ -763,7 +760,7 @@ the graph and return more context:
         driver,
         INDEX_NAME,
         FULLTEXT_INDEX_NAME,
-        retrieval_query="MATCH (node)-[:AUTHORED_BY]->(author:Author)" "RETURN author.name"
+        retrieval_query="MATCH (node)-[:AUTHORED_BY]->(author:Author) RETURN author.name"
         embedder=embedder,
     )
 
diff --git a/examples/customize/retrievers/result_formatter_vector_cypher_retriever.py b/examples/customize/retrievers/result_formatter_vector_cypher_retriever.py
index 42dda21..f1919f4 100644
--- a/examples/customize/retrievers/result_formatter_vector_cypher_retriever.py
+++ b/examples/customize/retrievers/result_formatter_vector_cypher_retriever.py
@@ -21,7 +21,12 @@ INDEX_NAME = "moviePlotsEmbedding"
 
 # for each Movie node matched by the vector search, retrieve more context:
 # the name of all actors starring in that movie
-RETRIEVAL_QUERY = " MATCH (node)<-[:ACTED_IN]-(p:Person) RETURN node.title as movieTitle, node.plot as moviePlot, collect(p.name) as actors, score as similarityScore"
+RETRIEVAL_QUERY = """
+RETURN  node.title as movieTitle,
+        node.plot as moviePlot,
+        collect { MATCH (actor:Actor)-[:ACTED_IN]->(node) RETURN actor.name } AS actors,
+        score as similarityScore
+"""
 
 
 def my_result_formatter(record: neo4j.Record) -> RetrieverResultItem:
diff --git a/examples/retrieve/hybrid_cypher_retriever.py b/examples/retrieve/hybrid_cypher_retriever.py
index f29a214..1ce7515 100644
--- a/examples/retrieve/hybrid_cypher_retriever.py
+++ b/examples/retrieve/hybrid_cypher_retriever.py
@@ -22,7 +22,12 @@ FULLTEXT_INDEX_NAME = "movieFulltext"
 
 # for each Movie node matched by the vector search, retrieve more context:
 # the name of all actors starring in that movie
-RETRIEVAL_QUERY = " MATCH (node)<-[:ACTED_IN]-(p:Person) RETURN node.title as movieTitle, node.plot as moviePlot, collect(p.name) as actors, score as similarityScore"
+RETRIEVAL_QUERY = """
+RETURN  node.title as movieTitle,
+        node.plot as moviePlot,
+        collect { MATCH (actor:Actor)-[:ACTED_IN]->(node) RETURN actor.name } AS actors,
+        score as similarityScore
+"""
 
 with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
     # Initialize the retriever
diff --git a/examples/retrieve/vector_cypher_retriever.py b/examples/retrieve/vector_cypher_retriever.py
index e8dab57..5730236 100644
--- a/examples/retrieve/vector_cypher_retriever.py
+++ b/examples/retrieve/vector_cypher_retriever.py
@@ -20,7 +20,12 @@ INDEX_NAME = "moviePlotsEmbedding"
 
 # for each Movie node matched by the vector search, retrieve more context:
 # the name of all actors starring in that movie
-RETRIEVAL_QUERY = " MATCH (node)<-[:ACTED_IN]-(p:Person) RETURN node.title as movieTitle, node.plot as moviePlot, collect(p.name) as actors, score as similarityScore"
+RETRIEVAL_QUERY = """
+RETURN  node.title as movieTitle,
+        node.plot as moviePlot,
+        collect { MATCH (actor:Actor)-[:ACTED_IN]->(node) RETURN actor.name } AS actors,
+        score as similarityScore
+"""
 
 with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
     # Initialize the retriever
diff --git a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py
index 376971e..1d68612 100644
--- a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py
+++ b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py
@@ -225,7 +225,7 @@ class LLMEntityRelationExtractor(EntityRelationExtractor):
                 logger.debug(f"Invalid JSON: {llm_result.content}")
             result = {"nodes": [], "relationships": []}
         try:
-            chunk_graph = Neo4jGraph(**result)
+            chunk_graph = Neo4jGraph.model_validate(result)
         except ValidationError as e:
             if self.on_error == OnError.RAISE:
                 raise LLMGenerationError("LLM response has improper format") from e
diff --git a/src/neo4j_graphrag/retrievers/text2cypher.py b/src/neo4j_graphrag/retrievers/text2cypher.py
index fb0a352..f275b08 100644
--- a/src/neo4j_graphrag/retrievers/text2cypher.py
+++ b/src/neo4j_graphrag/retrievers/text2cypher.py
@@ -55,7 +55,7 @@ class Text2CypherRetriever(Retriever):
         llm (neo4j_graphrag.generation.llm.LLMInterface): LLM object to generate the Cypher query.
         neo4j_schema (Optional[str]): Neo4j schema used to generate the Cypher query.
         examples (Optional[list[str], optional): Optional user input/query pairs for the LLM to use as examples.
-        custom_prompt (Optional[str]): Optional custom prompt to use instead of auto generated prompt. Will not include the neo4j_schema or examples args, if provided.
+        custom_prompt (Optional[str]): Optional custom prompt to use instead of auto generated prompt. Will include the neo4j_schema for schema and examples for examples prompt parameters, if they are provided.
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -99,7 +99,13 @@ class Text2CypherRetriever(Retriever):
         self.result_formatter = validated_data.result_formatter
         self.custom_prompt = validated_data.custom_prompt
         if validated_data.custom_prompt:
-            neo4j_schema = ""
+            if (
+                validated_data.neo4j_schema_model
+                and validated_data.neo4j_schema_model.neo4j_schema
+            ):
+                neo4j_schema = validated_data.neo4j_schema_model.neo4j_schema
+            else:
+                neo4j_schema = ""
         else:
             if (
                 validated_data.neo4j_schema_model
@@ -124,7 +130,7 @@ class Text2CypherRetriever(Retriever):
 
         Args:
             query_text (str): The natural language query used to search the Neo4j database.
-            prompt_params (Dict[str, Any]): additional values to inject into the custom prompt, if it is provided. Example: {'schema': 'this is the graph schema'}
+            prompt_params (Dict[str, Any]): additional values to inject into the custom prompt, if it is provided. If the schema or examples parameter is specified, it will overwrite the corresponding value passed during initialization. Example: {'schema': 'this is the graph schema'}
 
         Raises:
             SearchValidationError: If validation of the input arguments fail.


Test patch:
diff --git a/tests/unit/experimental/components/test_entity_relation_extractor.py b/tests/unit/experimental/components/test_entity_relation_extractor.py
index a1371ec..f76ab5c 100644
--- a/tests/unit/experimental/components/test_entity_relation_extractor.py
+++ b/tests/unit/experimental/components/test_entity_relation_extractor.py
@@ -178,6 +178,24 @@ async def test_extractor_llm_invalid_json() -> None:
         await extractor.run(chunks=chunks)
 
 
+@pytest.mark.asyncio
+async def test_extractor_llm_invalid_json_is_a_list() -> None:
+    """Test what happens when the returned JSON is a valid JSON list,
+    but it does not match the expected Pydantic model"""
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        # missing "label" for entity
+        content='[{"nodes": [{"id": 0, "entity_type": "Person", "properties": {}}], "relationships": []}]'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    with pytest.raises(LLMGenerationError):
+        await extractor.run(chunks=chunks)
+
+
 @pytest.mark.asyncio
 async def test_extractor_llm_badly_formatted_json_gets_fixed() -> None:
     llm = MagicMock(spec=LLMInterface)
diff --git a/tests/unit/retrievers/test_text2cypher.py b/tests/unit/retrievers/test_text2cypher.py
index 885698a..c5b7614 100644
--- a/tests/unit/retrievers/test_text2cypher.py
+++ b/tests/unit/retrievers/test_text2cypher.py
@@ -254,6 +254,73 @@ def test_t2c_retriever_initialization_with_custom_prompt_and_schema_and_examples
     llm.invoke.assert_called_once_with("This is a custom prompt. test")
 
 
+@patch("neo4j_graphrag.retrievers.base.get_version")
+def test_t2c_retriever_initialization_with_custom_prompt_and_schema_and_examples_for_prompt_params(
+    mock_get_version: MagicMock,
+    driver: MagicMock,
+    llm: MagicMock,
+    neo4j_record: MagicMock,
+) -> None:
+    mock_get_version.return_value = ((5, 23, 0), False, False)
+    prompt = "This is a custom prompt. {query_text} {schema} {examples}"
+    neo4j_schema = "dummy-schema"
+    examples = ["example-1", "example-2"]
+
+    retriever = Text2CypherRetriever(
+        driver=driver,
+        llm=llm,
+        custom_prompt=prompt,
+        neo4j_schema=neo4j_schema,
+        examples=examples,
+    )
+
+    driver.execute_query.return_value = (
+        [neo4j_record],
+        None,
+        None,
+    )
+    retriever.search(query_text="test")
+
+    llm.invoke.assert_called_once_with(
+        "This is a custom prompt. test dummy-schema example-1\nexample-2"
+    )
+
+
+@patch("neo4j_graphrag.retrievers.base.get_version")
+def test_t2c_retriever_initialization_with_custom_prompt_and_unused_schema_and_examples(
+    mock_get_version: MagicMock,
+    driver: MagicMock,
+    llm: MagicMock,
+    neo4j_record: MagicMock,
+) -> None:
+    mock_get_version.return_value = ((5, 23, 0), False, False)
+    prompt = "This is a custom prompt. {query_text} {schema} {examples}"
+    neo4j_schema = "dummy-schema"
+    examples = ["example-1", "example-2"]
+
+    retriever = Text2CypherRetriever(
+        driver=driver,
+        llm=llm,
+        custom_prompt=prompt,
+        neo4j_schema=neo4j_schema,
+        examples=examples,
+    )
+
+    driver.execute_query.return_value = (
+        [neo4j_record],
+        None,
+        None,
+    )
+    retriever.search(
+        query_text="test",
+        prompt_params={"schema": "another-dummy-schema", "examples": "another-example"},
+    )
+
+    llm.invoke.assert_called_once_with(
+        "This is a custom prompt. test another-dummy-schema another-example"
+    )
+
+
 @patch("neo4j_graphrag.retrievers.base.get_version")
 def test_t2c_retriever_invalid_custom_prompt_type(
     mock_get_version: MagicMock, driver: MagicMock, llm: MagicMock