Update app.py

app.py

@@ -39,7 +39,6 @@ MODELS = {
     }
 }
 
-# Define all invoice extraction functions first
 def get_api_key(model_choice):
     """Get the appropriate API key based on model choice"""
     api_key_env = MODELS[model_choice]["api_key_env"]
@@ -49,44 +48,6 @@ def get_api_key(model_choice):
         st.stop()
     return api_key
 
-def clean_json_response(text):
-    """Improved JSON extraction from API response with better error handling"""
-    # First try to parse directly as JSON
-    try:
-        return json.loads(text)
-    except json.JSONDecodeError:
-        pass
-    
-    # Try to extract JSON from markdown code blocks
-    json_match = re.search(r'```(?:json)?\n({.*?})\n```', text, re.DOTALL)
-    if json_match:
-        try:
-            return json.loads(json_match.group(1))
-        except json.JSONDecodeError:
-            pass
-    
-    # Try to extract any JSON-like content
-    json_match = re.search(r'\{.*\}', text, re.DOTALL)
-    if json_match:
-        try:
-            return json.loads(json_match.group(0))
-        except json.JSONDecodeError:
-            pass
-    
-    # Fallback to simple key-value parsing
-    try:
-        data = {}
-        for line in text.split('\n'):
-            if ':' in line:
-                parts = line.split(':', 1)
-                if len(parts) == 2:
-                    key, val = parts
-                    key = key.strip().strip('"').lower().replace(' ', '_')
-                    data[key] = val.strip().strip('"')
-        return data if data else None
-    except Exception:
-        return None
-
 def query_llm(model_choice, prompt):
     """Call the appropriate API based on model choice"""
     config = MODELS[model_choice]
@@ -122,6 +83,14 @@ def query_llm(model_choice, prompt):
                 content = response.json()["choices"][0]["message"]["content"]
                 st.session_state.last_api_response = content
                 st.session_state.last_api_response_raw = response.text
+                
+                # Special handling for Llama 4 Mavericks incomplete responses
+                if model_choice == "Llama 4 Mavericks" and content.count('{') != content.count('}'):
+                    st.warning("⚠️ Received incomplete JSON response. Trying to fix...")
+                    # Try to complete the JSON by adding missing closing braces
+                    missing_braces = content.count('{') - content.count('}')
+                    content += '}' * missing_braces + ']' * (content.count('[') - content.count(']'))
+                
                 return content
             except KeyError as e:
                 st.error(f"KeyError in response: {e}\nFull response: {response.json()}")
@@ -131,6 +100,53 @@ def query_llm(model_choice, prompt):
         st.error(f"🌐 Connection Failed: {str(e)}")
         return None
 
+def clean_json_response(text):
+    """Improved JSON extraction from API response with better error handling"""
+    if not text:
+        return None
+        
+    # Special handling for Llama 4 Mavericks incomplete responses
+    if 'line_items":' in text and ']' not in text.split('line_items":')[-1]:
+        # Try to complete the line items array
+        text = text.split('line_items":')[0] + 'line_items": []}'
+    
+    # First try to parse directly as JSON
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError as e:
+        st.warning(f"First JSON parse attempt failed: {str(e)}")
+    
+    # Try to extract JSON from markdown code blocks
+    json_match = re.search(r'```(?:json)?\n({.*?})\n```', text, re.DOTALL)
+    if json_match:
+        try:
+            return json.loads(json_match.group(1))
+        except json.JSONDecodeError as e:
+            st.warning(f"Markdown JSON parse failed: {str(e)}")
+    
+    # Try to extract any JSON-like content
+    json_match = re.search(r'\{.*\}', text, re.DOTALL)
+    if json_match:
+        try:
+            return json.loads(json_match.group(0))
+        except json.JSONDecodeError as e:
+            st.warning(f"Loose JSON parse failed: {str(e)}")
+    
+    # Fallback to simple key-value parsing
+    try:
+        data = {}
+        for line in text.split('\n'):
+            if ':' in line:
+                parts = line.split(':', 1)
+                if len(parts) == 2:
+                    key, val = parts
+                    key = key.strip().strip('"').lower().replace(' ', '_')
+                    data[key] = val.strip().strip('"')
+        return data if data else None
+    except Exception as e:
+        st.error(f"Final fallback parse failed: {str(e)}")
+        return None
+
 def get_extraction_prompt(model_choice, text):
     """Return the appropriate prompt based on model choice"""
     if model_choice == "DeepSeek v3":