initial commit

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.9-slim
+
+WORKDIR /app
+
+# Install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application files
+COPY . .
+
+# Create necessary directories
+RUN mkdir -p flask_session
+
+# Set environment variables
+ENV FLASK_APP=app.py
+ENV FLASK_ENV=production
+ENV PORT=7860
+
+# Expose the port Hugging Face Spaces uses
+EXPOSE 7860
+
+# Command to run the application
+CMD gunicorn --bind 0.0.0.0:7860 app:app

README.md

@@ -9,4 +9,60 @@ license: mit
 short_description: ipynb file generator using Gemini models to create notebooks
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# NoteGenie
+
+NoteGenie is an AI-powered Jupyter notebook generator that uses Google's Gemini models to create and edit notebooks from text prompts.
+
+## Features
+
+- Generate complete Jupyter notebooks with a text prompt
+- Edit existing notebooks with text instructions
+- Preview notebooks in the web interface
+- Download notebooks as .ipynb files
+- Streaming responses for a better user experience
+
+## Deployment on Hugging Face Spaces
+
+1. Create a new Space on Hugging Face
+2. Choose Docker template
+3. Upload this repository to the Space
+4. Set the following environment variables in your Space settings:
+   - SECRET_KEY: A secure random string for Flask sessions
+   - PORT: 7860 (default for Hugging Face Spaces)
+
+## Local Development
+
+1. Install dependencies:
+   ```
+   pip install -r requirements.txt
+   ```
+
+2. Run the application:
+   ```
+   python app.py
+   ```
+
+3. Open http://localhost:5000 in your browser
+
+## Using Docker
+
+Build and run the Docker container:
+
+```bash
+docker build -t notegenie .
+docker run -p 7860:7860 -e SECRET_KEY=your_secret_key notegenie
+```
+
+## Environment Variables
+
+- SECRET_KEY: Secret key for Flask session encryption
+- PORT: Port to run the application on (defaults to 5000 locally, 7860 for Hugging Face)
+- FLASK_ENV: Set to "development" for debug mode, "production" for production
+
+## API Key Setup
+
+NoteGenie requires a Google Gemini API key. Users can set their own API key in the web interface.
+
+## License
+
+[MIT License](LICENSE)

app.py

@@ -0,0 +1,200 @@
+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
+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
+
+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
+
+Session(app)
+
+# Map front-end model names to API model names
+MODEL_MAPPING = {
+    "gemini-2.0-pro": "gemini-2.0-pro-exp-02-05",
+    "gemini-2.0-flash": "gemini-2.0-flash",
+    "gemini-2.0-flash-thinking": "gemini-2.0-flash-thinking-exp-01-21"
+}
+
+def get_api_model_name(frontend_model_name):
+    return MODEL_MAPPING.get(frontend_model_name, frontend_model_name)
+
+@app.route("/", methods=["GET"])
+def index():
+    return render_template("index.html")
+
+@app.route("/set_api_key", methods=["POST"])
+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
+        session.permanent = True
+        session["api_key"] = api_key
+        
+        return jsonify({"success": True})
+    except Exception as 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:
+        return jsonify({"success": False, "message": "API key not set"}), 401
+    
+    # Handle both GET (for streaming) and POST requests
+    if request.method == "GET":
+        prompt = request.args.get("prompt")
+        model_name = request.args.get("model", "gemini-2.0-pro")
+        stream = request.args.get("stream", "false").lower() == "true"
+    else:
+        data = request.json
+        prompt = data.get("prompt")
+        model_name = data.get("model", "gemini-2.0-pro")
+        stream = data.get("stream", False)
+        format_only = data.get("format_only", False)
+    
+    if not prompt:
+        return jsonify({"success": False, "message": "Prompt is required"}), 400
+    
+    # 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"])
+    
+    try:
+        # 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,
+                "name": notebook_info["name"],
+                "description": notebook_info["description"]
+            })
+        elif stream:
+            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,
+                "name": notebook_info["name"],
+                "description": notebook_info["description"]
+            })
+    except Exception as 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:
+        return jsonify({"success": False, "message": "API key not set"}), 401
+    
+    data = request.json
+    notebook_json = data.get("notebook")
+    
+    if not notebook_json:
+        return jsonify({"success": False, "message": "Notebook content is required"}), 400
+    
+    # Store the notebook in the session for later access
+    session["current_notebook"] = notebook_json
+    
+    return jsonify({"success": True})
+
+@app.route("/edit_notebook", methods=["GET", "POST"])
+def edit_notebook_route():
+    if "api_key" not in session:
+        return jsonify({"success": False, "message": "API key not set"}), 401
+    
+    # Get edit prompt and current notebook
+    if request.method == "GET":
+        edit_prompt = request.args.get("edit_prompt")
+        model_name = request.args.get("model", "gemini-2.0-pro")
+        stream = request.args.get("stream", "true").lower() == "true"
+        # For GET streaming requests, get notebook from session
+        notebook_json = session.get("current_notebook")
+    else:
+        data = request.json
+        edit_prompt = data.get("edit_prompt")
+        notebook_json = data.get("notebook")
+        model_name = data.get("model", "gemini-2.0-pro")
+        stream = data.get("stream", False)
+    
+    if not edit_prompt:
+        return jsonify({"success": False, "message": "Edit prompt is required"}), 400
+    
+    if not notebook_json:
+        return jsonify({"success": False, "message": "No notebook available for editing. Please prepare the notebook first."}), 400
+    
+    # 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"])
+    
+    try:
+        if stream:
+            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,
+                "name": notebook_info["name"],
+                "description": notebook_info["description"]
+            })
+    except Exception as e:
+        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")
+    
+    if not notebook_json:
+        return jsonify({"success": False, "message": "Notebook content is required"}), 400
+    
+    if not filename.endswith(".ipynb"):
+        filename += ".ipynb"
+    
+    response = Response(
+        json.dumps(notebook_json, indent=2),
+        mimetype="application/json",
+        headers={"Content-Disposition": f"attachment;filename={filename}"}
+    )
+    
+    return response
+
+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"))

requirements.txt

@@ -0,0 +1,4 @@
+flask==2.0.1
+flask-session==0.4.0
+google-generativeai==0.3.1
+gunicorn==20.1.0

static/css/style.css

