src/search.py

"""
SearchClient + RAGPipeline - built-in live web knowledge integration.

Fully keyless by default - no API key, no quota, no cost ceiling. Uses
redundant public HTML endpoints so the pipeline keeps working when one is
rate-limited:

  Backends (tried in order until one returns results):
    1. DuckDuckGo HTML (html.duckduckgo.com)
    2. Bing HTML (www.bing.com/search)
    3. Mojeek HTML (www.mojeek.com/search) - scrape-friendly, no captcha
    4. Google CSE JSON API - ONLY if GOOGLE_CSE_API_KEY / GOOGLE_CSE_ID
       env vars are present (optional; not needed for normal use)

Every backend is scrape/HTML-based and needs no key. A per-backend cooldown
window suppresses a backend briefly after it fails so one bad endpoint cannot
stall the pipeline. RAGPipeline KV-caches results to avoid re-fetching.
"""

import json
import os
import re
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from typing import List, Optional

UA_POOL = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.4 Safari/605.1.15",
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0 Safari/537.36",
]

DEFAULT_CSE_ID = "f49a9160e6e4840d2"
_COOLDOWN_SECONDS = 90.0
_MAX_CONSECUTIVE_FAILS = 3


@dataclass
class SearchResult:
    title: str = ""
    url: str = ""
    snippet: str = ""

    def to_dict(self) -> dict:
        return {"title": self.title, "url": self.url, "snippet": self.snippet}


def _http_get(url: str, timeout: float = 8.0) -> Optional[str]:
    req = urllib.request.Request(url, headers={"User-Agent": UA_POOL[0]})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            body = resp.read().decode("utf-8", errors="ignore")
        # if an endpoint is returning an anti-bot stub, treat as failure
        if len(body) < 500 or re.search(r"captcha|unusual traffic|blocked|anomaly", body, re.I):
            return None
        return body
    except Exception:
        return None


def _unescape(text: str) -> str:
    return (text.replace("&amp;", "&").replace("&quot;", '"')
                .replace("&#x27;", "'").replace("&apos;", "'")
                .replace("&lt;", "<").replace("&gt;", ">")
                .replace("&nbsp;", " "))


def _strip_html(html: str) -> str:
    html = re.sub(r"<script.*?</script>", " ", html, flags=re.S | re.I)
    html = re.sub(r"<style.*?</style>", " ", html, flags=re.S | re.I)
    html = re.sub(r"<[^>]+>", " ", html)
    html = re.sub(r"\s+", " ", html)
    return html.strip()


