120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
import os
|
|
import urllib.request
|
|
import json
|
|
from typing import Dict
|
|
from datetime import datetime
|
|
from .base import ChatCompletionRequest
|
|
|
|
def get_current_time_info() -> str:
|
|
"""获取当前时间信息"""
|
|
now = datetime.now()
|
|
weekdays = ["星期一", "星期二", "星期三", "星期四", "星期五", "星期六", "星期日"]
|
|
return f"{now.strftime('%Y年%m月%d日 %H:%M:%S')} {weekdays[now.weekday()]}"
|
|
|
|
def get_web_search_mode(request: ChatCompletionRequest) -> str:
|
|
"""获取联网搜索模式"""
|
|
if getattr(request, 'deep_search', False):
|
|
return "deep"
|
|
elif getattr(request, 'web_search', False):
|
|
return "simple"
|
|
return ""
|
|
|
|
def execute_tavily_search(query: str, mode: str = "simple") -> str:
|
|
"""真实调用 Tavily 搜索 API"""
|
|
api_key = os.getenv("TAVILY_API_KEY")
|
|
if not api_key:
|
|
return "本地环境变量 TAVILY_API_KEY 未配置，无法进行搜索。"
|
|
|
|
url = "https://api.tavily.com/search"
|
|
headers = {"Content-Type": "application/json"}
|
|
data = {
|
|
"api_key": api_key,
|
|
"query": query,
|
|
"search_depth": "advanced" if mode == "deep" else "basic",
|
|
"include_answer": False,
|
|
"max_results": 5 if mode == "deep" else 3
|
|
}
|
|
|
|
req = urllib.request.Request(url, data=json.dumps(data).encode('utf-8'), headers=headers, method='POST')
|
|
try:
|
|
with urllib.request.urlopen(req) as response:
|
|
result = json.loads(response.read().decode('utf-8'))
|
|
results = result.get("results", [])
|
|
|
|
if not results:
|
|
return "搜索未返回结果。"
|
|
|
|
formatted_res = []
|
|
for i, res in enumerate(results):
|
|
formatted_res.append(f"[{i+1}] {res.get('title')}\n{res.get('content')}\n链接: {res.get('url')}")
|
|
|
|
return "\n\n".join(formatted_res)
|
|
except Exception as e:
|
|
return f"搜索请求失败，错误: {str(e)}"
|
|
|
|
|
|
|
|
def build_openai_search_tool(mode: str) -> Dict:
|
|
"""
|
|
构建兼容型联网搜索插件工具结构 (供 DeepSeek / OpenAI SDK 使用)
|
|
注意：此类提供标准的 Tool/Function Function calling 模板。
|
|
深度思考通常结合内置联网或者其他外挂流程实现。
|
|
"""
|
|
if mode == "deep":
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "web_search",
|
|
"description": "深度互联网搜索插件(查找并阅读网页内容)",
|
|
"parameters": {"type": "object", "properties": {"query": {"type": "string"}}},
|
|
}
|
|
}
|
|
else:
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "web_search",
|
|
"description": "进行互联网搜索并获取实时信息或资料以辅助回答。",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "要搜索的准确关键词或短语"
|
|
}
|
|
},
|
|
"required": ["query"]
|
|
},
|
|
}
|
|
}
|
|
|
|
def build_glm_search_tool(mode: str) -> Dict:
|
|
"""构建 GLM 联网搜索工具"""
|
|
today = get_current_time_info()
|
|
|
|
if mode == "deep":
|
|
# 深度搜索：返回搜索结果详情
|
|
return {
|
|
"type": "web_search",
|
|
"web_search": {
|
|
"enable": True,
|
|
"search_engine": "search_pro",
|
|
"search_result": True,
|
|
"search_prompt": f"你是一位智能助手。请用简洁的语言总结网络搜索{{search_result}}中的关键信息，按重要性排序并引用来源日期。今天的日期是{today}。",
|
|
"count": 5,
|
|
"search_recency_filter": "noLimit",
|
|
"content_size": "high",
|
|
},
|
|
}
|
|
else:
|
|
# 简单搜索
|
|
return {
|
|
"type": "web_search",
|
|
"web_search": {
|
|
"enable": True,
|
|
"search_engine": "search_pro",
|
|
"search_result": True,
|
|
"count": 5,
|
|
},
|
|
}
|