Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import google.generativeai as genai
 import json
+import time
 
 # Set up Gemini API
 api_key = "AIzaSyBbusXSbq7JTI8Sa7vjGqu-h2zluKhDpX8"
@@ -11,7 +12,7 @@ model = genai.GenerativeModel(
     model_name="gemini-1.5-pro-latest",
     generation_config=genai.GenerationConfig(
         temperature=0.2,
-        max_output_tokens=8192
+        max_output_tokens=2048
     )
 )
 
@@ -21,23 +22,50 @@ st.title("Lecture Notes Mindmap Generator")
 # Text area for input
 lecture_notes = st.text_area("Paste your lecture notes here:", height=300)
 
-# Function to generate mindmap data
-def generate_mindmap(notes):
+def generate_mindmap_step(notes, current_mindmap=None):
     prompt = f"""
-    Create a hierarchical mindmap structure from the following lecture notes. 
+    Create or extend a hierarchical mindmap structure from the following lecture notes. 
     The structure should be in JSON format, with each node having a 'name' and optional 'children' array.
-    Ensure all content from the notes is included in the mindmap.
-    Here are the lecture notes:
-
+    
+    If a current mindmap is provided, extend it. Otherwise, start a new one.
+    Focus on creating a comprehensive structure, but limit your response to about 1500 tokens.
+    
+    Current mindmap (if any):
+    {json.dumps(current_mindmap) if current_mindmap else "None"}
+    
+    Lecture notes:
     {notes}
-
+    
     Return only the JSON structure, without any additional text or explanation.
+    If you haven't covered all the content, include a node named "MORE_CONTENT_PENDING" at the end.
     """
     
-    response = model.generate_content(prompt)
-    return response.text
+    try:
+        response = model.generate_content(prompt)
+        mindmap_step = json.loads(response.text)
+        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)}")
+        return None, False
+
+def merge_mindmaps(base, extension):
+    if not base:
+        return extension
+    
+    def merge_nodes(node1, node2):
+        if 'children' not in node1:
+            node1['children'] = []
+        if 'children' in node2:
+            for child in node2['children']:
+                existing = next((n for n in node1['children'] if n['name'] == child['name']), None)
+                if existing:
+                    merge_nodes(existing, child)
+                else:
+                    node1['children'].append(child)
+    
+    merge_nodes(base, extension)
+    return base
 
-# Function to display mindmap
 def display_mindmap(data, indent=0):
     if isinstance(data, str):
         try:
@@ -45,7 +73,7 @@ def display_mindmap(data, indent=0):
         except json.JSONDecodeError:
             st.error("Failed to parse the generated mindmap data. Please try again.")
             return
-
+    
     if isinstance(data, dict):
         st.markdown('  ' * indent + f"- **{data['name']}**")
         if 'children' in data and isinstance(data['children'], list):
@@ -58,12 +86,28 @@ def display_mindmap(data, indent=0):
 # Button to generate mindmap
 if st.button("Generate Mindmap"):
     if lecture_notes:
-        with st.spinner("Generating mindmap..."):
-            try:
-                mindmap_data = generate_mindmap(lecture_notes)
-                st.session_state.mindmap_data = mindmap_data
-            except Exception as e:
-                st.error(f"An error occurred while generating the mindmap: {str(e)}")
+        with st.spinner("Generating mindmap... This may take a few minutes."):
+            full_mindmap = None
+            more_content = True
+            iteration = 0
+            
+            while more_content and iteration < 10:  # Limit to 10 iterations as a safeguard
+                iteration += 1
+                st.text(f"Processing iteration {iteration}...")
+                
+                step_mindmap, more_content = generate_mindmap_step(lecture_notes, full_mindmap)
+                
+                if step_mindmap:
+                    full_mindmap = merge_mindmaps(full_mindmap, step_mindmap)
+                    time.sleep(20)  # Respect API rate limit
+                else:
+                    break
+            
+            if full_mindmap:
+                st.session_state.mindmap_data = json.dumps(full_mindmap)
+                st.success("Mindmap generated successfully!")
+            else:
+                st.error("Failed to generate the complete mindmap. Please try again.")
     else:
         st.warning("Please enter your lecture notes first.")
 
@@ -76,7 +120,9 @@ st.write("""
 To use this app:
 1. Paste your lecture notes into the text area.
 2. Click the "Generate Mindmap" button.
-3. The generated mindmap will be displayed below.
+3. Wait for the mindmap to be generated (this may take a few minutes).
+4. The generated mindmap will be displayed below.
 
 This app uses Google's Gemini AI to structure your notes into a hierarchical mindmap format.
+It processes the notes in multiple steps to handle large inputs and complex structures.
 """)