#!/usr/bin/env python3 """重构 trade_capture: 最大容错,不限截图类型 核心思路: 1. LLM 提取截图中所有股票信息(不限类型、不限数量) 2. 代码逐只独立处理(买入/卖出/更新持仓/添加自选) 3. 每只股票独立校验,单只失败不影响其他 """ import os, sys, json, re, sqlite3 os.chdir("/home/hmo/MoFin/deploy/profile-scripts") os.environ["GIT_ALLOW_COMMIT"] = "1" DB = "/home/hmo/MoFin/data/mofin.db" # === 节标题解析器 === def _section_line(text, name): for line in text.split("\n"): if re.match(r'^\s*【' + name + r'】', line): return line return "" # === 解析多只股票的 LLM 输出 === def parse_screenshot_response(text): """从 LLM 输出中提取所有股票信息(不限数量)""" stocks = [] # 按【股票代码】分割,每段是一只股票的信息 parts = re.split(r'(?=【股票代码】)', text) for part in parts: code_line = _section_line(part, "股票代码") if not code_line: continue code = re.sub(r'^\s*【股票代码】\s*', '', code_line).strip() code = re.findall(r'\d{6}', code) if not code: continue code = code[0] stock = {"code": code} for field, label in [("name", "股票名称"), ("action", "交易动作"), ("shares", "交易数量"), ("price", "交易价格"), ("cost", "成本价"), ("pnl", "盈亏")]: line = _section_line(part, label) if line: val = re.sub(r'^\s*【' + label + r'】\s*', '', line).strip() stock[field] = val stocks.append(stock) return stocks # === 校验单只股票 === def validate_stock(stock): errors = [] code = stock.get("code", "") if not re.match(r'^\d{6}$', code): errors.append(f"代码'{code}'不是6位数字") else: conn = sqlite3.connect(DB, timeout=30) exists = conn.execute("SELECT 1 FROM stock_daily WHERE code=? LIMIT 1", (code,)).fetchone() conn.close() if not exists: errors.append(f"代码'{code}'在 stock_daily 中不存在") action = stock.get("action", "") if action and action not in ("买入", "卖出"): errors.append(f"交易动作'{action}'无效") shares = stock.get("shares", "") if shares: try: s = int(shares) if s <= 0: errors.append(f"数量{shares}必须>0") except ValueError: errors.append(f"数量'{shares}'不是整数") price = stock.get("price", "") if price: try: p = float(price) if p <= 0: errors.append(f"价格{price}必须>0") except ValueError: errors.append(f"价格'{price}'不是数字") return errors # === 执行单只股票操作 === def execute_stock(stock): code = stock["code"] action = stock.get("action", "") name = stock.get("name", "") conn = sqlite3.connect(DB, timeout=30) conn.execute("PRAGMA busy_timeout=30000") # 获取当前状态 old = conn.execute("SELECT shares, cost, name FROM holdings WHERE code=?", (code,)).fetchone() old_shares = old[0] if old else 0 old_cost = old[1] if old else 0 if not name and old: name = old[2] or "" result = {"code": code, "name": name} if action == "买入": shares = int(stock.get("shares", 0)) price = float(stock.get("price", 0)) new_shares = old_shares + shares new_cost = ((old_cost * old_shares) + (price * shares)) / new_shares if new_shares > 0 else price conn.execute(""" INSERT INTO holdings (code, name, shares, cost, price, is_active) VALUES (?, ?, ?, ?, ?, 1) ON CONFLICT(code) DO UPDATE SET shares=?, cost=?, price=?, is_active=1 """, (code, name, new_shares, round(new_cost, 4), price, new_shares, round(new_cost, 4), price)) result["action"] = "买入" result["new_shares"] = new_shares result["new_cost"] = round(new_cost, 4) elif action == "卖出": shares = int(stock.get("shares", 0)) price = float(stock.get("price", 0)) new_shares = max(0, old_shares - shares) if new_shares == 0: conn.execute("UPDATE holdings SET shares=0, price=?, is_active=0 WHERE code=?", (price, code)) else: conn.execute("UPDATE holdings SET shares=?, price=? WHERE code=?", (new_shares, price, code)) result["action"] = "卖出" result["new_shares"] = new_shares else: # 无交易动作(纯持仓图/自选股图)— 更新价格和持仓信息 if stock.get("shares"): new_shares = int(stock["shares"]) if new_shares != old_shares: conn.execute("UPDATE holdings SET shares=? WHERE code=?", (new_shares, code)) result["action"] = "更新持仓" result["new_shares"] = new_shares if stock.get("cost"): try: new_cost = float(stock["cost"]) conn.execute("UPDATE holdings SET cost=? WHERE code=?", (new_cost, code)) result["new_cost"] = new_cost except ValueError: pass if stock.get("price"): try: new_price = float(stock["price"]) conn.execute("UPDATE holdings SET price=? WHERE code=?", (new_price, code)) except ValueError: pass if not result.get("action"): result["action"] = "无变更" # 记录执行日志 from mofin_db import log_execution log_execution(conn, code, result.get("action",""), int(stock.get("shares",0) or 0), float(stock.get("price",0) or 0), source="trade_capture") conn.commit() conn.close() return result # === 完整流程 === def process_screenshot(ocr_text): """OCR文本 → LLM提取 → 逐只校验 → 逐只执行""" from llm_client import call_llm, REASSESS_MODEL prompt = f"""请从以下截图OCR文字中提取所有股票信息。不限类型(持仓/交易/自选都处理),不限数量。 对每只股票,判断是否有交易动作(买入/卖出),如果没有明确交易动作则作为持仓更新处理。 {ocr_text} 请严格按以下格式输出,每只股票一个区块: 【股票代码】6位数字 【股票名称】中文名称 【交易动作】买入/卖出/无(无明确交易时写"无") 【交易数量】正整数(有交易时填写) 【交易价格】数字(有交易时填写) 【成本价】数字(有持仓时填写) """ result = call_llm(prompt, model=REASSESS_MODEL, concurrent=True, max_tokens=None) if not result["ok"] or not result.get("content"): return {"ok": False, "error": "LLM调用失败"} stocks = parse_screenshot_response(result["content"]) if not stocks: return {"ok": False, "error": "未提取到任何股票信息", "raw": result["content"][:500]} # 逐只处理 results = [] for stock in stocks: errs = validate_stock(stock) if errs: results.append({"code": stock.get("code"), "ok": False, "errors": errs}) else: try: r = execute_stock(stock) results.append({"code": stock.get("code"), "ok": True, "result": r}) except Exception as e: results.append({"code": stock.get("code"), "ok": False, "errors": [str(e)]}) # clean_watchlist(持仓变动的副作用) try: from clean_watchlist import main as _cwl _cwl() except Exception: pass return {"ok": True, "results": results, "total": len(results), "success": sum(1 for r in results if r["ok"]), "failed": sum(1 for r in results if not r["ok"])} if __name__ == "__main__": print("trade_capture 模块已加载(最大容错版)")