Spaces:
Running
Running
File size: 4,732 Bytes
19ea30a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import json
from collections.abc import Generator
from typing import TYPE_CHECKING
import os
import gradio as gr
from datetime import datetime
from typing import Optional, Dict, List
import uuid
import asyncio
import time
from gradio.components import Component
from browser_use.browser.browser import Browser
from browser_use.browser.context import BrowserContext
from browser_use.agent.service import Agent
from src.browser.custom_browser import CustomBrowser
from src.browser.custom_context import CustomBrowserContext
from src.controller.custom_controller import CustomController
from src.agent.deep_research.deep_research_agent import DeepResearchAgent
class WebuiManager:
def __init__(self, settings_save_dir: str = "./tmp/webui_settings"):
self.id_to_component: dict[str, Component] = {}
self.component_to_id: dict[Component, str] = {}
self.settings_save_dir = settings_save_dir
os.makedirs(self.settings_save_dir, exist_ok=True)
def init_browser_use_agent(self) -> None:
"""
init browser use agent
"""
self.bu_agent: Optional[Agent] = None
self.bu_browser: Optional[CustomBrowser] = None
self.bu_browser_context: Optional[CustomBrowserContext] = None
self.bu_controller: Optional[CustomController] = None
self.bu_chat_history: List[Dict[str, Optional[str]]] = []
self.bu_response_event: Optional[asyncio.Event] = None
self.bu_user_help_response: Optional[str] = None
self.bu_current_task: Optional[asyncio.Task] = None
self.bu_agent_task_id: Optional[str] = None
def init_deep_research_agent(self) -> None:
"""
init deep research agent
"""
self.dr_agent: Optional[DeepResearchAgent] = None
self.dr_current_task = None
self.dr_agent_task_id: Optional[str] = None
self.dr_save_dir: Optional[str] = None
def add_components(self, tab_name: str, components_dict: dict[str, "Component"]) -> None:
"""
Add tab components
"""
for comp_name, component in components_dict.items():
comp_id = f"{tab_name}.{comp_name}"
self.id_to_component[comp_id] = component
self.component_to_id[component] = comp_id
def get_components(self) -> list["Component"]:
"""
Get all components
"""
return list(self.id_to_component.values())
def get_component_by_id(self, comp_id: str) -> "Component":
"""
Get component by id
"""
return self.id_to_component[comp_id]
def get_id_by_component(self, comp: "Component") -> str:
"""
Get id by component
"""
return self.component_to_id[comp]
def save_config(self, components: Dict["Component", str]) -> None:
"""
Save config
"""
cur_settings = {}
for comp in components:
if not isinstance(comp, gr.Button) and not isinstance(comp, gr.File) and str(
getattr(comp, "interactive", True)).lower() != "false":
comp_id = self.get_id_by_component(comp)
cur_settings[comp_id] = components[comp]
config_name = datetime.now().strftime("%Y%m%d-%H%M%S")
with open(os.path.join(self.settings_save_dir, f"{config_name}.json"), "w") as fw:
json.dump(cur_settings, fw, indent=4)
return os.path.join(self.settings_save_dir, f"{config_name}.json")
def load_config(self, config_path: str):
"""
Load config
"""
with open(config_path, "r") as fr:
ui_settings = json.load(fr)
update_components = {}
for comp_id, comp_val in ui_settings.items():
if comp_id in self.id_to_component:
comp = self.id_to_component[comp_id]
if comp.__class__.__name__ == "Chatbot":
update_components[comp] = comp.__class__(value=comp_val, type="messages")
else:
update_components[comp] = comp.__class__(value=comp_val)
if comp_id == "agent_settings.planner_llm_provider":
yield update_components # yield provider, let callback run
time.sleep(0.1) # wait for Gradio UI callback
config_status = self.id_to_component["load_save_config.config_status"]
update_components.update(
{
config_status: config_status.__class__(value=f"Successfully loaded config: {config_path}")
}
)
yield update_components
|