123 lines
4.8 KiB
Python
123 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
||
"""llm skill：New-API 统一调用入口（默认 MiniMax-M3）。
|
||
|
||
管线内所有内容生产型 LLM 调用（清洗探测/拆书抽取）必须经此入口：
|
||
- trust_env=False（本机代理环境变量会劫持内网直连，教训固化）；
|
||
- 超时 + 指数退避重试；<think> 剥离；JSON 三级容错提取（json-repair 兜底）；
|
||
- 每次调用向 stderr 打印 token 用量与耗时（成本审计），stdout 只出内容。
|
||
"""
|
||
import json
|
||
import pathlib
|
||
import re
|
||
import sys
|
||
import time
|
||
|
||
import click
|
||
import requests
|
||
|
||
BASE = "http://100.64.0.8:3000"
|
||
# New-API 普通令牌（仓库政策允许明文；严禁换管理令牌打 /v1）
|
||
TOKEN = "sk-DyVqO3lDmEvQZ3PqGpbNaaaHZHhbh0xaHRIiynhYSmVlLHl2"
|
||
DEFAULT_MODEL = "MiniMax-M3"
|
||
|
||
|
||
def chat(prompt, model=DEFAULT_MODEL, max_tokens=32000, temperature=0.2,
|
||
retries=2, timeout=900):
|
||
"""单轮对话，返回 (content, usage)。网络错/5xx/429 指数退避重试。
|
||
|
||
content 已剥离 <think>…</think>（推理模型可能把思考混进正文）。
|
||
"""
|
||
s = requests.Session()
|
||
s.trust_env = False # 本机代理 env 会劫持内网直连
|
||
payload = {
|
||
"model": model,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"max_tokens": max_tokens,
|
||
"temperature": temperature,
|
||
}
|
||
last_err = None
|
||
for attempt in range(retries + 1):
|
||
try:
|
||
t0 = time.time()
|
||
r = s.post(f"{BASE}/v1/chat/completions",
|
||
headers={"Authorization": f"Bearer {TOKEN}"},
|
||
json=payload, timeout=timeout)
|
||
# 429/5xx 属于可重试的服务端瞬时问题
|
||
if r.status_code in (429,) or r.status_code >= 500:
|
||
last_err = f"HTTP {r.status_code}: {r.text[:200]}"
|
||
raise requests.RequestException(last_err)
|
||
r.raise_for_status()
|
||
data = r.json()
|
||
content = data["choices"][0]["message"]["content"] or ""
|
||
content = re.sub(r"<think>.*?</think>", "", content, flags=re.S).strip()
|
||
usage = data.get("usage", {})
|
||
print(f"[llm] {model} in={usage.get('prompt_tokens', '?')} "
|
||
f"out={usage.get('completion_tokens', '?')} "
|
||
f"耗时{time.time() - t0:.0f}s finish={data['choices'][0].get('finish_reason')}",
|
||
file=sys.stderr)
|
||
return content, usage
|
||
except (requests.RequestException, KeyError, json.JSONDecodeError) as e:
|
||
last_err = str(e)
|
||
if attempt < retries:
|
||
wait = 8 * (2 ** attempt)
|
||
print(f"[llm] 第{attempt + 1}次失败({last_err[:120]})，{wait}s 后重试",
|
||
file=sys.stderr)
|
||
time.sleep(wait)
|
||
raise RuntimeError(f"LLM 调用重试耗尽: {last_err}")
|
||
|
||
|
||
def extract_json(text):
|
||
"""JSON 三级容错提取：直接解析 → 首尾括号截取 → json-repair 兜底。
|
||
|
||
opus 试拆实测过两类 JSON 病（中文引号、缺逗号）——任何模型都可能犯，统一在此兜住。
|
||
"""
|
||
try:
|
||
return json.loads(text)
|
||
except json.JSONDecodeError:
|
||
pass
|
||
# 剥 markdown 代码围栏后按最外层大括号/中括号截取
|
||
t = re.sub(r"^```(?:json)?\s*|\s*```$", "", text.strip(), flags=re.M)
|
||
for a, b in (("{", "}"), ("[", "]")):
|
||
i, j = t.find(a), t.rfind(b)
|
||
if i != -1 and j > i:
|
||
frag = t[i:j + 1]
|
||
try:
|
||
return json.loads(frag)
|
||
except json.JSONDecodeError:
|
||
import json_repair
|
||
return json_repair.loads(frag)
|
||
import json_repair
|
||
return json_repair.loads(t)
|
||
|
||
|
||
@click.group()
|
||
def cli():
|
||
pass
|
||
|
||
|
||
@cli.command("chat")
|
||
@click.option("--prompt-file", type=click.Path(exists=True), required=True)
|
||
@click.option("--model", default=DEFAULT_MODEL, show_default=True)
|
||
@click.option("--max-tokens", type=int, default=32000, show_default=True)
|
||
@click.option("--temperature", type=float, default=0.2, show_default=True)
|
||
@click.option("--out", type=click.Path(), help="内容写文件（不给则打 stdout）")
|
||
@click.option("--extract-json", "extract_", is_flag=True, help="容错提取 JSON 后输出")
|
||
def chat_cmd(prompt_file, model, max_tokens, temperature, out, extract_):
|
||
prompt = pathlib.Path(prompt_file).read_text()
|
||
content, _ = chat(prompt, model=model, max_tokens=max_tokens, temperature=temperature)
|
||
if extract_:
|
||
content = json.dumps(extract_json(content), ensure_ascii=False, indent=1)
|
||
if out:
|
||
pathlib.Path(out).write_text(content)
|
||
click.echo(f"已写 {out}（{len(content):,} 字符）", err=True)
|
||
else:
|
||
click.echo(content)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
cli()
|
||
except RuntimeError as e:
|
||
click.echo(f"[llm错误] {e}", err=True)
|
||
sys.exit(1)
|