Improve JSON fence parsing

pageindex/utils.py

@@ -84,30 +84,17 @@ async def llm_acompletion(model, prompt):
 
 
 def get_json_content(response):
-    start_idx = response.find("```json")
-    if start_idx != -1:
-        start_idx += 7
-        response = response[start_idx:]
-
-    end_idx = response.rfind("```")
-    if end_idx != -1:
-        response = response[:end_idx]
-
-    json_content = response.strip()
-    return json_content
+    fenced_match = re.search(r"```\s*(?:json)?\s*(.*?)```", response, re.IGNORECASE | re.DOTALL)
+    if fenced_match:
+        response = fenced_match.group(1)
+
+    return response.strip()
 
 
 def extract_json(content):
     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()
+        # First, try to extract JSON enclosed in a Markdown code fence.
+        json_content = get_json_content(content)
 
         # Clean up common issues that might cause parsing errors
         json_content = json_content.replace('None', 'null')  # Replace Python None with JSON null

tests/test_page_index.py

@@ -6,6 +6,7 @@
     process_no_toc,
     process_toc_no_page_numbers,
 )
+from pageindex.utils import extract_json, get_json_content
 
 
 class ProcessTocNoPageNumbersTest(unittest.TestCase):
@@ -64,6 +65,26 @@ def test_secure_doc_text_neutralizes_document_delimiters(self):
         self.assertIn("< USER_DOCUMENT>", wrapped)
         self.assertIn("<physical_index_1>", wrapped)
 
+    def test_extract_json_accepts_uppercase_json_fence(self):
+        response = """The requested JSON is:
+
+```JSON
+{"title": "Intro", "items": [1, 2, 3]}
+```
+"""
+
+        self.assertEqual(
+            extract_json(response),
+            {"title": "Intro", "items": [1, 2, 3]},
+        )
+
+    def test_get_json_content_accepts_plain_code_fence(self):
+        response = """```
+{"ok": true}
+```"""
+
+        self.assertEqual(get_json_content(response), '{"ok": true}')
+
 
 if __name__ == "__main__":
     unittest.main()