@@ -0,0 +1,761 @@
+:root {
+    /* Google Material Design colors */
+    --google-blue: #4285f4;
+    --google-red: #ea4335;
+    --google-yellow: #fbbc05;
+    --google-green: #34a853;
+    --google-gray: #5f6368;
+    --google-blue-hover: #3367d6;
+    --google-surface: #ffffff;
+    --google-background: #f8f9fa;
+    --google-border: #dadce0;
+    --google-text-primary: #202124;
+    --google-text-secondary: #5f6368;
+    --google-disabled: #e0e0e0;
+    
+    /* Spacing and dimensions */
+    --spacing-xs: 4px;
+    --spacing-sm: 8px;
+    --spacing-md: 16px;
+    --spacing-lg: 24px;
+    --spacing-xl: 32px;
+    --border-radius: 8px;
+    --shadow-1: 0 1px 2px 0 rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);
+    --shadow-2: 0 2px 6px 2px rgba(60, 64, 67, 0.15);
+    --shadow-3: 0 4px 8px 3px rgba(60, 64, 67, 0.15);
+}
+
+/* Base styles */
+body {
+    font-family: 'Roboto', sans-serif;
+    background-color: var(--google-background);
+    color: var(--google-text-primary);
+    margin: 0;
+    padding: 0;
+    height: 100vh;
+    display: flex;
+    flex-direction: column;
+    overflow: hidden;
+}
+
+/* Material typography */
+h1, h2, h3, h4, h5, h6 {
+    font-family: 'Google Sans', 'Roboto', sans-serif;
+    font-weight: 500;
+    margin: 0;
+    color: var(--google-text-primary);
+}
+
+.product-name {
+    font-size: 22px;
+    line-height: 28px;
+    font-weight: 500;
+}
+
+/* App container */
+.app-container {
+    display: flex;
+    flex-direction: column;
+    height: 100vh;
+    overflow: hidden;
+}
+
+/* Header styling */
+.app-header {
+    background-color: var(--google-surface);
+    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3);
+    padding: var(--spacing-md);
+    z-index: 10;
+}
+
+.header-content {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    max-width: 1600px;
+    margin: 0 auto;
+    width: 100%;
+}
+
+.logo-section {
+    display: flex;
+    align-items: center;
+    gap: var(--spacing-sm);
+}
+
+.logo-icon {
+    color: var(--google-blue);
+}
+
+.header-actions {
+    display: flex;
+    gap: var(--spacing-sm);
+}
+
+/* Main container */
+.main-container {
+    display: flex;
+    flex: 1;
+    overflow: hidden;
+    max-width: 1600px;
+    margin: 0 auto;
+    width: 100%;
+    height: calc(100vh - 64px);
+    padding: var(--spacing-md);
+    gap: var(--spacing-md);
+}
+
+/* Panels - update the width ratio */
+.panel {
+    display: flex;
+    flex-direction: column;
+    overflow: hidden;
+    background-color: var(--google-surface);
+    border-radius: var(--border-radius);
+    height: 100%;
+    box-shadow: var(--shadow-1);
+}
+
+.chat-panel {
+    flex: 0.4; /* Reduced from 1 to 0.4 (40% width) */
+    border-right: none;
+}
+
+.notebook-panel {
+    flex: 0.6; /* Increased from 1 to 0.6 (60% width) */
+    overflow: hidden;
+}
+
+/* Conversation container */
+.conversation-container {
+    flex: 1;
+    overflow-y: auto;
+    padding: var(--spacing-md);
+}
+
+/* Custom scrollbar styles */
+.conversation-container::-webkit-scrollbar,
+.notebook-preview::-webkit-scrollbar {
+    width: 8px;
+}
+
+.conversation-container::-webkit-scrollbar-track,
+.notebook-preview::-webkit-scrollbar-track {
+    background: transparent;
+}
+
+.conversation-container::-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 {
+    background-color: rgba(95, 99, 104, 0.5);
+}
+
+/* Input container */
+.input-container {
+    border-top: 1px solid var(--google-border);
+    padding: var(--spacing-md);
+    background-color: var(--google-surface);
+}
+
+.input-options {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    margin-bottom: var(--spacing-md);
+}
+
+.model-selector {
+    border: 1px solid var(--google-border);
+    border-radius: var(--border-radius);
+    padding: var(--spacing-xs) var(--spacing-md);
+    font-size: 14px;
+    color: var(--google-text-primary);
+    background-color: var(--google-surface);
+    width: auto;
+    outline: none;
+}
+
+.model-selector:focus {
+    border-color: var(--google-blue);
+}
+
+.mode-toggle {
+    display: flex;
+    border: 1px solid var(--google-blue);
+    border-radius: 4px;
+    overflow: hidden;
+}
+
+.toggle-button {
+    background: none;
+    border: none;
+    padding: 6px 12px;
+    font-size: 14px;
+    cursor: pointer;
+    color: var(--google-blue);
+}
+
+.toggle-button.active {
+    background-color: var(--google-blue);
+    color: white;
+}
+
+/* Fix for disabled toggle buttons */
+.toggle-button:disabled {
+    cursor: not-allowed;
+    opacity: 0.7;
+}
+
+/* Keep active styling even when disabled */
+.toggle-button.active:disabled {
+    background-color: var(--google-blue);
+    color: white;
+    opacity: 0.8;
+}
+
+/* Inactive button when disabled */
+.toggle-button:not(.active):disabled {
+    color: var(--google-text-secondary);
+    background-color: transparent;
+}
+
+.input-field-container {
+    position: relative;
+    margin-bottom: var(--spacing-sm);
+}
+
+.prompt-input {
+    width: 100%;
+    border: 1px solid var(--google-border);
+    border-radius: var(--border-radius);
+    padding: var(--spacing-md);
+    padding-right: 50px; /* Make room for the send button */
+    font-family: 'Roboto', sans-serif;
+    font-size: 16px;
+    resize: none;
+    outline: none;
+    transition: border-color 0.2s;
+}
+
+.prompt-input:focus {
+    border-color: var(--google-blue);
+    box-shadow: 0 0 0 1px var(--google-blue);
+}
+
+/* Send button inside input */
+.send-button {
+    position: absolute;
+    right: 12px;
+    top: 12px;
+    background-color: var(--google-blue);
+    color: white;
+    border: none;
+    border-radius: 50%;
+    width: 36px;
+    height: 36px;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    cursor: pointer;
+    transition: background-color 0.2s;
+}
+
+.send-button:hover {
+    background-color: var(--google-blue-hover);
+}
+
+.send-button:disabled {
+    background-color: var(--google-disabled);
+    cursor: not-allowed;
+}
+
+.send-button .material-icons {
+    font-size: 18px;
+}
+
+/* Hide the original action button container */
+.input-actions {
+    display: none;
+}
+
+/* Google Buttons */
+.google-button {
+    display: inline-flex;
+    align-items: center;
+    justify-content: center;
+    border: none;
+    border-radius: 4px;
+    font-family: 'Google Sans', 'Roboto', sans-serif;
+    font-size: 14px;
+    font-weight: 500;
+    padding: 8px 24px;
+    cursor: pointer;
+    transition: background-color 0.2s, box-shadow 0.2s;
+    white-space: nowrap;
+    gap: var(--spacing-sm);
+    height: 36px;
+    letter-spacing: 0.25px;
+    outline: none;
+    position: relative;
+}
+
+.google-button .material-icons {
+    font-size: 18px;
+}
+
+.google-button.primary {
+    background-color: var(--google-blue);
+    color: white;
+}
+
+.google-button.primary:hover {
+    background-color: var(--google-blue-hover);
+    box-shadow: var(--shadow-1);
+}
+
+.google-button.outlined {
+    background-color: transparent;
+    border: 1px solid var(--google-border);
+    color: var(--google-blue);
+}
+
+.google-button.outlined:hover {
+    background-color: rgba(66, 133, 244, 0.04);
+}
+
+.google-button.text {
+    background-color: transparent;
+    color: var(--google-blue);
+    padding: 8px 16px;
+}
+
+.google-button.text:hover {
+    background-color: rgba(66, 133, 244, 0.04);
+}
+
+.google-button:disabled {
+    background-color: var(--google-disabled);
+    color: var(--google-text-secondary);
+    cursor: not-allowed;
+    box-shadow: none;
+}
+
+/* Notebook panel */
+.notebook-panel {
+    overflow: hidden;
+}
+
+.notebook-header {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    padding: var(--spacing-md);
+    border-bottom: 1px solid var(--google-border);
+}
+
+.notebook-title {
+    font-size: 16px;
+    font-weight: 500;
+    color: var(--google-text-primary);
+}
+
+.notebook-preview {
+    flex: 1;
+    overflow-y: auto;
+    padding: var(--spacing-md);
+    background-color: var(--google-surface);
+}
+
+.placeholder-content {
+    display: flex;
+    flex-direction: column;
+    align-items: center;
+    justify-content: center;
+    height: 100%;
+    color: var(--google-text-secondary);
+}
+
+.large-icon {
+    font-size: 48px;
+    margin-bottom: var(--spacing-md);
+    color: var(--google-gray);
+}
+
+/* Welcome message */
+.welcome-message {
+    max-width: 600px;
+    margin: 0 auto;
+}
+
+.welcome-message h2 {
+    display: flex;
+    align-items: center;
+    gap: var(--spacing-sm);
+    margin-bottom: var(--spacing-md);
+    color: var(--google-blue);
+}
+
+.welcome-steps {
+    margin: var(--spacing-lg) 0;
+}
+
+.welcome-step {
+    display: flex;
+    align-items: center;
+    margin-bottom: var(--spacing-md);
+}
+
+.step-number {
+    width: 24px;
+    height: 24px;
+    border-radius: 50%;
+    background-color: var(--google-blue);
+    color: white;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    font-size: 14px;
+    margin-right: var(--spacing-md);
+}
+
+.step-text {
+    font-size: 16px;
+}
+
+.info-card {
+    display: flex;
+    background-color: #e8f0fe;
+    border-radius: var(--border-radius);
+    padding: var(--spacing-md);
+    margin: var(--spacing-md) 0;
+}
+
+.info-icon {
+    color: var(--google-blue);
+    margin-right: var(--spacing-md);
+}
+
+.info-content {
+    flex: 1;
+}
+
+.info-content p {
+    margin: 0;
+}
+
+.example-section {
+    margin-top: var(--spacing-lg);
+}
+
+.example-list {
+    list-style-type: none;
+    padding: 0;
+}
+
+.example-list li {
+    padding: var(--spacing-sm) 0;
+    border-left: 3px solid var(--google-blue);
+    padding-left: var(--spacing-md);
+    margin-bottom: var(--spacing-sm);
+    background-color: rgba(66, 133, 244, 0.05);
+    border-radius: 0 var(--border-radius) var(--border-radius) 0;
+}
+
+/* Messages */
+.message {
+    margin-bottom: var(--spacing-md);
+    max-width: 85%;
+    padding: var(--spacing-md);
+    border-radius: var(--border-radius);
+}
+
+.user-message {
+    background-color: #e8f0fe;
+    margin-left: auto;
+    position: relative;
+    box-shadow: var(--shadow-1);
+}
+
+.ai-message {
+    background-color: white;
+    border: 1px solid var(--google-border);
+    margin-right: auto;
+    position: relative;
+    box-shadow: var(--shadow-1);
+}
+
+/* Fix for user message formatting */
+.user-prefix {
+    display: inline;
+    font-weight: 500;
+}
+
+.message-content {
+    display: inline;
+}
+
+.message-content p {
+    display: inline;
+    margin: 0;
+}
+
+.message-content p:not(:first-child) {
+    display: block;
+    margin-top: var(--spacing-sm);
+}
+
+/* Typing indicator */
+.typing-indicator {
+    display: flex;
+    align-items: center;
+    gap: 4px;
+}
+
+.typing-dot {
+    background-color: var(--google-blue);
+    border-radius: 50%;
+    width: 8px;
+    height: 8px;
+    animation: typing-dot 1.4s infinite;
+}
+
+.typing-dot:nth-child(2) {
+    animation-delay: 0.2s;
+}
+
+.typing-dot:nth-child(3) {
+    animation-delay: 0.4s;
+}
+
+@keyframes typing-dot {
+    0%, 60%, 100% { transform: translateY(0); }
+    30% { transform: translateY(-6px); }
+}
+
+/* Notebook cell styling */
+.notebook-cell {
+    margin-bottom: var(--spacing-md);
+    border: 1px solid var(--google-border);
+    border-radius: var(--border-radius);
+    overflow: hidden;
+    background-color: var(--google-surface);
+    box-shadow: var(--shadow-1);
+}
+
+.cell-header {
+    padding: var(--spacing-sm) var(--spacing-md);
+    font-size: 13px;
+    font-weight: 500;
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    border-bottom: 1px solid var(--google-border);
+}
+
+.cell-type {
+    display: inline-flex;
+    align-items: center;
+    font-family: 'Google Sans', 'Roboto', sans-serif;
+    font-size: 13px;
+}
+
+.markdown-header {
+    background-color: #e8f0fe;
+    color: #000000;
+}
+
+/* Make code header match markdown header */
+.code-header {
+    background-color: #e8f0fe; /* Changed from #e6f4ea to match markdown-header */
+    color: #000000; /* Changed from #188038 to match markdown-header */
+}
+
+.cell-content {
+    padding: 0;
+}
+
+.cell-markdown .cell-content {
+    padding: var(--spacing-md);
+}
+
+.cell-code .cell-content {
+    padding: 0;
+}
+
+/* Code cell styling */
+.cell-code pre {
+    margin: 0 !important;
+    border-radius: 0;
+}
+
+.cell-code code {
+    padding: var(--spacing-md) var(--spacing-md) var(--spacing-md) 0 !important;
+    font-family: 'Roboto Mono', monospace;
+    font-size: 14px;
+}
+
+/* Modal styling */
+.google-modal {
+    border-radius: var(--border-radius);
+    border: none;
+    box-shadow: var(--shadow-3);
+}
+
+.google-modal .modal-header {
+    border-bottom: 1px solid var(--google-border);
+    padding: var(--spacing-md);
+}
+
+.google-modal .modal-title {
+    font-family: 'Google Sans', 'Roboto', sans-serif;
+    font-weight: 500;
+    font-size: 18px;
+    display: flex;
+    align-items: center;
+    gap: var(--spacing-sm);
+}
+
+.google-modal .modal-body {
+    padding: var(--spacing-md);
+}
+
+.google-modal .modal-footer {
+    border-top: 1px solid var(--google-border);
+    padding: var(--spacing-md);
+}
+
+/* Input fields in modal */
+.input-field {
+    margin-bottom: var(--spacing-md);
+}
+
+.input-label {
+    display: block;
+    margin-bottom: var(--spacing-sm);
+    font-size: 14px;
+    color: var(--google-text-secondary);
+}
+
+.text-input {
+    width: 100%;
+    border: 1px solid var(--google-border);
+    border-radius: 4px;
+    padding: 10px 12px;
+    font-family: 'Roboto', sans-serif;
+    font-size: 16px;
+    outline: none;
+    transition: border-color 0.2s;
+}
+
+.text-input:focus {
+    border-color: var(--google-blue);
+    box-shadow: 0 0 0 1px var(--google-blue);
+}
+
+.input-helper {
+    margin-top: 4px;
+    font-size: 12px;
+    color: var(--google-text-secondary);
+}
+
+/* Steps card */
+.steps-card {
+    background-color: var(--google-background);
+    border-radius: var(--border-radius);
+    overflow: hidden;
+    margin-bottom: var(--spacing-md);
+}
+
+.steps-header {
+    background-color: #e8f0fe;
+    padding: var(--spacing-sm) var(--spacing-md);
+    display: flex;
+    align-items: center;
+    gap: var(--spacing-sm);
+    color: var(--google-blue);
+}
+
+.steps-content {
+    padding: var(--spacing-md);
+}
+
+.steps-list {
+    margin: 0;
+    padding-left: 20px;
+}
+
+.steps-list li {
+    padding: 4px 0;
+}
+
+/* Floating scroll button */
+.scroll-button {
+    position: fixed;
+    bottom: 20px;
+    right: 20px;
+    width: 40px;
+    height: 40px;
+    border-radius: 50%;
+    background-color: white;
+    color: var(--google-text-primary);
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    border: 1px solid var(--google-border);
+    cursor: pointer;
+    box-shadow: var(--shadow-2);
+    opacity: 0;
+    transition: opacity 0.3s ease;
+    z-index: 1000;
+}
+
+.scroll-button.visible {
+    opacity: 1;
+}
+
+.scroll-button:hover {
+    background-color: #f8f9fa;
+}
+
+/* Add spinning animation for loading state */
+@keyframes spin-animation {
+    0% { transform: rotate(0deg); }
+    100% { transform: rotate(360deg); }
+}
+
+.material-icons.spinning {
+    animation: spin-animation 1.5s linear infinite;
+}
+
+/* Media queries for responsive design */
+@media (max-width: 768px) {
+    .main-container {
+        flex-direction: column;
+        padding: var(--spacing-sm);
+        gap: var(--spacing-sm);
+    }
+    
+    .chat-panel {
+        height: 40%; /* Reduced from 50% to 40% */
+        border-right: none;
+        border-bottom: none;
+        flex: none; /* Override the horizontal flex ratio when in column layout */
+    }
+    
+    .notebook-panel {
+        height: 60%; /* Increased from 50% to 60% */
+        flex: none; /* Override the horizontal flex ratio when in column layout */
+    }
+    
+    .send-button {
+        right: 8px;
+        top: 8px;
+        width: 32px;
+        height: 32px;
+    }
+}

