fateforge/llm.py

"""LLM 流式调用:与本地 llama-server 的 OpenAI 兼容 /v1/chat/completions 通信。"""
from __future__ import annotations

import json
from typing import AsyncIterator, Dict, List

import httpx

from .config import (
    LLM_BASE_URL,
    LLM_MAX_TOKENS,
    LLM_MODEL,
    LLM_TEMPERATURE,
    LLM_TIMEOUT_S,
)


async def stream_chat(
    system: str,
    user: str,
    *,
    max_tokens: int | None = None,
    temperature: float | None = None,
) -> AsyncIterator[str]:
    """流式调用 LLM,逐 chunk yield 文本。

    llama.cpp 的 OpenAI 兼容端点返回 SSE 格式:
        data: {"choices":[{"delta":{"content":"..."}}]}
        data: [DONE]
    """
    payload: Dict = {
        "model": LLM_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "max_tokens": max_tokens or LLM_MAX_TOKENS,
        "temperature": temperature if temperature is not None else LLM_TEMPERATURE,
        "stream": True,
    }
    headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}

    timeout = httpx.Timeout(
        connect=min(3.0, LLM_TIMEOUT_S),
        read=min(8.0, LLM_TIMEOUT_S),
        write=min(8.0, LLM_TIMEOUT_S),
        pool=min(3.0, LLM_TIMEOUT_S),
    )
    url = f"{LLM_BASE_URL.rstrip('/')}/v1/chat/completions"

    async with httpx.AsyncClient(timeout=timeout) as client:
        async with client.stream("POST", url, json=payload, headers=headers) as resp:
            resp.raise_for_status()
            async for raw in resp.aiter_lines():
                line = raw.strip()
                if not line.startswith("data:"):
                    continue
                data = line[5:].strip()
                if data == "[DONE]":
                    break
                try:
                    obj = json.loads(data)
                except json.JSONDecodeError:
                    continue
                try:
                    delta = obj["choices"][0].get("delta", {}) or {}
                    chunk = delta.get("content", "")
                except (KeyError, IndexError, TypeError):
                    chunk = ""
                if chunk:
                    yield chunk


async def stream_fate_reading(
    question: str,
    hex_info: dict,
    changed_info: dict | None,
    moving: List[int],
    level: str = "中平",
    reason: str = "",
    ui: dict | None = None,
    lang: str = "zh",
) -> AsyncIterator[str]:
    """组装 System+User Prompt,流式返回 LLM 文本。"""
    from .constants import build_system_prompt, build_user_prompt
    ui = ui or {"remedy": "自渡锦囊"}
    system_prompt = build_system_prompt(level, ui, lang=lang)
    user_prompt = build_user_prompt(question, hex_info, changed_info, moving, level, reason, lang=lang)
    async for chunk in stream_chat(system_prompt, user_prompt):
        yield chunk