fix: tolerate prose around model JSON

pageindex/utils.py

@@ -20,6 +20,8 @@
 # litellm is imported inside the functions that use it; eager import is slow
 # and fetches a remote model-cost map.
 
+_logger = logging.getLogger(__name__)
+
 # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY
 if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"):
     os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY")
@@ -154,39 +156,41 @@ def get_json_content(response):
     return json_content
 
 
-def extract_json(content):
+def _decode_embedded_json(content):
+    """Return the first JSON object or array embedded in a model response."""
+    decoder = json.JSONDecoder()
+    starts = [index for index in (content.find("{"), content.find("[")) if index >= 0]
+    if not starts:
+        return None
     try:
-        # First, try to extract JSON enclosed within ```json and ```
-        start_idx = content.find("```json")
-        if start_idx != -1:
-            start_idx += 7  # Adjust index to start after the delimiter
-            end_idx = content.rfind("```")
-            json_content = content[start_idx:end_idx].strip()
-        else:
-            # If no delimiters, assume entire content could be JSON
-            json_content = content.strip()
-
-        # Clean up common issues that might cause parsing errors
-        json_content = json_content.replace('None', 'null')  # Replace Python None with JSON null
-        json_content = json_content.replace('\n', ' ').replace('\r', ' ')  # Remove newlines
-        json_content = ' '.join(json_content.split())  # Normalize whitespace
-
-        # Attempt to parse and return the JSON object
-        return json.loads(json_content)
-    except json.JSONDecodeError as e:
-        logging.error(f"Failed to extract JSON: {e}")
-        # Try to clean up the content further if initial parsing fails
-        try:
-            # Remove any trailing commas before closing brackets/braces
-            json_content = json_content.replace(',]', ']').replace(',}', '}')
-            return json.loads(json_content)
-        except:
-            logging.error("Failed to parse JSON even after cleanup")
-            return {}
-    except Exception as e:
-        logging.error(f"Unexpected error while extracting JSON: {e}")
+        value, _ = decoder.raw_decode(content[min(starts):])
+    except json.JSONDecodeError:
+        return None
+    return value
+
+
+def extract_json(content):
+    if not isinstance(content, str):
+        _logger.error("Failed to extract JSON: response is not a string")
         return {}
 
+    # Prefer the fenced body, but also support providers that prepend or append prose.
+    json_content = get_json_content(content)
+    for candidate in (json_content, content):
+        parsed = _decode_embedded_json(candidate)
+        if parsed is not None:
+            return parsed
+
+    # Preserve the legacy repairs for models that emit Python's None or a trailing comma.
+    repaired = json_content.replace('None', 'null')
+    repaired = re.sub(r',\s*([}\]])', r'\1', repaired)
+    parsed = _decode_embedded_json(repaired)
+    if parsed is not None:
+        return parsed
+
+    _logger.error("Failed to parse JSON from model response")
+    return {}
+
 def write_node_id(data, node_id=0):
     if isinstance(data, dict):
         data['node_id'] = str(node_id).zfill(4)
@@ -974,4 +978,3 @@ def print_tree(tree, indent=0):
 def print_wrapped(text, width=100):
     for line in text.splitlines():
         print(textwrap.fill(line, width=width))
-

tests/test_page_index.py

@@ -6,6 +6,7 @@
     process_no_toc,
     process_toc_no_page_numbers,
 )
+from pageindex.utils import extract_json
 
 
 class ProcessTocNoPageNumbersTest(unittest.TestCase):
@@ -65,5 +66,40 @@ def test_secure_doc_text_neutralizes_document_delimiters(self):
         self.assertIn("<physical_index_1>", wrapped)
 
 
+class ExtractJsonTest(unittest.TestCase):
+    def test_extracts_object_wrapped_in_prose(self):
+        response = 'Here is the requested result: {"toc_detected": "yes"}. Hope that helps.'
+
+        self.assertEqual(extract_json(response), {"toc_detected": "yes"})
+
+    def test_extracts_fenced_json_with_braces_in_a_string(self):
+        response = '''The result is:
+```json
+{"thinking": "the source includes {braces}", "completed": "yes"}
+```
+'''
+
+        self.assertEqual(
+            extract_json(response),
+            {"thinking": "the source includes {braces}", "completed": "yes"},
+        )
+
+    def test_extracts_array_wrapped_in_prose(self):
+        response = 'Structured output follows: [{"title": "Introduction"}] Thanks.'
+
+        self.assertEqual(extract_json(response), [{"title": "Introduction"}])
+
+    def test_preserves_legacy_none_and_trailing_comma_repairs(self):
+        response = '{"toc_detected": None, "details": {"source": "model"},}'
+
+        self.assertEqual(
+            extract_json(response),
+            {"toc_detected": None, "details": {"source": "model"}},
+        )
+
+    def test_returns_empty_dict_without_json(self):
+        self.assertEqual(extract_json("I could not produce structured output."), {})
+
+
 if __name__ == "__main__":
     unittest.main()