Replace with updated NoteGenie version

.dockerignore

@@ -0,0 +1,14 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.env
+.venv
+env/
+venv/
+ENV/
+.idea/
+.vscode/
+.git/
+flask_session/
+*.log

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+# Ignore Flask session data
+/flask_session/
+
+# Ignore Python cache files
+/utils/__pycache__/

Dockerfile

@@ -12,8 +12,10 @@ COPY . .
 # Create necessary directories
 RUN mkdir -p flask_session
 
-# Set permissions for flask_session directory
-RUN chmod -R 777 flask_session
+# Set permissions for flask_session directory and touch API key file
+RUN chmod -R 777 flask_session && \
+    touch api_key.txt && \
+    chmod 666 api_key.txt
 
 # Set environment variables
 ENV FLASK_APP=app.py

README.md

@@ -1,6 +1,6 @@
 ---
 title: NoteGenie
-emoji: 📔
+emoji: 👁
 colorFrom: red
 colorTo: pink
 sdk: docker
@@ -30,17 +30,6 @@ NoteGenie is an AI-powered Jupyter notebook generator that uses Google's Gemini
    - SECRET_KEY: A secure random string for Flask sessions
    - PORT: 7860 (default for Hugging Face Spaces)
 
-### Troubleshooting on Hugging Face Spaces
-
-If you encounter issues when running NoteGenie on Hugging Face Spaces, try these steps:
-
-1. **API Key**: Ensure you've entered a valid Google Gemini API key
-2. **Browser Refresh**: Try completely refreshing the page
-3. **Shorter Prompts**: Use shorter, more concise prompts (Spaces may have connection timeouts)
-4. **Space Resources**: Check if your Space has enough resources allocated
-5. **Clear Cache**: Try clearing your browser cache or using an incognito window
-6. **Check Logs**: View the Space logs for detailed error information
-
 ## Local Development
 
 1. Install dependencies:
@@ -74,12 +63,6 @@ docker run -p 7860:7860 -e SECRET_KEY=your_secret_key notegenie
 
 NoteGenie requires a Google Gemini API key. Users can set their own API key in the web interface.
 
-## API Key Security
-
-* NoteGenie requires each user to provide their own Google Gemini API key
-* API keys are stored only in the user's browser session and are never shared between users
-* For security reasons, there is no "default" API key - each user must provide their own
-
 ## License
 
 [MIT License](LICENSE)

app.py

@@ -1,42 +1,36 @@
-from flask import Flask, request, jsonify, render_template, session, redirect, url_for, Response
+from flask import Flask, request, jsonify, render_template, session, redirect, url_for
+from flask_session import Session
 import google.generativeai as genai
 import json
 import uuid
 import os
 import logging
-import sys
 from utils.ai_helpers import generate_notebook, stream_notebook_generation, stream_notebook_edit, edit_notebook
 from utils.notebook_helpers import format_notebook, extract_notebook_info
-# Added for server-side sessions
-from flask_session import Session
 
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    stream=sys.stdout
+    format='%(asctime)s - notegenie - %(levelname)s - %(message)s'
 )
-logger = logging.getLogger('notegenie')
+logger = logging.getLogger()
 
 app = Flask(__name__)
 app.config["SECRET_KEY"] = os.environ.get("SECRET_KEY", "notegenie-secret-key-change-in-production")
+app.config["SESSION_TYPE"] = "filesystem"
+app.config["SESSION_PERMANENT"] = True
+app.config["SESSION_USE_SIGNER"] = True
+app.config["PERMANENT_SESSION_LIFETIME"] = 60 * 60 * 24 * 30  # 30 days
+app.config["SESSION_FILE_DIR"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), "flask_session")
+os.makedirs(app.config["SESSION_FILE_DIR"], exist_ok=True)  # Ensure directory exists
+
+# Set a more permissive file mode for session files to avoid permission issues
+app.config["SESSION_FILE_MODE"] = 0o666
 
