serp.py

import base64
from contextlib import asynccontextmanager
from typing import Optional
from duckduckgo_search import DDGS
import httpx
from pydantic import BaseModel, Field, field_validator
from playwright.async_api import Browser, BrowserContext, Page
from playwright.async_api import TimeoutError as PlaywrightTimeoutError
from urllib.parse import parse_qs, quote_plus, urlparse
import re
from lxml import etree
import asyncio
from asyncio import Semaphore

# Concurrency limit for Playwright browser contexts.
# This is to prevent too many concurrent browser contexts from being created,
PLAYWRIGHT_CONCURRENCY_LIMIT = 10


# One request must not be able to fan out without limit. Every search
# endpoint issues one outbound call per query (and `/serp/search` may try
# three backends for each), so the length of this list is a direct
# multiplier on outbound load - `n_results` was clamped from the start
# while the list holding the queries was not bounded at all.
MAX_QUERIES_PER_REQUEST = 50


class SerpQuery(BaseModel):
    queries: list[str] = Field(...,
                               min_length=1,
                               max_length=MAX_QUERIES_PER_REQUEST,
                               description="The list of queries to search for")
    n_results: int = Field(
        10, description="Number of results to return for each query. Valid values are 10, 25, 50 and 100")

    @field_validator("n_results")
    @classmethod
    def _clamp_n_results(cls, v: int) -> int:
        """Keep callers from passing a negative or absurdly large value straight into
        a page URL or loop bound (every backend but OPS did this unclamped)."""
        return max(1, min(v, 100))


class SerpResults(BaseModel):
    """Model for SERP scrapping results"""
    error: Optional[str]
    results: Optional[list[dict]]


class BraveSearchBlockedException(Exception):
    """Dummy exception to detect when the headless browser is flagged as suspicious."""

    def __init__(self, *args):
        super().__init__("Brave Search blocked the request, likely due to flagging browser as suspicious")
    pass


class GoogleScholarBlockedException(Exception):
    """Raised when Google Scholar serves its anti-bot interstitial.

    Scholar refuses datacenter IPs with a challenge page rather than
    results. That page has no `div.gs_ri`, so without this the scraper
    waits out the full selector timeout and then reports a bare "Timeout
    exceeded" - slow, and indistinguishable from a genuine selector
    regression.

    Carries the evidence that classified it. Without that, "blocked" is an
    assertion the reader has no way to check, and there is no way to tell
    which signal fired - which is exactly what made the first two attempts
    at this hard to debug from the deployment.
    """

    def __init__(self, evidence: str = ""):
        message = ("Google Scholar blocked the request (anti-bot interstitial served "
                   "instead of results); the deployment's IP is likely rate-limited.")
        if evidence:
            message = f"{message} Evidence: {evidence}"
        super().__init__(message)


class GoogleScholarUnavailableError(Exception):
    """Scholar returned neither results nor anything we can classify.

    Carries what the page actually was - final URL, title, a text excerpt -
    because the alternative is a bare selector timeout that says nothing.
    The first attempt at block detection here guessed at Google's challenge
    markup and matched none of it; the error was useless for working out
    why. Whatever we fail to classify, we report.
    """


class BrowserUnavailableError(Exception):
    """Raised when a Playwright-backed query is attempted but the browser failed to start."""

    def __init__(self, *args):
        super().__init__("Playwright browser is not available (it failed to start at startup).")


class EmptyResultsError(Exception):
    """A backend returned zero results.

    Deliberately distinct from a backend erroring, because zero results is
    ambiguous: it can mean the query genuinely has no hits, or that the
    backend served an anti-bot page that parsed to nothing. Callers (the
    `/serp/search` fallback chain in particular) rely on an exception to
    move on to the next backend, but should not treat this one as evidence
    that the backend itself is unhealthy.
    """


class DuckDuckGoBlockedException(EmptyResultsError):
    """Raised when DuckDuckGo returns zero results.

    The duckduckgo_search library doesn't always raise when its request is
    blocked/rate-limited — it can also just parse zero results from the
    anti-bot response page and return an empty list. A hard block usually
    does raise (and that still counts against the backend); this covers the
    soft case, which is indistinguishable from a genuinely empty search.
    """

    def __init__(self, *args):
        super().__init__("DuckDuckGo returned no results, likely blocked or rate-limited")


