google-mcp/lib/google_search.py

import asyncio
import logging
import re
import urllib.parse
from typing import Any, Dict, List

from lib.playwright_manager import PlaywrightManager
from lib.rate_limiter import RateLimiter

logger = logging.getLogger(__name__)

GOOGLE_URL = "https://www.google.com/search"


class GoogleSearch:
    """
    Search Google using Playwright headless browser with stealth mode.
    Respects rate limits and cleans up page resources after each search.
    """

    def __init__(self, min_interval: float = 5.0):
        self.rate_limiter = RateLimiter(min_interval_seconds=min_interval)

    async def search(
        self, query: str, num_results: int = 10, language: str = "en"
    ) -> List[Dict[str, Any]]:
        await self.rate_limiter.acquire()

        page = None
        try:
            page = await PlaywrightManager.get_page(stealth=True)

            url_params = (
                f"{GOOGLE_URL}?q={urllib.parse.quote(query)}"
                f"&num={min(num_results, 20)}&hl={language}"
            )
            logger.info(f"Navigating to Google search: {query}")

            await page.goto(url_params, wait_until="domcontentloaded", timeout=30000)
            await asyncio.sleep(3)

            html = await page.content()
            results = self._parse_results(html, num_results)
            logger.info(f"Google search returned {len(results)} results for: {query}")
            return results

        except Exception as e:
            logger.error(f"Google search failed for '{query}': {e}")
            raise RuntimeError(f"Google search failed: {e}") from e
        finally:
            if page:
                await PlaywrightManager.close_page(page)

    def _parse_results(
        self, html: str, max_results: int
    ) -> List[Dict[str, Any]]:
        results = []
        seen_urls = set()

        # Strategy 1: Parse /url?q= links (Google's redirect URLs)
        url_pattern = re.findall(
            r'<a[^>]*href="(/url\?q=([^&"]+)&[^"]*)"[^>]*>(.*?)</a>',
            html,
            re.DOTALL,
        )

        for _, raw_url, anchor_html in url_pattern:
            if len(results) >= max_results:
                break

            url = self._clean_url(raw_url)
            if not url or url in seen_urls or "google.com" in url:
                continue
            seen_urls.add(url)

            title = re.sub(r'<[^>]+>', "", anchor_html).strip()
            if not title or len(title) < 3:
                continue

            results.append({"title": title, "url": url, "snippet": ""})

        # Strategy 2: If we got results, try to pair them with nearby snippets
        if results:
            results = self._add_snippets(html, results)
        else:
            # Strategy 3: Parse data-href attributes (modern Google)
            data_hrefs = re.findall(r'data-href="(.*?)"', html)
            for raw_url in data_hrefs:
                if len(results) >= max_results:
                    break
                url = self._clean_url(raw_url)
                if not url or url in seen_urls or "google.com" in url:
                    continue
                seen_urls.add(url)
                results.append({"title": url, "url": url, "snippet": ""})

        # Strategy 4: Parse from <h3> tags and nearby <a> tags
        if not results:
            results = self._parse_from_h3(html, max_results)

        return results[:max_results]

    def _clean_url(self, raw_url: str) -> str:
        """Decode Google redirect URL."""
        raw_url = urllib.parse.unquote(raw_url)
        # Remove tracking parameters
        url = re.sub(r"&[a-z_]+=.*$", "", raw_url)
        url = re.sub(r"[?#].*$", "", url)
        if not url.startswith("http"):
            url = "https://" + url.lstrip("//")
        return url.strip()

    def _add_snippets(
        self, html: str, results: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Try to find snippets near result URLs."""
        snippet_blocks = re.findall(
            r'<span[^>]*>([\s\S]*?)</span>', html
        )
        for i, result in enumerate(results):
            for block in snippet_blocks:
                clean = re.sub(r'<[^>]+>', "", block).strip()
                if not 50 < len(clean) < 300:
                    continue
                pos = html.find(clean)
                if pos == -1:
                    continue
                window = html[max(0, pos - 500): pos + 500]
                if result["url"] in window:
                    result["snippet"] = clean[:200]
                    break
        return results

    def _parse_from_h3(
        self, html: str, max_results: int
    ) -> List[Dict[str, Any]]:
        """Parse results from h3 title tags and nearby links."""
        results = []
        h3_pattern = re.findall(
            r'<h3[^>]*>([\s\S]*?)</h3>', html, re.DOTALL
        )

        for h3_content in h3_pattern:
            if len(results) >= max_results:
                break

            # Extract title from h3
            title = re.sub(r'<[^>]+>', "", h3_content).strip()
            if not title or len(title) < 3:
                continue

            # Find the link inside the h3
            link_match = re.search(r'href="([^"]+)"', h3_content)
            url = ""
            if link_match:
                url = self._clean_url(link_match.group(1))

            if url and "google.com" not in url:
                results.append({"title": title, "url": url, "snippet": ""})
            elif title:
                results.append({"title": title, "url": url, "snippet": ""})

        return results