From 629a7b9b653f4ef6c64725ff076e7fdc925ba9d6 Mon Sep 17 00:00:00 2001
From: cnndabbler <didierlacroix@comcast.net>
Date: Fri, 12 Jun 2026 12:18:38 -0700
Subject: [PATCH] fix(extract_json): tolerate non-strict model JSON (e.g.
 DeepSeek)

extract_json() assumed the whole response is JSON and returned {} on any
parse failure, which then KeyError-crashed callers (toc_detector_single_page)
mid-index-build on models that wrap JSON in prose/fences. Add a balanced-brace
fallback that pulls the first {...}/[...] object out of the raw response, and
default toc_detector's key access so a single bad page can't abort the run.

Repros on deepseek/deepseek-v4-flash; OpenAI/glm happened to match the strict
path. Fixes intermittent 'Processing failed' on long PDFs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 pageindex/page_index.py |  4 ++--
 pageindex/utils.py      | 45 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/pageindex/page_index.py b/pageindex/page_index.py
index 9004309fb..fe462e231 100644
--- a/pageindex/page_index.py
+++ b/pageindex/page_index.py
@@ -118,8 +118,8 @@ def toc_detector_single_page(content, model=None):
 
     response = llm_completion(model=model, prompt=prompt)
     # print('response', response)
-    json_content = extract_json(response)    
-    return json_content['toc_detected']
+    json_content = extract_json(response)
+    return json_content.get('toc_detected', 'no')
 
 
 def check_if_toc_extraction_is_complete(content, toc, model=None):
diff --git a/pageindex/utils.py b/pageindex/utils.py
index f00ccf3a7..5e8bced6c 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -96,6 +96,44 @@ def get_json_content(response):
     return json_content
          
 
+def _extract_balanced_json(text):
+    """Find and parse the first balanced {...} or [...] object in text.
+
+    Robustness fallback for models (e.g. DeepSeek) that occasionally wrap the
+    JSON in prose or code fences instead of returning it bare. Returns the
+    parsed object, or None if nothing parseable is found.
+    """
+    for open_ch, close_ch in (('{', '}'), ('[', ']')):
+        start = text.find(open_ch)
+        if start == -1:
+            continue
+        depth = 0
+        in_str = False
+        esc = False
+        for i in range(start, len(text)):
+            ch = text[i]
+            if in_str:
+                if esc:
+                    esc = False
+                elif ch == '\\':
+                    esc = True
+                elif ch == '"':
+                    in_str = False
+                continue
+            if ch == '"':
+                in_str = True
+            elif ch == open_ch:
+                depth += 1
+            elif ch == close_ch:
+                depth -= 1
+                if depth == 0:
+                    try:
+                        return json.loads(text[start:i + 1])
+                    except json.JSONDecodeError:
+                        break
+    return None
+
+
 def extract_json(content):
     try:
         # First, try to extract JSON enclosed within ```json and ```
@@ -122,7 +160,12 @@ def extract_json(content):
             # Remove any trailing commas before closing brackets/braces
             json_content = json_content.replace(',]', ']').replace(',}', '}')
             return json.loads(json_content)
-        except:
+        except json.JSONDecodeError:
+            # Last resort: pull the first balanced JSON object out of the raw
+            # response (handles models that add prose/fences around the JSON).
+            obj = _extract_balanced_json(content)
+            if obj is not None:
+                return obj
             logging.error("Failed to parse JSON even after cleanup")
             return {}
     except Exception as e: