Files
tick-stock-panel/backend/app/strategy/custom_signals_ai.py
T
dev 8519a2bd16 feat(ai): AI 生成自定义信号条件
- 新增 /api/signals/ai/generate 接口:自然语言描述 → 结构化信号条件
- 新增 custom_signals_ai 模块:组装提示词 + 解析校验 AI 返回的 JSON
- 复用 custom_signals.validate() 白名单安全闸门
- 新增输出 token 上限和上下文窗口设置
- 实现 AI 请求 max_tokens 钳制和输入预算检查
- 自定义信号右值字段名容错处理
- 增强 JSON 解析容错(尾随逗号、垃圾字符)
- 前端自定义信号对话框接入 AI 生成
2026-08-21 09:55:24 +08:00

232 lines
8.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""AI 生成自定义信号 — 组装提示词 + 解析并校验 AI 返回的结构化条件。
职责:
- build_messages(description): 把用户一句描述 + 字段白名单/运算符/格式要求组装成 LLM 消息
- parse_and_validate(text): 把 LLM 返回的 JSON 解析为 {name, conditions},并复用
custom_signals.validate() 做白名单/运算符/偏移校验(安全闸门)
不知道: API、AI 调用、持久化。纯函数,无副作用。
"""
from __future__ import annotations
import json
import re
from app.indicators.pipeline import ENRICHED_COLUMNS, ENRICHED_COLUMNS_BY_CATEGORY
from app.strategy import custom_signals
# 分类 → 中文标签(与 /api/custom-signals/options 的分组一致)
_GROUP_LABELS = {
"basic": "基础", "ma": "均线 MA", "ema": "指数均线 EMA",
"macd": "MACD", "boll": "布林带 BOLL", "kdj": "KDJ",
"atr": "ATR", "volume": "量价", "extremes": "极值",
"momentum": "动量", "volatility": "波动率", "rsi": "RSI",
}
# 行情类字段不在 ENRICHED_COLUMNS_BY_CATEGORY 里, 单独归一组
_QUOTE_FIELDS = {
"open", "high", "low", "close", "volume", "amount", "turnover_rate",
"consecutive_limit_ups", "consecutive_limit_downs",
}
_FENCED_JSON_RE = re.compile(r"```(?:json)?\s*\n?(.*?)```", re.DOTALL)
def _format_fields() -> str:
"""按类别格式化白名单字段(key(中文标签)),供 LLM 参考。"""
allowed = custom_signals.ALLOWED_FIELDS
lines: list[str] = []
quote = sorted(f for f in _QUOTE_FIELDS if f in allowed)
lines.append(
"行情: " + ", ".join(f"{f}({ENRICHED_COLUMNS.get(f, f)})" for f in quote)
)
for cat, label in _GROUP_LABELS.items():
fields = [f for f in ENRICHED_COLUMNS_BY_CATEGORY.get(cat, []) if f in allowed]
if fields:
lines.append(
f"{label}: "
+ ", ".join(f"{f}({ENRICHED_COLUMNS.get(f, f)})" for f in fields)
)
return "\n".join(lines)
_SYSTEM_TEMPLATE = """你是A股量化信号设计专家。用户会描述一个信号思路,你要把它拆解为布尔条件组合(多条件之间是「且」关系,即同时满足),并输出结构化 JSON 供系统编译为选股/回测/监控信号。
可用字段(白名单,只能使用以下字段,禁止自造或使用白名单之外的字段):
{fields}
运算符(op):> >= < <= == !=
右值(right):
- 数字:写字符串形式,如 "2"、"3000"、"0.05"
- 另一字段:必须带 "field:" 前缀,如 "field:ma20";严禁裸写字段名,如 "macd_dea" 应写成 "field:macd_dea"
日期偏移(leftDays / rightDays):取 N 个交易日前的值,0 = 当日最新;范围 0~{max_days}。只有明确需要「前N日」时才使用偏移。
要求:
1. 只输出一个 JSON 对象,禁止 markdown 代码块、禁止任何解释或多余文字。
2. JSON 结构固定为:
{{"name": "简短中文信号名称(≤12字)", "conditions": [
{{"left": "字段", "op": "运算符", "right": "数字字符串或field:字段", "leftDays": 0, "rightDays": 0}}
]}}
示例(右值引用另一字段时必须带 field: 前缀,不能裸写字段名):
{{"name": "MACD金叉", "conditions": [
{{"left": "macd_dif", "op": ">", "right": "field:macd_dea", "leftDays": 0, "rightDays": 0}}
]}}
3. conditions 至少 1 个、最多 8 个;优先用最少的条件表达清晰的思路。
4. 多条件必须能同时满足,不要输出互相矛盾的条件。"""
def build_messages(description: str) -> list[dict]:
"""组装 LLM 消息:[system 提示词, user 描述]。"""
system = _SYSTEM_TEMPLATE.format(
fields=_format_fields(),
max_days=custom_signals.MAX_DAYS,
)
return [
{"role": "system", "content": system},
{"role": "user", "content": description},
]
def parse_and_validate(text: str) -> dict:
"""解析并校验 AI 返回的 JSON → {"name", "conditions"}。非法时抛 ValueError。
只取 name 与 conditionsid / kind 由用户在表单里填写,不信任 AI。
"""
raw = _extract_json_object(text)
if not isinstance(raw, dict):
raise ValueError("AI 返回的 JSON 不是对象")
name = raw.get("name")
if not isinstance(name, str) or not name.strip():
raise ValueError("AI 未返回信号名称 name")
name = name.strip()[:30]
conditions_raw = raw.get("conditions")
if not isinstance(conditions_raw, list) or not conditions_raw:
raise ValueError("AI 未返回任何条件 conditions")
conditions = [_normalize_condition(c) for c in conditions_raw]
# 复用现有白名单/运算符/偏移校验作为安全闸门(id/kind 用占位值)。
probe = {
"id": "aigenerated",
"name": name,
"kind": "both",
"conditions": conditions,
}
custom_signals.validate(probe)
return {"name": name, "conditions": conditions}
def _normalize_condition(c: object) -> dict:
if not isinstance(c, dict):
raise ValueError("条件的每一项必须是 JSON 对象")
left = c.get("left")
op = c.get("op")
right = c.get("right")
if right is None:
raise ValueError("条件缺少右值 right")
if isinstance(right, bool):
raise ValueError("右值不能是布尔值")
if isinstance(right, (int, float)):
right = _num_to_str(right)
if not isinstance(right, str) or not right.strip():
raise ValueError(f"右值非法: {right!r}")
right = right.strip()
# 兜底: AI 偶尔漏写 field: 前缀的裸字段名, 补全为规范形式
if not right.startswith("field:") and right in custom_signals.ALLOWED_FIELDS:
right = f"field:{right}"
return {
"left": str(left),
"op": str(op),
"right": right,
"leftDays": _norm_days(c.get("leftDays")),
"rightDays": _norm_days(c.get("rightDays")),
}
def _norm_days(value: object) -> object:
"""归一化日期偏移为 int(缺省 0);无法转换时原样返回,由 validate 报中文错误。"""
if value is None:
return 0
if isinstance(value, bool):
return value
try:
return int(value)
except (TypeError, ValueError):
return value
def _num_to_str(value) -> str:
if isinstance(value, int):
return str(value)
f = float(value)
return str(int(f)) if f.is_integer() else str(f)
def _extract_json_object(text: str) -> object:
"""从 LLM 文本提取 JSON 对象(多级容错)。
依次尝试: 整段 → markdown 围栏内 → 首个 {...} 平衡块;
每级再对 尾随垃圾 / 尾逗号 做轻量修复。全部失败才报错。
"""
source = text or ""
candidates: list[str] = []
stripped = source.strip()
if stripped:
candidates.append(stripped)
candidates.extend(
match.group(1).strip() for match in _FENCED_JSON_RE.finditer(source)
)
brace = _first_brace_block(source)
if brace and brace.strip() not in candidates:
candidates.append(brace)
last_error: Exception | None = None
for candidate in candidates:
parsed = _try_parse_json(candidate)
if parsed is not None:
return parsed
try:
json.loads(candidate)
except json.JSONDecodeError as e:
last_error = e
raise ValueError(f"AI 返回的不是合法 JSON: {last_error}")
def _try_parse_json(candidate: str) -> object | None:
"""尽力解析一段可能带尾随垃圾 / 尾逗号的 JSON;失败返 None。"""
variants = [candidate.strip()]
last = candidate.rfind("}")
if last >= 0 and last < len(candidate) - 1:
variants.append(candidate[:last + 1].strip())
for v in variants:
if not v:
continue
try:
return json.loads(v)
except json.JSONDecodeError:
pass
# 去掉数组/对象结尾的多余逗号 (AI 常见错误): `,}` / `,]`
cleaned = re.sub(r",\s*([}\]])", r"\1", v)
if cleaned != v:
try:
return json.loads(cleaned)
except json.JSONDecodeError:
pass
return None
def _first_brace_block(text: str) -> str:
"""括号配对截取首个 {...} 块(AI 偶尔混入前后解释文字时的兜底)。"""
start = text.find("{")
if start < 0:
return text
depth = 0
for i in range(start, len(text)):
if text[i] == "{":
depth += 1
elif text[i] == "}":
depth -= 1
if depth == 0:
return text[start:i + 1]
return text[start:]