MoviePilot/app/agent/tools/impl/search_web.py

"""
提供网络搜索及网站范围筛选，并将结果整理为有界摘要。
"""

import json
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Type
from urllib.parse import urlparse

from ddgs import DDGS
from pydantic import BaseModel, Field

from app.agent.tools.base import MoviePilotTool
from app.agent.tools.tags import ToolTag
from app.runtime.log import logger
from app.runtime.settings import get_runtime_setting

# 搜索超时时间（秒）
SEARCH_TIMEOUT = 20
# 单次搜索最多返回结果数
MAX_SEARCH_RESULTS = 20
# 默认搜索源
DEFAULT_SEARCH_ENGINE = "auto"
# 可显式调用的搜索引擎后端
SEARCH_ENGINE_BACKENDS = (
    "auto",
    "duckduckgo",
    "google",
    "brave",
    "yahoo",
    "wikipedia",
    "yandex",
    "mojeek",
)
SUPPORTED_SEARCH_ENGINES = SEARCH_ENGINE_BACKENDS
DDGS_AUTO_BACKEND = ",".join(
    backend for backend in SEARCH_ENGINE_BACKENDS if backend != DEFAULT_SEARCH_ENGINE
)
SITE_SEARCH_PATTERN = re.compile(r"\bsite:", re.IGNORECASE)


@dataclass(frozen=True)
class _SearchSiteFilter:
    """
    站点限定搜索参数

    :param domain: 规范化的网站域名，用于验证结果归属
    :param path: 限定的 URL 路径；不限制路径时为空字符串
    :param search_target: 加入 site: 查询条件的域名及可选路径
    """

    domain: str
    path: str
    search_target: str


class SearchWebInput(BaseModel):
    """
    搜索网络内容工具的输入参数模型
    """

    query: str = Field(
        ...,
        description="需要搜索的网络查询文本",
    )
    max_results: Optional[int] = Field(
        MAX_SEARCH_RESULTS,
        description="最多返回的搜索结果数，默认 20，上限 20",
    )
    search_engine: Optional[str] = Field(
        DEFAULT_SEARCH_ENGINE,
        description="搜索后端：auto、duckduckgo、google、brave、yahoo、wikipedia、yandex 或 mojeek；用户未指定时使用 auto",
    )
    site_url: Optional[str] = Field(
        None,
        description="限制搜索范围的可选网站、域名或 URL，例如 https://docs.python.org/3/ 或 github.com",
    )


