Update app.py

app.py

@@ -234,6 +234,24 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
     full_prompt = "".join(prompt_parts)
     return full_prompt, included_files
 
+def extract_json_from_text(text):
+    """
+    Attempts to extract a balanced JSON object from the given text.
+    It looks for the first '{' and returns the substring until the braces are balanced.
+    """
+    start = text.find('{')
+    if start == -1:
+        return None
+    count = 0
+    for i in range(start, len(text)):
+        if text[i] == '{':
+            count += 1
+        elif text[i] == '}':
+            count -= 1
+            if count == 0:
+                return text[start:i+1]
+    return None
+
 def call_gemini_api(prompt):
     """
     Calls the Gemini API using the provided prompt.
@@ -277,16 +295,20 @@ def call_gemini_api(prompt):
             time.sleep(1)
             api_status.empty()
             
-            # Improved JSON extraction using regex
-            json_response_text = response.text.strip()
-            json_response_text = json_response_text.replace("```json", "").replace("```", "")
+            # Remove markdown formatting if present
+            json_response_text = response.text.strip().replace("```json", "").replace("```", "")
+            # First try regex extraction
             match = re.search(r'({.*})', json_response_text, re.DOTALL)
             if match:
                 final_json_text = match.group(1)
+            else:
+                # Fallback: extract using our balanced-braces function
+                final_json_text = extract_json_from_text(json_response_text)
+            if final_json_text:
                 insights = json.loads(final_json_text)
                 return insights, None
             else:
-                st.warning("⚠️ Could not find valid JSON object.")
+                st.warning("⚠️ Could not extract a valid JSON object.")
                 return {"raw_response": response.text}, "AI response did not contain clear JSON object."
         except json.JSONDecodeError as json_err:
             st.error(f"🚨 Error parsing JSON: {json_err}")