Update app.py

app.py

@@ -95,32 +95,12 @@ append_answer_lock = threading.Lock()
 # Initialize browser
 browser = SimpleTextBrowser(request_kwargs={})
 
-# Initialize tools
-ti_tool = TextInspectorTool(None, 20000)  # Will be updated with session-specific model
-
 # Initialize vulnerability tools
 cvedb_tool = CVEDBTool()
 epss_tool = EpsTool()
 nvd_tool = NvdTool()
 kevin_tool = KevinTool()
 
-# Default tools (will be updated with session-specific models)
-WEB_TOOLS = [
-    web_search,  # duckduckgo
-    VisitTool(browser),
-    PageUpTool(browser),
-    PageDownTool(browser),
-    FinderTool(browser),
-    FindNextTool(browser),
-    ArchiveSearchTool(browser),
-] + ([ti_tool] if ti_tool else []) + [
-    cvedb_tool,  # CVEDB Tool
-    # report_generator,  # Report generation tool - COMMENTED: Only works locally
-    epss_tool,  # EPSS Tool
-    nvd_tool,  # NVD Tool
-    kevin_tool,  # KEVin Tool
-]
-
 def validate_hf_api_key(api_key: str) -> tuple[bool, str]:
     """Validate Hugging Face API key by making a test request."""
     if not api_key or not api_key.strip():
@@ -168,10 +148,19 @@ def create_model_with_api_key(hf_token: str, model_id: str = None) -> InferenceC
 
 def create_tools_with_model(model: InferenceClientModel):
     """Create tools with the provided model."""
+    # Verify the model was created correctly
+    if model is None:
+        raise ValueError("Model is None, cannot create TextInspectorTool")
+    
     # Update text inspector tool with the model
+    # 20000 = maximum characters to read from files (text_limit)
     ti_tool = TextInspectorTool(model, 20000)
     