_PLAYWRIGHT_CONCURRENCY_SEMAPHORE = Semaphore(PLAYWRIGHT_CONCURRENCY_LIMIT)


@asynccontextmanager
async def playwright_open_page(browser: Optional[Browser]):
    """Context manager for playwright pages"""
    if browser is None:
        raise BrowserUnavailableError()
    async with _PLAYWRIGHT_CONCURRENCY_SEMAPHORE:
        context: BrowserContext = await browser.new_context()
        page: Page = await context.new_page()
        try:
            yield page
        finally:
            # Close the context even if closing the page fails (a crashed
            # renderer), otherwise the context leaks for the process's life.
            try:
                await page.close()
            finally:
                await context.close()


async def _block_stylesheet_and_image_resources(route, request):
    """Shared `page.route("**/*", ...)` handler for every scraper below: skip
    fetching stylesheets/images, since only the page's text content matters
    here and this meaningfully speeds up navigation.
    """
    if request.resource_type in ["stylesheet", "image"]:
        await route.abort()
    else:
        await route.continue_()


# How long to wait for Scholar to show either results or a challenge.
SCHOLAR_SELECTOR_TIMEOUT_MS = 30_000

_SCHOLAR_RESULT_SELECTOR = "div.gs_ri"
# Elements Google's challenge pages have carried. Waiting for these
# *alongside* the results selector is what lets a block fail fast instead
# of costing a full timeout. Matched on presence rather than visibility:
# reCAPTCHA renders in an iframe and its container is commonly zero-height
# until it loads, so a challenge element can be attached but never visible.
_SCHOLAR_BLOCK_SELECTOR = (
    "form#captcha-form, #recaptcha, .g-recaptcha, #infoDiv, "
    "form[action*='sorry'], #gs_captcha_f, img[src*='sorry'], "
    "div#af-error-container, input[name='captcha']"
)
# Wording Google's interstitials have used. Only consulted once the
# selectors have failed - a legitimate Scholar search for "unusual traffic"
# returns papers whose snippets contain that phrase.
_SCHOLAR_BLOCK_TEXT_MARKERS = (
    "unusual traffic",
    "not a robot",
    "/sorry/index",
    "automated queries",
    "your computer network",
)
# A /sorry redirect is markup-independent evidence, unlike any selector.
_SCHOLAR_BLOCK_URL_MARKERS = ("/sorry", "/challenge")


def _looks_like_scholar_block(page_content: str) -> bool:
    lowered = page_content.lower()
    return any(marker in lowered for marker in _SCHOLAR_BLOCK_TEXT_MARKERS)


def _url_looks_blocked(url: Optional[str]) -> bool:
    return bool(url) and any(m in url.lower() for m in _SCHOLAR_BLOCK_URL_MARKERS)


async def _describe_page(page: Page, final_url: Optional[str]) -> str:
    """A short, human-readable summary of whatever we actually got.

    Deliberately best-effort: this runs on an error path, so a failure to
    read the page must not replace the original problem with a new one.
    """
    parts = [f"url={final_url or getattr(page, 'url', 'unknown')!r}"]
    try:
        parts.append(f"title={await page.title()!r}")
    except Exception:
        parts.append("title=<unreadable>")
    try:
        text = " ".join((await page.inner_text("body")).split())
        parts.append(f"text={text[:300]!r}")
    except Exception:
        parts.append("text=<unreadable>")
    return ", ".join(parts)


