427 lines
19 KiB
Python
427 lines
19 KiB
Python
|
||
import json
|
||
import hashlib
|
||
import re
|
||
import time
|
||
import httpx
|
||
from datetime import datetime
|
||
from typing import Any, Dict, Optional
|
||
from config.env import SearchConfig
|
||
from utils.log_util import logger
|
||
from fastapi.responses import StreamingResponse
|
||
from utils.response_util import ResponseUtil
|
||
from module_admin.entity.vo.ragflow_vo import ConverseWithChatAssistantModel
|
||
from module_admin.service.ragflow_service import RAGFlowService
|
||
|
||
def format_sse(data: dict, event: str | None = None) -> str:
|
||
payload = json.dumps(data, ensure_ascii=False)
|
||
prefix = f'event: {event}\n' if event else ''
|
||
return f'{prefix}data: {payload}\n\n'
|
||
|
||
class SearchServiceError(Exception):
|
||
"""搜索服务业务异常"""
|
||
|
||
class SearchService:
|
||
"""搜索服务 - 处理联网搜索请求"""
|
||
|
||
# 基础搜索关键词 - 用于判断是否应该使用搜索服务
|
||
SEARCH_KEYWORDS = (
|
||
'搜索', '查找', '查询', '搜一下', '帮我查', '帮我搜', '检索',
|
||
'最新', '最近', '今日', '今天', '现在', '当前', '昨日', '昨天',
|
||
'天气', '温度', '气温', '降水', '下雨', '下雪', '晴', '多云', '阴天', '雾霾',
|
||
'新闻', '时事', '资讯', '动态', '报道', '消息',
|
||
'发生了什么', '怎么回事', '事件', '事故', '活动',
|
||
'search', 'find', 'query', 'latest', 'recent', 'today', 'now', 'current',
|
||
'weather', 'temperature', 'news', 'events', 'what happened'
|
||
)
|
||
|
||
# 领域特定关键词
|
||
DOMAIN_KEYWORDS = {
|
||
'科技': ('AI', '人工智能', '区块链', '元宇宙', '量子计算', '5G', '芯片'),
|
||
'体育': ('足球', '篮球', '奥运会', '世界杯', 'NBA', '英超', '比分'),
|
||
'娱乐': ('电影', '电视剧', '明星', '演唱会', '综艺', '票房'),
|
||
'财经': ('股票', '基金', '汇率', '利率', '财经', '股市', '行情'),
|
||
'健康': ('疫情', '疫苗', '健康', '医疗', '疾病', '症状'),
|
||
'教育': ('考试', '考研', '高考', '录取', '学校', '教育'),
|
||
}
|
||
|
||
# 否定关键词 - 公司内部相关话题不应触发搜索
|
||
NEGATIVE_KEYWORDS = (
|
||
# 通用否定
|
||
'不要搜索', '不搜索', '无需搜索', '不用搜索', 'don\'t search', 'no search',
|
||
# 公司内部话题（基于 static_qa.json）
|
||
'厂区', '公司', '康达', '预约', '参观', '成立', '业务', '荣誉', '愿景',
|
||
'生产区', '生产基地', '子公司', '分支机构', '管理总部',
|
||
'高分子', '新材料', '胶粘剂', '复合材料',
|
||
'军工', '航空航天', '电子信息', '半导体',
|
||
'制造业', '制造业企业', '高新技术企业', '博士后科研工作站'
|
||
)
|
||
|
||
@classmethod
|
||
def should_handle(cls, question: Optional[str]) -> bool:
|
||
"""判断是否应该使用搜索服务处理该问题"""
|
||
if not question:
|
||
return False
|
||
|
||
normalized = question.strip().lower()
|
||
if any(neg_k in normalized for neg_k in cls.NEGATIVE_KEYWORDS):
|
||
return False
|
||
|
||
if any(k in question for k in cls.SEARCH_KEYWORDS) or any(k in normalized for k in cls.SEARCH_KEYWORDS):
|
||
return True
|
||
|
||
for domain, domain_terms in cls.DOMAIN_KEYWORDS.items():
|
||
has_domain_term = any(term in question for term in domain_terms) or any(term.lower() in normalized for term in domain_terms)
|
||
has_time_or_action = any(k in question for k in ('最新', '最近', '今日', '今天', '搜索', '查找', '查询')) or any(k in normalized for k in ('latest', 'recent', 'today', 'search', 'find', 'query'))
|
||
if has_domain_term and has_time_or_action:
|
||
return True
|
||
|
||
time_patterns = r'(\d{4}年|\d{1,2}月\d{1,2}日|昨天|昨日|今天|今日|最近\d+天|最近\d+周|最近\d+月)'
|
||
if re.search(time_patterns, question) and any(k in question for k in ('发生', '新闻', '事件', '动态')):
|
||
return True
|
||
|
||
return False
|
||
|
||
@staticmethod
|
||
def _simplify_text(text: str) -> str:
|
||
"""简化文本内容"""
|
||
if not text:
|
||
return ""
|
||
text = re.sub(r'\s+', ' ', text.strip())
|
||
text = re.sub(r'[，,、；;]+', '，', text)
|
||
text = re.sub(r'[。.!]+', '。', text)
|
||
max_length = 150
|
||
if len(text) > max_length:
|
||
matches = list(re.finditer(r'[。.!?]', text[:max_length]))
|
||
if matches:
|
||
last_end = matches[-1].end()
|
||
text = text[:last_end] + "..."
|
||
else:
|
||
text = text[:max_length] + "..."
|
||
return text
|
||
|
||
@classmethod
|
||
async def get_search_answer(cls, question: str, redis=None) -> Dict[str, Any]:
|
||
"""获取搜索结果 (Baidu AppBuilder API)"""
|
||
# if not SearchConfig.BAIDU_APPBUILDER_API_KEY:
|
||
# raise SearchServiceError('搜索服务未配置 BAIDU_APPBUILDER_API_KEY')
|
||
|
||
cache_key = None
|
||
if redis:
|
||
cache_key = cls._build_cache_key(question)
|
||
cached = await redis.get(cache_key)
|
||
if cached:
|
||
try:
|
||
logger.info('搜索结果命中缓存: query=%s', question)
|
||
return json.loads(cached)
|
||
except (TypeError, json.JSONDecodeError):
|
||
pass
|
||
|
||
if not SearchConfig.BAIDU_APPBUILDER_API_KEY:
|
||
logger.error("BAIDU_APPBUILDER_API_KEY is missing in configuration")
|
||
raise SearchServiceError('搜索服务未配置 BAIDU_APPBUILDER_API_KEY')
|
||
|
||
try:
|
||
# Baidu API Request Logic
|
||
headers = {
|
||
"X-Appbuilder-Authorization": f"Bearer {SearchConfig.BAIDU_APPBUILDER_API_KEY}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
payload_data = {
|
||
"messages": [{"role": "user", "content": question}],
|
||
"search_source": "baidu_search_v1", # Can be parameterized if needed
|
||
"stream": False,
|
||
# Other optional parameters can be added here
|
||
}
|
||
|
||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||
response = await client.post(
|
||
SearchConfig.SEARCH_API_BASE,
|
||
headers=headers,
|
||
json=payload_data
|
||
)
|
||
|
||
if response.status_code != 200:
|
||
error_msg = f"HTTP {response.status_code}: {response.text}"
|
||
logger.error(f"Baidu API error: {error_msg}")
|
||
raise SearchServiceError(f"Baidu API error: {error_msg}")
|
||
|
||
result_json = response.json()
|
||
# Handle potential non-200 OK business errors if any specific structure exists, though usually status code covers it or there's an error field.
|
||
|
||
# Format Response
|
||
formatted = cls._format_baidu_response(question, result_json)
|
||
|
||
# Cache
|
||
if redis and cache_key:
|
||
try:
|
||
await redis.set(
|
||
cache_key,
|
||
json.dumps(formatted, ensure_ascii=False),
|
||
ex=SearchConfig.SEARCH_CACHE_TTL
|
||
)
|
||
except Exception as ex:
|
||
logger.debug('缓存搜索结果失败: %s', ex)
|
||
|
||
return formatted
|
||
|
||
except Exception as exc:
|
||
logger.exception('Baidu搜索失败')
|
||
raise SearchServiceError(str(exc)) from exc
|
||
|
||
|
||
@classmethod
|
||
async def handle_search_chat(cls, converse_params: ConverseWithChatAssistantModel, redis=None) -> Any:
|
||
"""Handle chat request using Search Service + RAGFlow"""
|
||
start_time = time.perf_counter()
|
||
|
||
try:
|
||
# 1. Check Stream
|
||
if converse_params.stream:
|
||
return StreamingResponse(
|
||
cls._stream_baidu_search(converse_params.question),
|
||
media_type='text/event-stream'
|
||
)
|
||
|
||
# 2. Non-stream: Get Search Answer directly (Bypass RAGFlow)
|
||
search_payload = await cls.get_search_answer(converse_params.question, redis=redis)
|
||
|
||
# Use the "message" or "answer" from Baidu as the final response
|
||
# Format to match RAGFlow response structure if needed, or just return the helpful message
|
||
response_data = {
|
||
"answer": search_payload.get('message', ''),
|
||
"references": search_payload.get('results', [])
|
||
}
|
||
|
||
logger.info('搜索处理耗时 %.3fs', time.perf_counter() - start_time)
|
||
return ResponseUtil.success(data=response_data)
|
||
|
||
except SearchServiceError as exc:
|
||
logger.warning('搜索服务失败，降级到RAGFlow: %s', exc)
|
||
# Fallback to normal flow
|
||
# 修复：直接调用同步方法，不再使用await
|
||
result = RAGFlowService.converse_with_chat_assistant_services(converse_params)
|
||
|
||
if converse_params.stream:
|
||
async def stream_response():
|
||
try:
|
||
async for chunk in result:
|
||
payload = chunk.get('data') if isinstance(chunk, dict) else chunk
|
||
if not payload:
|
||
continue
|
||
body = payload if isinstance(payload, dict) else {'data': payload}
|
||
yield format_sse(body)
|
||
yield format_sse({'status': 'completed'}, event='end')
|
||
except Exception as exc:
|
||
logger.exception('ragflow流式对话异常: %s', exc)
|
||
yield format_sse({'message': str(exc)}, event='error')
|
||
finally:
|
||
logger.info('ragflow流式对话耗时 %.3fs', time.perf_counter() - start_time)
|
||
|
||
return StreamingResponse(stream_response(), media_type='text/event-stream')
|
||
|
||
code = result.get('code', 0)
|
||
if code != 0:
|
||
msg = result.get('message') or result.get('msg') or '接口异常'
|
||
return ResponseUtil.error(msg=msg, data=result.get('data', None))
|
||
return ResponseUtil.success(data=result.get('data', None))
|
||
|
||
@classmethod
|
||
async def _stream_baidu_search(cls, question: str):
|
||
"""Stream Search + Answer using ZhipuAI (GLM-4)"""
|
||
# Note: Method name kept for compatibility, but logic switched to ZhipuAI
|
||
|
||
# 先回一条提示，避免联网搜索阶段用户感知“卡住”
|
||
yield format_sse({'data': '稍等片刻，达达检索调取中'})
|
||
|
||
if not SearchConfig.ZHIPUAI_API_KEY:
|
||
yield format_sse({'message': 'Configuration Error: Missing ZHIPUAI_API_KEY. Please set valid API key.'}, event='error')
|
||
return
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {SearchConfig.ZHIPUAI_API_KEY}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
# ZhipuAI / GLM-4 Payload with Web Search Tool
|
||
payload = {
|
||
"model": "glm-4-flash", # Use Flash model for faster speed
|
||
"messages": [
|
||
{"role": "user", "content": f"请回答用户问题：{question}\n要求：\n1. 务必保持简练，字数严格控制在100字以内！\n2. 直接输出答案。"}
|
||
],
|
||
"tools": [{
|
||
"type": "web_search",
|
||
"web_search": {
|
||
"enable": True,
|
||
"search_query": question,
|
||
"search_result": True,
|
||
"count": 3 # Limit to 3 results for speed (optional, default is 10)
|
||
}
|
||
}],
|
||
"stream": True
|
||
}
|
||
|
||
try:
|
||
logger.info(f"ZhipuAI Search Request: {question}")
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
async with client.stream('POST', SearchConfig.ZHIPUAI_API_BASE, headers=headers, json=payload) as response:
|
||
if response.status_code != 200:
|
||
error_msg = await response.read()
|
||
logger.error(f"ZhipuAI Error: {error_msg}")
|
||
yield format_sse({'message': f"Search Error: {response.status_code}"}, event='error')
|
||
return
|
||
|
||
async for line in response.aiter_lines():
|
||
if not line.strip():
|
||
continue
|
||
data_str = line[6:].strip() if line.startswith('data: ') else line
|
||
if not data_str or data_str == '[DONE]':
|
||
continue
|
||
|
||
try:
|
||
chunk = json.loads(data_str)
|
||
choices = chunk.get('choices', [])
|
||
if choices:
|
||
delta = choices[0].get('delta', {})
|
||
content = delta.get('content', '')
|
||
|
||
# Check for tool calls or references if needed, but GLM-4 usually streams content directly after search
|
||
if content:
|
||
yield format_sse({'data': content})
|
||
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
yield format_sse({'status': 'completed'}, event='end')
|
||
|
||
except Exception as exc:
|
||
logger.exception("ZhipuAI Streaming Failed")
|
||
yield format_sse({'message': str(exc)}, event='error')
|
||
|
||
@classmethod
|
||
async def classify_intent(cls, question: str) -> str:
|
||
"""
|
||
Classify user intent: 'SEARCH' (Public Internet) vs 'RAG' (Internal KB)
|
||
Uses glm-4-flash for <300ms latency.
|
||
"""
|
||
# 0) 确定性兜底：明确公司内部话题一律走 RAG（避免 LLM 误判导致“走搜索”）
|
||
normalized = (question or '').strip().lower()
|
||
if any(neg_k in normalized for neg_k in cls.NEGATIVE_KEYWORDS) or any(k in normalized for k in ('康达新材', '康达')):
|
||
return 'RAG'
|
||
|
||
if not SearchConfig.ZHIPUAI_API_KEY:
|
||
# Fallback to simple keyword check if no key
|
||
return 'SEARCH' if cls.should_handle(question) else 'RAG'
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {SearchConfig.ZHIPUAI_API_KEY}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
prompt = (
|
||
f"判断用户问题意图：\n{question}\n\n"
|
||
"背景信息：我们是'康达新材'（简称'康达'）的内部助手。\n"
|
||
"分类规则：\n"
|
||
"1. RAG (内部知识): \n"
|
||
" - 涉及'康达'、'康达新材'、公司历史、架构、产品、内部规章的问题。\n"
|
||
" - 询问特定的产品手册、错误代码、内部文档。\n"
|
||
"2. SEARCH (公网搜索): \n"
|
||
" - 询问天气、今日新闻、股票大盘、通用百科知识。\n"
|
||
" - 与公司无关的外部实时资讯。\n"
|
||
"只输出分类标签 (SEARCH 或 RAG)，不要其他内容。"
|
||
)
|
||
|
||
payload = {
|
||
"model": "glm-4-flash",
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"temperature": 0.1, # Low temp for deterministic output
|
||
"max_tokens": 10
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=3.0) as client: # Fast timeout
|
||
response = await client.post(
|
||
SearchConfig.ZHIPUAI_API_BASE,
|
||
headers=headers,
|
||
json=payload
|
||
)
|
||
if response.status_code == 200:
|
||
data = response.json()
|
||
choices = data.get('choices', [])
|
||
if choices:
|
||
content = choices[0]['message']['content'].strip().upper()
|
||
# Simple heuristic cleanup
|
||
if 'SEARCH' in content:
|
||
return 'SEARCH'
|
||
return 'RAG'
|
||
except Exception:
|
||
pass # Fail safe to RAG (or logic below)
|
||
|
||
# Fallback to keyword match if LLM fails
|
||
return 'SEARCH' if cls.should_handle(question) else 'RAG'
|
||
|
||
@classmethod
|
||
def _format_baidu_response(cls, query: str, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""Format Baidu AppBuilder response to match expected structure"""
|
||
# The Baidu response structure based on docs example:
|
||
# { "choices": [ { "message": { "content": "..." } } ], "references": [ { "content": "...", "title": "...", "url": "..." } ] }
|
||
|
||
choices = payload.get('choices', [])
|
||
message_content = ""
|
||
if choices:
|
||
message_content = choices[0].get('message', {}).get('content', '')
|
||
|
||
references = payload.get('references', [])
|
||
|
||
results = []
|
||
for ref in references:
|
||
results.append({
|
||
'title': ref.get('title', ''),
|
||
'link': ref.get('url', ''),
|
||
'snippet': ref.get('content', ''), # Use content as snippet
|
||
'source': 'baidu'
|
||
})
|
||
|
||
# Build context for LLM (similar to original logic)
|
||
context_parts = []
|
||
if message_content:
|
||
context_parts.append(f"智能总结：{message_content}")
|
||
|
||
if results:
|
||
context_parts.append("\n相关信息：")
|
||
for idx, item in enumerate(results[:5], 1): # Limit to 5
|
||
title = item['title']
|
||
snippet = cls._simplify_text(item['snippet'])
|
||
context_parts.append(f"{idx}. {title}\n {snippet}")
|
||
|
||
context_for_llm = "\n".join(context_parts)
|
||
|
||
# Message for display (similar to original logic)
|
||
message_parts = [f'关于"{query}"的搜索结果：\n']
|
||
if message_content:
|
||
message_parts.append(f"【智能回答】{message_content}\n")
|
||
|
||
if results:
|
||
message_parts.append('\n【参考来源】')
|
||
for idx, item in enumerate(results[:3], 1):
|
||
title = item['title']
|
||
snippet = cls._simplify_text(item['snippet'])
|
||
message_parts.append(f'\n{idx}. {title}')
|
||
else:
|
||
message_parts.append('\n未找到相关结果。')
|
||
|
||
return {
|
||
'type': 'search',
|
||
'query': query,
|
||
'results': results, # Normalized results list
|
||
'message': '\n'.join(message_parts), # Display friendly message
|
||
'context_for_llm': context_for_llm, # Context for RAG extraction
|
||
'raw': payload
|
||
}
|
||
|
||
@staticmethod
|
||
def _build_cache_key(query: str) -> str:
|
||
query_hash = hashlib.md5(query.encode('utf-8')).hexdigest()[:16]
|
||
today = datetime.utcnow().strftime('%Y%m%d')
|
||
return f'search:baidu:{query_hash}:{today}'
|