npmaker/Final_Assignment

1from selenium import webdriver2from selenium.webdriver.firefox.options import Options3from selenium.webdriver.firefox.service import Service4from webdriver_manager.firefox import GeckoDriverManager5from selenium.webdriver.common.by import By6from selenium.webdriver.support.ui import WebDriverWait7from selenium.webdriver.support import expected_conditions as EC8import time9import re10from smolagents import Tool11import json12 13class WikipediaSeleniumTool(Tool):14    name = "wikipediaSeleniumTool"15    #description = "Tool that uses Selenium with Firefox to scrape Wikipedia pages efficiently."16    description = """17        This tool extracts content from Wikipedia pages using Selenium with Firefox.18        It's optimized for token efficiency and works well for parsing Wikipedia articles.19        Input should be a task containing a Wikipedia URL (e.g., 'Extract information from https://en.wikipedia.org/wiki/Artificial_intelligence').20        """21    inputs = {22        'task': {23            'type': 'string',24            'description': 'A task description that includes a Wikipedia URL'25        }26    }27    output_type = "string"28    is_initialized = False29 30    def __init__(self):31        """Initialize the Wikipedia Selenium Tool.32        """33        self.max_output_length = 100000034        self.headless = True35        36    def _setup_driver(self):37        """Set up and return a Firefox WebDriver."""38        print("setup driver")39        options = Options()40        if self.headless:41            options.add_argument('--headless')42        43        # Add more options for stability in containerized environments44        options.add_argument('--width=1920')45        options.add_argument('--height=1080')46        47        # Use GeckoDriverManager to automatically handle the driver48        service = Service(GeckoDriverManager().install())49        return webdriver.Firefox(service=service, options=options)50    51    def _is_wikipedia_url(self, url):52        """Check if the URL is a Wikipedia URL."""53        return bool(re.match(r'^https?://([\w-]+\.)?wikipedia\.org', url))54    55    def _extract_wikipedia_content(self, driver):56        """Extract content from Wikipedia page, focusing on the main content area."""57        print("_extract_wikipedia_content")58        try:59            # Wait for the content to load60            WebDriverWait(driver, 10).until(61                EC.presence_of_element_located((By.ID, "content"))62            )63            64            # Get the main content area65            content_div = driver.find_element(By.ID, "mw-content-text")66            67            # Extract title68            title_element = driver.find_element(By.ID, "firstHeading")69            title = title_element.text if title_element else "Unknown Title"70            71            # Extract only the paragraphs and headings from the content72            paragraphs = content_div.find_elements(By.TAG_NAME, "p")73            #print(f"paragraphs: {paragraphs}")74            headings = content_div.find_elements(By.XPATH, ".//h2|.//h3|.//h4|.//h5|.//h6")75            #print(f"headings: {headings}")76            #tr = content_div.find_elements(By.TAG_NAME,"tr")77            td = content_div.find_elements(By.TAG_NAME,"td")78            #print(f"tables: {td}")79            80            # Combine all texts81            text_parts = [f"# {title}\n\n"]82            83            # Process headings and paragraphs84            all_elements = []85            for heading in headings:86                all_elements.append({"type": "heading", "element": heading, "position": heading.location['y']})87            for para in paragraphs:88                if para.text.strip():  # Only include non-empty paragraphs89                    all_elements.append({"type": "paragraph", "element": para, "position": para.location['y']})90            for td in td:91                all_elements.append({"type": "td", "element": td, "position": td.location['y']})92            93            # Sort elements by their Y position94            all_elements.sort(key=lambda x: x["position"])95            96            # Extract and format text97            for item in all_elements:98                if item["type"] == "heading":99                    level = item["element"].tag_name[1]  # Get the heading level (2-6)100                    heading_text = item["element"].text.strip()101                    if heading_text and not heading_text.lower() == "contents":102                        text_parts.append(f"\n{'#' * int(level)} {heading_text}\n")103                elif item["type"] == "td":104                    td_text = re.sub('<[^<]+?>', '', item["element"].text)105                    if td_text:106                        text_parts.append(f"{td_text}\n")107                else:  # paragraph108                    para_text = item["element"].text.strip()109                    if para_text:110                        text_parts.append(f"{para_text}\n\n")111            112            full_content = "".join(text_parts)113            print(f"full_content: {full_content}")114            # Truncate if necessary115            if len(full_content) > self.max_output_length:116                full_content = full_content[:self.max_output_length] + "...[content truncated]"117                118            return full_content119        120        except Exception as e:121            return f"Error extracting Wikipedia content: {str(e)}"122    123    def forward(self, task) -> str:124        """Process Wikipedia URLs to extract content in a token-efficient way.125        126        Args:127            task: A task description that includes a Wikipedia URL128        129        Returns:130            The extracted content from the Wikipedia page131        """132        print("extract url from task")133        # Extract URL from task134        url_match = re.search(r'https?://[^\s"\'<>]+', task)135        if not url_match:136            return "No URL found in the task. Please provide a Wikipedia URL."137        138        url = url_match.group(0)139        140        # Check if it's a Wikipedia URL141        if not self._is_wikipedia_url(url):142            return f"The URL {url} is not a Wikipedia URL. This tool only handles Wikipedia pages."143        144        # Set up the driver145        driver = self._setup_driver()146        147        try:148            driver.get(url)149            time.sleep(2)  # Brief pause to ensure page loads150            151            # Extract the content152            content = self._extract_wikipedia_content(driver)153            154            return content155        156        except Exception as e:157            return f"Error processing Wikipedia page: {str(e)}"158        159        finally:160            driver.quit()161