feat:联网搜索和自动回复

plugins/TavilySearch/main.py

@@ -0,0 +1,202 @@
+"""
+TavilySearch 联网搜索插件
+
+基于 Tavily API 的联网搜索功能，仅作为 LLM Tool 供 AIChat 调用
+支持多 API Key 轮询，搜索结果返回给 AIChat 的 AI 处理（带上下文和人设）
+"""
+
+import tomllib
+import aiohttp
+from pathlib import Path
+from typing import List, Optional
+from loguru import logger
+from utils.plugin_base import PluginBase
+
+
+class TavilySearch(PluginBase):
+    """Tavily 联网搜索插件 - 仅作为 LLM Tool"""
+
+    description = "Tavily 联网搜索 - 支持多 Key 轮询的搜索工具"
+    author = "Assistant"
+    version = "1.0.0"
+
+    def __init__(self):
+        super().__init__()
+        self.config = None
+        self.api_keys = []
+        self.current_key_index = 0
+
+    async def async_init(self):
+        """异步初始化"""
+        try:
+            config_path = Path(__file__).parent / "config.toml"
+            if not config_path.exists():
+                logger.error(f"TavilySearch 配置文件不存在: {config_path}")
+                return
+
+            with open(config_path, "rb") as f:
+                self.config = tomllib.load(f)
+
+            self.api_keys = [k for k in self.config["tavily"]["api_keys"] if k and not k.startswith("#")]
+            if not self.api_keys:
+                logger.warning("TavilySearch: 未配置有效的 API Key")
+            else:
+                logger.success(f"TavilySearch 已加载，共 {len(self.api_keys)} 个 API Key")
+
+        except Exception as e:
+            logger.error(f"TavilySearch 初始化失败: {e}")
+            self.config = None
+
+    def _get_next_api_key(self) -> str:
+        """轮询获取下一个 API Key"""
+        if not self.api_keys:
+            return ""
+        key = self.api_keys[self.current_key_index]
+        self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
+        return key
+
+    async def _search_tavily(self, query: str) -> Optional[dict]:
+        """调用 Tavily API 进行搜索"""
+        api_key = self._get_next_api_key()
+        if not api_key:
+            logger.error("没有可用的 Tavily API Key")
+            return None
+
+        tavily_config = self.config["tavily"]
+        proxy_config = self.config.get("proxy", {})
+
+        payload = {
+            "api_key": api_key,
+            "query": query,
+            "search_depth": tavily_config.get("search_depth", "basic"),
+            "max_results": tavily_config.get("max_results", 5),
+            "include_raw_content": tavily_config.get("include_raw_content", False),
+            "include_images": tavily_config.get("include_images", False),
+        }
+
+        proxy = None
+        if proxy_config.get("enabled", False):
+            proxy_type = proxy_config.get("type", "http")
+            proxy_host = proxy_config.get("host", "127.0.0.1")
+            proxy_port = proxy_config.get("port", 7890)
+            proxy = f"{proxy_type}://{proxy_host}:{proxy_port}"
+
+        try:
+            import ssl
+            timeout = aiohttp.ClientTimeout(total=30)
+
+            # SSL 配置
+            ssl_config = self.config.get("ssl", {})
+            ssl_verify = ssl_config.get("verify", True)
+
+            connector = None
+            if not ssl_verify:
+                # 跳过 SSL 验证
+                ssl_context = ssl.create_default_context()
+                ssl_context.check_hostname = False
+                ssl_context.verify_mode = ssl.CERT_NONE
+                connector = aiohttp.TCPConnector(ssl=ssl_context)
+
+            async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
+                async with session.post(
+                    "https://api.tavily.com/search",
+                    json=payload,
+                    proxy=proxy
+                ) as resp:
+                    if resp.status == 200:
+                        result = await resp.json()
+                        logger.info(f"Tavily 搜索成功: {query[:30]}...")
+                        logger.info(f"Tavily 原始返回: {result}")
+                        return result
+                    else:
+                        error_text = await resp.text()
+                        logger.error(f"Tavily API 错误: {resp.status}, {error_text}")
+                        return None
+
+        except Exception as e:
+            logger.error(f"Tavily 搜索失败: {e}")
+            return None
+
+    def _format_search_results(self, results: dict) -> str:
+        """格式化搜索结果供 AI 处理"""
+        if not results or "results" not in results:
+            return "未找到相关搜索结果"
+
+        formatted = []
+        for i, item in enumerate(results["results"], 1):
+            title = item.get("title", "无标题")
+            content = item.get("content", "")
+            url = item.get("url", "")
+            formatted.append(f"【结果 {i}】\n标题: {title}\n内容: {content}\n来源: {url}\n")
+
+        return "\n".join(formatted)
+
+    def get_llm_tools(self) -> List[dict]:
+        """返回 LLM 工具定义"""
+        if not self.config or not self.config["behavior"]["enabled"]:
+            return []
+
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "tavily_web_search",
+                    "description": "使用 Tavily 进行联网搜索，获取最新的网络信息。适用于需要查询实时信息、新闻、知识等场景。",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string",
+                                "description": "搜索关键词或问题，建议使用简洁明确的搜索词"
+                            }
+                        },
+                        "required": ["query"]
+                    }
+                }
+            }
+        ]
+
+    async def execute_llm_tool(self, tool_name: str, arguments: dict, bot, from_wxid: str) -> dict:
+        """
+        执行 LLM 工具调用
+
+        只负责搜索，返回结果给 AIChat 的 AI 处理（带上下文和人设）
+        """
+        if tool_name != "tavily_web_search":
+            return None
+
+        if not self.config or not self.config["behavior"]["enabled"]:
+            return {"success": False, "message": "TavilySearch 插件未启用"}
+
+        if not self.api_keys:
+            return {"success": False, "message": "未配置 Tavily API Key"}
+
+        query = arguments.get("query", "")
+        if not query:
+            return {"success": False, "message": "搜索关键词不能为空"}
+
+        try:
+            logger.info(f"开始 Tavily 搜索: {query}")
+
+            # 调用 Tavily 搜索
+            search_results = await self._search_tavily(query)
+            if not search_results:
+                return {"success": False, "message": "搜索失败，请稍后重试"}
+
+            # 格式化搜索结果
+            formatted_results = self._format_search_results(search_results)
+
+            logger.success(f"Tavily 搜索完成: {query[:30]}...")
+
+            # 返回搜索结果，标记需要 AI 继续处理
+            return {
+                "success": True,
+                "message": formatted_results,
+                "need_ai_reply": True  # 标记需要 AI 基于此结果继续回复
+            }
+
+        except Exception as e:
+            logger.error(f"Tavily 搜索执行失败: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return {"success": False, "message": f"搜索失败: {str(e)}"}