251 lines
8.9 KiB
Python
251 lines
8.9 KiB
Python
import asyncio
|
||
import aiohttp
|
||
import json
|
||
import traceback
|
||
from filename_proceess_and_kgquery import get_entity
|
||
from config import SHIP_MODEL_NAME
|
||
|
||
# --- 配置区 ---
|
||
BASE_URL = "https://htknow.zkzdht.com/api/v1/knowledge"
|
||
ERROR_LOG_FILE = "error_log.txt"
|
||
|
||
HEADERS = {
|
||
"accept": "application/json",
|
||
"x-role": "admin",
|
||
"x-user-id": "1",
|
||
"x-user-name": "testuser"
|
||
}
|
||
|
||
# 全局 Semaphore，用于限制并发连接数，防止对服务器造成过大压力
|
||
# 根据实际情况调整数值
|
||
CONCURRENT_LIMIT = 10
|
||
semaphore = asyncio.Semaphore(CONCURRENT_LIMIT)
|
||
|
||
async def log_error(kb_id, file_id, error_msg):
|
||
"""
|
||
将错误信息追加写入日志文件
|
||
注意：文件 I/O 在 asyncio 中是阻塞的，如果日志量极大，可以考虑使用线程池执行器
|
||
"""
|
||
# 对于少量的日志写入，直接使用 await asyncio.to_thread 或直接同步写入通常足够
|
||
# 这里为了简化，直接使用同步写入（在日志量不大时影响较小）
|
||
# 如果需要完全异步，可以使用 aiofiles 库
|
||
try:
|
||
with open(ERROR_LOG_FILE, "a", encoding="utf-8") as f:
|
||
f.write(f"[ERROR] KnowledgeBaseID: {kb_id}, FileID: {file_id}, Message: {error_msg}\n")
|
||
except Exception as e:
|
||
print(f"写入错误日志时发生异常: {e}")
|
||
|
||
async def find_ship_info_by_hull(json_file_path, data):
|
||
"""从NER结果中提取舷号，再从JSON文件中查找对应的舰船信息"""
|
||
try:
|
||
# 同步文件读取，如果文件很大或数量很多，可以使用 aiofiles
|
||
with open(json_file_path, 'r', encoding='utf-8') as f:
|
||
ship_data = json.load(f)
|
||
|
||
hulls = [
|
||
entity['properties']['舷号']
|
||
for entity in data.get('entities', [])
|
||
if 'properties' in entity and '舷号' in entity['properties']
|
||
]
|
||
if not hulls:
|
||
return None
|
||
|
||
target_hull = hulls[0]
|
||
|
||
for item in ship_data:
|
||
if str(item.get('hull_number')) == str(target_hull):
|
||
return {
|
||
"model_name": item.get('model_name'),
|
||
"ship_name": item.get('ship_name')
|
||
}
|
||
return None
|
||
except Exception as e:
|
||
print(f"查找舰船信息时发生错误: {e}")
|
||
return None
|
||
|
||
def format_entity_text(data):
|
||
"""格式化实体文本 (纯计算逻辑，保持同步)"""
|
||
type_mapping = {
|
||
'舰艇': [('舰艇为', '名称'), ('，舷号', '舷号')],
|
||
'系统': [('，系统：', '名称')],
|
||
'子系统': [('，子系统：', '名称')],
|
||
'设备': [('，设备：', '名称')]
|
||
}
|
||
|
||
text_parts = []
|
||
for entity in data.get('entities', []):
|
||
entity_type = entity.get('type')
|
||
properties = entity.get('properties', {})
|
||
|
||
if entity_type in type_mapping:
|
||
entity_str = ""
|
||
for prefix, attr_key in type_mapping[entity_type]:
|
||
value = str(properties.get(attr_key, ''))
|
||
entity_str += f"{prefix}{value}"
|
||
text_parts.append(entity_str)
|
||
|
||
return ''.join(text_parts)
|
||
|
||
async def get_knowledge_base_files(session, kb_id):
|
||
"""
|
||
获取指定知识库的文件列表
|
||
"""
|
||
url = f"{BASE_URL}/knowledge_base/{kb_id}"
|
||
try:
|
||
async with session.get(url, headers=HEADERS, timeout=120) as response:
|
||
response.raise_for_status()
|
||
data = await response.json()
|
||
|
||
if 'files' in data and isinstance(data['files'], list):
|
||
return data['files']
|
||
else:
|
||
print(f"警告: 知识库 {kb_id} 中未找到 'files' 字段。")
|
||
return []
|
||
|
||
except Exception as e:
|
||
error_msg = f"获取知识库详情失败: {str(e)} | Traceback: {traceback.format_exc()}"
|
||
print(error_msg)
|
||
await log_error(kb_id, "N/A", error_msg)
|
||
return []
|
||
|
||
async def get_file_slices(session, file_id):
|
||
"""获取文件所有切片"""
|
||
url = f"{BASE_URL}/files/{file_id}/slices"
|
||
try:
|
||
async with session.get(url, headers=HEADERS, timeout=120) as response:
|
||
response.raise_for_status()
|
||
return await response.json()
|
||
except Exception as e:
|
||
error_msg = f"获取文件切片失败: {str(e)}"
|
||
print(error_msg)
|
||
# 错误会在调用处记录
|
||
raise Exception(error_msg)
|
||
|
||
async def update_slices(session, slices_data, file_id):
|
||
"""更新切片"""
|
||
url = f"{BASE_URL}/files/{file_id}/slices"
|
||
headers = {
|
||
**HEADERS,
|
||
"content-type": "application/json"
|
||
}
|
||
payload = {"slices": slices_data}
|
||
try:
|
||
async with session.put(url, headers=headers, json=payload, timeout=120) as response:
|
||
response.raise_for_status()
|
||
return await response.json()
|
||
except Exception as e:
|
||
error_msg = f"更新切片失败: {str(e)}"
|
||
print(error_msg)
|
||
raise Exception(error_msg)
|
||
|
||
async def process_single_file(session, kb_id, file):
|
||
"""
|
||
处理单个文件的逻辑，用于并发执行
|
||
"""
|
||
file_id = file.get('id')
|
||
filename = file.get('filename', 'Unknown')
|
||
|
||
print(f" 处理文件: {filename} (ID: {file_id})")
|
||
|
||
try:
|
||
# --- 步骤 A: 提取实体信息 ---
|
||
final_result = get_entity(filename)
|
||
|
||
# 处理舰船型号
|
||
xinghao = await find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
|
||
model_name = xinghao.get('model_name', '') if xinghao else ''
|
||
|
||
# 格式化其他实体文本
|
||
other_data = format_entity_text(final_result)
|
||
|
||
# --- 步骤 B: 获取切片 ---
|
||
slices = await get_file_slices(session, file_id)
|
||
|
||
# --- 步骤 C: 处理切片内容 ---
|
||
slices_v2 = []
|
||
info_suffix = f"({other_data}，型号为{model_name})" if model_name and other_data else ""
|
||
for ins in slices:
|
||
ins_id = ins.get('id')
|
||
content = ins.get('content', '')
|
||
|
||
# 查找第一行的结尾（换行符位置）
|
||
newline_idx = content.find('\n')
|
||
|
||
# 拼接新内容
|
||
if newline_idx == -1:
|
||
# 内容中没有换行符，直接在末尾添加
|
||
new_content = content + info_suffix
|
||
else:
|
||
# 在第一行末尾插入
|
||
new_content = content[:newline_idx] + info_suffix + content[newline_idx:]
|
||
|
||
slices_v2.append({"id": ins_id, "content": new_content})
|
||
|
||
# --- 步骤 D: 更新切片 ---
|
||
await update_slices(session, slices_v2, file_id)
|
||
print(f" 成功更新文件: {filename}")
|
||
|
||
except Exception as e:
|
||
# 捕获当前文件处理过程中的所有异常
|
||
error_msg = f"处理文件失败: {str(e)}"
|
||
print(error_msg)
|
||
# 记录错误日志
|
||
await log_error(kb_id, file_id, error_msg)
|
||
# 跳过当前文件
|
||
return False
|
||
return True
|
||
|
||
async def process_files_in_knowledge_base(session, kb_id):
|
||
"""
|
||
处理单个知识库中的所有文件
|
||
"""
|
||
print(f"\n=== 开始处理知识库 ID: {kb_id} ===")
|
||
|
||
# 1. 获取文件列表
|
||
files = await get_knowledge_base_files(session, kb_id)
|
||
if not files:
|
||
print(f"知识库 {kb_id} 中没有文件或获取文件列表失败，跳过。")
|
||
return
|
||
|
||
# 2. 并发处理每个文件
|
||
# 使用 asyncio.gather 来并发处理所有文件
|
||
# 如果文件数量巨大，可以分批处理以避免内存溢出
|
||
file_tasks = [
|
||
process_single_file(session, kb_id, file)
|
||
for file in files
|
||
]
|
||
|
||
# 并发执行所有文件任务
|
||
# return_exceptions=True 会捕获任务内部的异常并作为结果返回，而不是中断整个 gather
|
||
results = await asyncio.gather(*file_tasks, return_exceptions=True)
|
||
|
||
# 统计结果
|
||
success_count = sum(1 for r in results if r is True)
|
||
failed_count = len(results) - success_count
|
||
print(f"知识库 {kb_id} 处理完成。成功: {success_count}, 失败: {failed_count}")
|
||
|
||
async def main():
|
||
KNOWLEDGE_BASE_IDS = ['8']
|
||
print(f"开始批量处理 {len(KNOWLEDGE_BASE_IDS)} 个知识库...")
|
||
|
||
# 创建一个全局的 TCP 连接器，复用连接
|
||
connector = aiohttp.TCPConnector(limit=CONCURRENT_LIMIT, limit_per_host=5)
|
||
|
||
async with aiohttp.ClientSession(connector=connector) as session:
|
||
# 遍历所有知识库ID，并发处理
|
||
# 如果知识库数量很多，也可以在这里使用 asyncio.gather
|
||
for kb_id in KNOWLEDGE_BASE_IDS:
|
||
try:
|
||
await process_files_in_knowledge_base(session, kb_id)
|
||
except Exception as e:
|
||
error_msg = f"处理知识库顶层异常: {str(e)}"
|
||
print(error_msg)
|
||
await log_error(kb_id, "N/A", error_msg)
|
||
continue
|
||
|
||
print("所有知识库处理完成。")
|
||
|
||
if __name__ == "__main__":
|
||
# 运行异步主函数
|
||
asyncio.run(main())
|