changes

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, jsonify, render_template, session, redirect, url_for
+from flask import Flask, request, jsonify, render_template, session, redirect, url_for, Response, Response
 from flask_session import Session
 import google.generativeai as genai
 import json
@@ -67,9 +67,15 @@ def set_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
+        
+        # Also store in global variable as a backup
+        global global_api_key
+        global_api_key = api_key
+        
         logger.info("API key successfully set and validated")
         return jsonify({"success": True})
     except Exception as e:
@@ -78,9 +84,20 @@ def set_api_key():
 
 @app.route("/generate_notebook", methods=["GET", "POST"])
 def generate_notebook_route():
-    if "api_key" not in session:
+    # First check session, then fall back to global variable
+    api_key = None
+    if "api_key" in session:
+        api_key = session.get("api_key")
+    elif global_api_key is not None:
+        api_key = global_api_key
+        # Try to restore the session
+        session["api_key"] = global_api_key
+        logger.info("Restored API key from global variable")
+    
+    if not api_key:
         logger.warning("Generate notebook request without API key")
         return jsonify({"success": False, "message": "API key not set"}), 401
+    
     # Handle both GET (for streaming) and POST requests
     if request.method == "GET":
         prompt = request.args.get("prompt")
@@ -101,7 +118,9 @@ def generate_notebook_route():
     logger.info(f"Generate notebook with model: {api_model_name}, stream: {stream}")
     
     try:
-        genai.configure(api_key=session["api_key"])
+        # Configure with the api_key we retrieved
+        genai.configure(api_key=api_key)
+        
         # OPTIMIZATION: If format_only is True, skip the AI call and just format the provided content
         if request.method == "POST" and format_only:
             # Use client-provided content as is (it's already the AI response)
@@ -116,7 +135,7 @@ def generate_notebook_route():
             })
         elif stream:
             logger.info("Starting streaming notebook generation")
-            return stream_notebook_generation(prompt, api_model_name)
+            return stream_notebook_generation(prompt, api_model_name, api_key)
         else:
             notebook_content = generate_notebook(prompt, api_model_name)
             notebook_json = format_notebook(notebook_content)
@@ -150,7 +169,17 @@ def prepare_edit_notebook():
 
 @app.route("/edit_notebook", methods=["GET", "POST"])
 def edit_notebook_route():
-    if "api_key" not in session:
+    # First check session, then fall back to global variable
+    api_key = None
+    if "api_key" in session:
+        api_key = session.get("api_key")
+    elif global_api_key is not None:
+        api_key = global_api_key
+        # Try to restore the session
+        session["api_key"] = global_api_key
+        logger.info("Restored API key from global variable for edit")
+    
+    if not api_key:
         return jsonify({"success": False, "message": "API key not set"}), 401
     
     # Get edit prompt and current notebook
@@ -177,9 +206,10 @@ def edit_notebook_route():
     api_model_name = get_api_model_name(model_name)
     
     try:
-        genai.configure(api_key=session["api_key"])
+        # Use the api_key we retrieved, not session["api_key"]
+        genai.configure(api_key=api_key)
         if stream:
-            return stream_notebook_edit(edit_prompt, notebook_json, api_model_name)
+            return stream_notebook_edit(edit_prompt, notebook_json, api_model_name, api_key)
         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)
@@ -248,6 +278,41 @@ def session_test():
         "session_vars": session_vars
     })
 
+# Add a session check endpoint to debug session issues
+@app.route("/session_check", methods=["GET"])
+def session_check():
+    if "session_test" not in session:
+        session["session_test"] = True
+        is_new = True
+    else:
+        is_new = False
+    
+    # Log the current session state
+    session_vars = list(session.keys()) if session else []
+    logger.info(f"Session check - variables: {session_vars}")
+    
+    # Check if API key is available from URL parameter (fallback for HF)
+    api_key_param = request.args.get('api_key_param')
+    if api_key_param and "api_key" not in session:
+        global global_api_key
+        try:
+            # Validate it quickly before accepting
+            genai.configure(api_key=api_key_param)
+            model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
+            # If no exception, store it
+            session["api_key"] = api_key_param
+            global_api_key = api_key_param
+            logger.info("API key set from URL parameter")
+        except Exception as e:
+            logger.error(f"Invalid API key from URL parameter: {str(e)}")
+    
+    return jsonify({
+        "session_works": True,
+        "is_new_session": is_new,
+        "has_api_key": "api_key" in session,
+        "session_vars": session_vars
+    })
+
 if __name__ == "__main__":
     debug_mode = os.environ.get("FLASK_ENV") == "development"
     port = int(os.environ.get("PORT", 5000))