async def _extract_google_scholar_results(page: Page, n_results: int,
                                          timeout_ms: int = SCHOLAR_SELECTOR_TIMEOUT_MS,
                                          final_url: Optional[str] = None) -> list[dict]:
    """Extract results from an already-loaded Google Scholar results page.

    Raises GoogleScholarBlockedException when Google served an anti-bot
    interstitial instead of results, and GoogleScholarUnavailableError -
    carrying the page's URL, title and a text excerpt - when it served
    something we can't classify. A bare selector timeout is never the
    outcome, because it tells whoever reads the logs nothing about why.
    """
    url = final_url if final_url is not None else getattr(page, "url", None)

    async def _block_evidence() -> Optional[str]:
        """Which signal, if any, says this is a challenge page."""
        if _url_looks_blocked(url):
            return f"redirected to {url!r}"
        try:
            if await page.locator(_SCHOLAR_BLOCK_SELECTOR).count():
                return "challenge element present in the page"
        except Exception:
            pass
        if _looks_like_scholar_block(await page.content()):
            return "page wording matches Google's interstitial"
        return None

    # Classify before waiting on anything. A page carrying no results at
    # all is decidable the moment navigation completes - Scholar is
    # server-rendered, so a genuine results page already has div.gs_ri in
    # the DOM - and every signal (redirect URL, challenge element, Google's
    # wording) is available now. Doing this here rather than in the timeout
    # handler is what stops a block costing a full selector wait: measured
    # against the deployment, classification was already correct but
    # arrived 30s late for exactly this reason.
    #
    # The "no results" guard is also what makes the wording check safe this
    # early: a real search for "unusual traffic" returns papers, so it
    # never reaches here.
    if await page.locator(_SCHOLAR_RESULT_SELECTOR).count() == 0:
        early_evidence = await _block_evidence()
        if early_evidence:
            raise GoogleScholarBlockedException(early_evidence)

    try:
        await page.wait_for_selector(
            f"{_SCHOLAR_RESULT_SELECTOR}, {_SCHOLAR_BLOCK_SELECTOR}",
            state="attached", timeout=timeout_ms)
    except PlaywrightTimeoutError:
        evidence = await _block_evidence()
        if evidence:
            raise GoogleScholarBlockedException(evidence) from None
        raise GoogleScholarUnavailableError(
            "Google Scholar returned neither results nor a recognisable "
            f"challenge page ({await _describe_page(page, url)})") from None

    # Results win. Only a page with none of them gets classified, which
    # keeps the wording check from firing on a legitimate search whose
    # snippets happen to quote Google's interstitial.
    if await page.locator(_SCHOLAR_RESULT_SELECTOR).count() == 0:
        raise GoogleScholarBlockedException(
            await _block_evidence() or "no results on the page")

    items = await page.locator(_SCHOLAR_RESULT_SELECTOR).all()
    results = []
    for item in items[:n_results]:
        title = await item.locator("h3").inner_text(timeout=1000)
        body = await item.locator("div.gs_rs").inner_text(timeout=1000)
        href = await item.locator("h3 > a").get_attribute("href")

        results.append({
            "title": title,
            "body": body,
            "href": href
        })

    return results


def google_scholar_url(q: str, n_results: int) -> str:
    """The Google Scholar results URL for a query.

    Split out from the scraper so it can be asserted on directly: the
    Playwright tests replace `page.goto` and never see the URL.
    """
    return f"https://scholar.google.com/scholar?q={quote_plus(q)}&num={n_results}"


async def query_google_scholar(browser: Browser, q: str, n_results: int = 10):
    """Queries google scholar for the specified query and number of results. Returns relevant papers"""
    async with playwright_open_page(browser) as page:
        await page.route("**/*", _block_stylesheet_and_image_resources)

        await page.goto(google_scholar_url(q, n_results))

        return await _extract_google_scholar_results(page, n_results)


# A patent id's shape, e.g. "US11930446B2" or "EP4760514A1": a two-letter
# country code, at least 6 digits, and an optional kind-code suffix (a
# letter plus 0-2 digits). PATENT_ID_REGEX finds one inside free text (used
# to pull ids out of scraped search results); app.py reuses the same core
# pattern, anchored, to validate a user-supplied patent id - one definition
# so the two can't drift apart.
PATENT_ID_CORE = r"[A-Z]{2}\d{6,}(?:[A-Z]\d?)?"
PATENT_ID_REGEX = rf"\b{PATENT_ID_CORE}\b"


