nba-fan-hub/llm.py

# -*- coding: utf-8 -*-
"""DeepSeek 大模型调用层（OpenAI 兼容 /chat/completions，支持 function calling）"""
import json
import logging
import time

import requests

from config import LLM_BASE_URL, LLM_API_KEY, LLM_MODEL, LLM_TIMEOUT, LLM_MAX_TOKENS, LLM_TEMPERATURE

log = logging.getLogger("llm")

URL = f"{LLM_BASE_URL}/chat/completions"
HEADERS = {"Authorization": f"Bearer {LLM_API_KEY}", "Content-Type": "application/json"}


def _post(payload, timeout=LLM_TIMEOUT):
    resp = requests.post(URL, headers=HEADERS, json=payload, timeout=timeout)
    resp.raise_for_status()
    return resp.json()


def chat(messages, tools=None, temperature=LLM_TEMPERATURE, max_tokens=LLM_MAX_TOKENS):
    """基础对话。返回完整 OpenAI 响应 dict。
    tools 为 None 时不带工具；带工具时模型可能返回 tool_calls。"""
    payload = {
        "model": LLM_MODEL,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": False,
    }
    if tools:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"
    for attempt in range(3):
        try:
            return _post(payload)
        except requests.exceptions.HTTPError as e:
            body = e.response.text[:400] if e.response is not None else ""
            code = e.response.status_code if e.response is not None else 0
            if code in (400, 429, 500, 502, 503) and attempt < 2:
                log.warning("LLM HTTP %s(第%s次重试): %s", code, attempt + 1, body)
                time.sleep(2 * (attempt + 1))
                continue
            log.error("LLM HTTP %s: %s", code, body)
            raise
    raise RuntimeError("LLM 请求失败")


def parse_content(resp):
    """提取 assistant 文本内容"""
    try:
        return resp["choices"][0]["message"].get("content") or ""
    except Exception:
        return ""


def extract_tool_calls(resp):
    """提取 [{name, arguments(dict), id}]"""
    calls = []
    try:
        msg = resp["choices"][0]["message"]
        for tc in msg.get("tool_calls") or []:
            try:
                args = json.loads(tc["function"].get("arguments") or "{}")
            except json.JSONDecodeError:
                args = {}
            calls.append({"name": tc["function"]["name"], "arguments": args, "id": tc["id"]})
    except Exception:
        pass
    return calls


def extract_reasoning(resp):
    """提取思考模式下的 reasoning_content（回传时必须带上）"""
    try:
        return resp["choices"][0]["message"].get("reasoning_content") or ""
    except Exception:
        return ""


def count_tokens(messages):
    """粗略估算 token 数（中英混合，1字≈1token）"""
    total = 0
    for m in messages:
        total += len(m.get("content") or "") // 2 + 8
    return total