Update app.py

app.py

@@ -234,6 +234,14 @@ def clear_session_data(session_id: str):
             user_sessions[session_id]["agent"] = None
             logger.info(f"Session {session_id[:8]}... cleared")
 
+def clear_agent_only(session_id: str):
+    """Clear only the agent, keeping the API key for convenience."""
+    with session_lock:
+        if session_id in user_sessions:
+            if "agent" in user_sessions[session_id]:
+                del user_sessions[session_id]["agent"]
+            logger.info(f"Session {session_id[:8]}... agent cleared")
+
 def cleanup_old_sessions():
     """Clean up sessions older than 1 hour."""
     with session_lock:
@@ -308,7 +316,7 @@ class GradioUI:
                     logger.info(f"Using HF_TOKEN from .env file for session {session_id[:8]}...")
                 else:
                     logger.warning(f"No API key found for session {session_id[:8]}...")
-                    error_msg = "❌ No API key provided. Please enter your Hugging Face API key in the API Configuration section above or set HF_TOKEN in your .env file."
+                    error_msg = "❌ No API key configured for your session. Please enter your Hugging Face API key in the API Configuration section above and click 'Setup API Key'."
                     messages.append(gr.ChatMessage(role="assistant", content=error_msg))
                     yield messages
                     return
@@ -419,8 +427,10 @@ User Query: """
                 yield messages
             
             # Clear sensitive data from session after interaction (AUTOMATIC)
-            clear_session_data(session_id)
-            logger.info(f"Session {session_id[:8]}... automatically cleared after interaction")
+            # Note: We clear the agent but keep the API key for convenience
+            if "agent" in session_data:
+                del session_data["agent"]
+            logger.info(f"Session {session_id[:8]}... agent cleared after interaction")
             
             yield messages
         except Exception as e:
@@ -666,7 +676,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             # API Key Configuration Section
                             with gr.Accordion("🔑 API Configuration", open=False):
                                 gr.Markdown("**Configure your Hugging Face API Key**")
-                                gr.Markdown("🔒 **Security**: Each user has an independent session. API keys are automatically cleared after each interaction for your privacy.")
+                                gr.Markdown("🔒 **Security**: Each user has an independent session. Your API key is stored securely and automatically cleared after 1 hour of inactivity.")
                                 gr.Markdown("Get your API key from: https://huggingface.co/settings/tokens")
                                 
                                 api_key_input = gr.Textbox(
@@ -677,7 +687,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                                 )
                                 api_key_status = gr.Textbox(
                                     label="Status",
-                                    value="✅ HF_TOKEN found in .env file. To use a different key, enter it above and click 'Setup API Key'." if os.getenv("HF_TOKEN") else "⚠️ Please enter your Hugging Face API key or set HF_TOKEN in your .env file.",
+                                    value="✅ HF_TOKEN found in .env file. To use a different key, enter it above and click 'Setup API Key'." if os.getenv("HF_TOKEN") else "⚠️ Please enter your Hugging Face API key above and click 'Setup API Key' to start using the application.",
                                     interactive=False
                                 )
                                 
@@ -814,7 +824,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                         # API Key Configuration Section for Mobile
                         with gr.Accordion("🔑 API Configuration", open=False):
                             gr.Markdown("**Configure your Hugging Face API Key**")
-                            gr.Markdown("🔒 **Security**: Each user has an independent session. API keys are automatically cleared after each interaction for your privacy.")
+                            gr.Markdown("🔒 **Security**: Each user has an independent session. Your API key is stored securely and automatically cleared after 1 hour of inactivity.")
                             gr.Markdown("Get your API key from: https://huggingface.co/settings/tokens")
                             
                             mobile_api_key_input = gr.Textbox(
@@ -825,7 +835,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             )
                             mobile_api_key_status = gr.Textbox(
                                 label="Status",
-                                value="✅ HF_TOKEN found in .env file. To use a different key, enter it above and click 'Setup API Key'." if os.getenv("HF_TOKEN") else "⚠️ Please enter your Hugging Face API key or set HF_TOKEN in your .env file.",
+                                value="✅ HF_TOKEN found in .env file. To use a different key, enter it above and click 'Setup API Key'." if os.getenv("HF_TOKEN") else "⚠️ Please enter your Hugging Face API key above and click 'Setup API Key' to start using the application.",
                                 interactive=False
                             )