async def _extract_google_patents_results(page: Page, n_results: int) -> list[dict]:
    """Extract results from an already-loaded Google Patents search page."""
    # Wait for at least one search result item to appear
    # This ensures the page has loaded enough to start scraping
    await page.wait_for_function(
        "() => document.querySelectorAll('search-result-item').length >= 1",
        timeout=30_000
    )

    items = await page.locator("search-result-item").all()
    results = []
    for item in items:
        text = " ".join(await item.locator("span").all_inner_texts())
        match = re.search(PATENT_ID_REGEX, text)
        if not match:
            continue

        patent_id = match.group()

        try:
            title = await item.locator("h3, h4").first.inner_text(timeout=1000)
            body = await item.locator("div.abstract, div.result-snippet, .snippet, .result-text").first.inner_text(timeout=1000)
        except Exception:
            continue  # If we can't get title or body, skip this item

        results.append({
            "id": patent_id,
            "href": f"https://patents.google.com/patent/{patent_id}/en",
            "title": title,
            "body": body
        })

    return results[:n_results]


def google_patents_search_url(q: str, n_results: int) -> str:
    """The Google Patents search URL for a query."""
    return f"https://patents.google.com/?q={quote_plus(q)}&num={n_results}"


async def query_google_patents(browser: Browser, q: str, n_results: int = 10):
    """Queries google patents for the specified query and number of results. Returns relevant patents"""
    async with playwright_open_page(browser) as page:
        await page.route("**/*", _block_stylesheet_and_image_resources)

        await page.goto(google_patents_search_url(q, n_results))

        return await _extract_google_patents_results(page, n_results)


async def _extract_brave_results(page: Page, n_results: int) -> list[dict]:
    """Extract results from an already-loaded Brave Search results page.

    Raises BraveSearchBlockedException if the page looks like Brave's
    anti-bot interstitial rather than a results page.
    """
    results_cards = await page.locator('.snippet').all()

    if len(results_cards) == 0:
        page_content = await page.content()

        if "suspicious" in page_content:
            raise BraveSearchBlockedException()

    results = []

    for result in results_cards:
        title = await result.locator('.title').all_inner_texts()
        description = await result.locator('.snippet-description').all_inner_texts()
        url = await result.locator('a').nth(0).get_attribute('href')

        # Filter out results with no URL or brave-specific URLs
        if url is None or url.startswith('/'):
            continue

        results.append({
            "title": title[0] if title else "",
            "body": description[0] if description else "",
            "href": url
        })

        if len(results) >= n_results:
            break

    return results


def brave_search_url(q: str, n_results: int = 10) -> str:
    """The Brave Search results URL for a query.

    Brave paginates rather than taking a result count, so `n_results` is
    applied while extracting rather than in the URL; it is accepted here to
    keep the builders interchangeable.
    """
    return f"https://search.brave.com/search?q={quote_plus(q)}"


async def query_brave_search(browser: Browser, q: str, n_results: int = 10):
    """Queries Brave Search for the specified query."""
    async with playwright_open_page(browser) as page:
        await page.route("**/*", _block_stylesheet_and_image_resources)

        await page.goto(brave_search_url(q, n_results))

        return await _extract_brave_results(page, n_results)


def decode_bing_redirect(url: Optional[str]) -> Optional[str]:
    """Resolve a Bing `/ck/a` tracking link to the page it points at.

    Bing doesn't put the destination in a result's href; it wraps every link
    in `https://www.bing.com/ck/a?...&u=a1<base64url>&ntb=1`. Handing that
    back means callers get a tracking link instead of the page - useless for
    citation, deduplication, or deciding whether a result is worth fetching.

    Anything that isn't a decodable Bing redirect is returned unchanged: a
    link we can't decode is still a working link, and losing it would be
    worse than leaving it wrapped. The decoded value comes from a scraped
    page, so only http(s) destinations are let out.
    """
    if not url or "bing.com/ck/a" not in url:
        return url

    encoded = parse_qs(urlparse(url).query).get("u", [""])[0]
    if not encoded.startswith("a1"):
        return url

    payload = encoded[2:]
    try:
        # Bing strips base64 padding; restore however much is missing.
        decoded = base64.urlsafe_b64decode(
            payload + "=" * (-len(payload) % 4)).decode("utf-8")
    except (ValueError, UnicodeDecodeError):
        return url

    if not decoded.startswith(("http://", "https://")):
        return url
    return decoded


