Update backend/app/openrouter_client.py

backend/app/openrouter_client.py

@@ -1,6 +1,7 @@
 import os
 import base64
 import json
+import re
 from typing import Any, Dict
 
 import httpx
@@ -109,4 +110,33 @@ async def extract_fields_from_document(
         text = content
 
     # Try to parse JSON from the model output
-    return json.loads(text)
+    # The model might return JSON wrapped in markdown code blocks or with extra text
+    try:
+        # First, try direct JSON parsing
+        return json.loads(text)
+    except json.JSONDecodeError:
+        # Try to extract JSON from markdown code blocks
+        json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
+        if json_match:
+            try:
+                return json.loads(json_match.group(1))
+            except json.JSONDecodeError:
+                pass
+        
+        # Try to find JSON object in the text (look for {...})
+        json_match = re.search(r'\{.*\}', text, re.DOTALL)
+        if json_match:
+            try:
+                return json.loads(json_match.group(0))
+            except json.JSONDecodeError:
+                pass
+        
+        # If all parsing fails, return a default structure with the raw text
+        return {
+            "doc_type": "other",
+            "confidence": 50.0,
+            "fields": {
+                "raw_text": text[:500],  # First 500 chars
+                "error": "Could not parse JSON from model response"
+            }
+        }