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())