-    return [
+    # Verify the tool was created correctly
+    if ti_tool is None:
+        raise ValueError("Failed to create TextInspectorTool")
+    
+    tools = [
         web_search,  # duckduckgo
         VisitTool(browser),
         PageUpTool(browser),
@@ -179,13 +168,15 @@ def create_tools_with_model(model: InferenceClientModel):
         FinderTool(browser),
         FindNextTool(browser),
         ArchiveSearchTool(browser),
-    ] + ([ti_tool] if ti_tool else []) + [
+        ti_tool,  # TextInspectorTool - always available
         cvedb_tool,  # CVEDB Tool
         # report_generator,  # Report generation tool - COMMENTED: Only works locally
         epss_tool,  # EPSS Tool
         nvd_tool,  # NVD Tool
         kevin_tool,  # KEVin Tool
     ]
+    
+    return tools
 
 # Agent creation in a factory function
 def create_agent(hf_token: str = None, model_id: str = None, max_steps: int = 10):
@@ -199,6 +190,11 @@ def create_agent(hf_token: str = None, model_id: str = None, max_steps: int = 10
     model = create_model_with_api_key(hf_token, model_id)
     tools = create_tools_with_model(model)
     
+    # Verify that TextInspectorTool is in the tools list
+    has_text_inspector = any(hasattr(tool, 'name') and tool.name == 'inspect_file_as_text' for tool in tools)
+    if not has_text_inspector:
+        raise ValueError("TextInspectorTool not found in tools list")
+    
     agent = CodeAgent(
         model=model,
         tools=[visualizer] + tools,
@@ -211,9 +207,6 @@ def create_agent(hf_token: str = None, model_id: str = None, max_steps: int = 10
     logger.info("Agent created successfully")
     return agent
 
-# 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:
@@ -486,7 +479,12 @@ User Query: """
             # Combine system prompt with user message
             full_prompt = system_prompt + prompt
 
-            messages.append(gr.ChatMessage(role="user", content=prompt))
+            # Extract clean message for display (remove internal context)
+            display_message = prompt
+            if "[INTERNAL CONTEXT:" in prompt:
+                display_message = prompt.split("[INTERNAL CONTEXT:")[0].strip()
+
+            messages.append(gr.ChatMessage(role="user", content=display_message))
             yield messages
 
             logger.info(f"Starting agent interaction for session {session_id[:8]}...")
@@ -614,13 +612,18 @@ User Query: """
         ), file_uploads_log + [file_path]
 
     def log_user_message(self, text_input, file_uploads_log):
+        # Create the user message for display (clean, without file info)
+        display_message = text_input
+        
+        # Create the internal message for the agent (with file context)
+        internal_message = text_input
+        if len(file_uploads_log) > 0:
+            file_names = [os.path.basename(f) for f in file_uploads_log]
+            file_paths = [f for f in file_uploads_log]  # Full paths
+            internal_message += f"\n\n[Uploaded files available: {', '.join(file_names)}. Use inspect_file_as_text(file_path='uploads/filename') to analyze them. Use the content as plain text context if readable, no need to parse or create complex functions.]"
+        
         return (
-            text_input
-            + (
-                f"\nYou have been provided with these files, which might be helpful or not: {file_uploads_log}"
-                if len(file_uploads_log) > 0
-                else ""
-            ),
+            internal_message,  # This goes to the agent (with file context)
             gr.Textbox(
                 value="",
                 interactive=False,
@@ -821,18 +824,20 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                                 setup_api_btn = gr.Button("Setup API Key", variant="secondary")
 
                             # If an upload folder is provided, enable the upload feature
-                            if self.file_upload_folder is not None:
-                                upload_file = gr.File(label="Upload a file")
-                                upload_status = gr.Textbox(
-                                    label="Upload Status",
-                                    interactive=False,
-                                    visible=False,
-                                )
-                                upload_file.change(
-                                    self.upload_file,
-                                    [upload_file, file_uploads_log],
-                                    [upload_status, file_uploads_log],
-                                )
+                            # COMMENTED: File upload feature temporarily disabled - works but consumes too many steps for parsing
+                            # TODO: Re-enable after optimizing TextInspectorTool to use fewer steps
+                            # if self.file_upload_folder is not None:
+                            #     upload_file = gr.File(label="Upload a file")
+                            #     upload_status = gr.Textbox(
+                            #         label="Upload Status",
+                            #         interactive=False,
+                            #         visible=False,
+                            #     )
+                            #     upload_file.change(
+                            #         self.upload_file,
+                            #         [upload_file, file_uploads_log],
+                            #         [upload_status, file_uploads_log],
+                            #     )
 
                             # Powered by smolagents
                             with gr.Row():
@@ -1000,6 +1005,22 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             variant="primary",
                         )
 
+                        # File upload section for mobile (simple)
+                        # COMMENTED: File upload feature temporarily disabled - works but consumes too many steps for parsing
+                        # TODO: Re-enable after optimizing
+                        # if self.file_upload_folder is not None:
+                        #     mobile_upload_file = gr.File(label="📎 Upload PDF/TXT file (optional)")
+                        #     mobile_upload_status = gr.Textbox(
+                        #         label="Upload Status",
+                        #         interactive=False,
+                        #         visible=False,
+                        #     )
+                        #     mobile_upload_file.change(
+                        #         self.upload_file,
+                        #         [mobile_upload_file, file_uploads_log],
+                        #         [mobile_upload_status, file_uploads_log],
+                        #     )
+
                         # Examples Section for Mobile
                         with gr.Accordion("💡 Example Prompts", open=False):
                             gr.Markdown("**Click any example below to populate your request field:**")
@@ -1129,7 +1150,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
         
         demo.launch(
             debug=True,
-            server_name="0.0.0.0",
+            server_name="localhost",
             server_port=7860,
             share=False,
             **kwargs
@@ -1137,6 +1158,6 @@ This AI agent specializes in automated vulnerability research and analysis, buil
 
 # can this fix ctrl-c no response? no
 try:
-    GradioUI().launch(mcp_server=True)
+    GradioUI(file_upload_folder="uploads").launch()
 except KeyboardInterrupt:
     ...