No handling of JSON Array LLM output

[MathiasToftas]

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:

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 :(

[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).