changes

README.md

@@ -30,6 +30,17 @@ 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:

app.py

@@ -4,9 +4,19 @@ 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
 
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    stream=sys.stdout
+)
+logger = logging.getLogger('notegenie')
+
 app = Flask(__name__)
 app.config["SECRET_KEY"] = os.environ.get("SECRET_KEY", "notegenie-secret-key-change-in-production")
 app.config["SESSION_TYPE"] = "filesystem"
@@ -16,6 +26,15 @@ 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
 
+# 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"] = True
+    app.config["SESSION_COOKIE_HTTPONLY"] = True
+    app.config["SESSION_COOKIE_SAMESITE"] = "Lax"
+
 Session(app)
 
 # Map front-end model names to API model names
@@ -47,14 +66,17 @@ def set_api_key():
         # Store API key in session with permanent flag
         session.permanent = True
         session["api_key"] = api_key
+        logger.info("API key successfully set and validated")
         
         return jsonify({"success": True})
     except Exception as e:
+        logger.error(f"API key validation error: {str(e)}")
         return jsonify({"success": False, "message": str(e)}), 400
 
 @app.route("/generate_notebook", methods=["GET", "POST"])
 def generate_notebook_route():
     if "api_key" not in session:
+        logger.warning("Generate notebook request without API key")
         return jsonify({"success": False, "message": "API key not set"}), 401
     
     # Handle both GET (for streaming) and POST requests
@@ -74,10 +96,11 @@ def generate_notebook_route():
     
     # Map the frontend model name to the API model name
     api_model_name = get_api_model_name(model_name)
-    
-    genai.configure(api_key=session["api_key"])
+    logger.info(f"Generate notebook with model: {api_model_name}, stream: {stream}")
     
     try:
+        genai.configure(api_key=session["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)
@@ -92,6 +115,7 @@ def generate_notebook_route():
                 "description": notebook_info["description"]
             })
         elif stream:
+            logger.info("Starting streaming notebook generation")
             return stream_notebook_generation(prompt, api_model_name)
         else:
             notebook_content = generate_notebook(prompt, api_model_name)
@@ -105,6 +129,7 @@ def generate_notebook_route():
                 "description": notebook_info["description"]
             })
     except Exception as e:
+        logger.error(f"Error generating notebook: {str(e)}", exc_info=True)
         return jsonify({"success": False, "message": str(e)}), 500
 
 @app.route("/prepare_edit_notebook", methods=["POST"])
@@ -197,4 +222,5 @@ def download_notebook():
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 5000))
-    app.run(host="0.0.0.0", port=port, debug=(os.environ.get("FLASK_ENV") == "development"))
+    debug_mode = os.environ.get("FLASK_ENV") == "development"
+    app.run(host="0.0.0.0", port=port, debug=debug_mode)

static/js/main.js

@@ -295,105 +295,152 @@ document.addEventListener('DOMContentLoaded', function() {
             }
         }, 5000); // Check every 5 seconds
         
-        // Create a new event source
-        eventSource = new EventSource(`/generate_notebook?${new URLSearchParams({
-            prompt: prompt,
-            model: modelName,
-            stream: true
-        }).toString()}`);
-        
-        eventSource.onmessage = function(event) {
-            // Update our last-activity timestamp
-            lastChunkTime = Date.now();
+        // Add a retry mechanism
+        let retryCount = 0;
+        const MAX_RETRIES = 2;
+        
+        function createEventSource() {
+            // Create a new event source with a unique timestamp to prevent caching
+            const timestamp = Date.now();
+            eventSource = new EventSource(`/generate_notebook?${new URLSearchParams({
+                prompt: prompt,
+                model: modelName,
+                stream: true,
+                t: timestamp  // Add timestamp to prevent caching
+            }).toString()}`);
             
-            try {
-                const data = JSON.parse(event.data);
+            eventSource.onmessage = function(event) {
+                // Update our last-activity timestamp
+                lastChunkTime = Date.now();
                 
-                // Handle errors sent from the server
-                if (data.error) {
-                    console.error("Server error:", data.error);
-                    updateAiMessage(aiMessageId, `**Error:** ${data.error}`);
+                try {
+                    const data = JSON.parse(event.data);
                     
-                    // Try to salvage what we have so far
-                    if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
-                        processNotebookResponse(aiResponseText);
+                    // Handle heartbeat messages (for Hugging Face)
+                    if (data.heartbeat !== undefined) {
+                        console.log(`Heartbeat received: ${data.heartbeat}`);
+                        return; // Just a keepalive, no content to process
                     }
                     
-                    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();
+                    // 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;
                     }
                     
-                    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.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);
+                            }
+                        }
                     }
                     
-                    // Update preview periodically during streaming
-                    const now = Date.now();
-                    if (now - lastPreviewUpdate > PREVIEW_UPDATE_INTERVAL) {
-                        lastPreviewUpdate = now;
+                    if (data.done) {
+                        eventSource.close();
+                        eventSource = null;
+                        clearInterval(connectionTimer);
                         
-                        // Only try to update preview if we have meaningful content
-                        if (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL')) {
-                            updateNotebookPreviewDuringStream(aiResponseText);
-                        }
+                        // Process the complete response for final rendering
+                        processNotebookResponse(aiResponseText);
+                        setGeneratingState(false);
                     }
+                } catch (error) {
+                    console.error('Error parsing event data:', error, event.data);
                 }
+            };
+            
+            eventSource.onerror = function(err) {
+                console.error('EventSource error:', err);
                 
-                if (data.done) {
+                // Don't retry if we're already disconnected
+                if (eventSource.readyState === 2) {
                     eventSource.close();
                     eventSource = null;
                     clearInterval(connectionTimer);
                     
-                    // Process the complete response for final rendering
-                    processNotebookResponse(aiResponseText);
+                    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) {
+                        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 {
+                            // 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);
+                        }
+                    }
+                    
                     setGeneratingState(false);
                 }
-            } catch (error) {
-                console.error('Error parsing event data:', error, event.data);
-            }
-        };
+            };
+        }
         
