custom_tools_selenium.py

# custom_tools_selenium.py - FIXED VERSION
import time
from typing import Optional, Type
import os

# ✅ SỬA: Dùng CrewAI BaseTool thay vì LangChain
from crewai.tools import BaseTool
from pydantic import BaseModel, Field

# --- Selenium & WebDriver Manager ---
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager

# --- Làm sạch HTML ---
from bs4 import BeautifulSoup

UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/127.0.0.0 Safari/537.36"
)

def _make_chrome(headless: bool = True) -> webdriver.Chrome:
    """Tạo Chrome driver với các options tối ưu"""
    opts = Options()
    if headless:
        opts.add_argument("--headless=new")
    opts.add_argument("--disable-gpu")
    opts.add_argument("--no-sandbox")
    opts.add_argument("--disable-dev-shm-usage")
    opts.add_argument("--window-size=1366,2000")
    opts.add_argument("--disable-blink-features=AutomationControlled")
    opts.add_argument("--disable-extensions")
    opts.add_argument("--disable-logging")
    opts.add_argument(f"user-agent={UA}")

    try:
        service = Service(ChromeDriverManager().install())
        driver = webdriver.Chrome(service=service, options=opts)
        driver.set_page_load_timeout(45)
        return driver
    except Exception as e:
        print(f"❌ Chrome driver setup error: {e}")
        raise

def fetch_page_html(
    url: str,
    headless: bool = True,
    wait_css: Optional[str] = None,
    max_scroll: int = 8,
    pause: float = 0.8
) -> str:
    """Lấy HTML content từ trang web với Selenium"""
    driver = _make_chrome(headless=headless)
    try:
        print(f"🔄 Loading URL: {url}")
        driver.get(url)
        
        # Wait for page load
        WebDriverWait(driver, 30).until(
            lambda d: d.execute_script("return document.readyState") == "complete"
        )
        
        # Wait for specific element if provided
        if wait_css:
            WebDriverWait(driver, 30).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, wait_css))
            )
        
        # Auto scroll to load dynamic content
        last_height = 0
        for i in range(max_scroll):
            driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            time.sleep(pause)
            new_height = driver.execute_script("return document.body.scrollHeight")
            if new_height == last_height:
                break
            last_height = new_height
        
        print("✅ Page loaded successfully")
        return driver.page_source
        
    except Exception as e:
        print(f"❌ Error fetching page: {e}")
        raise
    finally:
        driver.quit()

def html_to_clean_text(html: str) -> str:
    """Chuyển HTML thành text sạch"""
    soup = BeautifulSoup(html, "html.parser")
    
    # Remove unwanted tags
    for tag in soup(["script", "style", "noscript", "nav", "header", "footer"]):
        tag.decompose()
    
    text = soup.get_text(separator="\n")
    lines = [ln.strip() for ln in text.splitlines()]
    lines = [ln for ln in lines if ln]
    return "\n".join(lines)

def skip_head_lines(text: str, n: int = 120) -> str:
    """Bỏ n dòng đầu (thường là navigation/header)"""
    lines = text.splitlines()
    if len(lines) <= n:
        return ""
    
    remain = "\n".join(lines[n:])
    compact, last_blank = [], False
    
    for ln in remain.splitlines():
        ln = ln.strip()
        if not ln:
            if not last_blank:
                compact.append("")
            last_blank = True
        else:
            compact.append(ln)
            last_blank = False
    
    return "\n".join(compact)

def trim_to_token_limit(text: str, max_tokens: int = 4000, model: str = "gpt-4o-mini") -> str:
    """Cắt text theo giới hạn token"""
    try:
        import tiktoken
        enc = tiktoken.encoding_for_model(model)
        tokens = enc.encode(text)
        if len(tokens) <= max_tokens:
            return text
        return enc.decode(tokens[:max_tokens])
    except Exception:
        # Fallback: ước lượng 1 token ~ 4 ký tự
        max_chars = max_tokens * 4
        return text[:max_chars]

# ✅ PYDANTIC MODEL CHO INPUT VALIDATION
class WebsiteContentInput(BaseModel):
    """Input schema for website content reading tool"""
    website_url: str = Field(..., description="URL của trang web cần đọc")

# ✅ CREWAI TOOL CLASS
class ReadWebsiteContentTool(BaseTool):
    name: str = "Read website content"
    description: str = (
        "Đọc nội dung trang web bằng Selenium, xử lý JavaScript, "
        "trả về text sạch đã được làm sạch và cắt theo token limit. "
        "Đặc biệt hữu ích cho job posting và trang động."
    )
    args_schema: Type[BaseModel] = WebsiteContentInput

    def _run(self, website_url: str) -> str:
        """Thực thi tool - đọc nội dung website"""
        try:
            # Validate URL
            if not (website_url and website_url.startswith("http")):
                return "❌ Error: URL không hợp lệ. Cần bắt đầu bằng http/https"

            print(f"🔍 Reading website: {website_url}")
            
            # Fetch HTML
            html = fetch_page_html(website_url, headless=True, wait_css=None)
            
            # Process content
            clean_text = html_to_clean_text(html)
            clean_text = skip_head_lines(clean_text, n=120)
            
            # Trim to token limit
            model_name = os.getenv("OPENAI_MODEL_NAME", "gpt-4o-mini")
            final_text = trim_to_token_limit(clean_text, max_tokens=4000, model=model_name)
            
            print(f"✅ Content extracted: {len(final_text)} characters")
            return final_text
            
        except Exception as e:
            error_msg = f"❌ Error reading website: {str(e)}"
            print(error_msg)
            return error_msg

# ✅ EXPORT TOOL INSTANCE
read_website_content = ReadWebsiteContentTool()