Update app.py

app.py

@@ -107,15 +107,17 @@ def clean_json_response(text, model_choice):
 
     # Handle Mistral's markdown response
     if model_choice == "Mistral Small":
-        json_match = re.search(r'```(?:json)?\n({.*?})\n```', text, re.DOTALL)
-        if json_match:
-            text = json_match.group(1)
-
+        # Remove all markdown formatting
+        text = text.replace('```json', '').replace('```', '').strip()
+    
     # Handle Llama's truncated response
     if model_choice == "Llama 4 Mavericks":
         if '"line_items":' in text and not text.strip().endswith('}}'):
-            text = text.split('"line_items":')[0] + '"line_items": []}}'
-
+            # Check if we have at least a complete header
+            if '"invoice_header":' in text:
+                # Return with empty line items
+                text = text.split('"line_items":')[0] + '"line_items": []}'
+    
     # Try parsing the cleaned JSON
     try:
         data = json.loads(text)
@@ -129,6 +131,16 @@ def clean_json_response(text, model_choice):
         
         return data
     except json.JSONDecodeError as e:
+        # Try one more time with strict=False for Llama
+        if model_choice == "Llama 4 Mavericks":
+            try:
+                # Find the last complete JSON object
+                end_pos = text.rfind('}')
+                if end_pos != -1:
+                    return json.loads(text[:end_pos+1])
+            except:
+                pass
+        
         st.warning(f"JSON parsing failed: {str(e)}")
         return None