MoFin/deploy/profile-scripts/llm_client.py

#!/usr/bin/env python3
"""llm_client.py — 共享 LLM 客户端（重试 + gateway 预检）

所有重评脚本统一通过此模块调用 LLM gateway，避免重复的 HTTP/重试逻辑。

用法:
    from llm_client import call_llm, REASSESS_MODEL, gateway_alive
    if not gateway_alive():
        print("Gateway 不可用，退出")
        sys.exit(1)
    result = call_llm(prompt)
    if result["ok"]:
        full_text = result["content"]
"""

import json
import time
import urllib.request
import urllib.error

# ── 常量：所有重评调用统一使用 ──
REASSESS_MODEL = "deepseek-v4-pro"
GATEWAY = "http://127.0.0.1:8643/v1/chat/completions"
GATEWAY_BASE = "http://127.0.0.1:8643/v1/models"
AUTH = "Bearer hermes123"


def gateway_alive(timeout=5):
    """快速预检 gateway 是否存活。失败立刻返回 False，不阻塞。"""
    try:
        req = urllib.request.Request(GATEWAY_BASE, headers={"Authorization": AUTH})
        urllib.request.build_opener(urllib.request.ProxyHandler({})).open(req, timeout=timeout)
        return True
    except Exception:
        return False


def call_llm(prompt, model=None, max_tokens=4096, timeout=150,
             retries=1, backoff=20, system=None):
    """调用 LLM gateway，带重试和结构化日志。

    Args:
        prompt: 用户消息内容
        model: 模型名（默认 REASSESS_MODEL）
        max_tokens: 最大输出 token 数
        timeout: 单次调用超时（秒）
        retries: 超时/5xx/连接错误时的重试次数
        backoff: 重试间隔（秒）
        system: 可选 system message

    Returns:
        {ok: bool, content: str, error: str|None, model: str,
         elapsed: float, attempts: int}
        永远不抛异常到调用方。
    """
    model_name = model or REASSESS_MODEL
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    payload = json.dumps({
        "model": model_name,
        "messages": messages,
        "max_tokens": max_tokens,
    }).encode()

    for attempt in range(retries + 1):
        t0 = time.monotonic()
        try:
            req = urllib.request.Request(
                GATEWAY,
                data=payload,
                headers={
                    "Content-Type": "application/json",
                    "Authorization": AUTH,
                }
            )
            opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
            resp = opener.open(req, timeout=timeout)
            elapsed = time.monotonic() - t0

            body = json.loads(resp.read().decode())
            if "choices" not in body:
                msg = f"API响应无choices字段: {str(body)[:200]}"
                print(f"  [LLM] 尝试{attempt+1}/{retries+1} 失败({elapsed:.1f}s): {msg}", flush=True)
                if attempt < retries:
                    time.sleep(backoff)
                continue

            content = body["choices"][0]["message"]["content"]
            print(f"  [LLM] 尝试{attempt+1}/{retries+1} 成功, {elapsed:.1f}s, "
                  f"输出{len(content)}字, model={model_name}", flush=True)
            return {
                "ok": True,
                "content": content,
                "error": None,
                "model": model_name,
                "elapsed": elapsed,
                "attempts": attempt + 1,
            }

        except urllib.error.URLError as e:
            elapsed = time.monotonic() - t0
            err_msg = str(e)
            print(f"  [LLM] 尝试{attempt+1}/{retries+1} 连接失败({elapsed:.1f}s): {err_msg[:120]}", flush=True)
            if attempt < retries:
                print(f"  [LLM] 等待{backoff}s后重试...", flush=True)
                time.sleep(backoff)
            else:
                return {
                    "ok": False, "content": "", "error": f"连接失败(重试{retries}次): {err_msg[:200]}",
                    "model": model_name, "elapsed": elapsed, "attempts": attempt + 1,
                }

        except Exception as e:
            elapsed = time.monotonic() - t0
            err_msg = str(e)
            print(f"  [LLM] 尝试{attempt+1}/{retries+1} 失败({elapsed:.1f}s): {err_msg[:120]}", flush=True)
            if attempt < retries:
                print(f"  [LLM] 等待{backoff}s后重试...", flush=True)
                time.sleep(backoff)
            else:
                return {
                    "ok": False, "content": "", "error": f"调用失败(重试{retries}次): {err_msg[:200]}",
                    "model": model_name, "elapsed": elapsed, "attempts": attempt + 1,
                }

    # Unreachable
    return {
        "ok": False, "content": "", "error": "未预期的调用结束",
        "model": model_name, "elapsed": 0, "attempts": retries + 1,
    }