refactor(trade_capture): 最大容错版——不限截图类型/不限股票数量/逐只独立处理
This commit is contained in:
xxm
@@ -1,12 +1,10 @@
|
||||
#!/usr/bin/env python3
|
||||
"""重构截图识别: 遵循 LLM 执行协议
|
||||
"""重构 trade_capture: 最大容错，不限截图类型
|
||||
|
||||
协议要求:
|
||||
1. LLM 输出节标题格式(不用 JSON)
|
||||
2. parse_response 按节标题精确匹配
|
||||
3. 多层门禁校验
|
||||
4. 写入前快照
|
||||
5. 确认反馈
|
||||
核心思路:
|
||||
1. LLM 提取截图中所有股票信息（不限类型、不限数量）
|
||||
2. 代码逐只独立处理（买入/卖出/更新持仓/添加自选）
|
||||
3. 每只股票独立校验，单只失败不影响其他
|
||||
"""
|
||||
import os, sys, json, re, sqlite3
|
||||
os.chdir("/home/hmo/MoFin/deploy/profile-scripts")
|
||||
@@ -14,112 +12,98 @@ os.environ["GIT_ALLOW_COMMIT"] = "1"
|
||||
|
||||
DB = "/home/hmo/MoFin/data/mofin.db"
|
||||
|
||||
# === 1. 定义节标题 schema ===
|
||||
TRADE_SECTIONS = [
|
||||
{"name": "交易动作", "field": "action", "type": "enum", "values": ["买入", "卖出"]},
|
||||
{"name": "股票代码", "field": "code", "type": "string", "constraint": "6位数字"},
|
||||
{"name": "股票名称", "field": "name", "type": "string"},
|
||||
{"name": "交易数量", "field": "shares", "type": "integer", "constraint": "必须>0"},
|
||||
{"name": "交易价格", "field": "price", "type": "number", "constraint": "必须>0"},
|
||||
]
|
||||
|
||||
# === 2. 节标题解析器(复用已有) ===
|
||||
# === 节标题解析器 ===
|
||||
def _section_line(text, name):
|
||||
for line in text.split("\n"):
|
||||
if re.match(r'^\s*【' + name + r'】', line):
|
||||
return line
|
||||
return ""
|
||||
|
||||
# === 3. 解析 LLM 输出 ===
|
||||
def parse_trade_response(text):
|
||||
result = {}
|
||||
for sec in TRADE_SECTIONS:
|
||||
line = _section_line(text, sec["name"])
|
||||
if not line:
|
||||
# === 解析多只股票的 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
|
||||
val = re.sub(r'^\s*【' + sec["name"] + r'】\s*', '', line).strip()
|
||||
code = re.sub(r'^\s*【股票代码】\s*', '', code_line).strip()
|
||||
code = re.findall(r'\d{6}', code)
|
||||
if not code:
|
||||
continue
|
||||
code = code[0]
|
||||
|
||||
if sec["type"] == "enum":
|
||||
if val in sec["values"]:
|
||||
result[sec["field"]] = val
|
||||
elif sec["type"] == "integer":
|
||||
nums = re.findall(r'\d+', val)
|
||||
if nums:
|
||||
result[sec["field"]] = int(nums[0])
|
||||
elif sec["type"] == "number":
|
||||
nums = re.findall(r'\d+\.?\d*', val)
|
||||
if nums:
|
||||
result[sec["field"]] = float(nums[0])
|
||||
elif sec["type"] == "string":
|
||||
result[sec["field"]] = val
|
||||
return result
|
||||
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
|
||||
|
||||
# === 4. 门禁校验 ===
|
||||
def validate_trade(parsed, current_holdings=None):
|
||||
# === 校验单只股票 ===
|
||||
def validate_stock(stock):
|
||||
errors = []
|
||||
|
||||
# 动作
|
||||
if parsed.get("action") not in ("买入", "卖出"):
|
||||
errors.append("交易动作必须是'买入'或'卖出'")
|
||||
|
||||
# 代码
|
||||
code = parsed.get("code", "")
|
||||
code = stock.get("code", "")
|
||||
if not re.match(r'^\d{6}$', code):
|
||||
errors.append(f"股票代码'{code}'不是6位数字")
|
||||
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 中不存在")
|
||||
errors.append(f"代码'{code}'在 stock_daily 中不存在")
|
||||
|
||||
# 数量
|
||||
shares = parsed.get("shares", 0)
|
||||
if not isinstance(shares, int) or shares <= 0:
|
||||
errors.append(f"交易数量'{shares}'必须是正整数")
|
||||
action = stock.get("action", "")
|
||||
if action and action not in ("买入", "卖出"):
|
||||
errors.append(f"交易动作'{action}'无效")
|
||||
|
||||
# 价格
|
||||
price = parsed.get("price", 0)
|
||||
if not isinstance(price, (int, float)) or price <= 0:
|
||||
errors.append(f"交易价格'{price}'必须是正数")
|
||||
elif price > 0:
|
||||
# 与现价偏差检查
|
||||
conn = sqlite3.connect(DB, timeout=30)
|
||||
row = conn.execute("SELECT close FROM stock_daily WHERE code=? ORDER BY date DESC LIMIT 1", (code,)).fetchone()
|
||||
conn.close()
|
||||
if row and row[0]:
|
||||
current_price = float(row[0])
|
||||
if current_price > 0:
|
||||
deviation = abs(price - current_price) / current_price
|
||||
if deviation > 0.1:
|
||||
errors.append(f"交易价格{price}与现价{current_price}偏差{deviation:.1%}超过10%")
|
||||
shares = stock.get("shares", "")
|
||||
if shares:
|
||||
try:
|
||||
s = int(shares)
|
||||
if s <= 0:
|
||||
errors.append(f"数量{shares}必须>0")
|
||||
except ValueError:
|
||||
errors.append(f"数量'{shares}'不是整数")
|
||||
|
||||
# 卖出数量检查
|
||||
if parsed.get("action") == "卖出" and current_holdings:
|
||||
current_shares = current_holdings.get(code, {}).get("shares", 0)
|
||||
if shares > current_shares:
|
||||
errors.append(f"卖出数量{shares}超过当前持仓{current_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
|
||||
|
||||
# === 5. 写入 holdings ===
|
||||
def execute_trade(parsed):
|
||||
code = parsed["code"]
|
||||
action = parsed["action"]
|
||||
shares = parsed["shares"]
|
||||
price = parsed["price"]
|
||||
name = parsed.get("name", "")
|
||||
# === 执行单只股票操作 ===
|
||||
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 FROM holdings WHERE code=?", (code,)).fetchone()
|
||||
# 获取当前状态
|
||||
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("""
|
||||
@@ -128,65 +112,101 @@ def execute_trade(parsed):
|
||||
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 == "卖出":
|
||||
new_shares = old_shares - shares
|
||||
if new_shares <= 0:
|
||||
new_shares = 0
|
||||
conn.execute("UPDATE holdings SET shares=0, is_active=0 WHERE code=?", (code,))
|
||||
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"] = "无变更"
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
print(f" ✅ {action} {name}({code}) {shares}股 @{price} → 持仓{new_shares}股")
|
||||
return {"ok": True, "new_shares": new_shares}
|
||||
return result
|
||||
|
||||
# === 6. 完整流程 ===
|
||||
# === 完整流程 ===
|
||||
def process_screenshot(ocr_text):
|
||||
"""OCR文本 → 结构化提取 → 校验 → 写入"""
|
||||
"""OCR文本 → LLM提取 → 逐只校验 → 逐只执行"""
|
||||
from llm_client import call_llm, REASSESS_MODEL
|
||||
|
||||
prompt = f"""请从以下OCR文字中提取交易信息，严格按节标题格式输出：
|
||||
prompt = f"""请从以下截图OCR文字中提取所有股票信息。不限类型（持仓/交易/自选都处理），不限数量。
|
||||
|
||||
对每只股票，判断是否有交易动作（买入/卖出），如果没有明确交易动作则作为持仓更新处理。
|
||||
|
||||
{ocr_text}
|
||||
|
||||
请严格按以下格式输出（每个字段一个节标题）：
|
||||
请严格按以下格式输出，每只股票一个区块：
|
||||
|
||||
【交易动作】买入 或 卖出
|
||||
【股票代码】6位数字
|
||||
【股票名称】中文名称
|
||||
【交易数量】正整数
|
||||
【交易价格】数字
|
||||
【交易动作】买入/卖出/无（无明确交易时写"无"）
|
||||
【交易数量】正整数（有交易时填写）
|
||||
【交易价格】数字（有交易时填写）
|
||||
【成本价】数字（有持仓时填写）
|
||||
"""
|
||||
|
||||
result = call_llm(prompt, model=REASSESS_MODEL, max_tokens=None)
|
||||
if not result["ok"] or not result.get("content"):
|
||||
return {"ok": False, "error": "LLM调用失败"}
|
||||
|
||||
parsed = parse_trade_response(result["content"])
|
||||
stocks = parse_screenshot_response(result["content"])
|
||||
if not stocks:
|
||||
return {"ok": False, "error": "未提取到任何股票信息", "raw": result["content"][:500]}
|
||||
|
||||
# 门禁校验
|
||||
errors = validate_trade(parsed)
|
||||
if errors:
|
||||
return {"ok": False, "errors": errors, "parsed": parsed}
|
||||
# 逐只处理
|
||||
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)]})
|
||||
|
||||
# 获取当前持仓(卖出检查用)
|
||||
conn = sqlite3.connect(DB, timeout=30)
|
||||
current = {}
|
||||
for row in conn.execute("SELECT code, shares, cost FROM holdings WHERE is_active=1"):
|
||||
current[row[0]] = {"shares": row[1], "cost": row[2]}
|
||||
conn.close()
|
||||
# clean_watchlist（持仓变动的副作用）
|
||||
try:
|
||||
from clean_watchlist import main as _cwl
|
||||
_cwl()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 二次校验(卖出数量)
|
||||
errors = validate_trade(parsed, current)
|
||||
if errors:
|
||||
return {"ok": False, "errors": errors, "parsed": parsed}
|
||||
|
||||
# 执行
|
||||
exec_result = execute_trade(parsed)
|
||||
return {"ok": True, "parsed": parsed, "exec": exec_result}
|
||||
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 模块已加载(协议版)")
|
||||
print(f"节标题: {[s['name'] for s in TRADE_SECTIONS]}")
|
||||
print("trade_capture 模块已加载(最大容错版)")
|
||||
|
||||
Reference in New Issue
Block a user