Update app.py

app.py

@@ -19,36 +19,41 @@ model = genai.GenerativeModel(
 # Streamlit UI
 st.title("Lecture Notes Mindmap Generator")
 
+# Initialize session state variables
+if 'lecture_notes' not in st.session_state:
+    st.session_state.lecture_notes = ""
+if 'current_mindmap' not in st.session_state:
+    st.session_state.current_mindmap = None
+if 'processed_chunks' not in st.session_state:
+    st.session_state.processed_chunks = 0
+
 # Text area for input
-lecture_notes = st.text_area("Paste your lecture notes here:", height=300)
+lecture_notes = st.text_area("Paste your lecture notes here:", height=300, value=st.session_state.lecture_notes)
+
+def chunk_notes(notes, chunk_size=500):
+    words = notes.split()
+    return [' '.join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
 
-def generate_mindmap_step(notes, current_mindmap=None):
+def generate_mindmap_step(notes_chunk, current_mindmap=None):
     prompt = f"""
-    Create or extend a hierarchical mindmap structure from the following lecture notes. 
+    Extend the hierarchical mindmap structure based on the following chunk of lecture notes. 
     The structure should be in JSON format, with each node having a 'name' and optional 'children' array.
     
-    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"}
+    Current mindmap:
+    {json.dumps(current_mindmap) if current_mindmap else "{}"}
     
-    Lecture notes:
-    {notes}
+    New lecture notes chunk:
+    {notes_chunk}
     
-    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.
+    Return only the JSON structure of the entire updated mindmap, without any additional text or explanation.
+    Ensure the JSON is valid and complete.
     """
     
     try:
         response = model.generate_content(prompt)
         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 to parse JSON, handling potential leading/trailing characters
         try:
             mindmap_step = json.loads(response_text)
         except json.JSONDecodeError:
@@ -61,35 +66,17 @@ def generate_mindmap_step(notes, current_mindmap=None):
             else:
                 raise ValueError("Unable to extract valid JSON from the response")
         
-        return mindmap_step, "MORE_CONTENT_PENDING" in str(mindmap_step)
+        return 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
+        return None
 
 def display_mindmap(data, indent=0):
     if isinstance(data, str):
         try:
             data = json.loads(data)
         except json.JSONDecodeError:
-            st.error("Failed to parse the generated mindmap data. Please try again.")
+            st.error("Failed to parse the mindmap data. Please try again.")
             return
     
     if isinstance(data, dict):
@@ -101,46 +88,52 @@ def display_mindmap(data, indent=0):
         for item in data:
             display_mindmap(item, indent)
 
-# Button to generate mindmap
-if st.button("Generate Mindmap"):
+# Button to generate or continue mindmap
+if st.button("Generate/Continue Mindmap"):
     if lecture_notes:
+        st.session_state.lecture_notes = lecture_notes
+        chunks = chunk_notes(lecture_notes)
+        
         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}...")
+            for i in range(st.session_state.processed_chunks, len(chunks)):
+                st.text(f"Processing chunk {i+1} of {len(chunks)}...")
                 
-                step_mindmap, more_content = generate_mindmap_step(lecture_notes, full_mindmap)
+                step_mindmap = generate_mindmap_step(chunks[i], st.session_state.current_mindmap)
                 
                 if step_mindmap:
-                    full_mindmap = merge_mindmaps(full_mindmap, step_mindmap)
+                    st.session_state.current_mindmap = step_mindmap
+                    st.session_state.processed_chunks = i + 1
                     time.sleep(20)  # Respect API rate limit
                 else:
-                    st.warning(f"Iteration {iteration} failed to generate a valid mindmap step. Continuing to next iteration.")
-                    continue
+                    st.warning(f"Failed to process chunk {i+1}. You can try continuing from this point.")
+                    break
             
-            if full_mindmap:
-                st.session_state.mindmap_data = json.dumps(full_mindmap)
-                st.success("Mindmap generated successfully!")
+            if st.session_state.current_mindmap:
+                st.success(f"Mindmap generated successfully! Processed {st.session_state.processed_chunks} out of {len(chunks)} chunks.")
             else:
-                st.error("Failed to generate the complete mindmap. Please try again.")
+                st.error("Failed to generate the mindmap. Please try again.")
     else:
         st.warning("Please enter your lecture notes first.")
 
-# Display the mindmap if data is available
-if 'mindmap_data' in st.session_state:
-    st.subheader("Generated Mindmap")
-    display_mindmap(st.session_state.mindmap_data)
+# Button to clear and restart
+if st.button("Clear and Restart"):
+    st.session_state.lecture_notes = ""
+    st.session_state.current_mindmap = None
+    st.session_state.processed_chunks = 0
+    st.success("Cleared all data. You can start a new mindmap generation.")
+
+# Display the current state of the mindmap
+if st.session_state.current_mindmap:
+    st.subheader("Current Mindmap")
+    display_mindmap(st.session_state.current_mindmap)
 
 st.write("""
 To use this app:
 1. Paste your lecture notes into the text area.
-2. Click the "Generate Mindmap" button.
-3. Wait for the mindmap to be generated (this may take a few minutes).
-4. The generated mindmap will be displayed below.
+2. Click the "Generate/Continue Mindmap" button to start or continue the process.
+3. The mindmap will be generated in chunks, and you can see the progress.
+4. If the process is interrupted, you can continue from where it left off.
+5. Use the "Clear and Restart" button to start over with new notes.
 
 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.