llm.py

"""
deploy/llm.py — LLM 流式/非流式调用

从 config 导入 API 配置，提供 stream_llm() 和 call_llm() 接口。
"""
import json
import time
import urllib.request
import urllib.error

from config import API_URL, MODEL_NAME, API_KEY, DEFAULT_TIMEOUT


def _estimate_tokens(text):
    """Rough token estimate: ~1 token per 1.5 chars for mixed CJK + English."""
    return max(1, int(len(text) / 1.5))


def _format_speed(start_time, reason_buf, content_buf, reason_start, content_start):
    """Format live speed stats for the UI."""
    elapsed = time.time() - start_time
    parts = []

    if reason_buf:
        reason_tokens = _estimate_tokens(reason_buf)
        reason_elapsed = time.time() - reason_start if reason_start else elapsed
        if reason_elapsed > 0:
            reason_speed = reason_tokens / reason_elapsed
            parts.append(f"🧠 {reason_speed:.0f} tok/s")

    if content_buf:
        content_tokens = _estimate_tokens(content_buf)
        content_elapsed = time.time() - content_start if content_start else elapsed
        if content_elapsed > 0:
            content_speed = content_tokens / content_elapsed
            parts.append(f"💬 {content_speed:.0f} tok/s")

    total = _estimate_tokens(reason_buf) + _estimate_tokens(content_buf)
    if total > 1:
        parts.append(f"📊 {total} tok")

    if not parts:
        return "⏱️ 思考中..."

    return "  ·  ".join(parts)


def stream_llm(messages, max_tokens=4096, temperature=0.7):
    """Stream the LLM API (Modal OpenAI-compatible endpoint).

    Yields (content, reasoning, speed) tuples.
    Uses a 50ms debounce buffer: accumulates chunks within 50ms windows
    before yielding, reducing Gradio frontend repaint frequency.
    """
    if not API_URL:
        yield "❌ API URL 未配置", "", "⏱️ --"
        return

    payload = json.dumps({
        "model": MODEL_NAME,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": True
    }).encode()

    headers = {
        "Content-Type": "application/json",
        "Accept": "text/event-stream"
    }
    if API_KEY:
        headers["Authorization"] = f"Bearer {API_KEY}"

    req = urllib.request.Request(API_URL, data=payload, headers=headers)

    DEBOUNCE_MS = 50  # batch chunks within this window

    reason_buf = ""
    content_buf = ""
    start_time = time.time()
    reason_start = None
    content_start = None

    try:
        with urllib.request.urlopen(req, timeout=DEFAULT_TIMEOUT) as resp:
            last_yield_time = time.time()
            done = False

            for line in resp:
                line = line.decode("utf-8").strip()
                if not line or not line.startswith("data: "):
                    continue

                data = line[6:]
                if data == "[DONE]":
                    done = True
                    # Final yield with remaining buffer
                    break

                try:
                    chunk = json.loads(data)
                    delta = chunk.get("choices", [{}])[0].get("delta", {})

                    if "reasoning_content" in delta and delta["reasoning_content"]:
                        if reason_start is None:
                            reason_start = time.time()
                        reason_buf += delta["reasoning_content"]

                    if "content" in delta and delta["content"]:
                        if content_start is None:
                            content_start = time.time()
                        content_buf += delta["content"]

                    # Debounce: only yield if enough time has passed
                    now = time.time()
                    if (now - last_yield_time) * 1000 >= DEBOUNCE_MS:
                        speed = _format_speed(start_time, reason_buf, content_buf, reason_start, content_start)
                        yield content_buf or "🧠 思考中...", reason_buf, speed
                        last_yield_time = now

                except json.JSONDecodeError:
                    continue

            # Always yield final state after stream ends
            speed = _format_speed(start_time, reason_buf, content_buf, reason_start, content_start)
            yield content_buf or "🧠 思考中...", reason_buf, speed

    except urllib.error.URLError as e:
        yield f"❌ 网络错误: {e}", reason_buf, "⏱️ --"
    except Exception as e:
        yield f"❌ API error: {e}", reason_buf, "⏱️ --"


def call_llm(messages, max_tokens=4096, temperature=0.7):
    """Call the LLM API and return the final (content, reasoning) tuple."""
    for content, reasoning, _ in stream_llm(messages, max_tokens, temperature):
        pass
    return content, reasoning