-        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);
-        };
+        // Initial event source creation
+        createEventSource();
     }
     
     function handleEditStreamingResponse(editPrompt, notebook, modelName, aiMessageId) {

utils/ai_helpers.py

@@ -2,6 +2,14 @@ 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."""
@@ -90,9 +98,13 @@ 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)
     
-    return response.text
+    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
 
 def edit_notebook(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05"):
     """Edit an existing notebook based on user request."""
@@ -110,25 +122,39 @@ def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05
     
     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
-                        print("Warning: Empty chunk received (no parts)")
+                        logger.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):
+                    except (AttributeError, IndexError) as e:
+                        logger.warning(f"Error accessing text property: {e}")
                         # If accessing text property fails, we'll try extracting from parts
                         pass
                     
@@ -137,29 +163,50 @@ 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
-                    print(f"Error processing chunk: {e}, chunk structure: {repr(chunk)[:200]}")
+                    logger.error(f"Error processing chunk: {e}, chunk structure: {repr(chunk)[:200]}")
                     continue
                 
-                # Briefly pause to prevent overwhelming the client
+                # Very brief 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)}"
-            print(error_message)
+            logger.error(error_message, exc_info=True)
             yield f"data: {json.dumps({'error': error_message})}\n\n"
             yield f"data: {json.dumps({'done': True})}\n\n"
     
-    return Response(stream_with_context(generate()), content_type="text/event-stream")
+    # 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
+    )
 
 def stream_notebook_edit(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05"):
     """Stream notebook editing responses from Gemini API."""
@@ -168,53 +215,83 @@ def stream_notebook_edit(edit_request, notebook_json, model_name="gemini-2.0-pro
     
     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:
-                        # Skip this empty chunk and continue
-                        print("Warning: Empty chunk received (no parts)")
+                        logger.warning("Empty chunk received (no parts)")
                         continue
                     
-                    # First try the standard text property
+                    # Try standard text property first
                     try:
                         if hasattr(chunk, 'text') and chunk.text:
                             yield f"data: {json.dumps({'chunk': chunk.text})}\n\n"
-                            continue  # If we successfully got text, continue to next chunk
-                    except (AttributeError, IndexError):
-                        # If accessing text property fails, we'll try extracting from parts
+                            last_send_time = time.time()
+                            continue
+                    except (AttributeError, IndexError) as e:
+                        logger.warning(f"Error accessing text property: {e}")
                         pass
                     
-                    # If we're here, we couldn't get text directly, try to extract from parts
+                    # Try parts extraction
                     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:
-                    # Log the error but continue - don't break the stream
-                    print(f"Error processing chunk: {e}, chunk structure: {repr(chunk)[:200]}")
+                    logger.error(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)}"
-            print(error_message)
+            logger.error(error_message, exc_info=True)
             yield f"data: {json.dumps({'error': error_message})}\n\n"
             yield f"data: {json.dumps({'done': True})}\n\n"
     
-    return Response(stream_with_context(generate()), content_type="text/event-stream")
+    # 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
+    )