Update app.py

app.py

@@ -4,6 +4,7 @@ import os
 import re
 import shutil
 import threading
+import uuid
 from typing import Optional
 from loguru import logger
 from datetime import datetime
@@ -85,65 +86,20 @@ if os.getenv("HF_TOKEN"):
         logger.warning(f"Failed to login with HF_TOKEN from environment: {e}")
         logger.info("You can still use the application by providing a valid API key in the interface")
 
-append_answer_lock = threading.Lock()
-
-custom_role_conversions = {"tool-call": "assistant", "tool-response": "user"}
-
-user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
-
-BROWSER_CONFIG = {
-    "viewport_size": 1024 * 5,
-    "downloads_folder": "downloads_folder",
-    "request_kwargs": {
-        "headers": {"User-Agent": user_agent},
-        "timeout": 300,
-    },
-    "serpapi_key": os.getenv("SERPAPI_API_KEY"),
-}
-
-os.makedirs(f"./{BROWSER_CONFIG['downloads_folder']}", exist_ok=True)
+# Global session storage for independent user sessions
+user_sessions = {}
+session_lock = threading.Lock()
 
-# Default Hugging Face model configuration
-model_id = os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct")
-logger.info(f"Default Hugging Face model: {model_id}")
+append_answer_lock = threading.Lock()
 
-# Define text_limit before using it
-text_limit = 20000
+# Initialize browser
+browser = SimpleTextBrowser()
 
-# Default model (will be overridden by session-specific models)
-# Note: This model may not work without a valid API key
-default_model = None
-ti_tool = None
+# Initialize tools
+ti_tool = TextInspectorTool(None, 20000)  # Will be updated with session-specific model
 
-# Only try to create default model if we have a valid token
-if os.getenv("HF_TOKEN"):
-    try:
-        # Test if the token is valid
-        api = HfApi(token=os.getenv("HF_TOKEN"))
-        api.whoami()  # This will raise an exception if token is invalid
-        # If we get here, token is valid
-        default_model = InferenceClientModel(
-            model_id,
-            custom_role_conversions={
-                "tool-call": "assistant",
-                "tool-response": "user"
-            },
-            token=os.getenv("HF_TOKEN")
-        )
-        ti_tool = TextInspectorTool(default_model, text_limit)
-        logger.info("Default model created successfully with valid token")
-    except Exception as e:
-        logger.warning(f"Failed to create default model: {e}")
-        default_model = None
-        ti_tool = None
-else:
-    logger.info("No HF_TOKEN provided, default model will be created when user provides API key")
-
-browser = SimpleTextBrowser(**BROWSER_CONFIG)
-
-# Tool configuration
+# Initialize vulnerability tools
 cvedb_tool = CVEDBTool()
-report_generator = ReportGeneratorTool()
 epss_tool = EpsTool()
 nvd_tool = NvdTool()
 kevin_tool = KevinTool()
@@ -186,50 +142,22 @@ def validate_hf_api_key(api_key: str) -> tuple[bool, str]:
         return False, f"❌ Invalid API key: {str(e)}"
 
 def create_model_with_api_key(hf_token: str, model_id: str = None) -> InferenceClientModel:
-    """Create a new InferenceClientModel instance with the provided HF_TOKEN."""
+    """Create a model instance with the provided API key."""
     if not model_id:
-        model_id = os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct")
-    
-    logger.info(f"Creating model {model_id} with token: {hf_token[:10]}...")
+        model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
     