-# Detect if running on Hugging Face Spaces
-IS_HUGGINGFACE = os.environ.get('SPACE_ID') is not None
-if IS_HUGGINGFACE:
-    logger.info("Running on Hugging Face Spaces environment")
-    # Ensure sessions work properly on Hugging Face
-    app.config["SESSION_COOKIE_SECURE"] = False  # Changed from True to False for HTTP connections
-    app.config["SESSION_COOKIE_HTTPONLY"] = True
-    app.config["SESSION_COOKIE_SAMESITE"] = "Lax"
-    # Add more durable storage for Hugging Face
-    app.config["SESSION_FILE_THRESHOLD"] = 10  # Low threshold to ensure writes
-    # Important: Don't use large session lifetime on Hugging Face
-    app.config["PERMANENT_SESSION_LIFETIME"] = 60 * 60 * 24  # 1 day only
-    # Added: Use server-side sessions to persist API keys across workers
-    app.config["SESSION_TYPE"] = "filesystem"
-    app.config["SESSION_FILE_DIR"] = os.path.join(os.getcwd(), "flask_session")
-    Session(app)
+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 = {
@@ -48,8 +42,35 @@ 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
+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")
+    
+    return api_key
+
 @app.route("/", methods=["GET"])
 def index():
+    # Test session functionality
+    session["session_test"] = True
+    logger.info(f"Session check - variables: {list(session.keys())}")
     return render_template("index.html")
 
 @app.route("/set_api_key", methods=["POST"])
@@ -57,16 +78,25 @@ def set_api_key():
     api_key = request.form.get("api_key")
     if not api_key:
         return jsonify({"success": False, "message": "API key is required"}), 400
+    
     try:
         # Test the API key
         genai.configure(api_key=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 with permanent flag
+        # Store API key in session
         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:
@@ -75,24 +105,13 @@ def set_api_key():
 
 @app.route("/generate_notebook", methods=["GET", "POST"])
 def generate_notebook_route():
-    # SECURITY FIX: Only use session API key, never a global one
-    api_key = session.get("api_key") if "api_key" in session else None
-
-    # Fallback: check URL parameter for API key if not set in session
-    if not api_key:
-        api_key_param = request.args.get('api_key_param')
-        if api_key_param:
-            try:
-                genai.configure(api_key=api_key_param)
-                model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
-                session["api_key"] = api_key_param
-                api_key = api_key_param
-                logger.info("API key set from URL parameter in generate_notebook_route")
-            except Exception as e:
-                logger.error(f"API key validation error from URL parameter in generate_notebook_route: {str(e)}")
+    api_key = get_api_key()
     if not api_key:
         logger.warning("Generate notebook request without API key")
-        return jsonify({"success": False, "message": "API key not set - please set your API key in the settings"}), 401
+        return jsonify({"success": False, "message": "API key not set"}), 401
+    
+    # Always configure genai with the API key for each request
+    genai.configure(api_key=api_key)
     
     # Handle both GET (for streaming) and POST requests
     if request.method == "GET":
@@ -111,18 +130,15 @@ def generate_notebook_route():
     
     # Map the frontend model name to the API model name
     api_model_name = get_api_model_name(model_name)
-    logger.info(f"Generate notebook with model: {api_model_name}, stream: {stream}")
     
     try:
-        # Configure with the api_key we retrieved
-        genai.configure(api_key=api_key)
-        
         # OPTIMIZATION: If format_only is True, skip the AI call and just format the provided content
         if request.method == "POST" and format_only:
             # Use client-provided content as is (it's already the AI response)
             notebook_content = prompt
             notebook_json = format_notebook(notebook_content)
             notebook_info = extract_notebook_info(notebook_content)
+            
             return jsonify({
                 "success": True,
                 "notebook": notebook_json,
@@ -130,12 +146,12 @@ def generate_notebook_route():
                 "description": notebook_info["description"]
             })
         elif stream:
-            logger.info("Starting streaming notebook generation")
-            return stream_notebook_generation(prompt, api_model_name, api_key)
+            return stream_notebook_generation(prompt, api_model_name)
         else:
             notebook_content = generate_notebook(prompt, api_model_name)
             notebook_json = format_notebook(notebook_content)
             notebook_info = extract_notebook_info(notebook_content)
+            
             return jsonify({
                 "success": True,
                 "notebook": notebook_json,
@@ -143,13 +159,14 @@ def generate_notebook_route():
                 "description": notebook_info["description"]
             })
     except Exception as e:
-        logger.error(f"Error generating notebook: {str(e)}", exc_info=True)
+        logger.error(f"Error generating notebook: {str(e)}")
         return jsonify({"success": False, "message": str(e)}), 500
 
 @app.route("/prepare_edit_notebook", methods=["POST"])
 def prepare_edit_notebook():
     """Store the notebook in the session for editing."""
-    if "api_key" not in session:
+    api_key = get_api_key()
+    if not api_key:
         return jsonify({"success": False, "message": "API key not set"}), 401
     
     data = request.json
@@ -165,11 +182,12 @@ def prepare_edit_notebook():
 
 @app.route("/edit_notebook", methods=["GET", "POST"])
 def edit_notebook_route():
-    # SECURITY FIX: Only use session API key, never a global one
-    api_key = session.get("api_key") if "api_key" in session else None
-    
+    api_key = get_api_key()
     if not api_key:
-        return jsonify({"success": False, "message": "API key not set - please set your API key in the settings"}), 401
+        return jsonify({"success": False, "message": "API key not set"}), 401
+    
+    # Always configure genai with the API key for each request
+    genai.configure(api_key=api_key)
     
     # Get edit prompt and current notebook
     if request.method == "GET":
@@ -195,15 +213,14 @@ def edit_notebook_route():
     api_model_name = get_api_model_name(model_name)
     
     try:
-        # Use the api_key we retrieved, not session["api_key"]
-        genai.configure(api_key=api_key)
         if stream:
-            return stream_notebook_edit(edit_prompt, notebook_json, api_model_name, api_key)
+            return stream_notebook_edit(edit_prompt, notebook_json, api_model_name)
         else:
             # Non-streaming path (not used in current UI but kept for API completeness)
             edited_content = edit_notebook(edit_prompt, notebook_json, api_model_name)
             notebook_json = format_notebook(edited_content)
             notebook_info = extract_notebook_info(edited_content)
+            
             return jsonify({
                 "success": True,
                 "notebook": notebook_json,
@@ -211,11 +228,13 @@ def edit_notebook_route():
                 "description": notebook_info["description"]
             })
     except Exception as e:
-        logger.error(f"Error editing notebook: {str(e)}", exc_info=True)
+        app.logger.error(f"Error editing notebook: {str(e)}")
         return jsonify({"success": False, "message": str(e)}), 500
 
 @app.route("/download_notebook", methods=["POST"])
 def download_notebook():
+    from flask import Response
+    
     data = request.json
     notebook_json = data.get("notebook")
     filename = data.get("filename", f"notebook_{uuid.uuid4()}.ipynb")
@@ -231,75 +250,30 @@ def download_notebook():
         mimetype="application/json",
         headers={"Content-Disposition": f"attachment;filename={filename}"}
     )
+    
     return response
 
-@app.route("/session_test", methods=["GET"])
-def session_test():
-    if "session_test" not in session:
-        session["session_test"] = True
-        is_new = True
-    else:
-        is_new = False
-    
-    # Log the current session state
-    session_vars = list(session.keys()) if session else []
-    logger.info(f"Session check - variables: {session_vars}")
-    
-    # Check if API key is available from URL parameter (fallback for HF)
-    api_key_param = request.args.get('api_key_param')
-    if api_key_param and "api_key" not in session:
-        try:
-            # Validate it quickly before accepting
-            genai.configure(api_key=api_key_param)
-            model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
-            # If no exception, store it
-            session["api_key"] = api_key_param
-            logger.info("API key set from URL parameter")
-        except Exception as e:
-            logger.error(f"Invalid API key from URL parameter: {str(e)}")
-    
-    return jsonify({
-        "session_works": True,
-        "is_new_session": is_new,
+# Add a session diagnostic endpoint
+@app.route("/check_session", methods=["GET"])
+def check_session():
+    # For debugging only - would be disabled in production
+    session_data = {
         "has_api_key": "api_key" in session,
-        "session_vars": session_vars
-    })
-
-# Add a session check endpoint to debug session issues
-@app.route("/session_check", methods=["GET"])
-def session_check():
-    if "session_test" not in session:
-        session["session_test"] = True
-        is_new = True
-    else:
-        is_new = False
+        "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),
+    }
     
-    # Log the current session state
-    session_vars = list(session.keys()) if session else []
-    logger.info(f"Session check - variables: {session_vars}")
+    # Check if running on Hugging Face Spaces
+    is_hf_space = "SPACE_ID" in os.environ
+    session_data["is_huggingface_space"] = is_hf_space
     
-    # Check if API key is available from URL parameter (fallback for HF)
-    # SECURITY FIX: Only set API key in the current user's session, never globally
-    api_key_param = request.args.get('api_key_param')
-    if api_key_param and "api_key" not in session:
-        try:
-            # Validate it quickly before accepting
-            genai.configure(api_key=api_key_param)
-            model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
-            # If no exception, store it in session only
-            session["api_key"] = api_key_param
-            logger.info("API key set from URL parameter")
-        except Exception as e:
-            logger.error(f"Invalid API key from URL parameter: {str(e)}")
+    if is_hf_space:
+        logger.info("Running on Hugging Face Spaces environment")
     
-    return jsonify({
-        "session_works": True,
-        "is_new_session": is_new,
-        "has_api_key": "api_key" in session,
-        "session_vars": session_vars
-    })
+    return jsonify(session_data)
 
 if __name__ == "__main__":
-    debug_mode = os.environ.get("FLASK_ENV") == "development"
     port = int(os.environ.get("PORT", 5000))
-    app.run(host="0.0.0.0", port=port, debug=debug_mode)
+    app.run(host="0.0.0.0", port=port, debug=(os.environ.get("FLASK_ENV") == "development"))

static/css/style.css

@@ -135,27 +135,23 @@ h1, h2, h3, h4, h5, h6 {
 
 /* Custom scrollbar styles */
 .conversation-container::-webkit-scrollbar,
-.notebook-preview::-webkit-scrollbar,
-.prompt-input::-webkit-scrollbar {
+.notebook-preview::-webkit-scrollbar {
     width: 8px;
 }
 
 .conversation-container::-webkit-scrollbar-track,
-.notebook-preview::-webkit-scrollbar-track,
-.prompt-input::-webkit-scrollbar-track {
+.notebook-preview::-webkit-scrollbar-track {
     background: transparent;
 }
 
 .conversation-container::-webkit-scrollbar-thumb,
-.notebook-preview::-webkit-scrollbar-thumb,
-.prompt-input::-webkit-scrollbar-thumb {
+.notebook-preview::-webkit-scrollbar-thumb {
     background-color: rgba(95, 99, 104, 0.3);
     border-radius: 20px;
 }
 
 .conversation-container::-webkit-scrollbar-thumb:hover,
-.notebook-preview::-webkit-scrollbar-thumb:hover,
-.prompt-input::-webkit-scrollbar-thumb:hover {
+.notebook-preview::-webkit-scrollbar-thumb:hover {
     background-color: rgba(95, 99, 104, 0.5);
 }
 

static/js/main.js

@@ -166,9 +166,6 @@ document.addEventListener('DOMContentLoaded', function() {
                 // Store API key in localStorage as backup
                 localStorage.setItem('notegenie_api_key', apiKey);
                 
-                // For Hugging Face, add a flag that API key was set in this session
-                sessionStorage.setItem('api_key_set', 'true');
-                
                 if (document.querySelector('#apiKeyModal.show')) {
                     showApiKeyFeedback('API key saved successfully!', 'success');
                     setTimeout(() => {
@@ -298,167 +295,116 @@ document.addEventListener('DOMContentLoaded', function() {
             }
         }, 5000); // Check every 5 seconds
         
-        // Add a retry mechanism
-        let retryCount = 0;
-        const MAX_RETRIES = 2;
+        // Get API key from localStorage as a backup
+        const backupApiKey = localStorage.getItem('notegenie_api_key');
         
-        function createEventSource() {
-            // Create a new event source with a unique timestamp to prevent caching
-            const timestamp = Date.now();
-            
-            // On Hugging Face, add API key as URL parameter if stored (fallback mechanism)
-            let apiKeyParam = '';
-            if (window.location.hostname.includes('huggingface.co') || 
-                window.location.hostname.includes('hf.space')) {
-                const storedApiKey = localStorage.getItem('notegenie_api_key');
-                if (storedApiKey && !sessionStorage.getItem('api_key_tried_in_url')) {
-                    apiKeyParam = `&api_key_param=${encodeURIComponent(storedApiKey)}`;
-                    sessionStorage.setItem('api_key_tried_in_url', 'true');
-                }
-            }
-            
-            eventSource = new EventSource(`/generate_notebook?${new URLSearchParams({
-                prompt: prompt,
-                model: modelName,
-                stream: true,
-                t: timestamp  // Add timestamp to prevent caching
-            }).toString()}${apiKeyParam}`);
+        // Create URL parameters including API key as fallback
+        const urlParams = new URLSearchParams({
+            prompt: prompt,
+            model: modelName,
+            stream: true
+        });
+        
+        // Add API key to URL params if available from localStorage (as a fallback)
+        if (backupApiKey) {
+            urlParams.append('api_key', backupApiKey);
+        }
+        
+        // Create a new event source with API key included
+        eventSource = new EventSource(`/generate_notebook?${urlParams.toString()}`);
+        
+        eventSource.onmessage = function(event) {
+            // Update our last-activity timestamp
+            lastChunkTime = Date.now();
             
-            eventSource.onmessage = function(event) {
-                // Update our last-activity timestamp
-                lastChunkTime = Date.now();
+            try {
+                const data = JSON.parse(event.data);
                 
-                try {
-                    const data = JSON.parse(event.data);
+                // Handle errors sent from the server
+                if (data.error) {
+                    console.error("Server error:", data.error);
+                    updateAiMessage(aiMessageId, `**Error:** ${data.error}`);
                     
-                    // Handle heartbeat messages (for Hugging Face)
-                    if (data.heartbeat !== undefined) {
-                        console.log(`Heartbeat received: ${data.heartbeat}`);
-                        return; // Just a keepalive, no content to process
+                    // Try to salvage what we have so far
+                    if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
+                        processNotebookResponse(aiResponseText);
                     }
                     
-                    // Handle errors sent from the server
-                    if (data.error) {
-                        console.error("Server error:", data.error);
-                        updateAiMessage(aiMessageId, `**Error:** ${data.error}`);
-                        
-                        // Try to salvage what we have so far
-                        if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
-                            processNotebookResponse(aiResponseText);
-                        }
-                        
-                        eventSource.close();
-                        eventSource = null;
-                        setGeneratingState(false);
-                        clearInterval(connectionTimer);
-                        return;
+                    eventSource.close();
+                    eventSource = null;
+                    setGeneratingState(false);
+                    clearInterval(connectionTimer);
+                    return;
+                }
+                
+                if (data.chunk) {
+                    aiResponseText += data.chunk;
+                    
+                    // Extract notebook info as soon as it's available
+                    const nameMatch = aiResponseText.match(/NOTEBOOK_NAME:?\s*(.+?)(?:\n|$)/);
+                    const descMatch = aiResponseText.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?:\n|$)/);
+                    
+                    if (nameMatch && nameMatch[1].trim()) {
+                        // Update notebook title immediately when found
+                        notebookTitleEl.textContent = nameMatch[1].trim();
                     }
                     
-                    if (data.chunk) {
-                        aiResponseText += data.chunk;
-                        
-                        // Extract notebook info as soon as it's available
-                        const nameMatch = aiResponseText.match(/NOTEBOOK_NAME:?\s*(.+?)(?:\n|$)/);
-                        const descMatch = aiResponseText.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?:\n|$)/);
-                        
-                        if (nameMatch && nameMatch[1].trim()) {
-                            // Update notebook title immediately when found
-                            notebookTitleEl.textContent = nameMatch[1].trim();
-                        }
-                        
-                        if (descMatch && descMatch[1].trim()) {
-                            // Update AI message to only display the description, not the full response
-                            updateAiMessage(aiMessageId, `**NoteGenie:** ${descMatch[1].trim()}`);
-                        } else {
-                            // Show a simple generating message while waiting for description
-                            updateAiMessage(aiMessageId, "**NoteGenie:** Generating your notebook...");
-                        }
-                        
-                        // Update preview periodically during streaming
-                        const now = Date.now();
-                        if (now - lastPreviewUpdate > PREVIEW_UPDATE_INTERVAL) {
-                            lastPreviewUpdate = now;
-                            
-                            // Only try to update preview if we have meaningful content
-                            if (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL')) {
-                                updateNotebookPreviewDuringStream(aiResponseText);
-                            }
-                        }
+                    if (descMatch && descMatch[1].trim()) {
+                        // Update AI message to only display the description, not the full response
+                        updateAiMessage(aiMessageId, `**NoteGenie:** ${descMatch[1].trim()}`);
+                    } else {
+                        // Show a simple generating message while waiting for description
+                        updateAiMessage(aiMessageId, "**NoteGenie:** Generating your notebook...");
                     }
                     
-                    if (data.done) {
-                        eventSource.close();
-                        eventSource = null;
-                        clearInterval(connectionTimer);
+                    // Update preview periodically during streaming
+                    const now = Date.now();
+                    if (now - lastPreviewUpdate > PREVIEW_UPDATE_INTERVAL) {
+                        lastPreviewUpdate = now;
                         
-                        // Process the complete response for final rendering
-                        processNotebookResponse(aiResponseText);
-                        setGeneratingState(false);
+                        // Only try to update preview if we have meaningful content
+                        if (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL')) {
+                            updateNotebookPreviewDuringStream(aiResponseText);
+                        }
                     }
-                } catch (error) {
-                    console.error('Error parsing event data:', error, event.data);
                 }
-            };
-            
-            eventSource.onerror = function(err) {
-                console.error('EventSource error:', err);
                 
-                // Don't retry if we're already disconnected
-                if (eventSource.readyState === 2) {
+                if (data.done) {
                     eventSource.close();
                     eventSource = null;
                     clearInterval(connectionTimer);
                     
-                    if (retryCount < MAX_RETRIES) {
-                        retryCount++;
-                        console.log(`Retrying connection (${retryCount}/${MAX_RETRIES})...`);
-                        updateAiMessage(aiMessageId, `**NoteGenie:** Connection issue, retrying... (${retryCount}/${MAX_RETRIES})`);
-                        
-                        // Wait a moment before retrying
-                        setTimeout(() => {
-                            createEventSource();
-                        }, 2000);
-                        return;
-                    }
-                    
-                    // Check if it's an auth error (most likely API key not set)
-                    if (err.status === 401) {
-                        // Updated error message to be more explicit
-                        updateAiMessage(aiMessageId, '**Error: API key not set or invalid.** \n\nYou must set your own Google Gemini API key to use NoteGenie. Please click the API Key button in the top right corner.');
-                        
-                        // Automatically show the API key modal
-                        showApiKeyModal();
-                    } else {
-                        // Try to salvage what we have so far
-                        if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
-                            updateAiMessage(aiMessageId, '**Warning:** Connection issue occurred but I\'ll try to process what I received so far.');
-                            processNotebookResponse(aiResponseText);
-                        } else {
-                            // Provide more detailed error message for Hugging Face users
-                            let errorMsg = '**Error:** Failed to generate notebook.';
-                            
-                            // Check if we're on Hugging Face (URL check)
-                            if (window.location.hostname.includes('huggingface.co') || 
-                                window.location.hostname.includes('hf.space')) {
-                                errorMsg += ' This might be due to Hugging Face Space limitations. Try: \n\n' +
-                                    '1. Refreshing the page \n' +
-                                    '2. Using a shorter prompt \n' +
-                                    '3. Re-entering your API key';
-                            } else {
-                                errorMsg += ' Please try again.';
-                            }
-                            
-                            updateAiMessage(aiMessageId, errorMsg);
-                        }
-                    }
-                    
+                    // Process the complete response for final rendering
+                    processNotebookResponse(aiResponseText);
                     setGeneratingState(false);
                 }
-            };
-        }
+            } catch (error) {
+                console.error('Error parsing event data:', error, event.data);
+            }
+        };
         
-        // Initial event source creation
-        createEventSource();
+        eventSource.onerror = function(err) {
+            console.error('EventSource error:', err);
+            eventSource.close();
+            eventSource = null;
+            clearInterval(connectionTimer);
+            
+            // Check if it's an auth error (most likely API key not set)
+            if (err.status === 401) {
+                updateAiMessage(aiMessageId, '**Error: API key not set or invalid.** \n\nPlease click the API Key button in the top right corner to set your Google Gemini API key.');
+                showApiKeyModal();
+            } else {
+                // Try to salvage what we have so far
+                if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
+                    updateAiMessage(aiMessageId, '**Warning:** Connection issue occurred but I\'ll try to process what I received so far.');
+                    processNotebookResponse(aiResponseText);
+                } else {
+                    updateAiMessage(aiMessageId, '**Error:** Failed to generate notebook. Please try again.');
+                }
+            }
+            
+            setGeneratingState(false);
+        };
     }
     
     function handleEditStreamingResponse(editPrompt, notebook, modelName, aiMessageId) {
@@ -481,6 +427,8 @@ document.addEventListener('DOMContentLoaded', function() {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/json',
+                // Add API key as header if available
+                ...(backupApiKey && {'X-API-Key': backupApiKey})
             },
             body: JSON.stringify({
                 notebook: notebook
@@ -494,12 +442,20 @@ document.addEventListener('DOMContentLoaded', function() {
         })
         .then(data => {
             if (data.success) {
-                // Now create a new event source for editing - the notebook is now in the session
-                eventSource = new EventSource(`/edit_notebook?${new URLSearchParams({
+                // Create URL parameters including API key as fallback
+                const urlParams = new URLSearchParams({
                     edit_prompt: editPrompt,
                     model: modelName,
                     stream: true
-                }).toString()}`);
+                });
+                
+                // Add API key to URL params if available from localStorage (as a fallback)
+                if (backupApiKey) {
+                    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()}`);
                 
                 // Track the last time we got a chunk - for timeout detection
                 let lastChunkTime = Date.now();
@@ -633,15 +589,7 @@ document.addEventListener('DOMContentLoaded', function() {
         })
         .catch(error => {
             console.error('Error preparing notebook for edit:', error);
-            
-            // Check if it's an auth error and handle API key issues more clearly
-            if (error.message && error.message.includes("API key not set")) {
-                updateAiMessage(aiMessageId, '**Error: API key not set or invalid.** \n\nYou must set your own Google Gemini API key to use NoteGenie. Please click the API Key button in the top right corner.');
-                showApiKeyModal();
-            } else {
-                updateAiMessage(aiMessageId, '**Error:** ' + error.message);
-            }
-            
+            updateAiMessage(aiMessageId, '**Error:** ' + error.message);
             setGeneratingState(false);
         });
     }

templates/index.html

@@ -93,7 +93,6 @@
                             <span class="material-icons info-icon">lightbulb</span>
                             <div class="info-content">
                                 <p><strong>Tip:</strong> Be specific in your request for best results. Include the topic, intended audience, and desired level of detail.</p>
-                                <p><strong>Important:</strong> NoteGenie requires you to provide your own Google AI Studio API key. Which is free.</p>
                             </div>
                         </div>
                         <div class="example-section">
@@ -254,28 +253,5 @@
     <script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/plugins/autoloader/prism-autoloader.min.js"></script>
     <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
     <script src="{{ url_for('static', filename='js/main.js') }}"></script>
-    
-    <!-- Hugging Face Spaces Session Check -->
-    <script>
-    if (window.location.hostname.includes('huggingface.co') || window.location.hostname.includes('hf.space')) {
-        console.log("Running on Hugging Face Spaces");
-        // Check if cookies are enabled
-        if (navigator.cookieEnabled) {
-            console.log("Cookies are enabled");
-        } else {
-            console.warn("Cookies are disabled - sessions won't work!");
-        }
-        
-        // Add a hidden diagonstic endpoint
-        fetch('/session_check')
-            .then(response => response.json())
-            .then(data => {
-                console.log("Session check:", data);
-            })
-            .catch(err => {
-                console.error("Session check failed:", err);
-            });
-    }
-    </script>
 </body>
 </html>

utils/ai_helpers.py

@@ -2,14 +2,6 @@ import google.generativeai as genai
 from flask import Response, stream_with_context
 import json
 import time
-import logging
-import os
-
-# Get logger from app
-logger = logging.getLogger('notegenie')
-
-# Check if we're running on Hugging Face Spaces
-IS_HUGGINGFACE = os.environ.get('SPACE_ID') is not None
 
 def craft_notebook_prompt(user_prompt):
     """Enhance the user prompt with instructions for generating a well-structured Jupyter notebook."""
@@ -98,13 +90,9 @@ def generate_notebook(user_prompt, model_name="gemini-2.0-pro-exp-02-05"):
     model = genai.GenerativeModel(model_name)
     
     enhanced_prompt = craft_notebook_prompt(user_prompt)
+    response = model.generate_content(enhanced_prompt)
     
-    try:
-        response = model.generate_content(enhanced_prompt)
-        return response.text
-    except Exception as e:
-        logger.error(f"Error generating notebook content: {str(e)}", exc_info=True)
-        raise
+    return response.text
 
 def edit_notebook(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05"):
     """Edit an existing notebook based on user request."""
@@ -115,50 +103,32 @@ def edit_notebook(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02
     
     return response.text
 
-def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05", api_key=None):
+def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05"):
     """Stream notebook generation responses from Gemini API."""
-    # Configure the API with the key if provided
-    if api_key:
-        genai.configure(api_key=api_key)
-        
     model = genai.GenerativeModel(model_name)
     enhanced_prompt = craft_notebook_prompt(user_prompt)
     
     def generate():
         try:
-            logger.info(f"Starting streaming generation with model {model_name}")
             response = model.generate_content(enhanced_prompt, stream=True)
             
             # Send a notification that streaming has started
             yield f"data: {json.dumps({'chunk': 'Starting notebook generation...'})}\n\n"
             
-            # Heartbeat counter for Hugging Face compatibility
-            heartbeat_counter = 0
-            last_send_time = time.time()
-            
             for chunk in response:
-                # Send heartbeat on Hugging Face to prevent connection timeout
-                current_time = time.time()
-                if IS_HUGGINGFACE and (current_time - last_send_time > 10):  # Send heartbeat every 10 seconds
-                    yield f"data: {json.dumps({'heartbeat': heartbeat_counter})}\n\n"
-                    heartbeat_counter += 1
-                    last_send_time = current_time
-                
                 try:
                     # More robust empty chunk detection
                     if not hasattr(chunk, 'parts') or not chunk.parts:
                         # Skip this empty chunk and continue
-                        logger.warning("Empty chunk received (no parts)")
+                        print("Warning: Empty chunk received (no parts)")
                         continue
                     
                     # First try the standard text property
                     try:
                         if hasattr(chunk, 'text') and chunk.text:
                             yield f"data: {json.dumps({'chunk': chunk.text})}\n\n"
-                            last_send_time = time.time()
                             continue  # If we successfully got text, continue to next chunk
-                    except (AttributeError, IndexError) as e:
-                        logger.warning(f"Error accessing text property: {e}")
+                    except (AttributeError, IndexError):
                         # If accessing text property fails, we'll try extracting from parts
                         pass
                     
@@ -167,139 +137,84 @@ def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05
                         # Extract text from part using different approaches
                         if hasattr(part, 'text') and part.text:
                             yield f"data: {json.dumps({'chunk': part.text})}\n\n"
-                            last_send_time = time.time()
                         elif isinstance(part, dict) and 'text' in part:
                             yield f"data: {json.dumps({'chunk': part['text']})}\n\n"
-                            last_send_time = time.time()
                         elif hasattr(part, 'string_value'):
                             yield f"data: {json.dumps({'chunk': part.string_value})}\n\n"
-                            last_send_time = time.time()
                 
                 except (AttributeError, IndexError, TypeError) as e:
                     # Log the error but continue - don't break the stream
-                    logger.error(f"Error processing chunk: {e}, chunk structure: {repr(chunk)[:200]}")
+                    print(f"Error processing chunk: {e}, chunk structure: {repr(chunk)[:200]}")
                     continue
                 
-                # Very brief pause to prevent overwhelming the client
+                # Briefly pause to prevent overwhelming the client
                 time.sleep(0.01)
             
-            # Final heartbeat for Hugging Face before completion
-            if IS_HUGGINGFACE:
-                yield f"data: {json.dumps({'heartbeat': 'final'})}\n\n"
-                
             yield f"data: {json.dumps({'done': True})}\n\n"
-            logger.info("Streaming generation completed successfully")
             
         except Exception as e:
             # Send error to client and close stream
             error_message = f"Error generating notebook: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            print(error_message)
             yield f"data: {json.dumps({'error': error_message})}\n\n"
             yield f"data: {json.dumps({'done': True})}\n\n"
     
-    # Set appropriate headers for Hugging Face compatibility
-    headers = None
-    if IS_HUGGINGFACE:
-        headers = {
-            'X-Accel-Buffering': 'no',  # Disable proxy buffering
-            'Cache-Control': 'no-cache',
-            'Connection': 'keep-alive'
-        }
-    
-    return Response(
-        stream_with_context(generate()), 
-        content_type="text/event-stream",
-        headers=headers
-    )
+    return Response(stream_with_context(generate()), content_type="text/event-stream")
 
-def stream_notebook_edit(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05", api_key=None):
+def stream_notebook_edit(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05"):
     """Stream notebook editing responses from Gemini API."""
-    # Configure the API with the key if provided
-    if api_key:
-        genai.configure(api_key=api_key)
-        
     model = genai.GenerativeModel(model_name)
     enhanced_prompt = craft_edit_prompt(edit_request, notebook_json)
     
     def generate():
         try:
-            logger.info(f"Starting streaming edit with model {model_name}")
             response = model.generate_content(enhanced_prompt, stream=True)
             
             # Send a notification that editing has started
             yield f"data: {json.dumps({'chunk': 'Starting notebook edit...'})}\n\n"
             
-            # Heartbeat counter for Hugging Face compatibility
-            heartbeat_counter = 0
-            last_send_time = time.time()
-            
             for chunk in response:
-                # Send heartbeat on Hugging Face to prevent connection timeout
-                current_time = time.time()
-                if IS_HUGGINGFACE and (current_time - last_send_time > 10):  # Send heartbeat every 10 seconds
-                    yield f"data: {json.dumps({'heartbeat': heartbeat_counter})}\n\n"
-                    heartbeat_counter += 1
-                    last_send_time = current_time
-                
-                # Process chunk similar to generation function
                 try:
                     # More robust empty chunk detection
                     if not hasattr(chunk, 'parts') or not chunk.parts:
-                        logger.warning("Empty chunk received (no parts)")
+                        # Skip this empty chunk and continue
+                        print("Warning: Empty chunk received (no parts)")
                         continue
                     
-                    # Try standard text property first
+                    # First try the standard text property
                     try:
                         if hasattr(chunk, 'text') and chunk.text:
                             yield f"data: {json.dumps({'chunk': chunk.text})}\n\n"
-                            last_send_time = time.time()
-                            continue
-                    except (AttributeError, IndexError) as e:
-                        logger.warning(f"Error accessing text property: {e}")
+                            continue  # If we successfully got text, continue to next chunk
+                    except (AttributeError, IndexError):
+                        # If accessing text property fails, we'll try extracting from parts
                         pass
                     
-                    # Try parts extraction
+                    # If we're here, we couldn't get text directly, try to extract from parts
                     for part in chunk.parts:
+                        # Extract text from part using different approaches
                         if hasattr(part, 'text') and part.text:
                             yield f"data: {json.dumps({'chunk': part.text})}\n\n"
-                            last_send_time = time.time()
                         elif isinstance(part, dict) and 'text' in part:
                             yield f"data: {json.dumps({'chunk': part['text']})}\n\n"
-                            last_send_time = time.time()
                         elif hasattr(part, 'string_value'):
                             yield f"data: {json.dumps({'chunk': part.string_value})}\n\n"
-                            last_send_time = time.time()
                 
                 except (AttributeError, IndexError, TypeError) as e:
-                    logger.error(f"Error processing chunk: {e}, chunk structure: {repr(chunk)[:200]}")
+                    # Log the error but continue - don't break the stream
+                    print(f"Error processing chunk: {e}, chunk structure: {repr(chunk)[:200]}")
                     continue
                 
+                # Briefly pause to prevent overwhelming the client
                 time.sleep(0.01)
             
-            # Final heartbeat for Hugging Face before completion
-            if IS_HUGGINGFACE:
-                yield f"data: {json.dumps({'heartbeat': 'final'})}\n\n"
-                
             yield f"data: {json.dumps({'done': True})}\n\n"
-            logger.info("Streaming edit completed successfully")
             
         except Exception as e:
+            # Send error to client and close stream
             error_message = f"Error editing notebook: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            print(error_message)
             yield f"data: {json.dumps({'error': error_message})}\n\n"
             yield f"data: {json.dumps({'done': True})}\n\n"
     
-    # Set appropriate headers for Hugging Face compatibility
-    headers = None
-    if IS_HUGGINGFACE:
-        headers = {
-            'X-Accel-Buffering': 'no',  # Disable proxy buffering
-            'Cache-Control': 'no-cache',
-            'Connection': 'keep-alive'
-        }
-    
-    return Response(
-        stream_with_context(generate()), 
-        content_type="text/event-stream",
-        headers=headers
-    )
+    return Response(stream_with_context(generate()), content_type="text/event-stream")