class SearchWebTool(MoviePilotTool):
    """
    网络搜索工具，支持 DDGS 搜索引擎和指定站点限定搜索。
    """

    name: str = Field(
        default="search_web",
        description="供模型调用的稳定工具名称",
    )
    tags: list[str] = Field(
        default=[ToolTag.Read, ToolTag.Web],
        description="工具的能力和权限分类标签，初始化时会统一归一化",
    )
    description: str = Field(
        default="需要最新信息、事实或参考来源时搜索网络。支持选择 DDGS 搜索后端、自动回退及通过 site_url 限定网站或 URL。默认使用系统代理，返回标题、摘要和链接。",
        description="供模型理解工具用途、输入约束和调用边界的说明",
    )
    args_schema: Type[BaseModel] = Field(
        default=SearchWebInput,
        description="工具输入参数的验证模型或 JSON Schema",
    )

    def get_tool_message(self, **kwargs: Any) -> Optional[str]:
        """
        根据搜索参数生成友好的提示消息
        """
        query = kwargs.get("query", "")
        max_results = kwargs.get("max_results", MAX_SEARCH_RESULTS)
        search_engine = self._normalize_search_engine(kwargs.get("search_engine"))
        site_url = kwargs.get("site_url")
        message = f"搜索网络内容: {query} (最多返回 {max_results} 条结果"
        if search_engine != DEFAULT_SEARCH_ENGINE:
            message += f"，搜索源: {search_engine}"
        if site_url:
            message += f"，限定站点: {site_url}"
        return f"{message})"

    async def run(  # type: ignore[override]
        self,
        query: str,
        max_results: Optional[int] = MAX_SEARCH_RESULTS,
        search_engine: Optional[str] = DEFAULT_SEARCH_ENGINE,
        site_url: Optional[str] = None,
        **kwargs: Any,
    ) -> str:
        """
        执行网络搜索。

        :param query: 搜索关键词
        :param max_results: 最大返回结果数
        :param search_engine: 指定搜索源，默认自动选择
        :param site_url: 指定站点或网址，传入时只返回该范围内的搜索结果
        :return: JSON格式的搜索结果或错误信息
        """
        search_engine = self._normalize_search_engine(search_engine)
        if search_engine not in SUPPORTED_SEARCH_ENGINES:
            supported = ", ".join(SUPPORTED_SEARCH_ENGINES)
            return f"错误: 不支持的搜索源 '{search_engine}'，支持的搜索源: {supported}"

        site_filter = self._normalize_site_filter(site_url)
        if site_url and not site_filter:
            return f"错误: site_url 无效，无法限定搜索范围: {site_url}"

        search_query = self._build_search_query(query=query, site_filter=site_filter)
        if not search_query:
            return "错误: query 不能为空"

        logger.info(
            f"执行工具: {self.name}, 参数: query={query}, "
            f"max_results={max_results}, search_engine={search_engine}, site_url={site_url}"
        )

        try:
            # 限制最大结果数
            max_results = min(
                max(1, max_results or MAX_SEARCH_RESULTS),
                MAX_SEARCH_RESULTS,
            )
            results: List[Dict[str, Any]] = []

            for engine in self._get_search_plan(search_engine):
                results = await self._search_with_backend(
                    engine=engine,
                    query=search_query,
                    max_results=max_results,
                    site_filter=site_filter,
                )
                if results:
                    break

            if not results:
                return f"未找到与 '{search_query}' 相关的搜索结果"

            # 格式化并裁剪结果
            formatted_results = self._format_and_truncate_results(results, max_results)
            return json.dumps(formatted_results, ensure_ascii=False, indent=2)

        except Exception as e:
            error_message = f"搜索网络内容失败: {str(e)}"
            logger.error(f"搜索网络内容失败: {e}", exc_info=True)
            return error_message

    @staticmethod
    def _normalize_search_engine(search_engine: Optional[str]) -> str:
        """
        规范化搜索源参数
        """
        engine = (search_engine or DEFAULT_SEARCH_ENGINE).strip().lower()
        aliases = {
            "ddgs": DEFAULT_SEARCH_ENGINE,
            "ddg": "duckduckgo",
            "duck": "duckduckgo",
            "search": DEFAULT_SEARCH_ENGINE,
            "search_engine": DEFAULT_SEARCH_ENGINE,
        }
        return aliases.get(engine, engine)

    @staticmethod
    def _get_search_plan(search_engine: str) -> List[str]:
        """
        根据搜索源配置生成兜底搜索顺序
        """
        if search_engine != DEFAULT_SEARCH_ENGINE:
            return [search_engine]
        return [DEFAULT_SEARCH_ENGINE]

    async def _search_with_backend(
        self,
        engine: str,
        query: str,
        max_results: int,
        site_filter: Optional[_SearchSiteFilter],
    ) -> List[Dict[str, Any]]:
        """
        使用指定后端执行搜索。

        :param engine: 搜索后端名称
        :param query: 已加工的搜索关键词
        :param max_results: 最大结果数
        :param site_filter: 站点限定条件
        :return: 搜索结果列表
        """
        logger.info(f"使用 DDGS 搜索后端 {self._get_ddgs_backend(engine)} 进行搜索...")
        return await self._search_ddgs(query, max_results, engine, site_filter)

    @staticmethod
    def _get_ddgs_backend(search_engine: str) -> str:
        """
        获取实际传给 DDGS 的搜索后端。

        :param search_engine: 用户指定的搜索源
        :return: DDGS 后端名称或逗号分隔的后端列表
        """
        if search_engine == DEFAULT_SEARCH_ENGINE:
            return DDGS_AUTO_BACKEND
        return search_engine

    @staticmethod
    def _normalize_site_filter(site_url: Optional[str]) -> Optional[_SearchSiteFilter]:
        """
        将用户传入的网址转换为搜索引擎 site 过滤条件。

        :param site_url: 用户传入的站点、域名或完整URL
        :return: 站点过滤条件，无法解析时返回 None
        """
        if not site_url:
            return None

        raw_site_url = site_url.strip()
        if not raw_site_url:
            return None

        parse_target = raw_site_url
        if not re.match(r"^https?://", raw_site_url, re.IGNORECASE):
            parse_target = f"https://{raw_site_url}"

        parsed = urlparse(parse_target)
        domain = (parsed.hostname or "").lower()
        if not domain:
            return None

        path = re.sub(r"/+", "/", parsed.path or "").rstrip("/")
        search_target = f"{domain}{path}" if path else domain
        return _SearchSiteFilter(domain=domain, path=path, search_target=search_target)

    @staticmethod
    def _build_search_query(
        query: str,
        site_filter: Optional[_SearchSiteFilter],
    ) -> str:
        """
        生成实际发送给搜索后端的搜索关键词。

        :param query: 原始搜索关键词
        :param site_filter: 站点限定条件
        :return: 加入 site 过滤后的关键词
        """
        search_query = (query or "").strip()
        if not site_filter or SITE_SEARCH_PATTERN.search(search_query):
            return search_query
        if not search_query:
            return f"site:{site_filter.search_target}"
        return f"{search_query} site:{site_filter.search_target}"

    @staticmethod
    def _filter_results_by_site(
        results: List[Dict[str, Any]],
        site_filter: Optional[_SearchSiteFilter],
    ) -> List[Dict[str, Any]]:
        """
        根据指定站点过滤搜索结果。

        :param results: 原始搜索结果
        :param site_filter: 站点限定条件
        :return: 站点范围内的搜索结果
        """
        if not site_filter:
            return results
        return [
            result
            for result in results
            if SearchWebTool._result_matches_site(result.get("url", ""), site_filter)
        ]

    @staticmethod
    def _result_matches_site(url: str, site_filter: _SearchSiteFilter) -> bool:
        """
        判断搜索结果 URL 是否属于指定站点。

        :param url: 搜索结果 URL
        :param site_filter: 站点限定条件
        :return: URL 属于指定站点时返回 True
        """
        if not url:
            return False

        parse_target = url
        if not re.match(r"^https?://", url, re.IGNORECASE):
            parse_target = f"https://{url}"

        parsed = urlparse(parse_target)
        result_host = SearchWebTool._normalize_host(parsed.hostname or "")
        target_host = SearchWebTool._normalize_host(site_filter.domain)
        if not result_host or not target_host:
            return False
        if result_host != target_host and not result_host.endswith(f".{target_host}"):
            return False
        if not site_filter.path:
            return True

        result_path = re.sub(r"/+", "/", parsed.path or "").rstrip("/")
        return result_path == site_filter.path or result_path.startswith(
            f"{site_filter.path}/"
        )

    @staticmethod
    def _normalize_host(host: str) -> str:
        """
        标准化域名以便比较。

        :param host: 原始域名
        :return: 去掉常见 www 前缀后的域名
        """
        normalized_host = (host or "").lower()
        if normalized_host.startswith("www."):
            return normalized_host[4:]
        return normalized_host

    @staticmethod
    def _source_label(search_engine: str) -> str:
        """
        将搜索源标识转换为结果中的展示名称。

        :param search_engine: 搜索源标识
        :return: 展示名称
        """
        labels = {
            "auto": "DDGS",
            "duckduckgo": "DuckDuckGo",
            "google": "Google",
            "brave": "Brave",
            "yahoo": "Yahoo",
            "wikipedia": "Wikipedia",
            "yandex": "Yandex",
            "mojeek": "Mojeek",
        }
        return labels.get(
            search_engine or DEFAULT_SEARCH_ENGINE,
            search_engine or "SearchEngine",
        )

    @staticmethod
    def _extract_result_url(result: Dict[str, Any]) -> str:
        """
        从不同搜索引擎结果结构中提取 URL。

        :param result: 搜索引擎返回的单条结果
        :return: URL 字符串
        """
        value = result.get("href") or result.get("url") or ""
        return str(value) if value else ""

    @staticmethod
    def _extract_result_snippet(result: Dict[str, Any]) -> str:
        """
        从不同搜索引擎结果结构中提取摘要。

        :param result: 搜索引擎返回的单条结果
        :return: 摘要字符串
        """
        value = (
            result.get("body")
            or result.get("snippet")
            or result.get("content")
            or ""
        )
        return str(value) if value else ""

    @staticmethod
    def _get_proxy_url(proxy_setting: Any) -> Optional[str]:
        """
        从代理设置中提取代理URL
        """
        if not proxy_setting:
            return None
        if isinstance(proxy_setting, dict):
            value = proxy_setting.get("http") or proxy_setting.get("https")
            return str(value) if value else None
        return str(proxy_setting) if isinstance(proxy_setting, str) else None

    async def _search_ddgs(
        self,
        query: str,
        max_results: int,
        search_engine: str = DEFAULT_SEARCH_ENGINE,
        site_filter: Optional[_SearchSiteFilter] = None,
    ) -> List[Dict[str, Any]]:
        """
        使用 DDGS 搜索引擎后端进行搜索。

        :param query: 搜索关键词
        :param max_results: 最大结果数
        :param search_engine: DDGS搜索后端
        :param site_filter: 站点限定条件
        :return: 搜索结果列表
        """
        try:

            def sync_search() -> list[dict[str, Any]]:
                """
                在线程中执行同步搜索
                """
                results: list[dict[str, Any]] = []
                ddgs_kwargs: dict[str, Any] = {"timeout": SEARCH_TIMEOUT}
                proxy_url = self._get_proxy_url(get_runtime_setting('PROXY'))
                if proxy_url:
                    ddgs_kwargs["proxy"] = proxy_url

                try:
                    with DDGS(**ddgs_kwargs) as ddgs:
                        ddgs_results = ddgs.text(
                            query,
                            max_results=max_results,
                            backend=self._get_ddgs_backend(search_engine),
                        )
                        if ddgs_results:
                            for result in ddgs_results:
                                source = (
                                    DEFAULT_SEARCH_ENGINE
                                    if search_engine == DEFAULT_SEARCH_ENGINE
                                    else search_engine
                                )
                                results.append(
                                    {
                                        "title": result.get("title", ""),
                                        "snippet": self._extract_result_snippet(result),
                                        "url": self._extract_result_url(result),
                                        "source": self._source_label(source),
                                    }
                                )
                except Exception as err:
                    logger.warning(f"搜索引擎搜索进程失败: {err}")
                return results

            raw_results = await self.run_blocking("web", sync_search)
            results = raw_results if isinstance(raw_results, list) else []
            return self._filter_results_by_site(results, site_filter)

        except Exception as e:
            logger.warning(f"搜索引擎搜索失败: {e}")
            return []

    @staticmethod
    def _format_and_truncate_results(
            results: List[Dict[str, Any]], max_results: int
    ) -> Dict[str, Any]:
        """
        格式化并裁剪搜索结果
        """
        formatted: dict[str, Any] = {"total_results": len(results)}
        formatted_results: list[dict[str, Any]] = []

        for idx, result in enumerate(results[:max_results], 1):
            title = result.get("title", "")[:200]
            snippet = result.get("snippet", "")
            url = result.get("url", "")
            source = result.get("source", "Unknown")

            # 裁剪摘要
            max_snippet_length = 1000
            if len(snippet) > max_snippet_length:
                snippet = snippet[:max_snippet_length] + "..."

            # 清理文本
            snippet = re.sub(r"\s+", " ", snippet).strip()

            formatted_results.append(
                {
                    "rank": idx,
                    "title": title,
                    "snippet": snippet,
                    "url": url,
                    "source": source,
                }
            )

        formatted["results"] = formatted_results
        if len(results) > max_results:
            formatted["note"] = f"仅显示前 {max_results} 条结果。"

        return formatted