-    # First, try to login with the token to ensure it's properly set
-    try:
-        login(hf_token)
-        logger.info("Successfully logged in with token")
-    except Exception as e:
-        logger.warning(f"Login failed: {e}")
-    
-    # Create the model with explicit token
-    model = InferenceClientModel(
-        model_id,
-        custom_role_conversions={
-            "tool-call": "assistant",
-            "tool-response": "user"
-        },
-        token=hf_token
+    return InferenceClientModel(
+        model_id=model_id,
+        hf_token=hf_token,
+        temperature=0.1,
+        max_tokens=4000,
     )
-    
-    # Verify the token is set correctly
-    if hasattr(model, 'token'):
-        logger.info(f"Model token attribute: {model.token[:10] if model.token else 'None'}...")
-    else:
-        logger.warning("Model does not have token attribute")
-    
-    # Test the model with a simple request to verify token works
-    try:
-        logger.info("Testing model with simple request...")
-        # This is a simple test to see if the model can be accessed
-        test_response = model.generate("Hello", max_new_tokens=5)
-        logger.info("Model test successful")
-    except Exception as e:
-        logger.error(f"Model test failed: {e}")
-        # Don't raise the exception, just log it
-    
-    logger.info(f"Model created successfully with token")
-    return model
 
-def create_tools_with_model(model: InferenceClientModel) -> list:
-    """Create tools list with the provided model."""
+def create_tools_with_model(model: InferenceClientModel):
+    """Create tools with the provided model."""
+    # Update text inspector tool with the model
+    ti_tool = TextInspectorTool(model, 20000)
+    
     return [
         web_search,  # duckduckgo
         VisitTool(browser),
@@ -238,7 +166,7 @@ def create_tools_with_model(model: InferenceClientModel) -> list:
         FinderTool(browser),
         FindNextTool(browser),
         ArchiveSearchTool(browser),
-        TextInspectorTool(model, text_limit),
+    ] + ([ti_tool] if ti_tool else []) + [
         cvedb_tool,  # CVEDB Tool
         # report_generator,  # Report generation tool - COMMENTED: Only works locally
         epss_tool,  # EPSS Tool
@@ -270,11 +198,60 @@ def create_agent(hf_token: str = None, model_id: str = None, max_steps: int = 10
     logger.info("Agent created successfully")
     return agent
 
-# Only create document_inspection_tool if default_model is available
-if default_model:
-    document_inspection_tool = TextInspectorTool(default_model, 20000)
-else:
-    document_inspection_tool = None
+# Document inspection tool will be created per session
+document_inspection_tool = None
+
+def get_user_session(request: gr.Request) -> str:
+    """Get or create a unique session ID for the user."""
+    if not request:
+        return str(uuid.uuid4())
+    
+    # Try to get session from headers or create new one
+    session_id = request.headers.get("x-session-id")
+    if not session_id:
+        session_id = str(uuid.uuid4())
+    
+    return session_id
+
+def get_session_data(session_id: str) -> dict:
+    """Get session data for a specific user."""
+    with session_lock:
+        if session_id not in user_sessions:
+            user_sessions[session_id] = {
+                "hf_token": None,
+                "agent": None,
+                "max_steps": 10,
+                "created_at": datetime.now()
+            }
+        return user_sessions[session_id]
+
+def clear_session_data(session_id: str):
+    """Clear session data for a specific user."""
+    with session_lock:
+        if session_id in user_sessions:
+            # Clear sensitive data
+            user_sessions[session_id]["hf_token"] = None
+            user_sessions[session_id]["agent"] = None
+            logger.info(f"Session {session_id[:8]}... cleared")
+
+def cleanup_old_sessions():
+    """Clean up sessions older than 1 hour."""
+    with session_lock:
+        current_time = datetime.now()
+        sessions_to_remove = []
+        
+        for session_id, session_data in user_sessions.items():
+            if session_data.get("created_at"):
+                time_diff = current_time - session_data["created_at"]
+                if time_diff.total_seconds() > 3600:  # 1 hour
+                    sessions_to_remove.append(session_id)
+        
+        for session_id in sessions_to_remove:
+            del user_sessions[session_id]
+            logger.info(f"Removed old session {session_id[:8]}...")
+        
+        if sessions_to_remove:
+            logger.info(f"Cleaned up {len(sessions_to_remove)} old sessions")
 
 class GradioUI:
     """A one-line interface to launch your agent in Gradio"""
@@ -308,59 +285,60 @@ class GradioUI:
         else:
             return message, "error"
 
-    def interact_with_agent(self, prompt, messages, session_state):
-        # Clear any stale session data at the beginning
-        if "hf_token" in session_state and not session_state.get("hf_token"):
-            del session_state["hf_token"]
-        if "agent" in session_state and not session_state.get("agent"):
-            del session_state["agent"]
+    def interact_with_agent(self, prompt, messages, request: gr.Request):
+        """Handle agent interaction with proper session management."""
+        # Get unique session ID for this user
+        session_id = get_user_session(request)
+        session_data = get_session_data(session_id)
+        
+        logger.info(f"Processing request for session {session_id[:8]}...")
         
-        # Get or create session-specific agent
-        if "agent" not in session_state:
+        # Check if we have a valid agent for this session
+        if not session_data.get("agent"):
             # Check if we have a valid HF_TOKEN in session
-            hf_token = session_state.get("hf_token")
+            hf_token = session_data.get("hf_token")
             
             # If no token in session, try to get it from .env file
             if not hf_token:
                 env_token = os.getenv("HF_TOKEN")
                 if env_token:
                     hf_token = env_token
-                    session_state["hf_token"] = env_token
-                    session_state["max_steps"] = 10  # Default max_steps
-                    logger.info(f"Using HF_TOKEN from .env file: {env_token[:10]}...")
+                    session_data["hf_token"] = env_token
+                    session_data["max_steps"] = 10  # Default max_steps
+                    logger.info(f"Using HF_TOKEN from .env file for session {session_id[:8]}...")
                 else:
-                    logger.warning("No API key found in session state or .env file")
+                    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."
                     messages.append(gr.ChatMessage(role="assistant", content=error_msg))
                     yield messages
                     return
             
-            logger.info(f"Agent not in session, checking for token: {hf_token[:10] if hf_token else 'None'}...")
+            logger.info(f"Creating agent for session {session_id[:8]}...")
             
             if hf_token:
                 try:
-                    max_steps = session_state.get("max_steps", 10)
-                    session_state["agent"] = create_agent(hf_token, max_steps=max_steps)
-                    logger.info("Agent created successfully in interact_with_agent")
+                    max_steps = session_data.get("max_steps", 10)
+                    session_data["agent"] = create_agent(hf_token, max_steps=max_steps)
+                    logger.info(f"Agent created successfully for session {session_id[:8]}...")
                 except Exception as e:
-                    logger.error(f"Failed to create agent in interact_with_agent: {e}")
+                    logger.error(f"Failed to create agent for session {session_id[:8]}: {e}")
                     error_msg = f"❌ Failed to create agent with provided API key: {str(e)}"
                     messages.append(gr.ChatMessage(role="assistant", content=error_msg))
                     yield messages
                     return
         else:
-            logger.info("Agent already exists in session")
+            logger.info(f"Agent already exists for session {session_id[:8]}...")
         
         # Adding monitoring
         try:
             # log the existence of agent memory
-            has_memory = hasattr(session_state["agent"], "memory")
+            has_memory = hasattr(session_data["agent"], "memory")
             print(f"Agent has memory: {has_memory}")
             if has_memory:
-                print(f"Memory type: {type(session_state['agent'].memory)}")
+                print(f"Memory type: {type(session_data['agent'].memory)}")
 
             # Prepare the system prompt
-            system_prompt = f"""You are a Vulnerability Intelligence Analyst. Complete the user request in {session_state.get('max_steps', 10)} steps maximum.
+            system_prompt = f"""You are a Vulnerability Intelligence Analyst. Complete the user request in {session_data.get('max_steps', 10)} steps maximum.
 
 AVAILABLE TOOLS: nvd_search, web_search, cvedb_search, kevin_search, epss_search
 
@@ -428,9 +406,9 @@ User Query: """
             messages.append(gr.ChatMessage(role="user", content=prompt))
             yield messages
 
-            logger.info("Starting agent interaction...")
+            logger.info(f"Starting agent interaction for session {session_id[:8]}...")
             for msg in stream_to_gradio(
-                session_state["agent"], task=full_prompt, reset_agent_memory=False
+                session_data["agent"], task=full_prompt, reset_agent_memory=False
             ):
                 # If the message contains an HTML report, we save it and update the message
                 if isinstance(msg.content, str) and msg.content.startswith("<!DOCTYPE html>"):
@@ -441,24 +419,25 @@ User Query: """
                 yield messages
             
             # Clear sensitive data from session after interaction (AUTOMATIC)
-            if "hf_token" in session_state:
-                del session_state["hf_token"]
-            if "agent" in session_state:
-                del session_state["agent"]
-            if "HF_TOKEN" in os.environ:
-                del os.environ["HF_TOKEN"]
-            logger.info("Session automatically cleared after interaction")
+            clear_session_data(session_id)
+            logger.info(f"Session {session_id[:8]}... automatically cleared after interaction")
             
             yield messages
         except Exception as e:
-            logger.error(f"Error in interaction: {str(e)}")
+            logger.error(f"Error in interaction for session {session_id[:8]}: {str(e)}")
             print(f"Error in interaction: {str(e)}")
             error_msg = f"❌ Error during interaction: {str(e)}"
             messages.append(gr.ChatMessage(role="assistant", content=error_msg))
             yield messages
 
-    def setup_api_key(self, api_key: str, max_steps: int, session_state) -> str:
-        """Setup API key for the session."""
+    def setup_api_key(self, api_key: str, max_steps: int, request: gr.Request) -> str:
+        """Setup API key for the user's session."""
+        # Get unique session ID for this user
+        session_id = get_user_session(request)
+        session_data = get_session_data(session_id)
+        
+        logger.info(f"Setting up API key for session {session_id[:8]}...")
+        
         # Check if API key is provided from interface
         if api_key and api_key.strip():
             # Use the API key from interface
@@ -477,30 +456,24 @@ User Query: """
         is_valid, message = validate_hf_api_key(token_to_use)
         
         if is_valid:
-            # Store HF_TOKEN in session state (but will be cleared after use)
-            session_state["hf_token"] = token_to_use
-            session_state["max_steps"] = max_steps
-            logger.info(f"API key stored in session from {source}: {token_to_use[:10]}...")
+            # Store HF_TOKEN in session data
+            session_data["hf_token"] = token_to_use
+            session_data["max_steps"] = max_steps
+            logger.info(f"API key stored in session {session_id[:8]}... from {source}")
             logger.info(f"Max steps set to: {max_steps}")
             
-            # Also set the environment variable for smolagents
-            os.environ["HF_TOKEN"] = token_to_use
-            logger.info("HF_TOKEN environment variable set")
-            
             # Create new agent with the HF_TOKEN and max_steps
             try:
-                session_state["agent"] = create_agent(token_to_use, max_steps=max_steps)
-                logger.info("Agent created successfully in setup_api_key")
+                session_data["agent"] = create_agent(token_to_use, max_steps=max_steps)
+                logger.info(f"Agent created successfully for session {session_id[:8]}...")
                 return f"✅ API key from {source} validated and agent created successfully! {message.split('!')[1] if '!' in message else ''}"
             except Exception as e:
-                logger.error(f"Failed to create agent in setup_api_key: {e}")
+                logger.error(f"Failed to create agent for session {session_id[:8]}: {e}")
                 return f"❌ Failed to create agent with API key from {source}: {str(e)}"
         else:
-            logger.warning(f"Invalid API key from {source}: {token_to_use[:10] if token_to_use else 'None'}...")
+            logger.warning(f"Invalid API key for session {session_id[:8]}... from {source}")
             return f"❌ Invalid API key from {source}: {message}"
 
-
-
     def upload_file(
         self,
         file,
@@ -693,7 +666,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**: API keys are automatically cleared after each interaction for your privacy.")
+                                gr.Markdown("🔒 **Security**: Each user has an independent session. API keys are automatically cleared after each interaction for your privacy.")
                                 gr.Markdown("Get your API key from: https://huggingface.co/settings/tokens")
                                 
                                 api_key_input = gr.Textbox(
@@ -742,10 +715,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                         <a target="_blank" href="https://github.com/huggingface/smolagents"><b>hf/smolagents</b></a>
                         </div>""")
 
-                        # Add session state to store session-specific data
-                        session_state = gr.State(
-                            {}
-                        )  # Initialize empty state for each session
+                        # Chat interface
                         stored_messages = gr.State([])
                         chatbot = gr.Chatbot(
                             label="open-Deep-Research",
@@ -765,7 +735,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                         # API Key setup event
                         setup_api_btn.click(
                             self.setup_api_key,
-                            [api_key_input, max_steps_slider, session_state],
+                            [api_key_input, max_steps_slider],
                             [api_key_status]
                         )
                         
@@ -777,8 +747,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             [stored_messages, text_input, launch_research_btn],
                         ).then(
                             self.interact_with_agent,
-                            # Include session_state in function calls
-                            [stored_messages, chatbot, session_state],
+                            [stored_messages, chatbot],
                             [chatbot],
                         ).then(
                             lambda: (
@@ -797,8 +766,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             [stored_messages, text_input, launch_research_btn],
                         ).then(
                             self.interact_with_agent,
-                            # Include session_state in function calls
-                            [stored_messages, chatbot, session_state],
+                            [stored_messages, chatbot],
                             [chatbot],
                         ).then(
                             lambda: (
@@ -846,7 +814,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**: API keys are automatically cleared after each interaction for your privacy.")
+                            gr.Markdown("🔒 **Security**: Each user has an independent session. API keys are automatically cleared after each interaction for your privacy.")
                             gr.Markdown("Get your API key from: https://huggingface.co/settings/tokens")
                             
                             mobile_api_key_input = gr.Textbox(
@@ -874,10 +842,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             
                             mobile_setup_api_btn = gr.Button("Setup API Key", variant="secondary")
                         
-                        # Add session state to store session-specific data
-                        session_state = gr.State(
-                            {}
-                        )  # Initialize empty state for each session
+                        # Chat interface for mobile
                         stored_messages = gr.State([])
                         file_uploads_log = gr.State([])
                         chatbot = gr.Chatbot(
@@ -894,7 +859,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                         # Mobile API Key setup event
                         mobile_setup_api_btn.click(
                             self.setup_api_key,
-                            [mobile_api_key_input, mobile_max_steps_slider, session_state],
+                            [mobile_api_key_input, mobile_max_steps_slider],
                             [mobile_api_key_status]
                         )
                         
@@ -977,8 +942,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             [stored_messages, text_input, launch_research_btn],
                         ).then(
                             self.interact_with_agent,
-                            # Include session_state in function calls
-                            [stored_messages, chatbot, session_state],
+                            [stored_messages, chatbot],
                             [chatbot],
                         ).then(
                             lambda: (
@@ -997,8 +961,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             [stored_messages, text_input, launch_research_btn],
                         ).then(
                             self.interact_with_agent,
-                            # Include session_state in function calls
-                            [stored_messages, chatbot, session_state],
+                            [stored_messages, chatbot],
                             [chatbot],
                         ).then(
                             lambda: (
@@ -1012,6 +975,20 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             [text_input, launch_research_btn],
                         )
 
+        # Start periodic cleanup of old sessions
+        def periodic_cleanup():
+            import time
+            while True:
+                try:
+                    cleanup_old_sessions()
+                    time.sleep(300)  # Run every 5 minutes
+                except Exception as e:
+                    logger.error(f"Error in periodic cleanup: {e}")
+                    time.sleep(300)
+        
+        cleanup_thread = threading.Thread(target=periodic_cleanup, daemon=True)
+        cleanup_thread.start()
+        
         demo.launch(debug=True, **kwargs)
 
 # can this fix ctrl-c no response? no