""" 故障记录数据库模块 用于在故障诊断智能体生成方案后自动记录故障信息到PostgreSQL数据库 记录内容：舷号、设备、故障、备件、系统 """ import json import re import asyncio from datetime import datetime from typing import Dict, Any, Optional, List from config import POSTGRES_CONNECTION_STRING async def init_fault_records_table(): """ 初始化 fault_records 表 """ try: from psycopg_pool import AsyncConnectionPool except ImportError: print("[fault_record_db] psycopg_pool 未安装，无法初始化表") return async with AsyncConnectionPool( POSTGRES_CONNECTION_STRING, kwargs={"autocommit": True} ) as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(""" CREATE TABLE IF NOT EXISTS fault_records ( id SERIAL PRIMARY KEY, ship_number TEXT NOT NULL DEFAULT '', device_name TEXT NOT NULL DEFAULT '', fault TEXT NOT NULL DEFAULT '', spare_parts JSONB NOT NULL DEFAULT '[]', system_name TEXT NOT NULL DEFAULT '', created_at TIMESTAMP NOT NULL DEFAULT NOW() ) """) await cur.execute(""" CREATE INDEX IF NOT EXISTS fault_records_ship_number_idx ON fault_records(ship_number) """) await cur.execute(""" CREATE INDEX IF NOT EXISTS fault_records_device_name_idx ON fault_records(device_name) """) await cur.execute(""" CREATE INDEX IF NOT EXISTS fault_records_system_name_idx ON fault_records(system_name) """) await cur.execute(""" CREATE INDEX IF NOT EXISTS fault_records_created_at_idx ON fault_records(created_at) """) print("[fault_record_db] fault_records 表初始化完成") async def extract_spare_parts_from_scheme(scheme_content: str) -> List[str]: """ 从方案内容中用大模型抽取备件名称列表 Args: scheme_content: 方案内容（Markdown格式） Returns: 备件名称列表 """ if not scheme_content or len(scheme_content.strip()) < 10: return [] prompt = f"""从以下维修方案内容中提取所有提到的备品备件名称。 方案内容： {scheme_content[:3000]} 提取规则： 1. 找到"备品备件"、"备件"等相关章节，提取其中的项目名称 2. 只提取名称，不要提取型号、规格、数量、材质、用途等 3. 如果没有专门的章节，也从全文中识别备件名称 以JSON数组格式返回： {{ "备品备件": [备件名称1, 备件名称2, ...] }} 如果没有提到备品备件，请返回空数组 []""" try: from modelsAPI.model_api import OpenaiAPI result_text = await OpenaiAPI.open_api_chat_without_thinking( query=prompt, json_output=True ) print("抽取出来的备品备件", result_text) result_text = result_text.strip() result_text = re.sub(r'^```json\s*', '', result_text, flags=re.IGNORECASE) result_text = re.sub(r'^```\s*', '', result_text) result_text = re.sub(r'```$', '', result_text) try: parsed = json.loads(result_text) except json.JSONDecodeError: json_match = re.search(r'\{.*\}', result_text, re.DOTALL) if json_match: parsed = json.loads(json_match.group(0)) else: return [] if isinstance(parsed, dict): if "备品备件" in parsed and isinstance(parsed["备品备件"], list): return [str(part).strip() for part in parsed["备品备件"] if part] all_parts = [] for v in parsed.values(): if isinstance(v, list): all_parts.extend([str(part).strip() for part in v if part]) return all_parts elif isinstance(parsed, list): return [str(part).strip() for part in parsed if part] return [] except Exception as e: print(f"[fault_record_db] 从方案提取备件失败: {str(e)}") return [] async def save_fault_record( ship_number: str, device_name: str, fault: str, spare_parts: List[str], system_name: str ) -> bool: """ 保存故障记录到数据库 Args: ship_number: 舷号 device_name: 设备名称 fault: 故障现象 spare_parts: 备件列表 system_name: 系统名称 Returns: 是否保存成功 """ if not ship_number and not device_name and not fault: print("[fault_record_db] 舷号、设备、故障均为空，跳过保存") return False try: from psycopg_pool import AsyncConnectionPool except ImportError: print("[fault_record_db] psycopg_pool 未安装，无法保存") return False try: async with AsyncConnectionPool( POSTGRES_CONNECTION_STRING, kwargs={"autocommit": True} ) as pool: async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ INSERT INTO fault_records (ship_number, device_name, fault, spare_parts, system_name, created_at) VALUES (%s, %s, %s, %s, %s, %s) """, ( ship_number or "", device_name or "", fault or "", json.dumps(spare_parts, ensure_ascii=False), system_name or "", datetime.now() ) ) print(f"[fault_record_db] 故障记录已保存: 舷号={ship_number}, 设备={device_name}, 故障={fault}, 系统={system_name}, 备件数={len(spare_parts)}") return True except Exception as e: print(f"[fault_record_db] 保存故障记录失败: {str(e)}") import traceback traceback.print_exc() return False async def record_fault_from_state(state: Dict[str, Any]) -> bool: """ 从工作流状态中提取故障信息并记录到数据库 在故障诊断智能体生成方案后调用 流程： 1. 从状态中获取舷号、设备、故障 2. 从方案内容中用大模型抽取备件 3. 通过设备名称反推系统 4. 保存到数据库 Args: state: 工作流状态字典 Returns: 是否记录成功 """ ship_number = state.get("ship_number", "") or "" device_name = state.get("device_name", "") or "" fault = state.get("fault", "") or "" if not device_name and not fault: print("[fault_record_db] 设备和故障均为空，跳过记录") return False scheme_content = state.get("last_generated_scheme", "") or state.get("response", "") or "" spare_parts = await extract_spare_parts_from_scheme(scheme_content) system_name = "其他" if device_name: from tools.fault_statistics import get_device_system_from_neo4j system_name = await get_device_system_from_neo4j(device_name, ship_number=ship_number) return await save_fault_record( ship_number=ship_number, device_name=device_name, fault=fault, spare_parts=spare_parts, system_name=system_name ) async def query_fault_records( ship_number: Optional[str] = None, device_name: Optional[str] = None, system_name: Optional[str] = None, start_date: Optional[str] = None, end_date: Optional[str] = None, limit: int = 1000 ) -> List[Dict[str, Any]]: """ 从数据库查询故障记录 Args: ship_number: 舷号过滤 device_name: 设备名称过滤 system_name: 系统名称过滤 start_date: 开始日期（格式：YYYY-MM-DD） end_date: 结束日期（格式：YYYY-MM-DD） limit: 返回记录数上限 Returns: 故障记录列表 """ try: from psycopg_pool import AsyncConnectionPool except ImportError: print("[fault_record_db] psycopg_pool 未安装，无法查询") return [] from datetime import timedelta conditions = [] params = [] idx = 1 if ship_number: conditions.append(f"ship_number = %s") params.append(ship_number) if device_name: conditions.append(f"device_name = %s") params.append(device_name) if system_name: conditions.append(f"system_name = %s") params.append(system_name) if start_date: try: start_dt = datetime.strptime(start_date, "%Y-%m-%d") conditions.append(f"created_at >= %s") params.append(start_dt) except ValueError: pass if end_date: try: end_dt = datetime.strptime(end_date, "%Y-%m-%d") + timedelta(days=1) - timedelta(seconds=1) conditions.append(f"created_at <= %s") params.append(end_dt) except ValueError: pass where_clause = " AND ".join(conditions) if conditions else "1=1" try: async with AsyncConnectionPool( POSTGRES_CONNECTION_STRING, kwargs={"autocommit": True} ) as pool: async with pool.connection() as conn: async with conn.cursor() as cur: sql = f""" SELECT id, ship_number, device_name, fault, spare_parts, system_name, created_at FROM fault_records WHERE {where_clause} ORDER BY created_at DESC LIMIT %s """ params.append(limit) await cur.execute(sql, params) rows = await cur.fetchall() columns = ["id", "ship_number", "device_name", "fault", "spare_parts", "system_name", "created_at"] results = [] for row in rows: record = dict(zip(columns, row)) if isinstance(record.get("spare_parts"), str): try: record["spare_parts"] = json.loads(record["spare_parts"]) except Exception: record["spare_parts"] = [] if isinstance(record.get("created_at"), datetime): record["created_at"] = record["created_at"].isoformat() results.append(record) return results except Exception as e: print(f"[fault_record_db] 查询故障记录失败: {str(e)}") import traceback traceback.print_exc() return []