tools/web_search.py

"""
tools/web_search.py —— 工具①：联网搜索

给 agent 一个"上网搜东西"的能力。
"""

import time

from langchain_core.tools import tool

from config import TAVILY_API_KEY


def _tavily(query: str):
    """用 Tavily 搜索。函数名以下划线开头，表示这是"内部辅助函数"，只给本文件自己调用。"""
    # 没有配密钥就直接返回 None（表示"用不了 Tavily"），让外层去用备胎 DuckDuckGo。
    if not TAVILY_API_KEY:
        return None
    import requests

    # 向 Tavily 的接口发一个请求，带上：密钥、搜索词、最多要 8 条结果、并让它顺便给个总结性答案。
    r = requests.post(
        "https://api.tavily.com/search",
        json={"api_key": TAVILY_API_KEY, "query": query, "max_results": 8, "include_answer": True},
        timeout=30,   # 最多等 30 秒，避免卡死
    )
    r.raise_for_status()   # 如果服务器返回错误状态码，这里会主动报错
    data = r.json()        # 把返回的 JSON 文本解析成 Python 能用的数据
    # 下面把搜索结果拼成一段整齐的文字返回（标题、网址、摘要），方便大模型阅读。
    parts = []
    if data.get("answer"):
        parts.append("Answer: " + data["answer"])
    for it in data.get("results", []):
        parts.append(f"[{it.get('title')}]({it.get('url')})\n{it.get('content', '')}")
    return "## Search Results\n\n" + "\n\n".join(parts) if parts else None


def _duckduckgo(query: str):
    """备用搜索引擎：DuckDuckGo。同样是内部辅助函数。"""
    # 不同版本的库名字不一样，try/except 是为了"哪个能导入就用哪个"，增强兼容性。
    try:
        from ddgs import DDGS
    except ImportError:
        from duckduckgo_search import DDGS

    results = DDGS().text(query, max_results=10)   # 搜索，最多取 10 条
    if not results:
        return None
    # 同样把结果拼成整齐文字返回。
    return "## Search Results\n\n" + "\n\n".join(
        f"[{r.get('title')}]({r.get('href') or r.get('url')})\n{r.get('body') or r.get('content', '')}"
        for r in results
    )


@tool
def web_search(query: str) -> str:
    """Search the web and return the top results (title, url, snippet). Use for general,
    up-to-date research; follow up with `visit_webpage` to read a result in full. For
    encyclopedic facts prefer `wikipedia_search`, which is more reliable."""
    # 第一步：先试 Tavily。成功拿到结果就直接返回。
    try:
        out = _tavily(query)
        if out:
            return out
    except Exception:
        pass   # Tavily 出错也不报错中断，继续往下用

    # 第二步：用 DuckDuckGo，并且最多重试 4 次（免费引擎容易临时失败，多试几次更稳）。
    last_err = None
    for attempt in range(4):
        try:
            out = _duckduckgo(query)
            if out:
                return out
        except Exception as e:
            last_err = e
        # 每次失败后等一会儿再试，且等待时间越来越长（1.5秒、3秒、4.5秒...），避免被对方当成攻击而封锁。
        time.sleep(1.5 * (attempt + 1))
    # 两套引擎都没搜到，就返回一句提示，告诉大模型"换个工具或换个说法再试"。
    return (
        f"Web search returned nothing (last error: {last_err}). "
        "Try `wikipedia_search`, rephrase the query, or `visit_webpage` on a known url."
    )