Building a Browser Use Agent with ADK and Selenium for Brand Search Optimization

import pandas as pd
from google.cloud import bigquery

# Sample data
data = [
    {
        "title": "Kids' Joggers",
        "description": "Comfortable joggers for active kids",
        "attributes": "Cotton blend, Elastic waistband, Pockets",
        "brand": "BSO Test Brand"
    },
    {
        "title": "Women's Active Shoes",
        "description": "Lightweight active shoes for all day comfort",
        "attributes": "Breathable mesh, Cushioned insole, Flexible sole",
        "brand": "BSO Test Brand"
    },
    {
        "title": "Men's Sneaker",
        "description": "Casual sneakers for everyday wear",
        "attributes": "Canvas upper, Rubber sole, Lace-up closure",
        "brand": "BSO Test Brand"
    }
]

# Create DataFrame
df = pd.DataFrame(data)

# Set up BigQuery client
client = bigquery.Client()
dataset_id = "your_dataset_id"
table_id = "product_data"

# Create dataset if it doesn't exist
dataset_ref = client.dataset(dataset_id)
try:
    client.get_dataset(dataset_ref)
    print(f"Dataset {dataset_id} already exists")
except Exception:
    dataset = bigquery.Dataset(dataset_ref)
    dataset = client.create_dataset(dataset)
    print(f"Dataset {dataset_id} created")

# Create table and upload data
table_ref = dataset_ref.table(table_id)
job_config = bigquery.LoadJobConfig()
job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
job_config.autodetect = True

job = client.load_table_from_dataframe(df, table_ref, job_config=job_config)
job.result()

print(f"Loaded {job.output_rows} rows into {dataset_id}.{table_id}")

from google.adk import Agent
from google.adk.agents import LLMAgent
from google.adk.tool_registry import StandardToolRegistry

from brand_search_optimization.keyword_finding_agent import KeywordFindingAgent
from brand_search_optimization.search_results_agent import SearchResultsAgent
from brand_search_optimization.comparison_agent import ComparisonRootAgent

class BrandSearchOptimizationAgent(Agent):
    """Root agent for Brand Search Optimization."""

    @classmethod
    def create(cls):
        # Create tool registry
        registry = StandardToolRegistry()
        
        # Register subagents
        keyword_finding = KeywordFindingAgent.create()
        search_results = SearchResultsAgent.create()
        comparison_root = ComparisonRootAgent.create()
        
        registry.register_tool(keyword_finding.as_tool())
        registry.register_tool(search_results.as_tool())
        registry.register_tool(comparison_root.as_tool())
        
        # Create the agent with the root prompt
        with open("brand_search_optimization/prompts/root_prompt.txt", "r") as f:
            prompt = f.read()
        
        return LLMAgent(
            registry=registry,
            prompt=prompt,
            name="Brand Search Optimization",
            description="A helpful assistant for Brand Search Optimization"
        )

from google.adk.agents import LLMAgent
from google.adk.tool_registry import StandardToolRegistry

from brand_search_optimization.tools.bigquery_tools import GetProductDetails

class KeywordFindingAgent(Agent):
    """Agent for finding keywords from product data."""

    @classmethod
    def create(cls):
        # Create tool registry
        registry = StandardToolRegistry()
        
        # Register tools
        registry.register_tool(GetProductDetails())
        
        # Create the agent with the keyword finding prompt
        with open("brand_search_optimization/prompts/keyword_finding_prompt.txt", "r") as f:
            prompt = f.read()
        
        return LLMAgent(
            registry=registry,
            prompt=prompt,
            name="Keyword Finding Agent",
            description="Finds relevant keywords for a brand based on product data"
        )

import os
from typing import Dict, Any
from google.cloud import bigquery
from google.adk.tool import Tool

class GetProductDetails(Tool):
    """Tool to retrieve product details from BigQuery."""
    
    description = "Get product details for a specific brand from the database"
    
    def __init__(self):
        self.client = bigquery.Client()
        self.project = os.getenv("PROJECT_ID")
        self.dataset = os.getenv("DATASET_ID")
        self.table = os.getenv("TABLE_ID")
    
    def __call__(self, brand_name: str) -> str:
        """
        Retrieves product details for the specified brand.
        
        Args:
            brand_name: The name of the brand to retrieve product details for
            
        Returns:
            Formatted markdown table with product details
        """
        query = f"""
        SELECT title, description, attributes
        FROM `{self.project}.{self.dataset}.{self.table}`
        WHERE brand = '{brand_name}'
        LIMIT 3
        """
        
        query_job = self.client.query(query)
        results = query_job.result()
        
        # Format results as markdown table
        markdown = "| Title | Description | Attributes |\n"
        markdown += "| ----- | ----------- | ---------- |\n"
        
        for row in results:
            markdown += f"| {row.title} | {row.description} | {row.attributes} |\n"
        
        return markdown

