57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
"""精简版 LLM 客户端 — 只保留智谱 GLM-5-turbo-nothinking,占坑期够用。"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import time
|
|
from typing import Any
|
|
|
|
from zai import ZhipuAiClient
|
|
|
|
logger = logging.getLogger("shagua.llm")
|
|
|
|
# 占坑期硬编码,与 pricebot-backend 同 key
|
|
ZHIPU_API_KEY = "c298ebdfda044e0387a3dc571f98ceed.QDY8lVCcGt04C6BD"
|
|
MODEL = "glm-5-turbo"
|
|
|
|
# 占坑期联调开关:DUOBIBI_MOCK_LLM=1 时各端点跳过真实 LLM 调用,改用本地确定性
|
|
# 算法造数据(不烧额度 / 不依赖外网,见 app/api/arena.py、app/api/worth_buy.py)。
|
|
MOCK_LLM = os.environ.get("DUOBIBI_MOCK_LLM", "").strip().lower() in ("1", "true", "yes", "on")
|
|
EXTRA_PARAMS: dict[str, Any] = {"thinking": {"type": "disabled"}}
|
|
|
|
DEFAULT_TEMPERATURE = 0.0
|
|
DEFAULT_MAX_TOKENS = 1024
|
|
|
|
_client = ZhipuAiClient(api_key=ZHIPU_API_KEY)
|
|
|
|
|
|
def chat(
|
|
messages: list[dict[str, Any]],
|
|
temperature: float = DEFAULT_TEMPERATURE,
|
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
) -> str:
|
|
start = time.time()
|
|
response = _client.chat.completions.create(
|
|
model=MODEL,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
**EXTRA_PARAMS,
|
|
)
|
|
latency_ms = int((time.time() - start) * 1000)
|
|
|
|
content = (response.choices[0].message.content or "").strip()
|
|
usage = response.usage
|
|
if usage is not None:
|
|
logger.info(
|
|
"[LLM] model=%s tokens in=%s out=%s total=%s latency=%sms",
|
|
MODEL,
|
|
usage.prompt_tokens,
|
|
usage.completion_tokens,
|
|
usage.total_tokens,
|
|
latency_ms,
|
|
)
|
|
else:
|
|
logger.info("[LLM] model=%s latency=%sms (no usage)", MODEL, latency_ms)
|
|
return content
|