Files
xiaozhi-esp32-server/main/xiaozhi-server/plugins_func/functions/search_from_ragflow.py
T
2025-11-08 16:00:37 +08:00

99 lines
3.5 KiB
Python

import requests
import sys
from config.logger import setup_logging
from plugins_func.register import register_function, ToolType, ActionResponse, Action
TAG = __name__
logger = setup_logging()
# 定义基础的函数描述模板
SEARCH_FROM_RAGFLOW_FUNCTION_DESC = {
"type": "function",
"function": {
"name": "search_from_ragflow",
"description": "从知识库中查询信息",
"parameters": {
"type": "object",
"properties": {"question": {"type": "string", "description": "查询的问题"}},
"required": ["question"],
},
},
}
@register_function(
"search_from_ragflow", SEARCH_FROM_RAGFLOW_FUNCTION_DESC, ToolType.SYSTEM_CTL
)
def search_from_ragflow(conn, question=None):
# 确保字符串参数正确处理编码
if question and isinstance(question, str):
# 确保问题参数是UTF-8编码的字符串
pass
else:
question = str(question) if question is not None else ""
base_url = conn.config["plugins"]["search_from_ragflow"].get("base_url", "")
api_key = conn.config["plugins"]["search_from_ragflow"].get("api_key", "")
dataset_ids = conn.config["plugins"]["search_from_ragflow"].get("dataset_ids", [])
url = base_url + "/api/v1/retrieval"
headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
# 确保payload中的字符串都是UTF-8编码
payload = {"question": question, "dataset_ids": dataset_ids}
try:
# 使用ensure_ascii=False确保JSON序列化时正确处理中文
response = requests.post(
url,
json=payload,
headers=headers,
timeout=5,
verify=False,
)
# 显式设置响应的编码为utf-8
response.encoding = "utf-8"
response.raise_for_status()
# 先获取文本内容,然后手动处理JSON解码
response_text = response.text
import json
result = json.loads(response_text)
if result.get("code") != 0:
error_detail = response.get("error", {}).get("detail", "")
# 安全地记录错误信息
logger.bind(tag=TAG).error(
"从RAGflow获取信息失败,原因:%s", str(error_detail)
)
return ActionResponse(Action.RESPONSE, None, "RAG接口返回异常")
chunks = result.get("data", {}).get("chunks", [])
contents = []
for chunk in chunks:
content = chunk.get("content", "")
if content:
# 安全地处理内容字符串
if isinstance(content, str):
contents.append(content)
elif isinstance(content, bytes):
contents.append(content.decode("utf-8", errors="replace"))
else:
contents.append(str(content))
# 构建适合大模型的上下文内容(每段前加编号,段间两个换行)
context_text = "\n\n".join(
f"{i+1}. {c.strip()}" for i, c in enumerate(contents[:5])
)
if not context_text:
context_text = "根据知识库查询结果,没有相关信息。"
return ActionResponse(Action.REQLLM, context_text, None)
except Exception as e:
# 使用安全的方式记录异常,避免编码问题
logger.bind(tag=TAG).error("从RAGflow获取信息失败,原因:%s", str(e))
return ActionResponse(Action.RESPONSE, None, "RAG接口返回异常")