56 lines
2.2 KiB
Python
56 lines
2.2 KiB
Python
|
|
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import List
|
|
from urllib.parse import quote
|
|
from server.agent.tools.search_tool import rag_search
|
|
from server.chat import utils
|
|
from server.knowledge_base.kb_doc_api import search_docs
|
|
|
|
|
|
def search_pic(query: str) -> List[str]:
|
|
try:
|
|
matches = re.findall(r'\{.*?\}', query)
|
|
if len(matches)>=2:
|
|
query = matches[0]
|
|
else:
|
|
return "<关键指令>不需要再调用该工具了</关键指令>"
|
|
uuid = json.loads(matches[1])["uuid"]
|
|
tip = utils.get_shared_variable(uuid)
|
|
# tip["END"] ="ok"
|
|
temp = {}
|
|
temp = json.loads(query)
|
|
res = search_docs(usr_query=temp["query"],fileName= [],top_k=10,score_threshold=0.9,query=temp["query"], knowledge_base_name="p_meiyupic")
|
|
if len(res)==0 and len(tip["source_docs"])==0:
|
|
utils.set_shared_variable(uuid,tip)
|
|
return "工具没有找到结果"
|
|
# 遍历 res 中的每个元素
|
|
result = ""
|
|
for item in res:
|
|
# 获取 source 的目录部分
|
|
source_dir = os.path.splitext(item.metadata['source'])[0]
|
|
# 获取 page_content
|
|
page_content = item.page_content
|
|
# 拼接字符串
|
|
if item.metadata['source'] in tip["source_docs"]:
|
|
continue
|
|
else:
|
|
tip["source_docs"].append(item.metadata['source'])
|
|
page_content = quote(page_content.replace("http://127.0.0.1:8099/chat_web_backend", "http://127.0.0.1:8099/chat_web_backend"),safe='/:?=&#+')
|
|
result += f'\n'
|
|
|
|
utils.set_shared_variable(uuid,tip)
|
|
if len(result)>0:
|
|
print(f"美术作品链接：{result}")
|
|
return f"注意：以下链接是图片不是参考文献，以下链接不要放到引文小标的格式输出而是以图片格式输出，禁止转义后面链接的编码，这个链接不能带中文。图片如下：{result}"
|
|
else:
|
|
return "<关键指令>不需要再调用该工具了</关键指令>"
|
|
|
|
|
|
|
|
except Exception as e:
|
|
return f"Failed to get picture.{e}"
|
|
|