omni-parse/omni/providers/searchers.py

"""Поисковые API и бесплатные движки: Tavily, Exa, You, Firecrawl, DDG, Mojeek, Bing."""
import html
import re
import urllib.parse

from .. import config, net
from ..model import Item

UA = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/141.0.0.0 Safari/537.36"}


def _mk(provider, title, url, text="", score=0.0, published="", extra=None):
    return Item(provider, "result", url=url, title=title, text=text, score=score,
                published=published, extra=extra or {})


ENV2SVC = {"FIRECRAWL_API_KEY": "firecrawl", "TAVILY_API_KEY": "tavily", "EXA_API_KEY": "exa",
           "SERPER_API_KEY": "serper", "JINA_API_KEY": "jina", "APIFY_TOKEN": "apify"}


def _key(name):
    """env первым, иначе ключ из vault (data/vault.json), иначе пусто."""
    v = config.KEYS.get(name)
    if v:
        return v
    from .. import vault
    svc = ENV2SVC.get(name)
    return vault.get(svc) if svc else ""


def available(name):
    need = {"tavily": "TAVILY_API_KEY", "exa": "EXA_API_KEY", "you": "YOU_API_KEY",
            "firecrawl": "FIRECRAWL_API_KEY"}.get(name)
    if not need:
        return True, "без ключей"
    return bool(_key(need)), ("готов" if _key(need) else f"нужен {need}")


def search_tavily(query, limit=15, depth="basic"):
    k = _key("TAVILY_API_KEY")
    if not k:
        return []
    r = net.post_json("https://api.tavily.com/search",
                      {"api_key": k, "query": query, "max_results": min(max(limit, 5), 20),
                       "search_depth": depth, "include_answer": False, "include_raw_content": False})
    d = r.json() or {}
    return [_mk("tavily", x.get("title", ""), x.get("url", ""), (x.get("content") or "")[:2000],
                float(x.get("score") or 0)) for x in (d.get("results") or [])]


def search_exa(query, limit=15):
    k = _key("EXA_API_KEY")
    if not k:
        return []
    r = net.post_json("https://api.exa.ai/search",
                      {"query": query, "numResults": min(max(limit, 5), 25),
                       "contents": {"text": {"maxCharacters": 1200}}},
                      headers={"x-api-key": k})
    d = r.json() or {}
    return [_mk("exa", x.get("title") or "", x.get("url", ""), (x.get("text") or "")[:2000],
                float(x.get("score") or 0), x.get("publishedDate") or "") for x in (d.get("results") or [])]


def search_you(query, limit=15):
    k = _key("YOU_API_KEY")
    if not k:
        return []
    url = f"https://ydc-index.io/v1/search?query={urllib.parse.quote(query)}&count={min(limit, 20)}"
    r = net.get(url, headers={**UA, "X-API-Key": k}, json_mode=True)
    d = r.json() or {}
    out = []
    for section in ("hits", "web"):
        for x in (d.get(section) or []):
            snips = x.get("snippets") or []
            txt = " ".join(snips) if isinstance(snips, list) else (x.get("description") or "")
            out.append(_mk("you", x.get("title") or "", x.get("url") or "", txt[:2000],
                           float(x.get("score") or 0), x.get("page_age") or ""))
        if out:
            break
    return out[:limit]


def search_firecrawl(query, limit=15):
    k = _key("FIRECRAWL_API_KEY")
    if not k:
        return []
    r = net.post_json("https://api.firecrawl.dev/v1/search",
                      {"query": query, "limit": min(max(limit, 3), 20)},
                      headers={"Authorization": f"Bearer {k}"})
    d = r.json() or {}
    return [_mk("firecrawl", x.get("title") or "", x.get("url", ""),
                (x.get("description") or "")[:2000]) for x in (d.get("data") or [])]


def search_ddg(query, limit=15):
    url = "https://html.duckduckgo.com/html/?q=" + urllib.parse.quote(query) + "&kl=ru-ru"
    r = net.get(url, headers=UA)
    if r.status != 200:
        return []
    out = []
    blocks = re.split(r'<div class="result results_links', r.text)[1:]
    for b in blocks:
        m = re.search(r'class="result__a" href="([^"]+)"[^>]*>(.*?)</a>', b, re.S)
        if not m:
            continue
        href = html.unescape(m.group(1))
        if "uddg=" in href:
            href = urllib.parse.unquote(href.split("uddg=")[1].split("&")[0])
        title = html.unescape(re.sub(r"<[^>]+>", "", m.group(2))).strip()
        sn = re.search(r'class="result__snippet"[^>]*>(.*?)</a>', b, re.S)
        snip = html.unescape(re.sub(r"<[^>]+>", "", sn.group(1))).strip() if sn else ""
        out.append(_mk("ddg", title, href, snip[:1500]))
        if len(out) >= limit:
            break
    return out


def search_mojeek(query, limit=15):
    """Mojeek: собственный индекс, html без капчи."""
    url = "https://www.mojeek.com/search?q=" + urllib.parse.quote(query)
    r = net.get(url, headers=UA)
    if r.status != 200:
        return []
    out = []
    for m in re.finditer(r'<a class="ob"[^>]*href="([^"]+)"[^>]*>(.*?)</a>.*?<p class="s">(.*?)</p>',
                         r.text, re.S):
        href, title, snip = m.group(1), m.group(2), m.group(3)
        out.append(_mk("mojeek", html.unescape(re.sub(r"<[^>]+>", "", title)).strip(),
                       html.unescape(href), html.unescape(re.sub(r"<[^>]+>", "", snip))[:1500]))
        if len(out) >= limit:
            break
    return out


def _bing_decode(href):
    """Bing прячет ссылку в ck/a?...&u=a1<base64url>."""
    if "u=a1" not in href:
        return href
    import base64
    raw = href.split("u=a1", 1)[1].split("&")[0]
    raw += "=" * (-len(raw) % 4)
    try:
        return base64.urlsafe_b64decode(raw).decode("utf-8", "ignore")
    except Exception:
        return href


def search_bing(query, limit=15):
    """Bing html: без ключа, первая страница. Ссылки приходят в base64."""
    url = "https://www.bing.com/search?q=" + urllib.parse.quote(query) + f"&count={min(limit, 30)}"
    r = net.get(url, headers={**UA, "Accept-Language": "ru,en;q=0.8"})
    if r.status != 200:
        return []
    out = []
    for m in re.finditer(r'<li class="b_algo".*?<h2>\s*<a[^>]+href="([^"]+)"[^>]*>(.*?)</a>', r.text, re.S):
        href = _bing_decode(html.unescape(m.group(1)))
        title = html.unescape(re.sub(r"<[^>]+>", "", m.group(2))).strip()
        if not href.startswith("http") or "bing.com/ck" in href:
            continue
        out.append(_mk("bing", title, href))
        if len(out) >= limit:
            break
    return out