xiaobaifupan/llm_strategy.py

from __future__ import annotations

import json
import time
import urllib.error
import urllib.request
from typing import Any

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 或模型未配置。")
    endpoint = f"{base_url.rstrip('/')}/chat/completions"
    payload = json.dumps(
        {
            "model": model,
            "messages": [{"role": "user", "content": "只回复 OK"}],
            "stream": False,
        },
        ensure_ascii=False,
    ).encode("utf-8")
    request = urllib.request.Request(
        endpoint,
        data=payload,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
            "User-Agent": "XiaobaiReviewWeb/0.5",
        },
        method="POST",
    )
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            result = json.loads(response.read().decode("utf-8"))
        reply = str(result["choices"][0]["message"]["content"]).strip()
    except urllib.error.HTTPError as exc:
        raise LLMCompilerError(_http_error_message(exc)) from exc
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, IndexError) as exc:
        raise LLMCompilerError(f"模型连接测试失败：{exc}") from exc
    return {
        "ok": True,
        "model": model,
        "reply": reply[:100],
        "latency_ms": round((time.perf_counter() - started) * 1000),
    }


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 或模型。")
    endpoint = f"{base_url.rstrip('/')}/chat/completions"
    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)}"
    )
    payload = json.dumps(
        {
            "model": model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt[:3000]},
            ],
            "stream": False,
        },
        ensure_ascii=False,
    ).encode("utf-8")
    request = urllib.request.Request(
        endpoint,
        data=payload,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
            "User-Agent": "XiaobaiReviewWeb/0.4",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            result = json.loads(response.read().decode("utf-8"))
        content = result["choices"][0]["message"]["content"].strip()
        if content.startswith("```"):
            content = content.strip("`")
            if content.startswith("json"):
                content = content[4:].strip()
        compiled = json.loads(content)
    except urllib.error.HTTPError as exc:
        raise LLMCompilerError(_http_error_message(exc).replace("模型连接测试", "LLM 策略编译")) from exc
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, IndexError) as exc:
        raise LLMCompilerError(f"LLM 策略编译失败：{exc}") from exc
    compiled["compiler"] = "llm"
    compiled["model"] = model
    return compiled


def _http_error_message(exc: urllib.error.HTTPError) -> str:
    detail = ""
    try:
        payload = json.loads(exc.read().decode("utf-8", errors="replace"))
        error = payload.get("error")
        if isinstance(error, dict):
            detail = str(error.get("message") or error.get("code") or "")
        elif error:
            detail = str(error)
        elif payload.get("message"):
            detail = str(payload["message"])
    except (json.JSONDecodeError, OSError):
        detail = ""
    suffix = f"：{detail[:300]}" if detail else ""
    return f"模型连接测试失败（HTTP {exc.code}）{suffix}"