class SearchClient:
    """Keyless multi-backend search. No API key required, no quota ceiling."""

    def __init__(self, google_api_key: Optional[str] = None, google_cse_id: Optional[str] = None):
        # Google CSE is an OPTIONAL 4th backend; only used when keys provided.
        self.api_key = google_api_key or os.environ.get("GOOGLE_CSE_API_KEY")
        self.cse_id = google_cse_id or os.environ.get("GOOGLE_CSE_ID") or DEFAULT_CSE_ID
        # per-backend cooldown tracking
        self._cooldowns: dict = {}
        self._fails: dict = {}

    def _cooled_down(self, name: str) -> bool:
        until = self._cooldowns.get(name, 0.0)
        return time.time() < until

    def _mark_fail(self, name: str):
        self._fails[name] = self._fails.get(name, 0) + 1
        if self._fails[name] >= _MAX_CONSECUTIVE_FAILS:
            self._cooldowns[name] = time.time() + _COOLDOWN_SECONDS
            self._fails[name] = 0

    def _mark_ok(self, name: str):
        self._fails[name] = 0

    def search(self, query: str, num: int = 5) -> List[SearchResult]:
        backends = [
            ("ddg", lambda: self._search_ddg(query, num)),
            ("bing", lambda: self._search_bing(query, num)),
            ("mojeek", lambda: self._search_mojeek(query, num)),
        ]
        if self.api_key:
            backends.append(("google", lambda: self._search_google(query, num)))

        for name, fn in backends:
            if self._cooled_down(name):
                continue
            try:
                results = fn()
            except Exception:
                results = []
            if results:
                self._mark_ok(name)
                return results
            self._mark_fail(name)
        return []

    def _search_google(self, query: str, num: int) -> List[SearchResult]:
        params = urllib.parse.urlencode({
            "key": self.api_key,
            "cx": self.cse_id,
            "q": query,
            "num": min(num, 10),
        })
        url = f"https://www.googleapis.com/customsearch/v1?{params}"
        body = _http_get(url)
        if not body:
            return []
        try:
            data = json.loads(body)
        except json.JSONDecodeError:
            return []
        return [SearchResult(
            title=item.get("title", ""),
            url=item.get("link", ""),
            snippet=item.get("snippet", ""),
        ) for item in data.get("items", [])]

    def _search_ddg(self, query: str, num: int) -> List[SearchResult]:
        url = "https://html.duckduckgo.com/html/?q=" + urllib.parse.quote(query)
        body = _http_get(url)
        if not body:
            return []
        results = []
        for m in re.finditer(r'<a[^>]+class="result__a"[^>]+href="([^"]+)"[^>]*>(.*?)</a>', body, re.S):
            href, title = m.group(1), re.sub(r"<[^>]+>", "", m.group(2))
            href = urllib.parse.unquote(href)
            m2 = re.search(r"uddg=([^&]+)", href)
            if m2:
                href = urllib.parse.unquote(m2.group(1))
            # skip DDG ads (y.js redirect wrapper)
            if "duckduckgo.com/y.js" in href or "ad_domain" in href or "ad_provider" in href:
                continue
            results.append(SearchResult(title=_unescape(title.strip()), url=href))
            if len(results) >= num:
                break
        for i, m in enumerate(re.finditer(r'class="result__snippet"[^>]*>(.*?)</a>', body, re.S)):
            if i < len(results):
                results[i].snippet = _unescape(re.sub(r"<[^>]+>", "", m.group(1)).strip())
        return results

    def _search_bing(self, query: str, num: int) -> List[SearchResult]:
        url = "https://www.bing.com/search?q=" + urllib.parse.quote(query) + "&count=" + str(num)
        body = _http_get(url)
        if not body:
            return []
        results = []
        # each result is an <li class="b_algo"> block; href may be a /ck/a redirect
        # with the real URL base64-encoded in the u= query param
        for m in re.finditer(r'<li class="b_algo".*?</li>', body, re.S):
            block = m.group(0)
            title = ""
            href = ""
            snippet = ""
            tm = re.search(r'<h2[^>]*><a[^>]+href="([^"]+)"[^>]*>(.*?)</a>', block, re.S)
            if not tm:
                continue
            href, title = tm.group(1), re.sub(r"<[^>]+>", "", tm.group(2)).strip()
            href = href.replace("&amp;", "&")
            if "bing.com/ck/a" in href:
                um = re.search(r"[?&]u=a1([A-Za-z0-9+/=]+)", href)
                if um:
                    try:
                        import base64
                        padded = um.group(1) + "=" * (-len(um.group(1)) % 4)
                        href = base64.b64decode(padded).decode("utf-8", errors="ignore")
                    except Exception:
                        pass
                else:
                    continue
            sm = re.search(r"<p[^>]*>(.*?)</p>", block, re.S)
            if sm:
                snippet = re.sub(r"<[^>]+>", "", sm.group(1)).strip()
            if href.startswith("http") and title:
                results.append(SearchResult(title=_unescape(title), url=href, snippet=_unescape(snippet)))
            if len(results) >= num:
                break
        return results

    def _search_mojeek(self, query: str, num: int) -> List[SearchResult]:
        url = "https://www.mojeek.com/search?q=" + urllib.parse.quote(query)
        body = _http_get(url)
        if not body:
            return []
        results = []
        # Mojeek result items: <a class="ob" href="...">title</a><p class="s">snippet</p>
        for m in re.finditer(r'<a class="ob" href="([^"]+)"[^>]*>(.*?)</a>', body, re.S):
            results.append(SearchResult(
                title=_unescape(re.sub(r"<[^>]+>", "", m.group(2)).strip()),
                url=m.group(1),
            ))
            if len(results) >= num:
                break
        for i, m in enumerate(re.finditer(r'<p class="s"[^>]*>(.*?)</p>', body, re.S)):
            if i < len(results):
                results[i].snippet = _unescape(re.sub(r"<[^>]+>", "", m.group(1)).strip())
        return results


class RAGPipeline:
    def __init__(self, client: Optional[SearchClient] = None, cache=None, max_context_chars: int = 6000):
        self.client = client or SearchClient()
        self.cache = cache
        self.max_context_chars = max_context_chars

    def retrieve(self, query: str, num: int = 5, top_k: int = 3, use_cache: bool = True) -> str:
        """Query -> search -> fetch -> chunk -> rank -> packed context string."""
        cache_key = f"rag:{query}"
        if use_cache and self.cache is not None:
            cached = self.cache.get(cache_key)
            if cached:
                return cached

        results = self.client.search(query, num=num)
        chunks = []
        for r in results:
            for c in self._fetch_and_chunk(r.url):
                chunks.append((r.title, c))

        if not chunks:
            context = self._fallback_context(results)
        else:
            ranked = self._rank(query, chunks, top_k)
            context = self._pack(ranked)

        if self.cache is not None:
            self.cache.set(cache_key, context, ttl=3600)
        return context

    def _fetch_and_chunk(self, url: str, max_chars: int = 3000) -> List[str]:
        body = _http_get(url, timeout=6.0)
        if not body:
            return []
        text = _strip_html(body)
        if not text:
            return []
        chunks = []
        start = 0
        while start < min(len(text), max_chars):
            end = min(start + 800, len(text))
            chunks.append(text[start:end])
            start = end
        return chunks

    def _rank(self, query: str, chunks: List[tuple], top_k: int) -> List[tuple]:
        q_tokens = set(re.findall(r"[a-z0-9]+", query.lower()))
        scored = []
        for title, chunk in chunks:
            c_tokens = set(re.findall(r"[a-z0-9]+", chunk.lower()))
            overlap = len(q_tokens & c_tokens) / max(1, len(q_tokens))
            scored.append((overlap, title, chunk))
        scored.sort(key=lambda t: -t[0])
        return scored[:top_k]

    def _pack(self, ranked: List[tuple]) -> str:
        parts = []
        total = 0
        for _, title, chunk in ranked:
            entry = f"[{title}]\n{chunk}"
            if total + len(entry) > self.max_context_chars:
                break
            parts.append(entry)
            total += len(entry)
        return "\n\n".join(parts)

    def _fallback_context(self, results: List[SearchResult]) -> str:
        parts = []
        for r in results[:3]:
            parts.append(f"[{r.title}]\n{r.url}\n{r.snippet}")
        return "\n\n".join(parts)