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

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

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

import time
from typing import Any, Callable, Optional

import httpx

from core.config import Settings
from core.jsonio import extract_json_object
from pipeline.tracing import TraceContext, TraceWriter, hash_prompt, redact_headers, timed_ms


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,
    trace_writer: TraceWriter | None = None,
    trace_context: TraceContext | None = None,
    trace_stage: str = "decode",
    trace_substage: str = "chat_json",
    prompt_name: str = "chat_json",
) -> 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):
        started = time.perf_counter()
        response_json: dict[str, Any] = {}
        request_payload = {
            "model": model,
            "messages": messages,
            "response_format": {"type": "json_object"},
        }
        try:
            resp = http_post(url, headers=headers, json=request_payload, timeout=timeout)
            resp.raise_for_status()
            response_json = resp.json()
            content = response_json["choices"][0]["message"]["content"]
        except httpx.HTTPError as exc:
            last_exc = exc
            if trace_writer is not None:
                trace_writer.llm_call(
                    context=trace_context or TraceContext(stage=trace_stage, substage=trace_substage),
                    stage=trace_stage,
                    substage=trace_substage,
                    provider="bailian",
                    model_name=model,
                    endpoint=url,
                    prompt_name=prompt_name,
                    prompt_hash=hash_prompt(system),
                    request_payload={"headers": redact_headers(headers), **request_payload},
                    status="failed",
                    error_message=str(exc),
                    latency_ms=timed_ms(started),
                    attempt_index=attempt + 1,
                )
            if attempt < 2:
                continue
            raise LLMError(f"llm_http_error: {exc}") from exc
        except (KeyError, IndexError, TypeError) as exc:
            if trace_writer is not None:
                trace_writer.llm_call(
                    context=trace_context or TraceContext(stage=trace_stage, substage=trace_substage),
                    stage=trace_stage,
                    substage=trace_substage,
                    provider="bailian",
                    model_name=model,
                    endpoint=url,
                    prompt_name=prompt_name,
                    prompt_hash=hash_prompt(system),
                    request_payload={"headers": redact_headers(headers), **request_payload},
                    response_payload=response_json,
                    status="failed",
                    error_message=str(exc),
                    latency_ms=timed_ms(started),
                    attempt_index=attempt + 1,
                )
            raise LLMError(f"llm_response_invalid: {exc}") from exc
        try:
            parsed = extract_json_object(content)
            if trace_writer is not None:
                trace_writer.llm_call(
                    context=trace_context or TraceContext(stage=trace_stage, substage=trace_substage),
                    stage=trace_stage,
                    substage=trace_substage,
                    provider="bailian",
                    model_name=model,
                    endpoint=url,
                    prompt_name=prompt_name,
                    prompt_hash=hash_prompt(system),
                    request_payload={"headers": redact_headers(headers), **request_payload},
                    response_payload=response_json,
                    parsed_payload=parsed,
                    status="done",
                    latency_ms=timed_ms(started),
                    attempt_index=attempt + 1,
                )
            return parsed
        except ValueError as exc:
            last_exc = exc
            if trace_writer is not None:
                trace_writer.llm_call(
                    context=trace_context or TraceContext(stage=trace_stage, substage=trace_substage),
                    stage=trace_stage,
                    substage=trace_substage,
                    provider="bailian",
                    model_name=model,
                    endpoint=url,
                    prompt_name=prompt_name,
                    prompt_hash=hash_prompt(system),
                    request_payload={"headers": redact_headers(headers), **request_payload},
                    response_payload={"content": content},
                    status="failed",
                    error_message=str(exc),
                    latency_ms=timed_ms(started),
                    attempt_index=attempt + 1,
                )
            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
    )