static/js/main.js

@@ -0,0 +1,1052 @@
+document.addEventListener('DOMContentLoaded', function() {
+    // DOM Elements
+    const conversationEl = document.getElementById('conversation');
+    const promptInputEl = document.getElementById('promptInput');
+    const actionBtnEl = document.getElementById('actionBtn');
+    const downloadBtnEl = document.getElementById('downloadBtn');
+    const notebookPreviewEl = document.getElementById('notebookPreview');
+    const notebookTitleEl = document.getElementById('notebookTitle');
+    const apiKeyInputEl = document.getElementById('apiKeyInput');
+    const saveApiKeyBtnEl = document.getElementById('saveApiKeyBtn');
+    const apiKeyFeedbackEl = document.getElementById('apiKeyFeedback');
+    const modelSelectEl = document.getElementById('modelSelect');
+    const generateModeBtnEl = document.getElementById('generateModeBtn');
+    const editModeBtnEl = document.getElementById('editModeBtn');
+    
+    // Create scroll-to-bottom button
+    const scrollToBottomBtn = document.createElement('button');
+    scrollToBottomBtn.className = 'scroll-to-bottom';
+    scrollToBottomBtn.innerHTML = '<i class="bi bi-arrow-down"></i>';
+    scrollToBottomBtn.title = 'Scroll to bottom';
+    document.body.appendChild(scrollToBottomBtn);
+    
+    // Scroll to bottom button event listener
+    scrollToBottomBtn.addEventListener('click', function() {
+        notebookPreviewEl.scrollTo({
+            top: notebookPreviewEl.scrollHeight,
+            behavior: 'smooth'
+        });
+    });
+    
+    // Show/hide scroll to bottom button based on scroll position
+    notebookPreviewEl.addEventListener('scroll', function() {
+        const scrollPosition = notebookPreviewEl.scrollTop + notebookPreviewEl.clientHeight;
+        const scrollThreshold = notebookPreviewEl.scrollHeight - 100;
+        
+        if (scrollPosition < scrollThreshold) {
+            scrollToBottomBtn.classList.add('visible');
+        } else {
+            scrollToBottomBtn.classList.remove('visible');
+        }
+    });
+    
+    // Global state
+    let currentNotebook = null;
+    let eventSource = null;
+    let aiResponseText = '';
+    let currentMode = 'generate'; // 'generate' or 'edit'
+    
+    // Add a debounce flag to prevent button spamming
+    let isActionInProgress = false;
+    
+    // Initialize tooltips
+    const tooltipTriggerList = [].slice.call(document.querySelectorAll('[data-bs-toggle="tooltip"]'));
+    tooltipTriggerList.map(function(tooltipTriggerEl) {
+        return new bootstrap.Tooltip(tooltipTriggerEl);
+    });
+    
+    // Check if we have an API key in localStorage on page load
+    const savedApiKey = localStorage.getItem('notegenie_api_key');
+    if (savedApiKey) {
+        // Try to set the API key automatically
+        apiKeyInputEl.value = savedApiKey;
+        checkAndSetApiKey();
+    }
+    
+    // Event Listeners
+    saveApiKeyBtnEl.addEventListener('click', saveApiKey);
+    actionBtnEl.addEventListener('click', handleAction);
+    downloadBtnEl.addEventListener('click', downloadNotebook);
+    generateModeBtnEl.addEventListener('click', () => setMode('generate'));
+    editModeBtnEl.addEventListener('click', () => setMode('edit'));
+    
+    // Allow Enter key to submit prompt (Shift+Enter for new line)
+    promptInputEl.addEventListener('keydown', function(e) {
+        if (e.key === 'Enter' && !e.shiftKey) {
+            e.preventDefault();
+            handleAction();
+        }
+    });
+    
+    // Set mode (generate or edit)
+    function setMode(mode) {
+        currentMode = mode;
+        
+        if (mode === 'generate') {
+            generateModeBtnEl.classList.add('active');
+            generateModeBtnEl.classList.remove('btn-outline-primary');
+            generateModeBtnEl.classList.add('btn-primary');
+            
+            editModeBtnEl.classList.remove('active');
+            editModeBtnEl.classList.remove('btn-primary');
+            editModeBtnEl.classList.add('btn-outline-primary');
+            
+            promptInputEl.placeholder = 'Describe the notebook you want...';
+        } else { // edit mode
+            editModeBtnEl.classList.add('active');
+            editModeBtnEl.classList.remove('btn-outline-primary');
+            editModeBtnEl.classList.add('btn-primary');
+            
+            generateModeBtnEl.classList.remove('active');
+            generateModeBtnEl.classList.remove('btn-primary');
+            generateModeBtnEl.classList.add('btn-outline-primary');
+            
+            promptInputEl.placeholder = 'Describe what changes you want to make to the notebook...';
+            promptInputEl.focus();
+        }
+    }
+    
+    // Handle action button click based on current mode
+    function handleAction() {
+        // Prevent multiple rapid clicks
+        if (isActionInProgress) return;
+        
+        // Get the prompt text and check if it's empty
+        const prompt = promptInputEl.value.trim();
+        
+        if (currentMode === 'generate') {
+            if (!prompt) {
+                addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please enter a prompt to generate a notebook.');
+                return;
+            }
+            
+            // Set the debounce flag and disable the button
+            isActionInProgress = true;
+            actionBtnEl.disabled = true;
+            
+            generateNotebook();
+        } else {
+            if (!prompt || !currentNotebook) {
+                if (!currentNotebook) {
+                    addSystemMessage('<i class="bi bi-exclamation-triangle"></i> No notebook to edit. Generate a notebook first.');
+                } else if (!prompt) {
+                    addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please describe the changes you want to make.');
+                }
+                return;
+            }
+            
+            // Set the debounce flag and disable the button
+            isActionInProgress = true;
+            actionBtnEl.disabled = true;
+            
+            editNotebook();
+        }
+    }
+    
+    // API Key handling
+    function checkAndSetApiKey() {
+        const apiKey = apiKeyInputEl.value.trim();
+        
+        if (!apiKey) {
+            return;
+        }
+        
+        fetch('/set_api_key', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/x-www-form-urlencoded',
+            },
+            body: new URLSearchParams({
+                'api_key': apiKey
+            })
+        })
+        .then(response => response.json())
+        .then(data => {
+            if (data.success) {
+                // Store API key in localStorage as backup
+                localStorage.setItem('notegenie_api_key', apiKey);
+                
+                if (document.querySelector('#apiKeyModal.show')) {
+                    showApiKeyFeedback('API key saved successfully!', 'success');
+                    setTimeout(() => {
+                        document.querySelector('#apiKeyModal .btn-close').click();
+                    }, 1500);
+                }
+            } else {
+                showApiKeyFeedback(`Error: ${data.message}`, 'danger');
+                localStorage.removeItem('notegenie_api_key'); // Remove invalid key
+            }
+        })
+        .catch(error => {
+            showApiKeyFeedback(`Error: ${error.message}`, 'danger');
+            localStorage.removeItem('notegenie_api_key'); // Remove key on error
+        });
+    }
+    
+    function saveApiKey() {
+        checkAndSetApiKey();
+    }
+    
+    function showApiKeyFeedback(message, type) {
+        apiKeyFeedbackEl.innerHTML = `<div class="alert alert-${type} mb-0">${message}</div>`;
+        setTimeout(() => {
+            apiKeyFeedbackEl.innerHTML = '';
+        }, 5000);
+    }
+    
+    // Notebook generation
+    function generateNotebook() {
+        const prompt = promptInputEl.value.trim();
+        
+        if (!prompt) {
+            addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please enter a prompt to generate a notebook.');
+            return;
+        }
+        
+        // Clear any existing notebook
+        notebookPreviewEl.innerHTML = '';
+        downloadBtnEl.disabled = true;
+        editModeBtnEl.disabled = true;
+        currentNotebook = null;
+        
+        // Add user message to conversation
+        addUserMessage(prompt);
+        
+        // Disable input during generation
+        setGeneratingState(true);
+        
+        const modelName = modelSelectEl.value;
+        
+        // Add AI typing indicator
+        const aiMessageId = 'ai-typing-' + Date.now();
+        addTypingIndicator(aiMessageId);
+        
+        // Always use streaming for better UX
+        handleStreamingResponse(prompt, modelName, aiMessageId);
+        
+        // Clear the input field after sending
+        promptInputEl.value = '';
+    }
+    
+    // Notebook editing
+    function editNotebook() {
+        const editPrompt = promptInputEl.value.trim();
+        
+        if (!editPrompt || !currentNotebook) {
+            if (!currentNotebook) {
+                addSystemMessage('<i class="bi bi-exclamation-triangle"></i> No notebook to edit. Generate a notebook first.');
+            } else {
+                addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please describe the changes you want to make.');
+            }
+            return;
+        }
+        
+        // Add user message to conversation
+        addUserMessage('Edit request: ' + editPrompt);
+        
+        // Disable input during editing
+        setGeneratingState(true);
+        
+        const modelName = modelSelectEl.value;
+        
+        // Add AI typing indicator
+        const aiMessageId = 'ai-typing-' + Date.now();
+        addTypingIndicator(aiMessageId);
+        
+        handleEditStreamingResponse(editPrompt, currentNotebook, modelName, aiMessageId);
+        
+        // Clear the input field after sending
+        promptInputEl.value = '';
+    }
+    
+    function handleStreamingResponse(prompt, modelName, aiMessageId) {
+        aiResponseText = '';
+        
+        // Clear any existing notebook preview and add loading indicator
+        notebookPreviewEl.innerHTML = '<div class="loading-preview-message text-center p-5"><div class="spinner-border text-primary" role="status"></div><p class="mt-3">Building notebook preview...</p></div>';
+        
+        // Close any existing event source
+        if (eventSource) {
+            eventSource.close();
+        }
+        
+        // Keep track of when we should update the preview (not on every tiny chunk)
+        let lastPreviewUpdate = 0;
+        const PREVIEW_UPDATE_INTERVAL = 1000; // Update preview every 1 second during streaming
+        
+        // Track the last time we got a chunk - for timeout detection
+        let lastChunkTime = Date.now();
+        const CONNECTION_TIMEOUT = 30000; // 30 seconds without data = timeout
+        
+        // Start a monitoring timer to detect stalled connections
+        const connectionTimer = setInterval(() => {
+            if (Date.now() - lastChunkTime > CONNECTION_TIMEOUT) {
+                clearInterval(connectionTimer);
+                if (eventSource) {
+                    console.log("Connection timed out - closing event source");
+                    eventSource.close();
+                    eventSource = null;
+                    
+                    // Process what we've got so far
+                    updateAiMessage(aiMessageId, "**NoteGenie:** Generation timed out, but I'll process what I've received so far.");
+                    processNotebookResponse(aiResponseText);
+                    setGeneratingState(false);
+                }
+            }
+        }, 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();
+            
+            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}`);
+                    
+                    // 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 (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 (data.done) {
+                    eventSource.close();
+                    eventSource = null;
+                    clearInterval(connectionTimer);
+                    
+                    // 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);
+            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) {
+        aiResponseText = '';
+        
+        // Clear any existing notebook preview and add loading indicator
+        notebookPreviewEl.innerHTML = '<div class="loading-preview-message text-center p-5"><div class="spinner-border text-primary" role="status"></div><p class="mt-3">Updating notebook based on your edit request...</p></div>';
+        
+        // Close any existing event source
+        if (eventSource) {
+            eventSource.close();
+        }
+        
+        // Keep track of when we should update the preview (not on every tiny chunk)
+        let lastPreviewUpdate = 0;
+        const PREVIEW_UPDATE_INTERVAL = 1000; // Update preview every 1 second during streaming, same as generate
+        
+        // First, send the notebook data to the server so it's available in the session
+        fetch('/prepare_edit_notebook', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+            },
+            body: JSON.stringify({
+                notebook: notebook
+            })
+        })
+        .then(response => {
+            if (!response.ok) {
+                throw new Error(`Server error: ${response.status}`);
+            }
+            return response.json();
+        })
+        .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({
+                    edit_prompt: editPrompt,
+                    model: modelName,
+                    stream: true
+                }).toString()}`);
+                
+                // Track the last time we got a chunk - for timeout detection
+                let lastChunkTime = Date.now();
+                const CONNECTION_TIMEOUT = 30000; // 30 seconds without data = timeout
+                
+                // Start a monitoring timer to detect stalled connections
+                const connectionTimer = setInterval(() => {
+                    if (Date.now() - lastChunkTime > CONNECTION_TIMEOUT) {
+                        clearInterval(connectionTimer);
+                        if (eventSource) {
+                            console.log("Connection timed out - closing event source");
+                            eventSource.close();
+                            eventSource = null;
+                            
+                            // Process what we've got so far
+                            updateAiMessage(aiMessageId, "**NoteGenie:** Edit processing timed out, but I'll use what I've received so far.");
+                            processNotebookResponse(aiResponseText);
+                            setGeneratingState(false);
+                            
+                            // Switch back to generate mode
+                            setMode('generate');
+                        }
+                    }
+                }, 5000); // Check every 5 seconds
+                
+                eventSource.onmessage = function(event) {
+                    // Update our last-activity timestamp
+                    lastChunkTime = Date.now();
+                    
+                    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}`);
+                            
+                            // 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);
+                            
+                            // Switch back to generate mode
+                            setMode('generate');
+                            return;
+                        }
+                        
+                        if (data.chunk) {
+                            aiResponseText += data.chunk;
+                            
+                            // Extract notebook info as soon as it's available, similar to generate function
+                            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 with the actual description from the edited notebook
+                                updateAiMessage(aiMessageId, `**NoteGenie:** ${descMatch[1].trim()}`);
+                            } else {
+                                // Show a simple editing message while waiting for description
+                                updateAiMessage(aiMessageId, "**NoteGenie:** Updating notebook based on your edit request...");
+                            }
+                            
+                            // Update preview periodically during streaming, just like in generate function
+                            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 (data.done) {
+                            eventSource.close();
+                            eventSource = null;
+                            clearInterval(connectionTimer);
+                            
+                            // Process the complete response for final rendering
+                            processNotebookResponse(aiResponseText);
+                            setGeneratingState(false);
+                            
+                            // Get the final description for the AI message
+                            const finalDescMatch = aiResponseText.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?=\n\s*---|\n\s*$|$)/);
+                            const finalDesc = finalDescMatch ? finalDescMatch[1].trim() : "I've updated the notebook based on your edit request.";
+                            
+                            // Update AI message with final description
+                            updateAiMessage(aiMessageId, `**NoteGenie:** ${finalDesc}`);
+                            
+                            // Switch back to generate mode
+                            setMode('generate');
+                        }
+                    } 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);
+                    
+                    // 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 edit notebook. ' + (err.message || 'Unknown error'));
+                    }
+                    
+                    setGeneratingState(false);
+                    
+                    // Switch back to generate mode
+                    setMode('generate');
+                };
+            } else {
+                throw new Error(data.message || 'Failed to prepare notebook for editing');
+            }
+        })
+        .catch(error => {
+            console.error('Error preparing notebook for edit:', error);
+            updateAiMessage(aiMessageId, '**Error:** ' + error.message);
+            setGeneratingState(false);
+        });
+    }
+    
+    // Show API key modal
+    function showApiKeyModal() {
+        const apiKeyModal = new bootstrap.Modal(document.getElementById('apiKeyModal'));
+        apiKeyModal.show();
+    }
+    
+    // Helper functions
+    function updateNotebookPreviewDuringStream(text) {
+        try {
+            // Extract notebook info for title
+            const nameMatch = text.match(/NOTEBOOK_NAME:?\s*(.+?)(?:\n|$)/);
+            const name = nameMatch ? nameMatch[1].trim() : 'Generating Notebook...';
+            
+            // Update notebook title
+            notebookTitleEl.textContent = name;
+            
+            // Find all cell markers in order of appearance
+            const cellMarkers = [...text.matchAll(/---\s*(MARKDOWN|CODE)\s*CELL\s*---/g)];
+            
+            // If no markers found, don't update preview yet
+            if (cellMarkers.length === 0) return;
+            
+            // Check if user is already near the bottom before deciding to auto-scroll
+            const isNearBottom = notebookPreviewEl.scrollHeight - notebookPreviewEl.scrollTop - notebookPreviewEl.clientHeight < 100;
+            
+            // Initialize cell counters
+            let markdownCellCount = 0;
+            let codeCellCount = 0;
+            
+            // Clear preview and prepare for new content
+            notebookPreviewEl.innerHTML = '';
+            
+            // Process each cell in order
+            for (let i = 0; i < cellMarkers.length; i++) {
+                const currentMarker = cellMarkers[i];
+                const nextMarker = cellMarkers[i + 1];
+                const cellType = currentMarker[1]; // MARKDOWN or CODE
+                
+                // Extract cell content between current marker and next marker (or end of text)
+                const markerEndIndex = currentMarker.index + currentMarker[0].length;
+                const contentEndIndex = nextMarker ? nextMarker.index : text.length;
+                let cellContent = text.substring(markerEndIndex, contentEndIndex).trim();
+                
+                if (cellType === 'MARKDOWN') {
+                    markdownCellCount++;
+                    createMarkdownCell(cellContent, markdownCellCount, notebookPreviewEl);
+                } else if (cellType === 'CODE') {
+                    // Extract Python code from between ```python and ```
+                    const codeMatch = cellContent.match(/```python\s*([\s\S]*?)```/);
+                    if (codeMatch) {
+                        codeCellCount++;
+                        createCodeCell(codeMatch[1].trim(), codeCellCount, notebookPreviewEl);
+                    }
+                }
+            }
+            
+            // Only auto-scroll if user was already near the bottom
+            if (isNearBottom) {
+                notebookPreviewEl.scrollTop = notebookPreviewEl.scrollHeight;
+            } else {
+                // Show the scroll to bottom button
+                scrollToBottomBtn.classList.add('visible');
+            }
+        } catch (error) {
+            console.error('Error updating preview during stream:', error);
+            // Don't update if there's an error - wait for complete response
+        }
+    }
+    
+    function createMarkdownCell(content, cellNumber, container) {
+        const cellDiv = document.createElement('div');
+        cellDiv.className = 'notebook-cell cell-markdown';
+        
+        // Add header for markdown cell
+        const headerDiv = document.createElement('div');
+        headerDiv.className = 'cell-header markdown-header';
+        headerDiv.innerHTML = `<span class="cell-type">Markdown [${cellNumber}]</span>`;
+        cellDiv.appendChild(headerDiv);
+        
+        // Add markdown content
+        const contentDiv = document.createElement('div');
+        contentDiv.className = 'cell-content';
+        contentDiv.innerHTML = marked.parse(content);
+        cellDiv.appendChild(contentDiv);
+        
+        container.appendChild(cellDiv);
+    }
+    
+    function createCodeCell(content, cellNumber, container) {
+        const cellDiv = document.createElement('div');
+        cellDiv.className = 'notebook-cell cell-code';
+        
+        // Add header for code cell
+        const headerDiv = document.createElement('div');
+        headerDiv.className = 'cell-header code-header';
+        headerDiv.innerHTML = `<span class="cell-type">Code [${cellNumber}]</span>`;
+        cellDiv.appendChild(headerDiv);
+        
+        // Add code content
+        const contentDiv = document.createElement('div');
+        contentDiv.className = 'cell-content';
+        const preEl = document.createElement('pre');
+        const codeEl = document.createElement('code');
+        codeEl.className = 'language-python';
+        codeEl.textContent = content;
+        preEl.appendChild(codeEl);
+        contentDiv.appendChild(preEl);
+        cellDiv.appendChild(contentDiv);
+        
+        container.appendChild(cellDiv);
+        
+        // Highlight syntax
+        Prism.highlightElement(codeEl);
+    }
+    
+    function processNotebookResponse(text) {
+        // Extract notebook info using improved regex patterns that handle markdown better
+        const nameMatch = text.match(/NOTEBOOK_NAME:?\s*(.+?)(?=\n\s*NOTEBOOK_DESCRIPTION|\n\s*---|\n\s*$|$)/);
+        const descMatch = text.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?=\n\s*---|\n\s*$|$)/);
+        
+        // Clean up the name and description
+        let name = nameMatch ? nameMatch[1].trim() : 'Generated Notebook';
+        let description = descMatch ? descMatch[1].trim() : '';
+        
+        // Remove markdown formatting
+        name = name.replace(/\*\*/g, '').replace(/\*/g, '');
+        description = description.replace(/\*\*/g, '').replace(/\*/g, '');
+        
+        // Update notebook title
+        notebookTitleEl.textContent = name;
+        
+        // Update the AI message with just the description
+        const chatMessageId = document.querySelector('.ai-message').id;
+        updateAiMessage(chatMessageId, `**NoteGenie:** ${description}`);
+        
+        // OPTIMIZATION: Process the notebook client-side if possible for small to medium notebooks
+        if (text.length < 50000) { // Only try client-side processing for reasonably sized notebooks
+            try {
+                // Try to process notebook directly in the client
+                const notebookJson = clientSideFormatNotebook(text);
+                renderNotebook(notebookJson);
+                currentNotebook = notebookJson;
+                downloadBtnEl.disabled = false;
+                editModeBtnEl.disabled = false;
+                return; // Exit if successful
+            } catch (error) {
+                console.log("Client-side formatting failed, falling back to server:", error);
+                // Fall through to server-side processing if client-side fails
+            }
+        }
+        
+        // OPTIMIZATION: Show loading indicator with progress message
+        notebookPreviewEl.innerHTML = '<div class="loading-preview-message text-center p-5"><div class="spinner-border text-primary" role="status"></div><p class="mt-3">Finalizing notebook format...</p></div>';
+        
+        // Make API call to convert text to notebook format with a timeout to prevent UI freezing
+        setTimeout(() => {
+            fetch('/generate_notebook', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({
+                    prompt: text,  // Send the full AI response as a prompt
+                    model: 'gemini-2.0-flash',  // Use a faster model for formatting
+                    stream: false,
+                    format_only: true
+                })
+            })
+            .then(response => response.json())
+            .then(data => {
+                if (data.success) {
+                    renderNotebook(data.notebook);
+                    currentNotebook = data.notebook;
+                    downloadBtnEl.disabled = false;
+                    editModeBtnEl.disabled = false;
+                }
+            })
+            .catch(error => {
+                console.error('Error formatting notebook:', error);
+                notebookPreviewEl.innerHTML = `<div class="alert alert-danger">Error formatting notebook: ${error.message}</div>`;
+            });
+        }, 10); // Small delay to allow UI to update
+    }
+    
+    // ADDED: Client-side notebook formatting function to reduce server dependency
+    function clientSideFormatNotebook(content) {
+        // Extract cells with improved regex that captures end of file properly
+        const markdownCells = content.match(/---\s*MARKDOWN\s*CELL\s*---\s*([\s\S]*?)(?=---\s*(?:MARKDOWN|CODE)\s*CELL\s*---|$)/g) || [];
+        const codeCells = content.match(/---\s*CODE\s*CELL\s*---\s*```python\s*([\s\S]*?)```/gs) || [];
+        
+        // Get cell order with improved markers
+        const cellMarkers = [...content.matchAll(/---\s*(MARKDOWN|CODE)\s*CELL\s*---/g)];
+        const cellTypes = cellMarkers.map(match => match[1]);
+        
+        // Initialize notebook structure
+        const cells = [];
+        let mdIdx = 0;
+        let codeIdx = 0;
+        
+        // Populate cells in the correct order
+        for (let i = 0; i < cellTypes.length; i++) {
+            const cellType = cellTypes[i];
+            
+            if (cellType === 'MARKDOWN' && mdIdx < markdownCells.length) {
+                // Extract markdown content with improved regex
+                let md = markdownCells[mdIdx].replace(/---\s*MARKDOWN\s*CELL\s*---/, '').trim();
+                cells.push({
+                    cell_type: 'markdown',
+                    metadata: {},
+                    source: md.split('\n')
+                });
+                mdIdx++;
+            } else if (cellType === 'CODE' && codeIdx < codeCells.length) {
+                // Extract code content with improved regex
+                let codeMatch = codeCells[codeIdx].match(/```python\s*([\s\S]*?)```/s);
+                let code = codeMatch ? codeMatch[1].trim() : '';
+                cells.push({
+                    cell_type: 'code',
+                    execution_count: null,
+                    metadata: {},
+                    outputs: [],
+                    source: code.split('\n')
+                });
+                codeIdx++;
+            }
+        }
+        
+        // Process the final cell if it wasn't captured above
+        if (cellMarkers.length > 0) {
+            const lastMarker = cellMarkers[cellMarkers.length - 1];
+            const lastMarkerEndIndex = lastMarker.index + lastMarker[0].length;
+            
+            // If there's content after the last marker, process it
+            if (lastMarkerEndIndex < content.length) {
+                const lastCellType = lastMarker[1]; // MARKDOWN or CODE
+                const lastCellContent = content.substring(lastMarkerEndIndex).trim();
+                
+                // Only add if not already captured
+                const isAlreadyCaptured = 
+                    (lastCellType === 'MARKDOWN' && mdIdx >= markdownCells.length) ||
+                    (lastCellType === 'CODE' && codeIdx >= codeCells.length);
+                    
+                if (!isAlreadyCaptured && lastCellContent) {
+                    if (lastCellType === 'MARKDOWN') {
+                        cells.push({
+                            cell_type: 'markdown',
+                            metadata: {},
+                            source: lastCellContent.split('\n')
+                        });
+                    } else if (lastCellType === 'CODE') {
+                        // Extract code if present
+                        const codeMatch = lastCellContent.match(/```python\s*([\s\S]*?)```/s);
+                        if (codeMatch) {
+                            cells.push({
+                                cell_type: 'code',
+                                execution_count: null,
+                                metadata: {},
+                                outputs: [],
+                                source: codeMatch[1].trim().split('\n')
+                            });
+                        }
+                    }
+                }
+            }
+        }
+        
+        // Fallback approach if no cells were found or if the extraction didn't work
+        // ...existing code...
+        
+        return {
+            cells: cells,
+            metadata: {
+                kernelspec: {
+                    display_name: 'Python 3',
+                    language: 'python',
+                    name: 'python3'
+                },
+                language_info: {
+                    name: 'python',
+                    version: '3.8.0'
+                }
+            },
+            nbformat: 4,
+            nbformat_minor: 4
+        };
+    }
+    
+    function downloadNotebook() {
+        if (!currentNotebook) {
+            return;
+        }
+        
+        // Improved filename cleaning - remove markdown formatting and unsafe filename characters
+        let cleanName = notebookTitleEl.textContent || 'generated_notebook';
+        // First remove markdown formatting characters
+        cleanName = cleanName.replace(/\*\*/g, '').replace(/\*/g, '').replace(/_/g, ' ');
+        // Then remove any characters that aren't safe for filenames
+        cleanName = cleanName.replace(/[^\w\s-]/gi, '').trim();
+        // Replace multiple spaces with single space
+        cleanName = cleanName.replace(/\s+/g, ' ');
+        
+        const filename = cleanName || 'generated_notebook';
+        
+        fetch('/download_notebook', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+            },
+            body: JSON.stringify({
+                notebook: currentNotebook,
+                filename: `${filename}.ipynb`
+            })
+        })
+        .then(response => response.blob())
+        .then(blob => {
+            const url = window.URL.createObjectURL(blob);
+            const a = document.createElement('a');
+            a.style.display = 'none';
+            a.href = url;
+            a.download = `${filename}.ipynb`;
+            document.body.appendChild(a);
+            a.click();
+            window.URL.revokeObjectURL(url);
+        })
+        .catch(error => {
+            console.error('Error downloading notebook:', error);
+            addSystemMessage('Error downloading notebook: ' + error.message);
+        });
+    }
+    
+    function renderNotebook(notebook) {
+        notebookPreviewEl.innerHTML = '';
+        
+        if (!notebook || !notebook.cells || !notebook.cells.length) {
+            notebookPreviewEl.innerHTML = '<div class="alert alert-warning">No notebook content available</div>';
+            return;
+        }
+        
+        let cellCount = 0;
+        
+        notebook.cells.forEach((cell, index) => {
+            cellCount++;
+            
+            if (cell.cell_type === 'markdown') {
+                const content = Array.isArray(cell.source) ? cell.source.join('\n') : cell.source;
+                createMarkdownCell(content, cellCount, notebookPreviewEl);
+            } else if (cell.cell_type === 'code') {
+                const codeContent = Array.isArray(cell.source) ? cell.source.join('\n') : cell.source;
+                createCodeCell(codeContent, cellCount, notebookPreviewEl);
+            }
+        });
+        
+        // Check if we should scroll to the bottom after rendering
+        const isNearBottom = notebookPreviewEl.scrollHeight - notebookPreviewEl.scrollTop - notebookPreviewEl.clientHeight < 100;
+        if (isNearBottom) {
+            notebookPreviewEl.scrollTop = notebookPreviewEl.scrollHeight;
+        } else {
+            // Show the scroll to bottom button
+            scrollToBottomBtn.classList.add('visible');
+        }
+        
+        // Enable edit and download buttons when notebook is available
+        downloadBtnEl.disabled = false;
+        editModeBtnEl.disabled = false;
+    }
+    
+    // UI Message handling
+    function addUserMessage(text) {
+        // Remove welcome message if present
+        const welcomeMessage = document.querySelector('.welcome-message');
+        if (welcomeMessage) {
+            welcomeMessage.remove();
+        }
+        
+        const messageDiv = document.createElement('div');
+        messageDiv.className = 'message user-message';
+        
+        // Fix: Create a span for the "You:" prefix and handle the text separately
+        // to prevent marked from adding paragraph tags that cause line breaks
+        const userPrefix = document.createElement('span');
+        userPrefix.className = 'user-prefix';
+        userPrefix.innerHTML = '<strong>You:</strong> ';
+        
+        const messageContent = document.createElement('span');
+        messageContent.className = 'message-content';
+        messageContent.innerHTML = marked.parse(text);
+        
+        // Remove any leading <p> and trailing </p> tags that marked adds
+        messageContent.innerHTML = messageContent.innerHTML
+            .replace(/^<p>/, '')
+            .replace(/<\/p>$/, '');
+        
+        messageDiv.appendChild(userPrefix);
+        messageDiv.appendChild(messageContent);
+        
+        conversationEl.appendChild(messageDiv);
+        conversationEl.scrollTop = conversationEl.scrollHeight;
+    }
+    
+    function addTypingIndicator(id) {
+        const messageDiv = document.createElement('div');
+        messageDiv.className = 'message ai-message';
+        messageDiv.id = id;
+        
+        const typingDiv = document.createElement('div');
+        typingDiv.className = 'typing-indicator';
+        for (let i = 0; i < 3; i++) {
+            const dot = document.createElement('span');
+            dot.className = 'typing-dot';
+            typingDiv.appendChild(dot);
+        }
+        
+        messageDiv.appendChild(typingDiv);
+        conversationEl.appendChild(messageDiv);
+        conversationEl.scrollTop = conversationEl.scrollHeight;
+    }
+    
+    function updateAiMessage(id, text) {
+        const messageDiv = document.getElementById(id);
+        if (messageDiv) {
+            if (text.startsWith('**NoteGenie:**')) {
+                text = text.replace('**NoteGenie:**', '**NoteGenie:**');
+            } else if (text.startsWith('**Error:**')) {
+                text = text.replace('**Error:**', '**<i class="bi bi-exclamation-triangle"></i> Error:**');
+            }
+            
+            messageDiv.innerHTML = marked.parse(text);
+            // Ensure auto-scrolling in the chat panel
+            conversationEl.scrollTop = conversationEl.scrollHeight;
+        }
+    }
+    
+    function addSystemMessage(text) {
+        const messageDiv = document.createElement('div');
+        messageDiv.className = 'info-card mb-3';
+        messageDiv.innerHTML = text;
+        conversationEl.appendChild(messageDiv);
+        conversationEl.scrollTop = conversationEl.scrollHeight;
+    }
+    
+    // UI state management
+    function setGeneratingState(isGenerating) {
+        if (isGenerating) {
+            actionBtnEl.disabled = true;
+            // Update loading state for the send button
+            actionBtnEl.innerHTML = '<span class="material-icons spinning">autorenew</span>';
+            promptInputEl.disabled = true;
+            
+            // Disable buttons but preserve their active/inactive visual state
+            generateModeBtnEl.disabled = true;
+            editModeBtnEl.disabled = true;
+            
+            // Make sure mode toggle visual state is preserved
+            if (currentMode === 'generate') {
+                generateModeBtnEl.classList.add('active');
+                editModeBtnEl.classList.remove('active');
+            } else {
+                editModeBtnEl.classList.add('active');
+                generateModeBtnEl.classList.remove('active');
+            }
+        } else {
+            // Reset the debounce flag when generation completes
+            isActionInProgress = false;
+            
+            actionBtnEl.disabled = false;
+            // Always use the send icon for the inline button
+            actionBtnEl.innerHTML = '<span class="material-icons">send</span>';
+            promptInputEl.disabled = false;
+            promptInputEl.focus();
+            
+            generateModeBtnEl.disabled = false;
+            
+            // Only enable edit mode button if we have a current notebook
+            editModeBtnEl.disabled = !currentNotebook;
+        }
+    }
+});

templates/index.html

@@ -0,0 +1,189 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>NoteGenie - AI-Powered Jupyter Notebook Generator</title>
+    <!-- Google fonts -->
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Google+Sans:wght@400;500;700&family=Roboto:wght@300;400;500&display=swap" rel="stylesheet">
+    <!-- Material Icons -->
+    <link href="https://fonts.googleapis.com/icon?family=Material+Icons" rel="stylesheet">
+    <!-- Bootstrap (still used for grid and components) -->
+    <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0-alpha1/dist/css/bootstrap.min.css" rel="stylesheet">
+    <!-- Prism for code highlighting -->
+    <link href="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/themes/prism.min.css" rel="stylesheet">
+    <!-- Custom styles -->
+    <link href="{{ url_for('static', filename='css/style.css') }}" rel="stylesheet">
+</head>
+<body>
+    <div class="app-container">
+        <!-- App Header -->
+        <header class="app-header">
+            <div class="header-content">
+                <div class="logo-section">
+                    <span class="material-icons logo-icon">auto_awesome</span>
+                    <h1 class="product-name">NoteGenie</h1>
+                </div>
+                <div class="header-actions">
+                    <button class="google-button outlined" data-bs-toggle="modal" data-bs-target="#apiKeyModal">
+                        <span class="material-icons">vpn_key</span>
+                        <span>API Key</span>
+                    </button>
+                </div>
+            </div>
+        </header>
+        
+        <!-- Main Content Area -->
+        <main class="main-container">
+            <!-- Left Panel - Chat Interface -->
+            <div class="panel chat-panel">
+                <div id="conversation" class="conversation-container">
+                    <div class="welcome-message">
+                        <h2><span class="material-icons">auto_awesome</span> Welcome to NoteGenie</h2>
+                        <p>Generate complete Jupyter notebooks from your text prompts using Google's Gemini AI.</p>
+                        <div class="welcome-steps">
+                            <div class="welcome-step">
+                                <div class="step-number">1</div>
+                                <div class="step-text">Set your Google Gemini API key</div>
+                            </div>
+                            <div class="welcome-step">
+                                <div class="step-number">2</div>
+                                <div class="step-text">Describe the notebook you want</div>
+                            </div>
+                            <div class="welcome-step">
+                                <div class="step-number">3</div>
+                                <div class="step-text">View and refine the generated notebook</div>
+                            </div>
+                            <div class="welcome-step">
+                                <div class="step-number">4</div>
+                                <div class="step-text">Download as .ipynb file</div>
+                            </div>
+                        </div>
+                        <div class="info-card">
+                            <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>
+                            </div>
+                        </div>
+                        <div class="example-section">
+                            <h5>Example prompts:</h5>
+                            <ul class="example-list">
+                                <li>"Create a notebook for data visualization with Plotly Express showing different chart types"</li>
+                                <li>"Make a machine learning notebook that shows text classification using BERT"</li>
+                                <li>"Build a beginner-friendly introduction to pandas with common data operations"</li>
+                            </ul>
+                        </div>
+                    </div>
+                </div>
+                
+                <div class="input-container">
+                    <div class="input-options">
+                        <select id="modelSelect" class="model-selector">
+                            <option value="gemini-2.0-pro">Gemini 2.0 Pro</option>
+                            <option value="gemini-2.0-flash">Gemini 2.0 Flash</option>
+                            <option value="gemini-2.0-flash-thinking">Gemini 2.0 Flash Thinking</option>
+                        </select>
+                        <div class="mode-toggle">
+                            <button type="button" class="toggle-button active" id="generateModeBtn">Generate</button>
+                            <button type="button" class="toggle-button" id="editModeBtn" disabled>Edit</button>
+                        </div>
+                    </div>
+                    <div class="input-field-container">
+                        <textarea id="promptInput" class="prompt-input" placeholder="Describe the notebook you want..." rows="3"></textarea>
+                        <button id="actionBtn" class="send-button">
+                            <span class="material-icons">send</span>
+                        </button>
+                    </div>
+                    <!-- Keep the original div but it will be hidden via CSS -->
+                    <div class="input-actions">
+                        <!-- Original button is now moved into the input field -->
+                    </div>
+                </div>
+            </div>
+            
+            <!-- Right Panel - Notebook Preview -->
+            <div class="panel notebook-panel">
+                <div class="notebook-header">
+                    <h3 id="notebookTitle" class="notebook-title"></h3>
+                    <div class="notebook-actions">
+                        <button id="downloadBtn" class="google-button" disabled>
+                            <span class="material-icons">download</span>
+                            <span>Download</span>
+                        </button>
+                    </div>
+                </div>
+                <div id="notebookPreview" class="notebook-preview">
+                    <div class="placeholder-content">
+                        <span class="material-icons large-icon">description</span>
+                        <p>Generate a notebook to see the preview here</p>
+                    </div>
+                </div>
+            </div>
+        </main>
+    </div>
+    
+    <!-- API Key Modal -->
+    <div class="modal fade" id="apiKeyModal" tabindex="-1" aria-hidden="true">
+        <div class="modal-dialog modal-dialog-centered">
+            <div class="modal-content google-modal">
+                <div class="modal-header">
+                    <h5 class="modal-title">
+                        <span class="material-icons">vpn_key</span>
+                        Set Google Gemini API Key
+                    </h5>
+                    <button type="button" class="btn-close" data-bs-dismiss="modal" aria-label="Close"></button>
+                </div>
+                <div class="modal-body">
+                    <div class="info-card">
+                        <span class="material-icons info-icon">info</span>
+                        <div class="info-content">
+                            <p><strong>You need to set your Google API key to use NoteGenie.</strong></p>
+                            <p>Your key will be remembered for future sessions on this device.</p>
+                        </div>
+                    </div>
+                    
+                    <div class="steps-card">
+                        <div class="steps-header">
+                            <span class="material-icons">key</span>
+                            <strong>How to get your API key:</strong>
+                        </div>
+                        <div class="steps-content">
+                            <ol class="steps-list">
+                                <li>Go to <a href="https://aistudio.google.com/" target="_blank">Google AI Studio</a></li>
+                                <li>Click on "Get API key" in the top left</li>
+                                <li>Click "Create API key"</li>
+                                <li>Copy the generated API key and paste it below</li>
+                            </ol>
+                        </div>
+                    </div>
+                    
+                    <div class="input-field">
+                        <label for="apiKeyInput" class="input-label">API Key</label>
+                        <input type="password" class="text-input" id="apiKeyInput" placeholder="Enter your Gemini API key">
+                        <div class="input-helper">Your API key is stored securely and never sent to third parties.</div>
+                    </div>
+                    <div id="apiKeyFeedback"></div>
+                </div>
+                <div class="modal-footer">
+                    <button type="button" class="google-button text" data-bs-dismiss="modal">Cancel</button>
+                    <button type="button" class="google-button primary" id="saveApiKeyBtn">Save</button>
+                </div>
+            </div>
+        </div>
+    </div>
+
+    <!-- Floating scroll button -->
+    <button class="scroll-button" id="scrollToBottomBtn">
+        <span class="material-icons">keyboard_arrow_down</span>
+    </button>
+
+    <!-- Scripts -->
+    <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0-alpha1/dist/js/bootstrap.bundle.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/components/prism-core.min.js"></script>
+    <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>
+</body>
+</html>

utils/ai_helpers.py

@@ -0,0 +1,220 @@
+import google.generativeai as genai
+from flask import Response, stream_with_context
+import json
+import time
+
+def craft_notebook_prompt(user_prompt):
+    """Enhance the user prompt with instructions for generating a well-structured Jupyter notebook."""
+    enhanced_prompt = f"""
+    Create a complete Jupyter notebook based on this request: "{user_prompt}"
+    
+    Please structure your response as follows:
+    
+    NOTEBOOK_NAME: [Short, descriptive name for the notebook with no formating "**"]
+    NOTEBOOK_DESCRIPTION: [a description of the notebook's purpose with no formating]
+    
+    Then provide the complete notebook with proper alternating Markdown and code cells. 
+    Format each cell as follows:
+    
+    --- MARKDOWN CELL ---
+    [Markdown content here]
+    
+    --- CODE CELL ---
+    ```python
+    [Python code here]
+    ```
+    
+    Important guidelines:
+    - Include comprehensive explanations in Markdown cells
+    - Ensure all code is executable and properly commented
+    - Include data loading, processing, and visualization where appropriate
+    - Add explanatory text before and after code sections
+    - Include example outputs or expected results when relevant
+    - Structure the notebook with clear section headers in Markdown
+    """
+    return enhanced_prompt
+
+def craft_edit_prompt(edit_request, notebook_json):
+    """Create a prompt for editing an existing notebook."""
+    # Extract cells from notebook JSON for context
+    cells = []
+    for i, cell in enumerate(notebook_json.get('cells', [])):
+        cell_type = cell.get('cell_type', 'unknown')
+        source = cell.get('source', [])
+        if isinstance(source, list):
+            source = '\n'.join(source)
+        
+        if cell_type == 'markdown':
+            cells.append(f"--- CELL {i+1} (MARKDOWN) ---\n{source}")
+        elif cell_type == 'code':
+            cells.append(f"--- CELL {i+1} (CODE) ---\n```python\n{source}\n```")
+    
+    notebook_content = '\n\n'.join(cells)
+    
+    # Create prompt with edit instructions
+    enhanced_prompt = f"""
+    I have a Jupyter notebook that I'd like you to modify based on this edit request: "{edit_request}"
+    
+    Here's the current notebook content:
+    
+    NOTEBOOK_STRUCTURE:
+    {notebook_content}
+    
+    Please provide the complete updated notebook with the requested changes, following the same format:
+    
+    NOTEBOOK_NAME: [Keep or update the notebook name]
+    NOTEBOOK_DESCRIPTION: [Keep or update the notebook description]
+    
+    Then provide the complete notebook with proper alternating Markdown and code cells. 
+    Format each cell as follows:
+    
+    --- MARKDOWN CELL ---
+    [Markdown content here]
+    
+    --- CODE CELL ---
+    ```python
+    [Python code here]
+    ```
+    
+    Important guidelines:
+    - Make only the changes requested in the edit request
+    - Preserve the overall structure of the notebook
+    - Keep all content from the original notebook that doesn't need modification
+    - Ensure all code remains executable and properly commented
+    - Feel free to reorganize, add, or remove cells as needed to fulfill the edit request
+    """
+    return enhanced_prompt
+
+def generate_notebook(user_prompt, model_name="gemini-2.0-pro-exp-02-05"):
+    """Generate a complete notebook using Gemini API."""
+    model = genai.GenerativeModel(model_name)
+    
+    enhanced_prompt = craft_notebook_prompt(user_prompt)
+    response = model.generate_content(enhanced_prompt)
+    
+    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."""
+    model = genai.GenerativeModel(model_name)
+    
+    enhanced_prompt = craft_edit_prompt(edit_request, notebook_json)
+    response = model.generate_content(enhanced_prompt)
+    
+    return response.text
+
+def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05"):
+    """Stream notebook generation responses from Gemini API."""
+    model = genai.GenerativeModel(model_name)
+    enhanced_prompt = craft_notebook_prompt(user_prompt)
+    
+    def generate():
+        try:
+            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"
+            
+            for chunk in response:
+                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)")
+                        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"
+                            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
+                    
+                    # 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"
+                        elif isinstance(part, dict) and 'text' in part:
+                            yield f"data: {json.dumps({'chunk': part['text']})}\n\n"
+                        elif hasattr(part, 'string_value'):
+                            yield f"data: {json.dumps({'chunk': part.string_value})}\n\n"
+                
+                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]}")
+                    continue
+                
+                # Briefly pause to prevent overwhelming the client
+                time.sleep(0.01)
+            
+            yield f"data: {json.dumps({'done': True})}\n\n"
+            
+        except Exception as e:
+            # Send error to client and close stream
+            error_message = f"Error generating notebook: {str(e)}"
+            print(error_message)
+            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")
+
+def stream_notebook_edit(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05"):
+    """Stream notebook editing responses from Gemini API."""
+    model = genai.GenerativeModel(model_name)
+    enhanced_prompt = craft_edit_prompt(edit_request, notebook_json)
+    
+    def generate():
+        try:
+            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"
+            
+            for chunk in response:
+                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)")
+                        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"
+                            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
+                    
+                    # 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"
+                        elif isinstance(part, dict) and 'text' in part:
+                            yield f"data: {json.dumps({'chunk': part['text']})}\n\n"
+                        elif hasattr(part, 'string_value'):
+                            yield f"data: {json.dumps({'chunk': part.string_value})}\n\n"
+                
+                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]}")
+                    continue
+                
+                # Briefly pause to prevent overwhelming the client
+                time.sleep(0.01)
+            
+            yield f"data: {json.dumps({'done': True})}\n\n"
+            
+        except Exception as e:
+            # Send error to client and close stream
+            error_message = f"Error editing notebook: {str(e)}"
+            print(error_message)
+            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")

utils/notebook_helpers.py

@@ -0,0 +1,289 @@
+import re
+import json
+import functools
+
+# Add a simple LRU cache for regex patterns
+def get_cached_pattern(pattern, flags=0):
+    """Cache compiled regex patterns for better performance."""
+    @functools.lru_cache(maxsize=32)
+    def _get_pattern(pattern_str, pattern_flags):
+        return re.compile(pattern_str, pattern_flags)
+    
+    return _get_pattern(pattern, flags)
+
+def extract_notebook_info(content):
+    """Extract notebook name and description from the AI response."""
+    # Improved regex pattern that handles multiline and markdown formatting better
+    name_match = re.search(r"NOTEBOOK_NAME:?\s*(.+?)(?=\n\s*NOTEBOOK_DESCRIPTION|\n\s*---|\n\s*$|$)", content, re.DOTALL)
+    desc_match = re.search(r"NOTEBOOK_DESCRIPTION:?\s*(.+?)(?=\n\s*---|\n\s*$|$)", content, re.DOTALL)
+    
+    # Extract and clean up potential markdown formatting
+    name = name_match.group(1).strip() if name_match else "Generated Notebook"
+    description = desc_match.group(1).strip() if desc_match else "Notebook generated using NoteGenie"
+    
+    # Remove markdown formatting from name and description
+    name = re.sub(r'\*\*(.*?)\*\*', r'\1', name)  # Remove bold formatting
+    name = re.sub(r'\*(.*?)\*', r'\1', name)      # Remove italic formatting
+    name = re.sub(r'_(.*?)_', r'\1', name)        # Remove underline formatting
+    
+    description = re.sub(r'\*\*(.*?)\*\*', r'\1', description)  # Remove bold formatting
+    description = re.sub(r'\*(.*?)\*', r'\1', description)      # Remove italic formatting
+    description = re.sub(r'_(.*?)_', r'\1', description)        # Remove underline formatting
+    
+    return {
+        "name": name,
+        "description": description
+    }
+
+def format_notebook(content):
+    """Convert the AI text response into a properly formatted Jupyter notebook JSON.
+    Optimized for performance with larger texts."""
+    # Use faster pattern matching approach with improved end-of-file handling
+    markdown_pattern = get_cached_pattern(r"---\s*MARKDOWN\s*CELL\s*---\s*([\s\S]*?)(?=---\s*(?:MARKDOWN|CODE)\s*CELL\s*---|$)", re.DOTALL)
+    code_pattern = get_cached_pattern(r"---\s*CODE\s*CELL\s*---\s*```python\s*([\s\S]*?)```", re.DOTALL)
+    cell_marker_pattern = get_cached_pattern(r"---\s*(MARKDOWN|CODE)\s*CELL\s*---", re.DOTALL)
+    
+    # OPTIMIZATION: Do a quick initial scan to determine notebook size and complexity
+    complexity = len(content) // 1000  # Rough estimate based on content length
+    cell_count = len(cell_marker_pattern.findall(content))
+    
+    # For very large notebooks, use a more memory-efficient but slower approach
+    if complexity > 200 or cell_count > 50:  # If over ~200KB or 50 cells
+        return format_large_notebook(content)
+    
+    # For regular notebooks, use the standard approach which is faster for medium-sized content
+    try:
+        # Extract cells from the content in a single pass if possible
+        markdown_cells = markdown_pattern.findall(content)
+        code_cells = code_pattern.findall(content)
+        
+        # If the AI didn't use the expected format, try alternate patterns
+        if not markdown_cells and not code_cells:
+            # Simplified handling for non-standard format
+            sections = re.split(r"```python|```", content)
+            cells = []
+            
+            for i, section in enumerate(sections):
+                section = section.strip()
+                if section and i % 2 == 0:
+                    # This is markdown content
+                    cells.append({"cell_type": "markdown", "source": section})
+                elif section:
+                    # This is code content
+                    cells.append({"cell_type": "code", "source": section})
+        else:
+            # Interleave markdown and code cells in the correct order
+            cells = []
+            
+            # Find overall ordering of cells
+            all_matches = list(cell_marker_pattern.finditer(content))
+            all_types = [m.group(1) for m in all_matches]
+            
+            md_idx = 0
+            code_idx = 0
+            
+            for i, cell_type in enumerate(all_types):
+                marker = all_matches[i]
+                marker_end = marker.end()
+                next_marker_start = all_matches[i+1].start() if i+1 < len(all_matches) else len(content)
+                cell_content = content[marker_end:next_marker_start].strip()
+                
+                if cell_type == "MARKDOWN":
+                    if md_idx < len(markdown_cells) or (i == len(all_types) - 1 and cell_content):
+                        if md_idx < len(markdown_cells):
+                            cell_source = markdown_cells[md_idx].strip()
+                            md_idx += 1
+                        else:
+                            # Handle the last markdown cell if it wasn't captured by the pattern
+                            cell_source = cell_content
+                            
+                        cells.append({
+                            "cell_type": "markdown",
+                            "source": cell_source
+                        })
+                elif cell_type == "CODE":
+                    if code_idx < len(code_cells) or (i == len(all_types) - 1 and "```python" in cell_content):
+                        if code_idx < len(code_cells):
+                            cell_source = code_cells[code_idx].strip()
+                            code_idx += 1
+                        else:
+                            # Handle the last code cell if it wasn't captured by the pattern
+                            code_match = re.search(r"```python\s*([\s\S]*?)```", cell_content, re.DOTALL)
+                            cell_source = code_match.group(1).strip() if code_match else ""
+                            
+                        cells.append({
+                            "cell_type": "code",
+                            "source": cell_source
+                        })
+        
+        # Ensure we have at least a title cell if nothing was extracted
+        if not cells:
+            notebook_info = extract_notebook_info(content)
+            cells.append({
+                "cell_type": "markdown",
+                "source": f"# {notebook_info['name']}\n\n{notebook_info['description']}"
+            })
+            
+            # Try to extract any code blocks that might be present - only if needed
+            code_blocks = re.findall(r"```python\s*(.*?)```", content, re.DOTALL)
+            for block in code_blocks:
+                cells.append({
+                    "cell_type": "code",
+                    "source": block.strip()
+                })
+        
+        # Format cells for Jupyter notebook structure - optimize by processing in chunks
+        formatted_cells = []
+        for cell in cells:
+            cell_source = cell["source"]
+            # Only split if it's a string, not if it's already a list
+            if isinstance(cell_source, str):
+                # OPTIMIZATION: For very large cells, process line by line to avoid memory issues
+                if len(cell_source) > 10000:  # If cell is over 10KB
+                    source_lines = []
+                    for line in cell_source.splitlines():
+                        source_lines.append(line)
+                else:
+                    source_lines = cell_source.split("\n")
+            else:
+                source_lines = cell_source
+                
+            formatted_cell = {
+                "cell_type": cell["cell_type"],
+                "metadata": {},
+                "source": source_lines
+            }
+            
+            if cell["cell_type"] == "code":
+                formatted_cell["execution_count"] = None
+                formatted_cell["outputs"] = []
+            
+            formatted_cells.append(formatted_cell)
+        
+        # Create the notebook structure
+        notebook = {
+            "cells": formatted_cells,
+            "metadata": {
+                "kernelspec": {
+                    "display_name": "Python 3",
+                    "language": "python",
+                    "name": "python3"
+                },
+                "language_info": {
+                    "name": "python",
+                    "version": "3.8.0"
+                }
+            },
+            "nbformat": 4,
+            "nbformat_minor": 4
+        }
+        
+        return notebook
+    except Exception as e:
+        # If standard approach fails, fall back to the more robust method
+        print(f"Error in standard format_notebook: {e}. Using fallback method.")
+        return format_large_notebook(content)
+
+def format_large_notebook(content):
+    """Memory-efficient formatter for very large notebooks.
+    Processes content in chunks to avoid memory issues."""
+    # Get notebook info
+    notebook_info = extract_notebook_info(content)
+    
+    # Initialize cells with the title
+    cells = [{
+        "cell_type": "markdown",
+        "metadata": {},
+        "source": [f"# {notebook_info['name']}", "", notebook_info['description']]
+    }]
+    
+    # Process content in chunks using incremental parsing
+    # Find cell markers and their positions
+    marker_positions = []
+    for match in re.finditer(r"---\s*(MARKDOWN|CODE)\s*CELL\s*---", content):
+        marker_positions.append((match.start(), match.end(), match.group(1)))
+    
+    # If no markers are found, try to extract code blocks directly
+    if not marker_positions:
+        # Just extract code blocks and treat everything else as markdown
+        remaining_text = content
+        last_end = 0
+        
+        for match in re.finditer(r"```python\s*(.*?)```", content, re.DOTALL):
+            # If there's text before this code block, add it as markdown
+            if match.start() > last_end:
+                markdown_text = content[last_end:match.start()].strip()
+                if markdown_text:
+                    cells.append({
+                        "cell_type": "markdown",
+                        "metadata": {},
+                        "source": markdown_text.split("\n")
+                    })
+            
+            # Add the code block
+            code_text = match.group(1).strip()
+            if code_text:
+                cells.append({
+                    "cell_type": "code",
+                    "metadata": {},
+                    "source": code_text.split("\n"),
+                    "execution_count": None,
+                    "outputs": []
+                })
+            
+            last_end = match.end()
+        
+        # If there's text after the last code block, add it as markdown
+        if last_end < len(content):
+            markdown_text = content[last_end:].strip()
+            if markdown_text:
+                cells.append({
+                    "cell_type": "markdown",
+                    "metadata": {},
+                    "source": markdown_text.split("\n")
+                })
+    else:
+        # Process each cell based on its markers
+        for i, (start, end, cell_type) in enumerate(marker_positions):
+            # Find the end of this cell (start of next cell or end of content)
+            cell_end = marker_positions[i+1][0] if i+1 < len(marker_positions) else len(content)
+            cell_content = content[end:cell_end].strip()
+            
+            if cell_type == "MARKDOWN":
+                cells.append({
+                    "cell_type": "markdown",
+                    "metadata": {},
+                    "source": cell_content.split("\n")
+                })
+            elif cell_type == "CODE":
+                # Extract code from between triple backticks
+                code_match = re.search(r"```python\s*(.*?)```", cell_content, re.DOTALL)
+                if code_match:
+                    code_text = code_match.group(1).strip()
+                    cells.append({
+                        "cell_type": "code",
+                        "metadata": {},
+                        "source": code_text.split("\n"),
+                        "execution_count": None,
+                        "outputs": []
+                    })
+    
+    # Create the notebook structure
+    notebook = {
+        "cells": cells,
+        "metadata": {
+            "kernelspec": {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "python3"
+            },
+            "language_info": {
+                "name": "python",
+                "version": "3.8.0"
+            }
+        },
+        "nbformat": 4,
+        "nbformat_minor": 4
+    }
+    
+    return notebook