projects/repobot/v2/llm.py

"""和大模型打交道的部分：客户端、计费、带重试的请求。和 v1 相比只是从 repobot.py 里挪了出来。"""
import os
import random
import sys
import time

import openai
from openai import OpenAI

MODEL = os.environ.get("LLM_MODEL", "deepseek-flash")

# 美元 / 每一百万词元，高峰价，截至 2026 年 9 月。用之前去 DeepSeek 的价格页面核对
PRICES = {"deepseek-flash": (0.006, 0.30, 1.20), "deepseek-v4-pro": (0.044, 1.32, 3.96)}

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.deepseek.com"),
    timeout=60,
    max_retries=0,  # 重试由下面的 with_retry 负责
)
RETRYABLE = (openai.RateLimitError, openai.APITimeoutError, openai.APIConnectionError, openai.InternalServerError)


def cost_usd(usage):
    if MODEL not in PRICES or usage is None:
        return 0.0
    hit_price, miss_price, out_price = PRICES[MODEL]
    hit = getattr(usage, "prompt_cache_hit_tokens", 0) or 0
    miss = usage.prompt_tokens - hit
    return (hit * hit_price + miss * miss_price + usage.completion_tokens * out_price) / 1_000_000


def with_retry(make_request, max_attempts=4):
    """执行一次请求，遇到可以重试的错误就按指数退避重试。流式请求只在建立连接时重试。"""
    for attempt in range(1, max_attempts + 1):
        try:
            return make_request()
        except RETRYABLE as e:
            if attempt == max_attempts:
                raise
            wait = 2 ** (attempt - 1) + random.random()
            print(f"\n[{type(e).__name__}，{wait:.1f} 秒后重试]", file=sys.stderr)
            time.sleep(wait)


def chat(messages, **kwargs):
    """普通的一次性调用，返回 (文本, usage)。"""
    response = with_retry(lambda: client.chat.completions.create(
        model=MODEL, messages=messages, extra_body={"thinking": {"type": "disabled"}}, **kwargs))
    return response.choices[0].message.content, response.usage


def open_stream(messages, thinking=False):
    return with_retry(lambda: client.chat.completions.create(
        model=MODEL,
        messages=messages,
        stream=True,
        stream_options={"include_usage": True},
        max_tokens=4000,
        extra_body={"thinking": {"type": "enabled" if thinking else "disabled"}},
    ))