data_stardard.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419
  1. import os
  2. import time
  3. import asyncio
  4. import io
  5. import csv
  6. import datetime
  7. import httpx
  8. import json
  9. import sqlite3
  10. import re
  11. # --- LangChain Imports ---
  12. from langchain_openai import ChatOpenAI
  13. from langchain_core.prompts import ChatPromptTemplate
  14. from langchain_core.output_parsers import JsonOutputParser
  15. from langchain_core.outputs import Generation
  16. # --- 保持工具类不变 ---
  17. class SafeJsonOutputParser(JsonOutputParser):
  18. def parse_result(self, result, *, partial: bool = False):
  19. if isinstance(result, list) and len(result) > 0:
  20. generation = result[0]
  21. elif isinstance(result, Generation):
  22. generation = result
  23. else:
  24. raise ValueError(f"Unexpected result type: {type(result)}")
  25. text = generation.text
  26. text = re.sub(r"<think>.*?</think>", "", text, flags=re.S).strip()
  27. text = re.sub(r"^```(?:json)?|```$", "", text, flags=re.I | re.M).strip()
  28. match = re.search(r"(\[\s*{.*}\s*\]|\{\s*\".*\"\s*\})", text, flags=re.S)
  29. if not match:
  30. # 兼容:有时候 LLM 可能直接返回 SQL 字符串而不是 JSON,这里做个简单的容错
  31. if "SELECT" in text.upper():
  32. return {"sql": text}
  33. raise ValueError(f"Invalid json output: {text[:200]}")
  34. json_text = match.group(1)
  35. return json.loads(json_text)
  36. class TransactionParserAgent:
  37. def __init__(self, api_key: str, multimodal_api_url: str, base_url: str = "https://api.deepseek.com"):
  38. self.llm = ChatOpenAI(
  39. model="deepseek-chat",
  40. api_key=api_key,
  41. base_url=base_url,
  42. temperature=0.0, # SQL生成需要极其精确
  43. max_retries=3,
  44. http_client=httpx.Client(timeout=60.0)
  45. )
  46. self.multimodal_api_url = multimodal_api_url
  47. self.parser = SafeJsonOutputParser()
  48. async def _invoke_miner_u(self, file_path: str) -> str:
  49. """调用 MinerU 并提取纯行数据 (保持 httpx 调用不变,因为这不是 LLM)"""
  50. miner_start_time = time.perf_counter()
  51. print("\n" + "=" * 40)
  52. print("📌 【步骤1 - 数据提取】 开始执行")
  53. dealRows = 0
  54. try:
  55. # MinerU 是独立服务,继续使用原生 httpx
  56. async with httpx.AsyncClient() as client:
  57. with open(file_path, 'rb') as f:
  58. files = {'file': (os.path.basename(file_path), f)}
  59. data = {'folderId': 'text'}
  60. print("🔄数据提取中...")
  61. response = await client.post(self.multimodal_api_url, files=files, data=data, timeout=120.0)
  62. if response.status_code == 200:
  63. res_json = response.json()
  64. full_md_list = []
  65. for element in res_json.get('convert_json', []):
  66. if 'md' in element:
  67. full_md_list.append(element['md'])
  68. if 'rows' in element:
  69. dealRows += len(element['rows'])
  70. print(f"📊 提取结果:共提取 {dealRows - 1} 条数据")
  71. return "\n\n".join(full_md_list)
  72. return ""
  73. except Exception as e:
  74. print(f"❌ MinerU 调用异常: {e}")
  75. return ""
  76. finally:
  77. print(f"✅ 【步骤1 - 数据提取】 执行完成")
  78. print(f"⏱️ 执行耗时:{time.perf_counter() - miner_start_time:.2f} 秒")
  79. # --- 🆕 核心逻辑:SQLite 转换引擎 ---
  80. def _init_sqlite_db(self, data_rows: list, header_line: str, delimiter='|') -> tuple:
  81. """
  82. 将 Markdown 行数据灌入 SQLite 内存数据库的通用宽表
  83. 返回: (conn, header_mapping_info)
  84. """
  85. # 1. 创建内存数据库
  86. conn = sqlite3.connect(":memory:")
  87. cursor = conn.cursor()
  88. header_fingerprint = "".join(header_line.strip().strip('|').split())
  89. header_added = False # 确保数据库里只进一个表头
  90. # 2. 分析最大列数,建立通用宽表 (row_id, c0, c1, ... c30)
  91. max_cols = 0
  92. parsed_rows = []
  93. # 预处理:清洗 Markdown 分隔符
  94. for row in data_rows:
  95. # 去除首尾的 |
  96. clean_row = row.strip().strip('|')
  97. # A. 过滤掉纯分割线(如 | --- | --- |)
  98. if not re.search(r'[\u4e00-\u9fa5a-zA-Z0-9]', clean_row):
  99. continue
  100. # B. 提取当前行的指纹
  101. current_fingerprint = "".join(clean_row.split())
  102. # C. 核心判断:
  103. if current_fingerprint == header_fingerprint:
  104. if not header_added:
  105. # 只有第一次见到表头指纹时,才放入数据库
  106. header_added = True
  107. else:
  108. # 之后再见到一模一样的表头,直接跳过
  109. continue
  110. # 分割
  111. parts = [p.strip() for p in clean_row.split(delimiter)]
  112. if len(parts) > max_cols:
  113. max_cols = len(parts)
  114. parsed_rows.append(parts)
  115. if max_cols == 0:
  116. return None, None
  117. # 动态建表语句
  118. cols_def = ", ".join([f"c{i} TEXT" for i in range(max_cols)])
  119. create_sql = f"CREATE TABLE temp_raw_data (row_id INTEGER PRIMARY KEY AUTOINCREMENT, {cols_def});"
  120. cursor.execute(create_sql)
  121. # 3. 批量插入数据
  122. insert_sql = f"INSERT INTO temp_raw_data ({', '.join([f'c{i}' for i in range(max_cols)])}) VALUES ({', '.join(['?' for _ in range(max_cols)])})"
  123. # 补全数据(如果某行比最长行短,补None)
  124. final_data = []
  125. for p in parsed_rows:
  126. padding = [None] * (max_cols - len(p))
  127. final_data.append(p + padding)
  128. cursor.executemany(insert_sql, final_data)
  129. conn.commit()
  130. return conn, max_cols
  131. def _get_sql_generation_prompt(self) -> ChatPromptTemplate:
  132. system_template = """
  133. # Role
  134. 你是一个 SQLite 专家。
  135. # Task
  136. 你有一个名为 `temp_raw_data` 的表,里面存储了 OCR 识别后的原始数据。
  137. 表的列名为 `c0`, `c1`, `c2`... `cN`。
  138. 请根据提供的【表头】和【数据样本】,编写一条 SQL 查询语句,将原始列映射为标准输出字段。
  139. # Target Schema (Output Columns)
  140. 你的 SQL 必须 `SELECT` 出以下字段(顺序不能变):
  141. 1. `txId`: 交易流水号。如果原始数据没有,使用 `row_id`。
  142. 2. `txDate`: 交易日期 (格式 YYYY-MM-DD)。
  143. 3. `txTime`: 交易时间 (格式 HH:mm:ss)。如果没有则返回 '00:00:00'。
  144. 4. `txAmount`: 交易金额 (绝对值数字,**必须去除逗号**,转为 REAL/FLOAT)。
  145. 5. `txDirection`: 交易方向 (必须经过逻辑判断输出 '收入' 或 '支出')。
  146. 6. `txBalance`: 余额 (去除逗号)。
  147. 7. `txSummary`: 摘要/用途。
  148. 8. `txCounterparty`: 对方账号/户名。
  149. # Logic Rules (Crucial!)
  150. 1. **Direction Logic**:
  151. - 如果有单独的借/贷列:通常 "借"=`支出`, "贷"=`收入`。
  152. - 如果有单独的收入/支出列:哪一列有值就是哪个方向。
  153. - 如果金额有正负号:负号通常是支出。
  154. - 请使用 SQL 的 `CASE WHEN ... THEN ... ELSE ... END` 语法处理。
  155. 2. **Data Cleaning**:
  156. - 金额字段必须处理千分位逗号:`CAST(REPLACE(c?, ',', '') AS REAL)`
  157. - 日期必须清洗。
  158. # Output JSON Format
  159. ```json
  160. {{
  161. "sql": "SELECT ... FROM temp_raw_data WHERE ..."
  162. }}
  163. """
  164. user_template = """
  165. # Table Info
  166. Max Columns: {max_cols} Generic Column Names: c0, c1, ... c{max_cols_minus_1}
  167. # Data Preview (Header + First 3 Rows)
  168. {data_preview}
  169. # Instruction
  170. 请编写 SQL 语句来提取并清洗数据。 注意:不要包含 Markdown 的 sql 标签,直接返回 JSON。 忽略表头行(通常 row_id = 1 是表头,所以 WHERE row_id > 1)。 """
  171. return ChatPromptTemplate.from_messages([("system", system_template), ("user", user_template)])
  172. async def _generate_transform_sql(self, header_row: str, sample_rows: list, max_cols: int) -> str:
  173. """让 LLM 编写 SQL"""
  174. # 构建预览数据,带上 c0, c1 这种列名提示,方便 LLM 对应
  175. preview_text = ""
  176. # 表头预览
  177. header_parts = [p.strip() for p in header_row.strip().strip('|').split('|')]
  178. header_map = " | ".join([f"c{i}({val})" for i, val in enumerate(header_parts)])
  179. preview_text += f"Mapping Hint: {header_map}\n"
  180. preview_text += "-" * 50 + "\n"
  181. # 数据预览
  182. for row in sample_rows:
  183. preview_text += row + "\n"
  184. prompt_params = {
  185. "max_cols": max_cols,
  186. "max_cols_minus_1": max_cols - 1,
  187. "data_preview": preview_text
  188. }
  189. chain = self._get_sql_generation_prompt() | self.llm | self.parser
  190. print(f"🧠 [LLM] 正在生成 SQL 清洗逻辑...")
  191. try:
  192. result = await chain.ainvoke(prompt_params)
  193. sql = result.get("sql")
  194. print(f"💡 [LLM] 生成 SQL:\n{sql}")
  195. return sql
  196. except Exception as e:
  197. print(f"❌ SQL 生成失败: {e}")
  198. return ""
  199. async def parse_to_csv(self, file_path: str) -> str:
  200. # 1. 获取 Markdown
  201. md_text = await self._invoke_miner_u(file_path)
  202. if not md_text: return ""
  203. # 记录开始时间(使用time.perf_counter获取高精度时间)
  204. start_time = time.perf_counter()
  205. print("\n" + "=" * 40)
  206. print("📌 【步骤2 - 标准化转换】 开始执行")
  207. # 2. 预处理数据行
  208. raw_lines = md_text.splitlines()
  209. clean_lines = [l.strip() for l in raw_lines if l.strip() and "|" in l]
  210. # 简单判定表头 (包含2个以上关键词)
  211. header_line = ""
  212. header_idx = 0
  213. keywords = ["日期", "金额", "余额", "摘要", "用途", "借", "贷"]
  214. for idx, line in enumerate(clean_lines):
  215. if sum(1 for k in keywords if k in line) >= 2:
  216. header_line = line
  217. header_idx = idx
  218. break
  219. if not header_line:
  220. header_line = clean_lines[0]
  221. # 数据行 (保留原始数据,之后灌入 DB)
  222. data_rows = clean_lines # 把表头也灌进去,通过 row_id > header_idx + 1 来过滤
  223. # 3. 灌入 SQLite
  224. conn, max_cols = self._init_sqlite_db(data_rows,header_line)
  225. if not conn:
  226. return ""
  227. try:
  228. # 4. LLM 生成 SQL
  229. # 取表头和前3条数据作为样本
  230. sample_data = clean_lines[header_idx:header_idx + 4]
  231. sql_query = await self._generate_transform_sql(header_line, sample_data, max_cols)
  232. if not sql_query:
  233. return ""
  234. # 5. 执行 SQL
  235. cursor = conn.cursor()
  236. # 为了安全,确保 SQL 只是 SELECT
  237. if "DROP" in sql_query.upper() or "DELETE" in sql_query.upper():
  238. raise ValueError("Unsafe SQL detected")
  239. # 有时候 LLM 忘记过滤表头,我们强制在 SQL 外层或提示中处理
  240. # 这里的简单做法是假设 SQL 正确,或者在 SQL 后追加 limit 测试
  241. print(f"🚀 [SQLite] 执行查询...")
  242. cursor.execute(sql_query)
  243. results = cursor.fetchall()
  244. print(f"✅ 提取成功,共 {len(results)} 条数据")
  245. # 6. 导出为 CSV 字符串
  246. output = io.StringIO()
  247. writer = csv.writer(output, quoting=csv.QUOTE_MINIMAL, lineterminator='\n')
  248. # 写入标准表头
  249. csv_header = ["txId", "txDate", "txTime", "txAmount", "txDirection", "txBalance", "txSummary",
  250. "txCounterparty", "createdAt"]
  251. writer.writerow(csv_header)
  252. created_at = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
  253. for row in results:
  254. # row 是元组 (id, date, time, amt, dir, bal, sum, counter)
  255. # 转换 tuple 为 list 并添加 createdAt
  256. row_list = list(row)
  257. # --- 🆕 新增:txAmount 取绝对值逻辑 ---
  258. try:
  259. raw_amount = str(row_list[3]).replace(',', '') # 再次确保去除逗号
  260. if raw_amount:
  261. # 转换为浮点数取绝对值,再转回字符串(或保持 float)
  262. row_list[3] = abs(float(raw_amount))
  263. except (ValueError, TypeError):
  264. # 如果转换失败(例如识别到了文字),保持原样或设为 0.0
  265. print(f"⚠️ 金额转换失败: {row_list[3]}")
  266. row_list[3] = 0.0
  267. # 安全性清洗:处理可能的 None
  268. row_list = [str(x) if x is not None else "" for x in row_list]
  269. # 确保只取前8个字段 (以防 LLM 多选了)
  270. final_row = row_list[:8] + [created_at]
  271. writer.writerow(final_row)
  272. return output.getvalue()
  273. except sqlite3.Error as e:
  274. print(f"❌ SQLite 执行错误: {e}")
  275. # 可以在这里做一个重试机制:把错误信息返给 LLM 让它修正 SQL
  276. return ""
  277. finally:
  278. conn.close()
  279. print(f"✅ 【步骤2 - 标准化转换】 执行完成")
  280. print(f"⏱️ 总耗时: {time.perf_counter() - start_time:.2f} 秒")
  281. # --- 流程入口 ---
  282. async def parse_and_save_to_file(self, file_path: str, output_dir: str = "output") -> str:
  283. current_script_path = os.path.abspath(__file__)
  284. current_dir = os.path.dirname(current_script_path)
  285. file_full_name = os.path.basename(file_path)
  286. file_name = os.path.splitext(file_full_name)[0]
  287. output_dir = os.path.normpath(os.path.join(current_dir, "..", "..", output_dir))
  288. os.makedirs(output_dir, exist_ok=True)
  289. timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
  290. file_name = f"{file_name}_data_standard_{timestamp}.csv"
  291. full_path = os.path.join(output_dir, file_name)
  292. csv_result = await self.parse_to_csv(file_path)
  293. if csv_result:
  294. with open(full_path, "w", encoding="utf-8") as f:
  295. f.write(csv_result)
  296. return full_path
  297. else:
  298. raise Exception("数据解析失败,未生成有效内容")
  299. async def run_workflow_task(self, input_file_path: str) -> dict:
  300. # 1. 记录开始时间(使用time.perf_counter获取高精度时间)
  301. start_time = time.perf_counter()
  302. print(f"BEGIN---数据标准化任务开始---")
  303. try:
  304. print(f"待执行标准化的文件:{input_file_path}")
  305. saved_path = await self.parse_and_save_to_file(input_file_path, "data_files")
  306. print(f"结果文件保存至:{saved_path}")
  307. return {
  308. "status": "success",
  309. "file_path": saved_path,
  310. "file_name": os.path.basename(saved_path),
  311. "timestamp": datetime.datetime.now().isoformat()
  312. }
  313. except Exception as e:
  314. return {"status": "error", "message": str(e)}
  315. finally:
  316. end_time = time.perf_counter()
  317. elapsed_time = end_time - start_time
  318. print(f"⏱️ 执行总耗时: {elapsed_time:.2f} 秒")
  319. print(f"END---数据标准化任务结束")
  320. async def data_standize(api_key: str, base_url: str, multimodal_api_url: str, input_file_path: str) -> dict:
  321. """
  322. 数据标准化入口方法
  323. """
  324. # 创建Agent
  325. agent = TransactionParserAgent(
  326. api_key=api_key,
  327. base_url=base_url,
  328. multimodal_api_url=multimodal_api_url
  329. )
  330. # 执行标准化处理
  331. return await agent.run_workflow_task(input_file_path)
  332. # --- 运行 ---
  333. async def main():
  334. agent = TransactionParserAgent(
  335. api_key="sk-8634dbc2866540c4b6003bb5733f23d8",
  336. multimodal_api_url="http://103.154.31.78:20012/api/file/read"
  337. )
  338. current_script_path = os.path.abspath(__file__)
  339. current_dir = os.path.dirname(current_script_path)
  340. # 模拟 Workflow 传入一个待处理文件
  341. input_pdf = "data_files/11111.png"
  342. filepath = os.path.normpath(os.path.join(current_dir, "..", "..", input_pdf))
  343. if not os.path.exists(filepath):
  344. print(f"{filepath}文件不存在")
  345. return
  346. result = await agent.run_workflow_task(filepath)
  347. if result["status"] == "success":
  348. print(f"🎯 【数据标准化】任务完成!")
  349. else:
  350. print(f"❌ 任务失败: {result['message']}")
  351. if __name__ == "__main__":
  352. asyncio.run(main())