Create-knowledge-find-decode/core/llm.py

"""文本 LLM：判断 / 拆分 / 解构用，走阿里云百炼 OpenAI-compatible /chat/completions。

只做一件事：给 system + user，返回解析好的 JSON dict。HTTP/鉴权风格同 extractor。
"""
from __future__ import annotations

from typing import Any, Callable, Optional

import httpx

from core.config import Settings
from core.jsonio import extract_json_object


class LLMError(RuntimeError):
    pass


def chat_json(
    system: str,
    user: str,
    *,
    model: Optional[str] = None,
    settings: Optional[Settings] = None,
    http_post: Callable[..., Any] = httpx.post,
    env_file: str = ".env",
    timeout: float = 60.0,
) -> dict:
    """调一次对话，强约束输出 JSON，返回解析后的 dict。带一次重试。"""
    settings = settings or Settings.from_env(env_file)
    api_key = settings.bailian_api_key
    if not api_key:
        raise LLMError("missing ALIYUN_BAILIAN_API_KEY")
    model = model or settings.llm_model
    messages = [
        {"role": "system", "content": system + (
            "\n只输出一个严格合法的 JSON 对象，不要解释或 markdown。"
            "字符串值要写在一行内，内部的换行写成 \\n、双引号写成 \\\"，不要出现裸换行或裸双引号。"
        )},
        {"role": "user", "content": user},
    ]
    url = f"{settings.bailian_base_url.rstrip('/')}/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    last_exc: Optional[Exception] = None
    # 最多 3 轮：HTTP 错误重试；JSON 不合法则把坏输出喂回去做 self-repair
    for attempt in range(3):
        try:
            resp = http_post(url, headers=headers, json={
                "model": model,
                "messages": messages,
                "response_format": {"type": "json_object"},
            }, timeout=timeout)
            resp.raise_for_status()
            content = resp.json()["choices"][0]["message"]["content"]
        except httpx.HTTPError as exc:
            last_exc = exc
            if attempt < 2:
                continue
            raise LLMError(f"llm_http_error: {exc}") from exc
        except (KeyError, IndexError, TypeError) as exc:
            raise LLMError(f"llm_response_invalid: {exc}") from exc
        try:
            return extract_json_object(content)
        except ValueError as exc:
            last_exc = exc
            messages = messages + [
                {"role": "assistant", "content": content},
                {"role": "user", "content": (
                    "上面的输出不是严格合法的 JSON。请只重新输出严格合法的 JSON，"
                    "字符串内的换行写成 \\n、双引号写成 \\\"，不要任何解释或 markdown。"
                )},
            ]
    raise LLMError(f"llm_json_unrepairable: {last_exc}")


# 默认对话器类型：(system, user) -> dict。stage 可注入假实现做离线测试。
ChatFn = Callable[[str, str], dict]


def default_chat(env_file: str = ".env", model: Optional[str] = None) -> ChatFn:
    settings = Settings.from_env(env_file)
    return lambda system, user: chat_json(
        system, user, model=model, settings=settings
    )