|
|
import pdb |
|
|
|
|
|
import pyperclip |
|
|
from typing import Optional, Type, Callable, Dict, Any, Union, Awaitable, TypeVar |
|
|
from pydantic import BaseModel |
|
|
from browser_use.agent.views import ActionResult |
|
|
from browser_use.browser.context import BrowserContext |
|
|
from browser_use.controller.service import Controller, DoneAction |
|
|
from browser_use.controller.registry.service import Registry, RegisteredAction |
|
|
from main_content_extractor import MainContentExtractor |
|
|
from browser_use.controller.views import ( |
|
|
ClickElementAction, |
|
|
DoneAction, |
|
|
ExtractPageContentAction, |
|
|
GoToUrlAction, |
|
|
InputTextAction, |
|
|
OpenTabAction, |
|
|
ScrollAction, |
|
|
SearchGoogleAction, |
|
|
SendKeysAction, |
|
|
SwitchTabAction, |
|
|
) |
|
|
import logging |
|
|
import inspect |
|
|
import asyncio |
|
|
import os |
|
|
from langchain_core.language_models.chat_models import BaseChatModel |
|
|
from browser_use.agent.views import ActionModel, ActionResult |
|
|
|
|
|
from src.utils.mcp_client import create_tool_param_model, setup_mcp_client_and_tools |
|
|
|
|
|
from browser_use.utils import time_execution_sync |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
Context = TypeVar('Context') |
|
|
|
|
|
|
|
|
class CustomController(Controller): |
|
|
def __init__(self, exclude_actions: list[str] = [], |
|
|
output_model: Optional[Type[BaseModel]] = None, |
|
|
ask_assistant_callback: Optional[Union[Callable[[str, BrowserContext], Dict[str, Any]], Callable[ |
|
|
[str, BrowserContext], Awaitable[Dict[str, Any]]]]] = None, |
|
|
): |
|
|
super().__init__(exclude_actions=exclude_actions, output_model=output_model) |
|
|
self._register_custom_actions() |
|
|
self.ask_assistant_callback = ask_assistant_callback |
|
|
self.mcp_client = None |
|
|
self.mcp_server_config = None |
|
|
|
|
|
def _register_custom_actions(self): |
|
|
"""Register all custom browser actions""" |
|
|
|
|
|
@self.registry.action( |
|
|
"When executing tasks, prioritize autonomous completion. However, if you encounter a definitive blocker " |
|
|
"that prevents you from proceeding independently – such as needing credentials you don't possess, " |
|
|
"requiring subjective human judgment, needing a physical action performed, encountering complex CAPTCHAs, " |
|
|
"or facing limitations in your capabilities – you must request human assistance." |
|
|
) |
|
|
async def ask_for_assistant(query: str, browser: BrowserContext): |
|
|
if self.ask_assistant_callback: |
|
|
if inspect.iscoroutinefunction(self.ask_assistant_callback): |
|
|
user_response = await self.ask_assistant_callback(query, browser) |
|
|
else: |
|
|
user_response = self.ask_assistant_callback(query, browser) |
|
|
msg = f"AI ask: {query}. User response: {user_response['response']}" |
|
|
logger.info(msg) |
|
|
return ActionResult(extracted_content=msg, include_in_memory=True) |
|
|
else: |
|
|
return ActionResult(extracted_content="Human cannot help you. Please try another way.", |
|
|
include_in_memory=True) |
|
|
|
|
|
@self.registry.action( |
|
|
'Upload file to interactive element with file path ', |
|
|
) |
|
|
async def upload_file(index: int, path: str, browser: BrowserContext, available_file_paths: list[str]): |
|
|
if path not in available_file_paths: |
|
|
return ActionResult(error=f'File path {path} is not available') |
|
|
|
|
|
if not os.path.exists(path): |
|
|
return ActionResult(error=f'File {path} does not exist') |
|
|
|
|
|
dom_el = await browser.get_dom_element_by_index(index) |
|
|
|
|
|
file_upload_dom_el = dom_el.get_file_upload_element() |
|
|
|
|
|
if file_upload_dom_el is None: |
|
|
msg = f'No file upload element found at index {index}' |
|
|
logger.info(msg) |
|
|
return ActionResult(error=msg) |
|
|
|
|
|
file_upload_el = await browser.get_locate_element(file_upload_dom_el) |
|
|
|
|
|
if file_upload_el is None: |
|
|
msg = f'No file upload element found at index {index}' |
|
|
logger.info(msg) |
|
|
return ActionResult(error=msg) |
|
|
|
|
|
try: |
|
|
await file_upload_el.set_input_files(path) |
|
|
msg = f'Successfully uploaded file to index {index}' |
|
|
logger.info(msg) |
|
|
return ActionResult(extracted_content=msg, include_in_memory=True) |
|
|
except Exception as e: |
|
|
msg = f'Failed to upload file to index {index}: {str(e)}' |
|
|
logger.info(msg) |
|
|
return ActionResult(error=msg) |
|
|
|
|
|
@time_execution_sync('--act') |
|
|
async def act( |
|
|
self, |
|
|
action: ActionModel, |
|
|
browser_context: Optional[BrowserContext] = None, |
|
|
|
|
|
page_extraction_llm: Optional[BaseChatModel] = None, |
|
|
sensitive_data: Optional[Dict[str, str]] = None, |
|
|
available_file_paths: Optional[list[str]] = None, |
|
|
|
|
|
context: Context | None = None, |
|
|
) -> ActionResult: |
|
|
"""Execute an action""" |
|
|
|
|
|
try: |
|
|
for action_name, params in action.model_dump(exclude_unset=True).items(): |
|
|
if params is not None: |
|
|
if action_name.startswith("mcp"): |
|
|
|
|
|
logger.debug(f"Invoke MCP tool: {action_name}") |
|
|
mcp_tool = self.registry.registry.actions.get(action_name).function |
|
|
result = await mcp_tool.ainvoke(params) |
|
|
else: |
|
|
result = await self.registry.execute_action( |
|
|
action_name, |
|
|
params, |
|
|
browser=browser_context, |
|
|
page_extraction_llm=page_extraction_llm, |
|
|
sensitive_data=sensitive_data, |
|
|
available_file_paths=available_file_paths, |
|
|
context=context, |
|
|
) |
|
|
|
|
|
if isinstance(result, str): |
|
|
return ActionResult(extracted_content=result) |
|
|
elif isinstance(result, ActionResult): |
|
|
return result |
|
|
elif result is None: |
|
|
return ActionResult() |
|
|
else: |
|
|
raise ValueError(f'Invalid action result type: {type(result)} of {result}') |
|
|
return ActionResult() |
|
|
except Exception as e: |
|
|
raise e |
|
|
|
|
|
async def setup_mcp_client(self, mcp_server_config: Optional[Dict[str, Any]] = None): |
|
|
self.mcp_server_config = mcp_server_config |
|
|
if self.mcp_server_config: |
|
|
self.mcp_client = await setup_mcp_client_and_tools(self.mcp_server_config) |
|
|
self.register_mcp_tools() |
|
|
|
|
|
def register_mcp_tools(self): |
|
|
""" |
|
|
Register the MCP tools used by this controller. |
|
|
""" |
|
|
if self.mcp_client: |
|
|
for server_name in self.mcp_client.server_name_to_tools: |
|
|
for tool in self.mcp_client.server_name_to_tools[server_name]: |
|
|
tool_name = f"mcp.{server_name}.{tool.name}" |
|
|
self.registry.registry.actions[tool_name] = RegisteredAction( |
|
|
name=tool_name, |
|
|
description=tool.description, |
|
|
function=tool, |
|
|
param_model=create_tool_param_model(tool), |
|
|
) |
|
|
logger.info(f"Add mcp tool: {tool_name}") |
|
|
logger.debug( |
|
|
f"Registered {len(self.mcp_client.server_name_to_tools[server_name])} mcp tools for {server_name}") |
|
|
else: |
|
|
logger.warning(f"MCP client not started.") |
|
|
|
|
|
async def close_mcp_client(self): |
|
|
if self.mcp_client: |
|
|
await self.mcp_client.__aexit__(None, None, None) |
|
|
|
