120 lines
5.4 KiB
Python
120 lines
5.4 KiB
Python
|
import asyncio
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
from typing import AsyncIterable, List, Optional
|
|||
|
|
from fastapi import Body, Request, HTTPException
|
|||
|
|
from fastapi.concurrency import run_in_threadpool
|
|||
|
|
from sse_starlette.sse import EventSourceResponse
|
|||
|
|
from server.knowledge_base.kb_doc_api import search_docs
|
|||
|
|
from configs import (
|
|||
|
|
TEMPERATURE,
|
|||
|
|
VECTOR_SEARCH_TOP_K,
|
|||
|
|
SCORE_THRESHOLD,
|
|||
|
|
LLM_MODELS,
|
|||
|
|
JOURNAL_KNOWLEDGE_BASE,
|
|||
|
|
MAX_TOKENS,
|
|||
|
|
)
|
|||
|
|
from server.utils import BaseResponse, get_prompt_template
|
|||
|
|
from server.chat.policy_fun_iast import get_llm_model_response
|
|||
|
|
from server.knowledge_base.kb_service.base import KBServiceFactory
|
|||
|
|
|
|||
|
|
async def abstract_search(
|
|||
|
|
query: str = Body(..., description="用户输入", examples=["人工智能领域相关的文献有啥"]),
|
|||
|
|
knowledge_base_name_list: List[str] = Body(..., description="知识库名称", examples=[["t_journal_article_bge_v1"]]),
|
|||
|
|
# top_k: int = Body(VECTOR_SEARCH_TOP_K, description="匹配向量数", ge=1),
|
|||
|
|
# fileName: List[str] = Body([], description="文件名称", examples=[["123.txt"]]),
|
|||
|
|
# stream: bool = Body(False, description="流式输出"),
|
|||
|
|
# model_name: str = Body(LLM_MODELS[0], description="LLM 模型名称。"),
|
|||
|
|
# temperature: float = Body(TEMPERATURE, description="LLM 采样温度", ge=0.0, le=1.0),
|
|||
|
|
# max_tokens: Optional[int] = Body(MAX_TOKENS,description="限制LLM生成Token数量，默认None代表模型最大值"),
|
|||
|
|
# score_threshold: float = Body(SCORE_THRESHOLD,
|
|||
|
|
# description=(
|
|||
|
|
# "知识库匹配相关度阈值，取值范围在0-1之间，"
|
|||
|
|
# "SCORE越小，相关度越高，取到1相当于不筛选，建议设置在0.5左右"
|
|||
|
|
# ),
|
|||
|
|
# ge=0.0,
|
|||
|
|
# le=1.0
|
|||
|
|
# ),
|
|||
|
|
# request: Request = None,
|
|||
|
|
# query_rewrite_model_name: str = Body(
|
|||
|
|
# 'Qwen2-72B-Instruct',
|
|||
|
|
# description="Query rewrite model name."
|
|||
|
|
# ),
|
|||
|
|
) -> EventSourceResponse:
|
|||
|
|
|
|||
|
|
async def search_iterator() -> AsyncIterable[str]:
|
|||
|
|
try:
|
|||
|
|
# 重写查询
|
|||
|
|
search_query_response = await run_in_threadpool(
|
|||
|
|
get_llm_model_response,
|
|||
|
|
strategy_name="query rewrite",
|
|||
|
|
llm_model_name=LLM_MODELS[0],
|
|||
|
|
template_prompt_name="query_rewrite",
|
|||
|
|
prompt_param_dict={"query": query, "history": []},
|
|||
|
|
temperature=0.01,
|
|||
|
|
max_tokens=512
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 解析重写后的查询
|
|||
|
|
try:
|
|||
|
|
data = json.loads(search_query_response.strip("```json\n").strip("```"))
|
|||
|
|
queries = data.get('query', [])
|
|||
|
|
search_query = ' '.join(queries)
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
search_query = query
|
|||
|
|
|
|||
|
|
logging.info(f'段落推荐输入问题: {search_query}')
|
|||
|
|
|
|||
|
|
# 执行文档搜索（这里只搜索期刊知识库）
|
|||
|
|
journaldocs = await run_in_threadpool(
|
|||
|
|
search_docs,
|
|||
|
|
query=search_query,
|
|||
|
|
fileName=[],
|
|||
|
|
knowledge_base_name=JOURNAL_KNOWLEDGE_BASE,
|
|||
|
|
top_k=VECTOR_SEARCH_TOP_K,
|
|||
|
|
score_threshold=SCORE_THRESHOLD
|
|||
|
|
)
|
|||
|
|
logging.info(f"段落推荐召回资料: {journaldocs}")
|
|||
|
|
# 处理搜索结果，去重并整理
|
|||
|
|
seen_ids = set()
|
|||
|
|
journallist = []
|
|||
|
|
# journaldocs = []
|
|||
|
|
for doc in journaldocs:
|
|||
|
|
doc_id = doc.metadata.get('ID')
|
|||
|
|
content = doc.metadata.get('abstract', '').strip()
|
|||
|
|
#处理content不是完整段落的情况
|
|||
|
|
if content:
|
|||
|
|
if not content.endswith('。'):
|
|||
|
|
last_period_index = content.rfind('。')
|
|||
|
|
if last_period_index != -1:
|
|||
|
|
content = content[:last_period_index + 1]
|
|||
|
|
else:
|
|||
|
|
content += '...'
|
|||
|
|
|
|||
|
|
if doc_id and content and doc_id not in seen_ids:
|
|||
|
|
journallist.append({
|
|||
|
|
'id': doc_id,
|
|||
|
|
'title': doc.metadata.get('title', ''),
|
|||
|
|
'content': content
|
|||
|
|
})
|
|||
|
|
seen_ids.add(doc_id)
|
|||
|
|
|
|||
|
|
search_results = {'journal_article': journallist}
|
|||
|
|
logging.info(f"段落推荐召回资料整理结果：{journallist}")
|
|||
|
|
|
|||
|
|
# 如果没有找到相关文档，返回默认响应
|
|||
|
|
# if not journallist:
|
|||
|
|
# search_results['journal_article'].append({
|
|||
|
|
# 'id': 'empty',
|
|||
|
|
# 'title': 'empty',
|
|||
|
|
# 'content': '暂未找到相关资料。'
|
|||
|
|
# })
|
|||
|
|
|
|||
|
|
yield json.dumps(search_results, ensure_ascii=False)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
# 捕获并返回错误信息
|
|||
|
|
error_response = BaseResponse(code=500, msg=str(e))
|
|||
|
|
yield json.dumps(error_response.dict(), ensure_ascii=False)
|
|||
|
|
|
|||
|
|
return EventSourceResponse(search_iterator())
|