275 lines
8.1 KiB
Python
275 lines
8.1 KiB
Python
"""
|
||
RAG搜索工具模块
|
||
包含知识库搜索、文件处理等功能
|
||
"""
|
||
import os
|
||
import json
|
||
import base64
|
||
import asyncio
|
||
import httpx
|
||
from typing import Optional, Dict, Any
|
||
from langchain_core.tools import tool
|
||
|
||
from config import RAG_CONFIG
|
||
from utils.function_tracker import track_function_calls
|
||
|
||
|
||
@track_function_calls
|
||
async def rag_search(
|
||
query: str,
|
||
top_k: Optional[int] = 5,
|
||
file_id: Optional[int] = None,
|
||
kb_id: Optional[str] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
知识库搜索工具
|
||
调用知识库搜索API，检索与query相关的文档内容
|
||
|
||
Args:
|
||
query: 搜索关键词（必填）
|
||
top_k: 返回结果数量（默认5）
|
||
file_id: 文件ID（可选）
|
||
kb_id: 知识库ID（可选）
|
||
|
||
Returns:
|
||
搜索结果，包含sourceCitation格式的数据
|
||
"""
|
||
|
||
if not query or not query.strip():
|
||
return {"success": False, "error": "查询文本不能为空"}
|
||
|
||
SEARCH_URL = RAG_CONFIG["search_endpoint"]
|
||
|
||
params = {"query": query.strip()}
|
||
|
||
if file_id is None:
|
||
config_file_id = RAG_CONFIG.get("file-id")
|
||
if config_file_id is not None:
|
||
try:
|
||
file_id = int(config_file_id) if isinstance(config_file_id, str) else config_file_id
|
||
except (ValueError, TypeError):
|
||
file_id = None
|
||
|
||
if kb_id is None:
|
||
config_kb_id = RAG_CONFIG.get("kb-id")
|
||
if config_kb_id is not None:
|
||
kb_id = str(config_kb_id) if config_kb_id is not None else None
|
||
|
||
if file_id is not None and file_id != "":
|
||
params["file_id"] = file_id
|
||
if kb_id is not None and kb_id != "":
|
||
params["kb_id"] = kb_id
|
||
|
||
print(111111111111111111)
|
||
print(RAG_CONFIG["x-user-id"])
|
||
print(RAG_CONFIG["x-user-name"])
|
||
print(RAG_CONFIG["x-role"])
|
||
|
||
headers = {
|
||
"accept": "application/json",
|
||
"x-user-id": RAG_CONFIG["x-user-id"],
|
||
"x-role": RAG_CONFIG["x-role"]
|
||
}
|
||
|
||
print(f"RAG params: {params}")
|
||
|
||
max_retries = 1
|
||
retry_delay = 2
|
||
|
||
response = None
|
||
last_exception = None
|
||
|
||
for attempt in range(max_retries + 1):
|
||
try:
|
||
async with httpx.AsyncClient(timeout=180.0, verify=False) as client:
|
||
response = await client.get(SEARCH_URL, params=params, headers=headers)
|
||
break
|
||
|
||
except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.RequestError) as e:
|
||
last_exception = e
|
||
if attempt < max_retries:
|
||
print(f"RAG请求超时，第 {attempt+1} 次重试，2秒后重新请求...")
|
||
await asyncio.sleep(retry_delay)
|
||
else:
|
||
return {"success": False, "error": f"请求超时: {str(e)}"}
|
||
|
||
try:
|
||
if response.status_code != 200:
|
||
return {"success": False, "error": f"HTTP {response.status_code}"}
|
||
|
||
result = response.json()
|
||
data = result.get("results", [])
|
||
|
||
filtered_data = [item for item in data if item.get("score", 0.0) >= 0.4]
|
||
|
||
sorted_data = sorted(
|
||
filtered_data,
|
||
key=lambda x: x.get("score", 0.0),
|
||
reverse=True
|
||
)[:top_k]
|
||
|
||
source_citation = {}
|
||
|
||
for idx, item in enumerate(sorted_data, start=1):
|
||
if not isinstance(item, dict):
|
||
continue
|
||
|
||
file_info = item.get("file")
|
||
kb_info = item.get("kb")
|
||
|
||
filename = file_info.get("filename", "") if file_info else ""
|
||
resource = filename or (kb_info.get("name", "未分组文档") if kb_info else "未分组文档")
|
||
|
||
positions = item.get("positions", [])
|
||
page_idx = positions[0].get("page_idx", 0) if positions else 0
|
||
|
||
citation_item = {
|
||
"index": idx,
|
||
"id": item.get("id", ""),
|
||
"text": item.get("content", ""),
|
||
"resource": resource,
|
||
"filename": filename,
|
||
"score": item.get("score", 0.0),
|
||
"page_idx": page_idx,
|
||
"file_id": item.get("file_id", ""),
|
||
"kb_id": kb_info.get("id", "") if kb_info else "",
|
||
"kb_name": kb_info.get("name", "") if kb_info else "",
|
||
"positions": positions
|
||
}
|
||
|
||
if resource not in source_citation:
|
||
source_citation[resource] = []
|
||
source_citation[resource].append(citation_item)
|
||
|
||
return {
|
||
"success": True,
|
||
"sourceCitation": source_citation,
|
||
"total": len(sorted_data),
|
||
"total_available": len(data)
|
||
}
|
||
|
||
except Exception as e:
|
||
return {"success": False, "error": str(e)}
|
||
|
||
|
||
@tool
|
||
async def rag_search_tool(query: str, top_k: Optional[int] = 5) -> str:
|
||
"""
|
||
RAG检索工具（供工作流调用）
|
||
在知识库中检索与query相关的文档内容
|
||
|
||
Args:
|
||
query: 检索查询
|
||
top_k: 返回结果数量
|
||
|
||
Returns:
|
||
检索结果文本
|
||
"""
|
||
result = await rag_search(query, top_k=top_k)
|
||
|
||
if not result.get("success", False):
|
||
return ""
|
||
|
||
source_citation = result.get("sourceCitation", {})
|
||
if not source_citation:
|
||
return ""
|
||
|
||
results_parts = []
|
||
for resource, items in source_citation.items():
|
||
for item in items:
|
||
text = item.get("text", "")
|
||
if text:
|
||
results_parts.append(f"【{resource}】{text}")
|
||
|
||
results_text = "\n\n".join(results_parts) if results_parts else ""
|
||
if len(results_text) > 2000:
|
||
results_text = results_text[:2000] + "..."
|
||
return results_text
|
||
|
||
|
||
@track_function_calls
|
||
async def file_to_text(file_path: str) -> Dict[str, Any]:
|
||
"""
|
||
文件转文本工具
|
||
支持多种文件格式：PDF、Word、TXT、图片等
|
||
|
||
Args:
|
||
file_path: 文件路径
|
||
|
||
Returns:
|
||
转换结果，包含文本内容
|
||
"""
|
||
if not file_path or not file_path.strip():
|
||
return {"success": False, "error": "文件路径不能为空"}
|
||
|
||
file_path = file_path.strip()
|
||
|
||
if not os.path.exists(file_path):
|
||
return {"success": False, "error": f"文件不存在: {file_path}"}
|
||
|
||
try:
|
||
file_ext = os.path.splitext(file_path)[1].lower()
|
||
|
||
if file_ext == ".txt":
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
text = f.read()
|
||
return {"success": True, "text": text}
|
||
|
||
elif file_ext == ".pdf":
|
||
import fitz
|
||
doc = fitz.open(file_path)
|
||
text_parts = []
|
||
for page in doc:
|
||
text_parts.append(page.get_text())
|
||
text = "\n".join(text_parts)
|
||
doc.close()
|
||
return {"success": True, "text": text}
|
||
|
||
elif file_ext in [".docx", ".doc"]:
|
||
from docx import Document
|
||
doc = Document(file_path)
|
||
text_parts = []
|
||
for para in doc.paragraphs:
|
||
text_parts.append(para.text)
|
||
text = "\n".join(text_parts)
|
||
return {"success": True, "text": text}
|
||
|
||
elif file_ext in [".jpg", ".jpeg", ".png", ".bmp", ".gif"]:
|
||
with open(file_path, "rb") as f:
|
||
image_data = f.read()
|
||
base64_image = base64.b64encode(image_data).decode("utf-8")
|
||
base64_url = f"data:image/{file_ext[1:]};base64,{base64_image}"
|
||
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
text = await OpenaiAPI.vlm_chat(
|
||
image_path=base64_url,
|
||
prompt="请描述这张图片的内容"
|
||
)
|
||
return {"success": True, "text": text}
|
||
|
||
else:
|
||
return {"success": False, "error": f"不支持的文件格式: {file_ext}"}
|
||
|
||
except Exception as e:
|
||
return {"success": False, "error": f"文件处理失败: {str(e)}"}
|
||
|
||
|
||
@tool
|
||
async def file_to_text_tool(file_path: str) -> str:
|
||
"""
|
||
文件转文本工具（供工作流调用）
|
||
支持多种文件格式：PDF、Word、TXT、图片等
|
||
|
||
Args:
|
||
file_path: 文件路径
|
||
|
||
Returns:
|
||
文件文本内容
|
||
"""
|
||
result = await file_to_text(file_path)
|
||
|
||
if not result.get("success", False):
|
||
return f"文件处理失败: {result.get('error', '未知错误')}"
|
||
|
||
return result.get("text", "")
|