ailine/backend/app/core/web_search.py

"""
联网搜索公共工具 - 无需 API Key，免费使用 DuckDuckGo
Web Search Public Utility - Free, no API Key, using DuckDuckGo
"""

from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from datetime import datetime


@dataclass
class SearchResult:
    """搜索结果数据类"""
    title: str
    url: str
    snippet: str
    source: str = "DuckDuckGo"
    timestamp: datetime = None

    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now()


class WebSearchTool:
    """联网搜索公共工具类"""

    def __init__(self, max_results: int = 5):
        self.max_results = max_results

    def search(self, query: str, max_results: Optional[int] = None) -> List[SearchResult]:
        """
        使用 DuckDuckGo 搜索
        
        Args:
            query: 搜索关键词
            max_results: 返回结果数量，默认使用初始化时的设置
            
        Returns:
            搜索结果列表
        """
        try:
            from duckduckgo_search import DDGS
            
            num_results = max_results or self.max_results
            
            with DDGS() as ddgs:
                results = ddgs.text(query, max_results=num_results)
                
                search_results = []
                for r in results:
                    search_results.append(SearchResult(
                        title=r.get("title", ""),
                        url=r.get("href", ""),
                        snippet=r.get("body", ""),
                        source="DuckDuckGo"
                    ))
                
                return search_results
                
        except ImportError:
            # 如果 duckduckgo-search 未安装，返回模拟数据
            return self._search_mock(query, max_results)
        except Exception as e:
            print(f"搜索出错：{e}")
            # 出错时返回模拟数据
            return self._search_mock(query, max_results)

    def _search_mock(self, query: str, max_results: Optional[int] = None) -> List[SearchResult]:
        """模拟搜索结果（兜底方案）"""
        mock_results = [
            SearchResult(
                title=f"{query} - 搜索结果 1",
                url="https://example.com/result1",
                snippet=f"这是关于 {query} 的模拟搜索结果，包含相关信息摘要...",
            ),
            SearchResult(
                title=f"{query} - 搜索结果 2",
                url="https://example.com/result2",
                snippet=f"更多关于 {query} 的内容，涵盖多个方面和细节...",
            ),
            SearchResult(
                title=f"{query} - 搜索结果 3",
                url="https://example.com/result3",
                snippet=f"深入分析 {query} 的各个维度，提供全面的视角...",
            ),
        ]
        
        num = max_results or self.max_results
        return mock_results[:num]

    def format_search_results(self, results: List[SearchResult]) -> str:
        """
        格式化搜索结果（带引用溯源）
        
        Args:
            results: 搜索结果列表
            
        Returns:
            格式化后的 Markdown 文本
        """
        if not results:
            return "未找到相关搜索结果"
        
        lines = []
        lines.append("## 🔍 联网搜索结果\n")
        
        for idx, result in enumerate(results, 1):
            lines.append(f"### [{idx}] {result.title}")
            lines.append(f"- 🔗 来源：[{result.url}]({result.url})")
            lines.append(f"- 📝 摘要：{result.snippet}")
            lines.append(f"- 📅 时间：{result.timestamp.strftime('%Y-%m-%d %H:%M:%S')}")
            lines.append("")
        
        # 添加引用溯源说明
        lines.append("---")
        lines.append("💡 **引用溯源说明**：")
        lines.append("- 以上搜索结果均标注了来源链接")
        lines.append("- 使用方括号数字标识引用（如 [1]、[2]）")
        lines.append("- 可通过链接追溯原始信息")
        
        return "\n".join(lines)


# 单例实例
_web_search_tool = None


def get_web_search_tool() -> WebSearchTool:
    """获取联网搜索工具单例"""
    global _web_search_tool
    if _web_search_tool is None:
        _web_search_tool = WebSearchTool()
    return _web_search_tool


def web_search(query: str, max_results: int = 5) -> str:
    """
    便捷函数：联网搜索并返回格式化结果
    
    Args:
        query: 搜索关键词
        max_results: 返回结果数量
        
    Returns:
        格式化后的搜索结果文本
    """
    tool = get_web_search_tool()
    results = tool.search(query, max_results)
    return tool.format_search_results(results)