async def _extract_bing_results(page: Page, n_results: int) -> list[dict]:
    """Extract results from an already-loaded Bing results page."""
    await page.wait_for_selector("li.b_algo")

    results = []

    items = await page.query_selector_all("li.b_algo")
    for item in items[:n_results]:
        title_el = await item.query_selector("h2 > a")
        url = await title_el.get_attribute("href") if title_el else None
        title = await title_el.inner_text() if title_el else ""

        snippet = ""

        # Try several fallback selectors
        for selector in [
            "div.b_caption p",  # typical snippet
            "div.b_caption",    # sometimes snippet is here
            "div.b_snippet",    # used in some result types
            "div.b_text",       # used in some panels
            "p"                 # fallback to any paragraph
        ]:
            snippet_el = await item.query_selector(selector)
            if snippet_el:
                snippet = await snippet_el.inner_text()
                if snippet.strip():
                    break

        if title and url:
            results.append({
                "title": title.strip(),
                "href": decode_bing_redirect(url.strip()),
                "body": snippet.strip()
            })

    return results


def bing_search_url(q: str, n_results: int = 10) -> str:
    """The Bing results URL for a query.

    Bing paginates rather than taking a result count, so `n_results` is
    applied while extracting rather than in the URL; it is accepted here to
    keep the builders interchangeable.
    """
    return f"https://www.bing.com/search?q={quote_plus(q)}"


async def query_bing_search(browser: Browser, q: str, n_results: int = 10):
    """Queries bing search for the specified query"""
    async with playwright_open_page(browser) as page:
        await page.route("**/*", _block_stylesheet_and_image_resources)

        await page.goto(bing_search_url(q, n_results))

        return await _extract_bing_results(page, n_results)


def _ddg_text_blocking(q: str, n_results: int) -> list[dict]:
    """The synchronous half of a DuckDuckGo query.

    `DDGS.text()` performs blocking HTTP I/O. It is deliberately isolated
    here so `query_ddg_search` can hand it to a worker thread rather than
    calling it on the event loop.
    """
    return list(DDGS().text(q, max_results=n_results))


async def query_ddg_search(q: str, n_results: int = 10):
    """Queries duckduckgo search for the specified query.

    The underlying library call is synchronous, so it runs in a worker
    thread: calling it directly would freeze the whole event loop for the
    duration of the request, and DuckDuckGo is the first backend tried for
    every query in `/serp/search`.
    """
    raw_results = await asyncio.to_thread(_ddg_text_blocking, q, n_results)

    results = [
        {"title": r["title"], "body": r["body"], "href": r["href"]}
        for r in raw_results
    ]

    if not results:
        raise DuckDuckGoBlockedException()

    return results


async def query_arxiv(client: httpx.AsyncClient, query: str, max_results: int = 3):
    """Searches arXiv for the specified query and returns a list of results with titles and PDF URLs."""
    ATOM_NAMESPACE = {'atom': 'http://www.w3.org/2005/Atom'}
    ARXIV_API_URL = 'https://export.arxiv.org/api/query?'

    search_params = {
        'search_query': query,
        'start': 0,
        'max_results': max_results
    }
    query_url = ARXIV_API_URL

    response = await client.get(query_url, params=search_params)
    response.raise_for_status()

    root = etree.fromstring(response.content)
    entries = root.findall('atom:entry', ATOM_NAMESPACE)

    results = []
    for entry in entries:
        title = entry.find(
            'atom:title', ATOM_NAMESPACE).text.strip().replace('\n', ' ')
        id = entry.find('atom:id', ATOM_NAMESPACE).text.strip()
        pdf_url = entry.find(
            'atom:id', ATOM_NAMESPACE).text.replace('/abs/', '/pdf/')
        summary = entry.find(
            'atom:summary', ATOM_NAMESPACE).text.strip()
        results.append({'title': title, 'href': pdf_url,
                       'body': summary, 'id': id})

    return results