Python/耐烧蚀合金智能体/skills/wiki-helper/scripts/search_materials.py

44 lines
1.3 KiB
Python

# 列举导入模块
import argparse
from pathlib import Path
import sys
sys.path.append(Path(__file__).resolve().parent.parent.parent.parent.parent.as_posix())
from utils.request import Request
parser = argparse.ArgumentParser()
parser.add_argument("--question", required=True, type=str, help="检索语句")
args = parser.parse_args()
# 实例请求客户端
request = Request(cache_enabled=True)
response = request.post(
url="http://127.0.0.1:19828/api/v1/projects/current/search", # LLM wiki 提供的搜索接口
headers={
"Content-Type": "application/json",
},
json={
"query": args.question,
"topK": 10,
"includeContent": True, # 包含文档内容
}, # 若客户端启用语义搜索则接口默认开启混合搜索
)
if not response.get("ok") or not response.get("results"):
print("未搜索到相关资料。")
materials = [] # 资料
for i, result in enumerate(response["results"], 1):
# 资料标题
title = result.get("title", f"资料{i}")
# 若资料内容为空则跳过
if not (content := result.get("content", "").strip()):
continue
materials.append(f"[{i}] {title}\n{content}")
try:
print("\n\n---\n\n".join(materials))
except UnicodeEncodeError:
sys.stdout.buffer.write(
("\n\n---\n\n".join(materials) + "\n").encode("utf-8", errors="replace")
)