320 lines
12 KiB
Python
320 lines
12 KiB
Python
from datetime import datetime
|
||
import hashlib
|
||
import re
|
||
from fastapi import Body, Request
|
||
from sse_starlette.sse import EventSourceResponse
|
||
from fastapi.concurrency import run_in_threadpool
|
||
from configs import (
|
||
LLM_MODELS,
|
||
VECTOR_SEARCH_TOP_K,
|
||
SCORE_THRESHOLD,
|
||
SELF_TOP_K,
|
||
SELF_SCORE_THRESHOLD,
|
||
TEMPERATURE,
|
||
SELF_TEMPERATURE,
|
||
USE_RERANKER,
|
||
RERANKER_MODEL,
|
||
RERANKER_MAX_LENGTH,
|
||
SELF_MAX_TOKENS,
|
||
SELF_USE_RERANKER,
|
||
MODEL_PATH
|
||
)
|
||
from server.chat.ZhipuSearchAPI import ZhipuSearchAPIWrapper
|
||
from server.chat.policy_fun_iast import get_llm_model_response
|
||
from server.knowledge_base.model.kb_document_model import DocumentWithVSId
|
||
from server.utils import wrap_done, get_ChatOpenAI
|
||
from server.utils import BaseResponse, get_prompt_template
|
||
from langchain.chains import LLMChain
|
||
from langchain.callbacks import AsyncIteratorCallbackHandler
|
||
from typing import AsyncIterable, List, Optional, Dict, Tuple, Union, AsyncGenerator
|
||
import asyncio
|
||
from langchain.prompts.chat import ChatPromptTemplate
|
||
from server.chat.utils import History, get_similar_documents, get_text_by_regex
|
||
from server.knowledge_base.kb_service.base import KBServiceFactory
|
||
import json
|
||
from server.knowledge_base.kb_doc_api import search_self_docs
|
||
from configs.basic_config import *
|
||
from langchain.chains.question_answering import load_qa_chain
|
||
from typing import Generator
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
async def word_explain(
|
||
query: str = Body(..., description="用户输入", examples=["智慧科协"]),
|
||
fileNames: List = Body([], description="文件名称", examples=[["孟庆海同志在“智慧科协2.0”5·30场景建设工作部署会议上的讲话.docx"]]),
|
||
knowledge_base_name_list: List[str] = Body(..., description="知识库列表", examples=[["p_cast0101011"]]),
|
||
history: List[History] = Body(
|
||
[],
|
||
description="历史对话",
|
||
examples=[[
|
||
{"role": "user", "content": "我们来玩成语接龙，我先来，生龙活虎"},
|
||
{"role": "assistant", "content": "虎头虎脑"}]
|
||
]
|
||
),
|
||
# stream: bool = Body(False, description="流式输出"),
|
||
) -> EventSourceResponse:
|
||
"""
|
||
名词解释api\n
|
||
入参：\n
|
||
query: str 用户输入\n
|
||
fileNames: List[str] 文件名称\n
|
||
knowledge_base_name_list: List[str] 知识库列表\n
|
||
history: List[History] 历史对话\n
|
||
"""
|
||
logger.info(
|
||
f"名词解释请求: query={query}, "
|
||
f"文件={fileNames}, "
|
||
f"知识库数={knowledge_base_name_list}, "
|
||
f"历史记录数={len(history)}, "
|
||
# f"流式模式={stream}"
|
||
)
|
||
|
||
try:
|
||
# 知识库存在性校验
|
||
for kb_name in knowledge_base_name_list:
|
||
if not KBServiceFactory.get_service_by_name(kb_name):
|
||
return BaseResponse(code=404, msg=f"知识库 {kb_name} 不存在")
|
||
|
||
# 名词检查
|
||
check_result = await run_in_threadpool(
|
||
get_llm_model_response,
|
||
strategy_name="word_check",
|
||
llm_model_name=LLM_MODELS[0],
|
||
template_prompt_name="word_check",
|
||
prompt_param_dict={"query": query},
|
||
temperature=0.01,
|
||
max_tokens=512
|
||
)
|
||
logger.info(f"名词检查结果: {check_result}")
|
||
docs = []
|
||
internet_context = ""
|
||
internet_doc = []
|
||
if '0' in check_result:
|
||
# logger.warning(f"非名词请求被拒绝: {query}")
|
||
# return BaseResponse(code=500, msg="请选择需要解释的单独名词。")
|
||
prompt_name = "word_explain_reject"
|
||
else:
|
||
prompt_name = "word_explain"
|
||
# 并发获取文档和网络结果
|
||
docs, internet_context, internet_doc = await retrieve_documents(
|
||
query=query,
|
||
fileNames=fileNames,
|
||
knowledge_base_names=knowledge_base_name_list,
|
||
top_k=SELF_TOP_K,
|
||
score_threshold=SELF_SCORE_THRESHOLD
|
||
)
|
||
|
||
return EventSourceResponse(
|
||
generate_chat_response(
|
||
query=query,
|
||
docs=docs,
|
||
internet_context=internet_context,
|
||
internet_docs=internet_doc,
|
||
fileNames=fileNames,
|
||
# stream=stream,
|
||
model_name=LLM_MODELS[0],
|
||
prompt_name=prompt_name
|
||
)
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"处理请求异常: {str(e)}", exc_info=True)
|
||
return BaseResponse(code=500, msg=f"处理请求时发生错误: {str(e)}")
|
||
|
||
|
||
async def retrieve_documents(
|
||
query: str,
|
||
fileNames: List[str],
|
||
knowledge_base_names: List[str],
|
||
top_k: int,
|
||
score_threshold: float
|
||
) -> Tuple[List[DocumentWithVSId], str, List[str]]:
|
||
"""
|
||
并发检索文档和网络信息
|
||
返回格式: (文档列表, 网络上下文, 名称+超链接格式列表)
|
||
"""
|
||
async def fetch_kb_docs(kb_name: str) -> List[DocumentWithVSId]:
|
||
"""获取知识库匹配结果"""
|
||
try:
|
||
return await run_in_threadpool(
|
||
search_self_docs,
|
||
query=query,
|
||
fileNames=fileNames,
|
||
knowledge_base_name=kb_name,
|
||
top_k=top_k,
|
||
score_threshold=score_threshold
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"知识库 {kb_name} 检索失败: {str(e)}")
|
||
return []
|
||
# TODO 更换duckduckgo
|
||
async def fetch_web_results() -> List[dict]:
|
||
"""获取网络搜索结果"""
|
||
try:
|
||
wrapper = ZhipuSearchAPIWrapper()
|
||
return await run_in_threadpool(
|
||
wrapper.zhipu_search,
|
||
origin_query=query
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"网络检索失败: {str(e)}")
|
||
return []
|
||
|
||
try:
|
||
# 并发执行所有检索任务
|
||
tasks = [
|
||
*[fetch_kb_docs(kb) for kb in knowledge_base_names],
|
||
fetch_web_results()
|
||
]
|
||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
# 处理结果
|
||
self_docs: List[DocumentWithVSId] = []
|
||
internet_context = ""
|
||
internet_docs = []
|
||
|
||
for result in results:
|
||
if isinstance(result, Exception):
|
||
# 记录并发任务中的异常
|
||
logger.error(f"并发任务出现异常: {str(result)}")
|
||
continue
|
||
elif isinstance(result, list) and all(isinstance(doc, DocumentWithVSId) for doc in result):
|
||
self_docs.extend(result)
|
||
elif isinstance(result, list) and all(isinstance(item, dict) for item in result):
|
||
# 相似度过滤互联网检索结果
|
||
if len(results[1]) > 0:
|
||
try:
|
||
sentences = [doc["title"] for doc in results[1]]
|
||
sentences_page_content = [str(i+1)+":【"+doc["title"]+doc["content"]+"】" for i,doc in enumerate(results[1])]
|
||
except Exception as e:
|
||
sentences = [doc["source"] for doc in results[1]]
|
||
sentences_page_content = [str(i+1)+":【"+doc["source"]+doc["content"]+"】" for i,doc in enumerate(results[1])]
|
||
res = get_llm_model_response(
|
||
strategy_name="default_similar",
|
||
llm_model_name=LLM_MODELS[0],
|
||
template_prompt_name="default_similar",
|
||
prompt_param_dict={"input": query, "title": f"{sentences_page_content}", "time": datetime.now().strftime("%Y%m%d")},
|
||
temperature=0.01,
|
||
max_tokens=512
|
||
)
|
||
try:
|
||
index =[]
|
||
if res == "无":
|
||
index = []
|
||
else:
|
||
index = res.split(",")
|
||
index = [int(i)-1 for i in index]
|
||
similar_document = get_similar_documents(index=index,sentences=sentences,query=query, docs=results[1], top_k=top_k)
|
||
except Exception as e:
|
||
print(e)
|
||
similar_document = get_similar_documents(index=[],sentences=sentences,query=query, docs=results[1], top_k=top_k)
|
||
# 提取所有 content 字段并拼接
|
||
for num,item in enumerate(similar_document, start=1):
|
||
content = item.get('content', '')
|
||
title = item.get('title', '')
|
||
cleaned_title = re.sub(r'（发布时间：.*?）', '', title)
|
||
link = item.get('link', '')
|
||
internet_docs.append(f"[{num}] [{cleaned_title}]({link})")
|
||
internet_context += content
|
||
|
||
return self_docs, internet_context, internet_docs
|
||
|
||
except Exception as e:
|
||
logger.error(f"文档检索异常: {str(e)}")
|
||
return [], "", []
|
||
|
||
|
||
def generate_chat_response(
|
||
query: str,
|
||
docs: List[Dict],
|
||
internet_context: str,
|
||
fileNames: List[str],
|
||
internet_docs: List[str],
|
||
# stream: bool,
|
||
model_name: str,
|
||
prompt_name: str
|
||
) -> Union[EventSourceResponse, BaseResponse]:
|
||
"""
|
||
生成聊天响应，整合本地文档和网络上下文
|
||
"""
|
||
callback = AsyncIteratorCallbackHandler()
|
||
model = get_ChatOpenAI(
|
||
model_name=model_name,
|
||
temperature=SELF_TEMPERATURE,
|
||
max_tokens=SELF_MAX_TOKENS,
|
||
callbacks=[callback],
|
||
)
|
||
|
||
# 构建提示模板
|
||
prompt_template = get_prompt_template("knowledge_base_chat", prompt_name)
|
||
chat_prompt = ChatPromptTemplate.from_messages([
|
||
History(role="system", content=prompt_template).to_msg_template(False)
|
||
])
|
||
|
||
# 配置问答链
|
||
qa_chain = load_qa_chain(
|
||
model,
|
||
chain_type="stuff",
|
||
prompt=chat_prompt,
|
||
verbose=True
|
||
)
|
||
|
||
async def generation_wrapper() -> AsyncGenerator[str, None]:
|
||
try:
|
||
# 预处理文件名
|
||
clean_files = [re.sub(r"\.\w+$", "", f) for f in fileNames]
|
||
|
||
# 构建上下文数据
|
||
context_data = {
|
||
"input_documents": docs,
|
||
"question": query,
|
||
"fileName": ", ".join(clean_files),
|
||
"internet_context": internet_context
|
||
}
|
||
|
||
task = asyncio.create_task(wrap_done(
|
||
qa_chain.acall(context_data),
|
||
callback.done
|
||
))
|
||
|
||
# 流式输出处理
|
||
async for token in callback.aiter():
|
||
yield json.dumps({"text": token}, ensure_ascii=False)
|
||
|
||
# 添加来源文档
|
||
if docs or internet_docs:
|
||
# response = {}
|
||
# source_documents = "\n".join(internet_docs)
|
||
yield json.dumps({"docs": internet_docs}, ensure_ascii=False)
|
||
|
||
await task
|
||
except Exception as e:
|
||
logger.error(f"生成响应异常: {str(e)}")
|
||
yield json.dumps({"error": "生成响应时发生错误"}, ensure_ascii=False)
|
||
|
||
return generation_wrapper()
|
||
|
||
|
||
def process_source_documents(docs: List[Dict]) -> List[str]:
|
||
"""
|
||
处理并格式化来源文档信息
|
||
"""
|
||
seen = set()
|
||
processed = []
|
||
|
||
for idx, doc in enumerate(docs, 1):
|
||
# 提取文档摘要
|
||
summary = doc.metadata.get("summary", "")
|
||
text = get_text_by_regex(summary) or doc.page_content[:200]
|
||
|
||
# 生成唯一标识
|
||
content_hash = hashlib.md5(text.encode()).hexdigest()
|
||
if content_hash in seen:
|
||
continue
|
||
|
||
seen.add(content_hash)
|
||
source = doc.metadata.get("source", "未知来源")
|
||
processed.append(f"[{idx}] {text[:10]}...（来源：{source}）")
|
||
|
||
logger.info(f"处理后的引用文档数: {len(processed)}")
|
||
return processed |