101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
from typing import Any
|
||
|
||
from backend.llm import transport as llm_transport
|
||
from screener import FACTOR_FIELDS, REGIMES
|
||
|
||
|
||
class LLMCompilerError(RuntimeError):
|
||
pass
|
||
|
||
|
||
def test_llm_connection(
|
||
api_key: str,
|
||
base_url: str,
|
||
model: str,
|
||
timeout: int = 30,
|
||
) -> dict[str, Any]:
|
||
if not api_key or not model:
|
||
raise LLMCompilerError("API Key 或模型未配置。")
|
||
try:
|
||
result = llm_transport.chat_completion(
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
model=model,
|
||
messages=[{"role": "user", "content": "只回复 OK"}],
|
||
timeout=timeout,
|
||
user_agent="XiaobaiReviewWeb/0.5",
|
||
)
|
||
reply = str(result.content).strip()
|
||
except llm_transport.OpenAIHTTPError as exc:
|
||
raise LLMCompilerError(exc.describe("模型连接测试失败")) from exc
|
||
except llm_transport.OpenAITransportError as exc:
|
||
raise LLMCompilerError(f"模型连接测试失败：{exc}") from exc
|
||
return {
|
||
"ok": True,
|
||
"model": model,
|
||
"reply": reply[:100],
|
||
"latency_ms": result.latency_ms,
|
||
}
|
||
|
||
|
||
def compile_strategy_with_llm(
|
||
prompt: str,
|
||
regime: str,
|
||
api_key: str,
|
||
base_url: str,
|
||
model: str,
|
||
timeout: int = 45,
|
||
) -> dict[str, Any]:
|
||
if not api_key or not model:
|
||
raise LLMCompilerError("尚未配置 LLM API Key 或模型。")
|
||
schema = {
|
||
"name": "策略名称",
|
||
"description": "策略说明",
|
||
"regimes": [regime],
|
||
"formula": {
|
||
"universe": {"exclude_st": True, "listed_days_min": 120},
|
||
"filters": [{"field": "return_5d", "op": ">=", "value": 0}],
|
||
"score": [{"field": "sector_strength", "weight": 0.3, "direction": "desc"}],
|
||
"limit": 15,
|
||
"min_score": 0.55,
|
||
},
|
||
}
|
||
system_prompt = (
|
||
"你是A股量化策略编译器。只输出JSON对象，不输出Markdown。"
|
||
"不得生成Python、SQL、网络请求或未提供的因子。"
|
||
f"当前市场阶段为{REGIMES.get(regime, regime)}。"
|
||
f"可用因子为：{json.dumps(FACTOR_FIELDS, ensure_ascii=False)}。"
|
||
"运算符只能使用 >, >=, <, <=, ==, !=, between, in。"
|
||
"score权重均大于0且不超过1，direction只能是asc或desc。"
|
||
"退潮和冰点策略必须提高门槛并允许结果为空。"
|
||
f"严格遵循以下结构：{json.dumps(schema, ensure_ascii=False)}"
|
||
)
|
||
try:
|
||
result = llm_transport.chat_completion(
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": prompt[:3000]},
|
||
],
|
||
timeout=timeout,
|
||
user_agent="XiaobaiReviewWeb/0.4",
|
||
)
|
||
content = result.content.strip()
|
||
if content.startswith("```"):
|
||
content = content.strip("`")
|
||
if content.startswith("json"):
|
||
content = content[4:].strip()
|
||
compiled = json.loads(content)
|
||
except llm_transport.OpenAIHTTPError as exc:
|
||
raise LLMCompilerError(exc.describe("LLM 策略编译失败")) from exc
|
||
except (llm_transport.OpenAITransportError, json.JSONDecodeError) as exc:
|
||
raise LLMCompilerError(f"LLM 策略编译失败：{exc}") from exc
|
||
compiled["compiler"] = "llm"
|
||
compiled["model"] = model
|
||
return compiled
|