Update app.py

app.py

@@ -42,7 +42,25 @@ def generate_mindmap_step(notes, current_mindmap=None):
     
     try:
         response = model.generate_content(prompt)
-        mindmap_step = json.loads(response.text)
+        response_text = response.text.strip()
+        
+        # Debug: Print raw response
+        st.text("Raw API Response:")
+        st.text(response_text)
+        
+        # Attempt to parse JSON, handling potential leading/trailing characters
+        try:
+            mindmap_step = json.loads(response_text)
+        except json.JSONDecodeError:
+            # Try to find and extract JSON content
+            start = response_text.find('{')
+            end = response_text.rfind('}') + 1
+            if start != -1 and end != -1:
+                json_content = response_text[start:end]
+                mindmap_step = json.loads(json_content)
+            else:
+                raise ValueError("Unable to extract valid JSON from the response")
+        
         return mindmap_step, "MORE_CONTENT_PENDING" in str(mindmap_step)
     except Exception as e:
         st.error(f"An error occurred while generating the mindmap: {str(e)}")
@@ -101,7 +119,8 @@ if st.button("Generate Mindmap"):
                     full_mindmap = merge_mindmaps(full_mindmap, step_mindmap)
                     time.sleep(20)  # Respect API rate limit
                 else:
-                    break
+                    st.warning(f"Iteration {iteration} failed to generate a valid mindmap step. Continuing to next iteration.")
+                    continue
             
             if full_mindmap:
                 st.session_state.mindmap_data = json.dumps(full_mindmap)