from google.adk.agents import LLMAgent
from google.adk.tool_registry import StandardToolRegistry

from brand_search_optimization.tools.selenium_tools import (
    GoToUrl, TakeScreenshot, FindElementWithText,
    ClickElementWithText, EnterTextIntoElement,
    ScrollDown, GetPageSource, AnalyzeWebpageAndDetermineAction
)

class SearchResultsAgent(Agent):
    """Agent for searching and retrieving top search results."""

    @classmethod
    def create(cls):
        # Create tool registry
        registry = StandardToolRegistry()
        
        # Register tools
        registry.register_tool(GoToUrl())
        registry.register_tool(TakeScreenshot())
        registry.register_tool(FindElementWithText())
        registry.register_tool(ClickElementWithText())
        registry.register_tool(EnterTextIntoElement())
        registry.register_tool(ScrollDown())
        registry.register_tool(GetPageSource())
        registry.register_tool(AnalyzeWebpageAndDetermineAction())
        
        # Create the agent with the search results prompt
        with open("brand_search_optimization/prompts/search_results_prompt.txt", "r") as f:
            prompt = f.read()
        
        return LLMAgent(
            registry=registry,
            prompt=prompt,
            name="Search Results Agent",
            description="Searches the web for keywords and retrieves top search results"
        )

import os
import time
from typing import Dict, Any, Optional
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from webdriver_manager.chrome import ChromeDriverManager
from google.adk.tool import Tool

# Initialize WebDriver
def get_driver():
    if os.getenv("DISABLE_WEB_DRIVER") == "1":
        return None
    
    chrome_options = Options()
    if os.getenv("HEADLESS", "0") == "1":
        chrome_options.add_argument("--headless")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    
    service = Service(ChromeDriverManager().install())
    driver = webdriver.Chrome(service=service, options=chrome_options)
    return driver

_driver = None

def driver():
    global _driver
    if _driver is None:
        _driver = get_driver()
    return _driver

class GoToUrl(Tool):
    """Tool to navigate to a URL."""
    
    description = "Navigate to a specified URL"
    
    def __call__(self, url: str) -> str:
        """Navigate to the specified URL."""
        if driver() is None:
            return "WebDriver is disabled."
        
        driver().get(url)
        return f"Navigated to {url}"

class TakeScreenshot(Tool):
    """Tool to take a screenshot of the current page."""
    
    description = "Take a screenshot of the current page"
    
    def __call__(self, tool_context: Dict[str, Any], file_name: str = "screenshot.png") -> str:
        """Take a screenshot and save it to the specified file."""
        if driver() is None:
            return "WebDriver is disabled."
        
        screenshot_path = file_name
        driver().save_screenshot(screenshot_path)
        
        # Register the screenshot as an artifact
        tool_context["artifact_paths"].append(screenshot_path)
        
        return f"Screenshot saved to {screenshot_path}"

class FindElementWithText(Tool):
    """Tool to find an element containing specific text."""
    
    description = "Find an element on the page containing specific text"
    
    def __call__(self, text_pattern: str) -> str:
        """Find an element containing the specified text pattern."""
        if driver() is None:
            return "WebDriver is disabled."
        
        try:
            element = driver().find_element(By.XPATH, f"//*[contains(text(),'{text_pattern}')]")
            return f"Found element with text: {text_pattern}"
        except Exception as e:
            return f"Element with text '{text_pattern}' not found: {str(e)}"

class ClickElementWithText(Tool):
    """Tool to click an element containing specific text."""
    
    description = "Click an element containing specific text"
    
    def __call__(self, text_pattern: str) -> str:
        """Click an element containing the specified text pattern."""
        if driver() is None:
            return "WebDriver is disabled."
        
        try:
            element = driver().find_element(By.XPATH, f"//*[contains(text(),'{text_pattern}')]")
            element.click()
            return f"Clicked element with text: {text_pattern}"
        except Exception as e:
            return f"Failed to click element with text '{text_pattern}': {str(e)}"

class EnterTextIntoElement(Tool):
    """Tool to enter text into an element."""
    
    description = "Enter text into an element identified by its ID"
    
    def __call__(self, text_to_enter: str, element_id: str, press_enter: bool = True) -> str:
        """Enter text into the specified element."""
        if driver() is None:
            return "WebDriver is disabled."
        
        try:
            element = driver().find_element(By.ID, element_id)
            element.clear()
            element.send_keys(text_to_enter)
            
            if press_enter:
                element.send_keys(Keys.RETURN)
                
            return f"Entered text '{text_to_enter}' into element with ID '{element_id}'"
        except Exception as e:
            return f"Failed to enter text into element with ID '{element_id}': {str(e)}"

