pipeline/extract.py

"""Step 4: fetch the top-ranked pages and extract clean article text."""
from __future__ import annotations

from typing import List

import trafilatura

from . import config


def extract_sources(results: List[dict]) -> List[dict]:
    """For each result, download + extract main text (capped). Skips pages that fail.

    Adds 'text' to each result dict and returns only those with usable content.
    """
    sources: List[dict] = []
    for r in results:
        url = r.get("url")
        if not url:
            continue
        text = _extract_one(url)
        if not text:
            # fall back to the search snippet so the source still contributes something
            text = (r.get("snippet") or "").strip()
        if not text:
            continue
        r = dict(r)
        r["text"] = text[: config.SOURCE_CHAR_CAP]
        sources.append(r)
    return sources


def _extract_one(url: str) -> str:
    try:
        downloaded = trafilatura.fetch_url(url)
        if not downloaded:
            return ""
        text = trafilatura.extract(
            downloaded,
            include_comments=False,
            include_tables=False,
            no_fallback=False,
            favor_precision=True,
        )
        return (text or "").strip()
    except Exception:
        return ""