108 lines
5.6 KiB
Python
108 lines
5.6 KiB
Python
|
import json
|
|||
|
|
import re
|
|||
|
|
import concurrent
|
|||
|
|
from fastapi.concurrency import run_in_threadpool
|
|||
|
|
from langchain.tools import YouTubeSearchTool
|
|||
|
|
from pydantic import BaseModel, Field
|
|||
|
|
from server.chat import utils
|
|||
|
|
from server.chat.utils import doc_to_list, solve_knowledge_map,solve_mental_data,shared_variable
|
|||
|
|
from server.knowledge_base.kb_doc_api import search_docs
|
|||
|
|
from server.knowledge_base.model.kb_document_model import DocumentWithVSId
|
|||
|
|
from configs import kb_config
|
|||
|
|
|
|||
|
|
|
|||
|
|
def rag_search1(query: str):
|
|||
|
|
"""
|
|||
|
|
根据用户输入的query，返回rag搜索结果
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
matches = re.findall(r'\{.*?\}', query)
|
|||
|
|
if len(matches)>=2:
|
|||
|
|
query = matches[0]
|
|||
|
|
else:
|
|||
|
|
return "<关键指令>不需要再调用该工具了</关键指令>"
|
|||
|
|
time_based_uuid = json.loads(matches[1])["uuid"]
|
|||
|
|
search = json.loads(query)
|
|||
|
|
search_query = search["query"]
|
|||
|
|
search_keywords = []
|
|||
|
|
search_text = f"{search_query}"
|
|||
|
|
if type(search["keywords"]) == list:
|
|||
|
|
search_keywords = search["keywords"]
|
|||
|
|
for keyword in search_keywords:
|
|||
|
|
search_text += f" {keyword}"
|
|||
|
|
else:
|
|||
|
|
search_keywords = search["keywords"].split(",")
|
|||
|
|
for keyword in search_keywords:
|
|||
|
|
search_text += f" {keyword}"
|
|||
|
|
result = []
|
|||
|
|
source_docs = {}
|
|||
|
|
knownledge_name = []
|
|||
|
|
if type(search["knowledge_name"]) == list:
|
|||
|
|
knownledge_name=search["knowledge_name"]
|
|||
|
|
else:
|
|||
|
|
knownledge_name=search["knowledge_name"].split(",")
|
|||
|
|
for knownledge in knownledge_name:
|
|||
|
|
if not knownledge in kb_config.CH_BASE_NAME:
|
|||
|
|
knownledge_name.remove(knownledge)
|
|||
|
|
if len(knownledge_name)==0:
|
|||
|
|
result.append(f"没有找到匹配的知识库,请必须更换联网思索搜索更多知识库内容")
|
|||
|
|
return result
|
|||
|
|
# knownledge_name=kb_config.CH_BASE_NAME
|
|||
|
|
|
|||
|
|
knownledge_name=solve_knowledge_map(knownledge_name)
|
|||
|
|
num = 0
|
|||
|
|
for knownledge in knownledge_name:
|
|||
|
|
source_docs[knownledge] = []
|
|||
|
|
seen_docs = set()
|
|||
|
|
duplicate_indices = []
|
|||
|
|
doc_list = search_docs(usr_query=search_text,fileName= [],top_k=5,score_threshold=0.9,query=search_text, knowledge_base_name=knownledge)
|
|||
|
|
|
|||
|
|
|
|||
|
|
for inum,doc in enumerate(doc_list):
|
|||
|
|
solve_mental_data(knownledge,doc_list,doc=doc,seen_docs=seen_docs,duplicate_indices=duplicate_indices,knowledge=result,inum=inum)
|
|||
|
|
|
|||
|
|
# 从policydocs中删除重复的文档（从后往前删除以防止索引错位）
|
|||
|
|
for index in sorted(duplicate_indices, reverse=True):
|
|||
|
|
del doc_list[index]
|
|||
|
|
# 处理原文来源进入数组。使用开关语句明确各个条件分支
|
|||
|
|
match knownledge:
|
|||
|
|
# 属于政策库分支，入参为中文政策库名称
|
|||
|
|
case kb_config.DEFAULT_POLICY_BASE:
|
|||
|
|
doc_to_list(num,kb_config.DEFAULT_POLICY_BASE_NAME,doc_list,source_docs[knownledge])
|
|||
|
|
# 属于期刊论文库分支，入参为期刊论文库的中文名称
|
|||
|
|
case kb_config.DEFAULT_JOURNAL_BASE:
|
|||
|
|
doc_to_list(num,kb_config.DEFAULT_JOURNAL_BASE_NAME,doc_list,source_docs[knownledge])
|
|||
|
|
# 属于报告库分支，入参为报告库中文名称
|
|||
|
|
case kb_config.DEFAULT_REPORT_BASE:
|
|||
|
|
doc_to_list(num,kb_config.DEFAULT_REPORT_BASE_NAME,doc_list,source_docs[knownledge])
|
|||
|
|
# 属于冶金行业新闻库分支，入参为冶金行业新闻库中文名称
|
|||
|
|
case kb_config.GY_NEWS_BASE:
|
|||
|
|
doc_to_list(num,kb_config.GY_NEWS_BASE_NAME,doc_list,source_docs[knownledge])
|
|||
|
|
# 属于冶金行业报告库分支，入参为冶金行业报告库中文名称
|
|||
|
|
case kb_config.GY_REPORT_BASE:
|
|||
|
|
doc_to_list(num,kb_config.GY_REPORT_BASE_NAME,doc_list,source_docs[knownledge])
|
|||
|
|
# 属于冶金专业知识库分支，入参为冶金专业知识库中文名称
|
|||
|
|
case kb_config.GY_JOURNAL_BASE:
|
|||
|
|
doc_to_list(num,kb_config.GY_JOURNAL_BASE_NAME,doc_list,source_docs[knownledge])
|
|||
|
|
case _:
|
|||
|
|
print(f"输入了没有的知识库名称")
|
|||
|
|
return("输入了没有的知识库名称")
|
|||
|
|
num += len(source_docs[knownledge])
|
|||
|
|
# 构建缓存对象用于h'per_query_cache'用于传递给其他方法使用uuid获取本轮对话的文献来源
|
|||
|
|
del num
|
|||
|
|
source = utils.get_shared_variable(time_based_uuid)
|
|||
|
|
print(utils.get_shared_variable(time_based_uuid))
|
|||
|
|
source["source_docs"]=source_docs
|
|||
|
|
utils.set_shared_variable(time_based_uuid,source)
|
|||
|
|
if 0<len(result)<3:
|
|||
|
|
return f"当前资料：{result}\n<关键指令>搜索结果较少，更换知识库或联网思索重新搜索！！！</关键指令>"
|
|||
|
|
if len(result)==0:
|
|||
|
|
return "注意：【指令：更换知识库或联网思索继续搜索！！！】"
|
|||
|
|
except:
|
|||
|
|
# return "入参格式需为{\"knowledge_name\":[\"XXX\",\"XXX\"],\"query\":\"XXX\",\"keywords\":[\"XXX\", \"XXX\", \"XXX\", \"XXX\"]}检查输入参数如果没有缺少必要值，当前工具异常请换其他工具"
|
|||
|
|
return "当前工具异常！请换其他工具"
|
|||
|
|
|
|||
|
|
return f"当前资料：{result}\n<关键指令>总结此内容！！！</关键指令>"
|
|||
|
|
|
|||
|
|
class RagSearchInput(BaseModel):
|
|||
|
|
query: str = Field(...,description="查询对象")
|