class ScrollDown(Tool):
    """Tool to scroll down the page."""
    
    description = "Scroll down the page"
    
    def __call__(self, amount: int = 500) -> str:
        """Scroll down by the specified amount of pixels."""
        if driver() is None:
            return "WebDriver is disabled."
        
        driver().execute_script(f"window.scrollBy(0, {amount});")
        return f"Scrolled down by {amount} pixels"

class GetPageSource(Tool):
    """Tool to get the page source."""
    
    description = "Get the HTML source of the current page"
    
    def __call__(self, limit: int = 10000) -> str:
        """Get the HTML source of the current page, limited to a specified number of characters."""
        if driver() is None:
            return "WebDriver is disabled."
        
        page_source = driver().page_source
        
        if len(page_source) > limit:
            return page_source[:limit] + "... (truncated)"
        
        return page_source

class AnalyzeWebpageAndDetermineAction(Tool):
    """Tool to analyze a webpage and determine actions to take."""
    
    description = "Analyze a webpage and determine actions to achieve a specific task"
    
    def __call__(self, user_task: str) -> str:
        """
        Analyze the current webpage and determine actions to achieve the specified task.
        
        Args:
            user_task: The task to achieve (e.g., "Find the top 3 products for running shoes")
            
        Returns:
            A set of recommended actions to take
        """
        if driver() is None:
            return "WebDriver is disabled."
        
        # Get the page source
        page_source = driver().page_source
        
        # You would typically use an LLM to analyze the page here
        # For this example, we'll return a simple response
        return f"Page analyzed. To achieve '{user_task}', I recommend the following actions:\n" + \
               "1. Identify main product elements on the page\n" + \
               "2. Extract product titles from the top results\n" + \
               "3. Return the top product titles in a structured format"

from google.adk.agents import LLMAgent
from google.adk.tool_registry import StandardToolRegistry

class GeneratorAgent(Agent):
    """Agent for generating the initial comparison report."""

    @classmethod
    def create(cls):
        # Create tool registry
        registry = StandardToolRegistry()
        
        # Create the agent with the generator prompt
        with open("brand_search_optimization/prompts/generator_prompt.txt", "r") as f:
            prompt = f.read()
        
        return LLMAgent(
            registry=registry,
            prompt=prompt,
            name="Generator Agent",
            description="Generates a comparison report between brand titles and top search results"
        )

class CritiqueAgent(Agent):
    """Agent for critiquing and improving the comparison report."""

    @classmethod
    def create(cls):
        # Create tool registry
        registry = StandardToolRegistry()
        
        # Create the agent with the critique prompt
        with open("brand_search_optimization/prompts/critique_prompt.txt", "r") as f:
            prompt = f.read()
        
        return LLMAgent(
            registry=registry,
            prompt=prompt,
            name="Critique Agent",
            description="Critiques and improves the comparison report"
        )

class ComparisonRootAgent(Agent):
    """Root agent for comparison that coordinates between generator and critique."""

    @classmethod
    def create(cls):
        # Create tool registry
        registry = StandardToolRegistry()
        
        # Register subagents
        generator = GeneratorAgent.create()
        critique = CritiqueAgent.create()
        
        registry.register_tool(generator.as_tool())
        registry.register_tool(critique.as_tool())
        
        # Create the agent with the comparison root prompt
        with open("brand_search_optimization/prompts/comparison_root_prompt.txt", "r") as f:
            prompt = f.read()
        
        return LLMAgent(
            registry=registry,
            prompt=prompt,
            name="Comparison Root Agent",
            description="Coordinates comparison and critique to generate optimization recommendations"
        )

class GetProductDetailsCustomDB(Tool):
    """Tool to retrieve product details from a custom database."""
    
    description = "Get product details for a specific brand from the database"
    
    def __call__(self, brand_name: str) -> str:
        """
        Retrieves product details for the specified brand from your custom database.
        
        Args:
            brand_name: The name of the brand to retrieve product details for
            
        Returns:
            Formatted markdown table with product details
        """
        # Replace this with your custom database logic
        # For example: query a MySQL database, read from CSV, etc.
        
        # Example with a simple dictionary
        products = [
            {
                "title": "Kids' Joggers",
                "description": "Comfortable joggers for active kids",
                "attributes": "Cotton blend, Elastic waistband, Pockets"
            },
            # Add more products...
        ]
        
        # Format results as markdown table
        markdown = "| Title | Description | Attributes |\n"
        markdown += "| ----- | ----------- | ---------- |\n"
        
        for product in products:
            markdown += f"| {product['title']} | {product['description']} | {product['attributes']} |\n"
        
        return markdown