More robust JSON parsing (#270)

* More robust JSON parsing

* Added missing dev dependency

* Update tests/test_paperqa.py

Co-authored-by: James Braza <[email protected]>

* Switched to pytest raise values

* Fixed more unit tests from PR review

---------

Co-authored-by: James Braza <[email protected]>

dev-requirements.txt

@@ -14,3 +14,6 @@ pre-commit
 pytest
 types-requests
 types-setuptools
+pytest-asyncio
+pytest-sugar
+pytest-timer

paperqa/docs.py

@@ -41,6 +41,7 @@
     gather_with_concurrency,
     get_loop,
     guess_is_4xx,
+    llm_read_json,
     maybe_is_html,
     maybe_is_pdf,
     maybe_is_text,
@@ -657,9 +658,7 @@ async def process(match):  # noqa: C901, PLR0912
                 success = True
                 if self.prompts.summary_json:
                     try:
-                        # fetch from markdown ```json if present
-                        context = context.split("```json")[-1].split("```")[0]
-                        result_data = json.loads(context)
+                        result_data = llm_read_json(context)
                     except json.decoder.JSONDecodeError:
                         # fallback to string
                         success = False

paperqa/utils.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import inspect
+import json
 import math
 import re
 import string
@@ -179,3 +180,14 @@ def is_coroutine_callable(obj):
     elif callable(obj):  # noqa: RET505
         return inspect.iscoroutinefunction(obj.__call__)
     return False
+
+
+def llm_read_json(text: str) -> dict:
+    """Read LLM output and extract JSON data from it."""
+    # fetch from markdown ```json if present
+    text = text.strip().split("```json")[-1].split("```")[0]
+    # split anything before the first {
+    text = "{" + text.split("{", 1)[-1]
+    # split anything after the last }
+    text = text.rsplit("}", 1)[0] + "}"
+    return json.loads(text)

pyproject.toml

@@ -40,7 +40,7 @@ name = "paper-qa"
 readme = "README.md"
 requires-python = ">=3.8"
 urls = {repository = "https://github.com/whitead/paper-qa"}
-version = "4.5.0"
+version = "4.5.1"
 
 [tool.codespell]
 check-filenames = true