Replace with updated NoteGenie version

app.py

@@ -29,9 +29,6 @@ app.config["SESSION_FILE_MODE"] = 0o666
 
 Session(app)
 
-# API key file storage path (as backup for session)
-API_KEY_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "api_key.txt")
-
 # Map front-end model names to API model names
 MODEL_MAPPING = {
     "gemini-2.0-pro": "gemini-2.0-pro-exp-02-05",
@@ -42,24 +39,11 @@ MODEL_MAPPING = {
 def get_api_model_name(frontend_model_name):
     return MODEL_MAPPING.get(frontend_model_name, frontend_model_name)
 
-# Function to get API key with multiple fallbacks
+# Function to get API key from different sources
 def get_api_key():
     # Try to get from session first
     api_key = session.get("api_key")
     
-    # If not in session, try to get from file
-    if not api_key:
-        try:
-            if os.path.exists(API_KEY_FILE):
-                with open(API_KEY_FILE, "r") as f:
-                    api_key = f.read().strip()
-                # Restore session if we found a key
-                if api_key:
-                    session["api_key"] = api_key
-                    logger.info("API key restored from backup file")
-        except Exception as e:
-            logger.error(f"Error reading API key file: {str(e)}")
-    
     # Try to get from request header or param (for direct API calls)
     if not api_key:
         api_key = request.headers.get("X-API-Key") or request.args.get("api_key")
@@ -85,18 +69,10 @@ def set_api_key():
         model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
         response = model.generate_content("Say 'API key is valid'")
         
-        # Store API key in session
+        # Store API key in session only
         session.permanent = True
         session["api_key"] = api_key
         
-        # Also store in backup file as failsafe
-        try:
-            with open(API_KEY_FILE, "w") as f:
-                f.write(api_key)
-            os.chmod(API_KEY_FILE, 0o666)  # Make readable/writable
-        except Exception as e:
-            logger.error(f"Failed to write API key to backup file: {str(e)}")
-        
         logger.info("API key successfully set and validated")
         return jsonify({"success": True})
     except Exception as e:
@@ -176,7 +152,7 @@ def prepare_edit_notebook():
         return jsonify({"success": False, "message": "Notebook content is required"}), 400
     
     # Store the notebook in the session for later access
-    session["current_notebook"] = notebook_json
+    session["current_notebook"] = json.dumps(notebook_json)  # Store as JSON string
     
     return jsonify({"success": True})
 
@@ -196,6 +172,8 @@ def edit_notebook_route():
         stream = request.args.get("stream", "true").lower() == "true"
         # For GET streaming requests, get notebook from session
         notebook_json = session.get("current_notebook")
+        if notebook_json:
+            notebook_json = json.loads(notebook_json)  # Parse JSON string back to dict
     else:
         data = request.json
         edit_prompt = data.get("edit_prompt")
@@ -262,7 +240,6 @@ def check_session():
         "session_vars": list(session.keys()),
         "session_file_dir_exists": os.path.exists(app.config["SESSION_FILE_DIR"]),
         "session_file_dir_writable": os.access(app.config["SESSION_FILE_DIR"], os.W_OK),
-        "api_key_file_exists": os.path.exists(API_KEY_FILE),
     }
     
     # Check if running on Hugging Face Spaces

static/js/main.js

@@ -111,6 +111,14 @@ document.addEventListener('DOMContentLoaded', function() {
         // Prevent multiple rapid clicks
         if (isActionInProgress) return;
         
+        // Check if API key is set
+        const apiKey = localStorage.getItem('notegenie_api_key');
+        if (!apiKey) {
+            addSystemMessage('<i class="bi bi-exclamation-triangle"></i> API key is not set. Please set your API key first.');
+            showApiKeyModal();
+            return;
+        }
+        
         // Get the prompt text and check if it's empty
         const prompt = promptInputEl.value.trim();
         
@@ -253,12 +261,26 @@ document.addEventListener('DOMContentLoaded', function() {
         const aiMessageId = 'ai-typing-' + Date.now();
         addTypingIndicator(aiMessageId);
         
+        // Update AI message to show we're starting the edit
+        updateAiMessage(aiMessageId, "**NoteGenie:** Starting to edit your notebook...");
+        
         handleEditStreamingResponse(editPrompt, currentNotebook, modelName, aiMessageId);
         
         // Clear the input field after sending
         promptInputEl.value = '';
     }
     
+    // Create a custom EventSource that supports headers 
+    function createEventSourceWithHeaders(url, headers) {
+        if (typeof EventSourcePolyfill !== 'undefined') {
+            return new EventSourcePolyfill(url, { headers: headers });
+        }
+        
+        // If native EventSource is all we have but we need headers, 
+        // we fall back to using URL parameters for authentication
+        return new EventSource(url);
+    }
+
     function handleStreamingResponse(prompt, modelName, aiMessageId) {
         aiResponseText = '';
         
@@ -310,8 +332,14 @@ document.addEventListener('DOMContentLoaded', function() {
             urlParams.append('api_key', backupApiKey);
         }
         
-        // Create a new event source with API key included
-        eventSource = new EventSource(`/generate_notebook?${urlParams.toString()}`);
+        // Set headers with API key if available
+        const headers = {};
+        if (backupApiKey) {
+            headers['X-API-Key'] = backupApiKey;
+        }
+        
+        // Create a new event source with API key included in both URL and headers
+        eventSource = createEventSourceWithHeaders(`/generate_notebook?${urlParams.toString()}`, headers);
         
         eventSource.onmessage = function(event) {
             // Update our last-activity timestamp
@@ -422,6 +450,14 @@ document.addEventListener('DOMContentLoaded', function() {
         let lastPreviewUpdate = 0;
         const PREVIEW_UPDATE_INTERVAL = 1000; // Update preview every 1 second during streaming, same as generate
         
+        // Get API key from localStorage as a backup
+        const backupApiKey = localStorage.getItem('notegenie_api_key');
+        
+        // Make sure notebook is in proper format before sending
+        const notebookToSend = typeof notebook === 'string' ? JSON.parse(notebook) : notebook;
+        
+        console.log('Sending notebook for editing:', notebookToSend); // Debug log
+        
         // First, send the notebook data to the server so it's available in the session
         fetch('/prepare_edit_notebook', {
             method: 'POST',
@@ -431,7 +467,7 @@ document.addEventListener('DOMContentLoaded', function() {
                 ...(backupApiKey && {'X-API-Key': backupApiKey})
             },
             body: JSON.stringify({
-                notebook: notebook
+                notebook: notebookToSend
             })
         })
         .then(response => {
@@ -442,6 +478,8 @@ document.addEventListener('DOMContentLoaded', function() {
         })
         .then(data => {
             if (data.success) {
+                updateAiMessage(aiMessageId, "**NoteGenie:** Processing your edit request...");
+                
                 // Create URL parameters including API key as fallback
                 const urlParams = new URLSearchParams({
                     edit_prompt: editPrompt,
@@ -454,8 +492,14 @@ document.addEventListener('DOMContentLoaded', function() {
                     urlParams.append('api_key', backupApiKey);
                 }
                 
-                // Now create a new event source for editing with the API key included
-                eventSource = new EventSource(`/edit_notebook?${urlParams.toString()}`);
+                // Set headers with API key if available
+                const headers = {};
+                if (backupApiKey) {
+                    headers['X-API-Key'] = backupApiKey;
+                }
+                
+                // Now create a new event source for editing with the API key included in both URL and headers
+                eventSource = createEventSourceWithHeaders(`/edit_notebook?${urlParams.toString()}`, headers);
                 
                 // Track the last time we got a chunk - for timeout detection
                 let lastChunkTime = Date.now();
@@ -590,6 +634,7 @@ document.addEventListener('DOMContentLoaded', function() {
         .catch(error => {
             console.error('Error preparing notebook for edit:', error);
             updateAiMessage(aiMessageId, '**Error:** ' + error.message);
+            notebookPreviewEl.innerHTML = '<div class="alert alert-danger">Error: Failed to prepare notebook for editing. Please try regenerating the notebook.</div>';
             setGeneratingState(false);
         });
     }