diff --git a/pageindex/utils.py b/pageindex/utils.py index 235dd09cc..fb6c1a990 100644 --- a/pageindex/utils.py +++ b/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 diff --git a/tests/test_page_index.py b/tests/test_page_index.py index 170d014ea..44ff8c60e 100644 --- a/tests/test_page_index.py +++ b/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("", 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()