Spaces:
Running
Running
from smolagents.tools import Tool | |
from helium import S | |
from selenium.webdriver.common.by import By | |
import json | |
import logging | |
logger = logging.getLogger(__name__) | |
class ScrapeTextTool(Tool): | |
name = "scrape_text" | |
description = "Scrapes text or table data from elements matching a CSS selector on the current page." | |
inputs = { | |
"selector": {"type": "string", "default": "p", "nullable": False, "description": "CSS selector to target elements"}, | |
"extract_table": {"type": "boolean", "default": False, "nullable": False, "description": "If True, extract table data as JSON"} | |
} | |
output_type = "string" | |
def __init__(self, driver): | |
super().__init__() | |
self.driver = driver | |
self.is_initialized = self.driver is not None | |
logger.debug(f"ScrapeTextTool initialized: is_initialized={self.is_initialized}") | |
def forward(self, selector="p", extract_table=False): | |
if not self.is_initialized: | |
return "Error: ScrapeTextTool is not initialized" | |
try: | |
if extract_table: | |
tables = self.driver.find_elements(By.CSS_SELECTOR, selector) | |
if not tables: | |
return "No tables found for selector" | |
table_data = [] | |
for table in tables: | |
rows = table.find_elements(By.TAG_NAME, "tr") | |
table_rows = [] | |
for row in rows: | |
cells = row.find_elements(By.TAG_NAME, "td") or row.find_elements(By.TAG_NAME, "th") | |
row_data = [cell.text.strip() for cell in cells if cell.text.strip()] | |
if row_data: | |
table_rows.append(row_data) | |
if table_rows: | |
table_data.append(table_rows) | |
return json.dumps(table_data) if table_data else "No table data found" | |
else: | |
elements = self.driver.find_elements(By.CSS_SELECTOR, selector) | |
text_list = [element.text.strip() for element in elements if element.text.strip()] | |
return "\n".join(text_list) if text_list else "No text found for selector" | |
except Exception as e: | |
return f"Failed to scrape with selector {selector}: {str(e)}" |