Merge pull request #460 from vvincent1234/fix/security

optmize webui settings and fix vulnerability

Author: warmshao

src/utils/agent_state.py

@@ -1,5 +1,6 @@
 import asyncio
 
+
 class AgentState:
     _instance = None
 
@@ -27,4 +28,4 @@ def set_last_valid_state(self, state):
         self.last_valid_state = state
 
     def get_last_valid_state(self):
-        return self.last_valid_state
+        return self.last_valid_state

src/utils/deep_research.py

@@ -19,7 +19,13 @@
 from browser_use.browser.context import BrowserContext
 from browser_use.controller.service import Controller, DoneAction
 from main_content_extractor import MainContentExtractor
-from langchain.schema import SystemMessage, HumanMessage
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    ToolMessage,
+    SystemMessage
+)
 from json_repair import repair_json
 from src.agent.custom_prompts import CustomSystemPrompt, CustomAgentMessagePrompt
 from src.controller.custom_controller import CustomController

src/utils/default_config_settings.py

@@ -4,6 +4,9 @@
 from pathlib import Path
 from typing import Dict, Optional
 import requests
+import json
+import gradio as gr
+import uuid
 
 from langchain_anthropic import ChatAnthropic
 from langchain_mistralai import ChatMistralAI
@@ -196,12 +199,14 @@ def update_model_dropdown(llm_provider, api_key=None, base_url=None):
     else:
         return gr.Dropdown(choices=[], value="", interactive=True, allow_custom_value=True)
 
+
 class MissingAPIKeyError(Exception):
     """Custom exception for missing API key."""
+
     def __init__(self, provider: str, env_var: str):
         provider_display = PROVIDER_DISPLAY_NAMES.get(provider, provider.upper())
         super().__init__(f"💥 {provider_display} API key not found! 🔑 Please set the "
-                        f"`{env_var}` environment variable or provide it in the UI.")
+                         f"`{env_var}` environment variable or provide it in the UI.")
 
 
 def encode_image(img_path):
@@ -270,3 +275,70 @@ async def capture_screenshot(browser_context):
         return encoded
     except Exception as e:
         return None
+
+
+class ConfigManager:
+    def __init__(self):
+        self.components = {}
+        self.component_order = []
+
+    def register_component(self, name: str, component):
+        """Register a gradio component for config management."""
+        self.components[name] = component
+        if name not in self.component_order:
+            self.component_order.append(name)
+        return component
+
+    def save_current_config(self):
+        """Save the current configuration of all registered components."""
+        current_config = {}
+        for name in self.component_order:
+            component = self.components[name]
+            # Get the current value from the component
+            current_config[name] = getattr(component, "value", None)
+
+        return save_config_to_file(current_config)
+
+    def update_ui_from_config(self, config_file):
+        """Update UI components from a loaded configuration file."""
+        if config_file is None:
+            return [gr.update() for _ in self.component_order] + ["No file selected."]
+
+        loaded_config = load_config_from_file(config_file.name)
+
+        if not isinstance(loaded_config, dict):
+            return [gr.update() for _ in self.component_order] + ["Error: Invalid configuration file."]
+
+        # Prepare updates for all components
+        updates = []
+        for name in self.component_order:
+            if name in loaded_config:
+                updates.append(gr.update(value=loaded_config[name]))
+            else:
+                updates.append(gr.update())
+
+        updates.append("Configuration loaded successfully.")
+        return updates
+
+    def get_all_components(self):
+        """Return all registered components in the order they were registered."""
+        return [self.components[name] for name in self.component_order]
+
+
+def load_config_from_file(config_file):
+    """Load settings from a config file (JSON format)."""
+    try:
+        with open(config_file, 'r') as f:
+            settings = json.load(f)
+        return settings
+    except Exception as e:
+        return f"Error loading configuration: {str(e)}"
+
+
+def save_config_to_file(settings, save_dir="./tmp/webui_settings"):
+    """Save the current settings to a UUID.json file with a UUID name."""
+    os.makedirs(save_dir, exist_ok=True)
+    config_file = os.path.join(save_dir, f"{uuid.uuid4()}.json")
+    with open(config_file, 'w') as f:
+        json.dump(settings, f, indent=2)
+    return f"Configuration saved to {config_file}"