news/llm/extractor.py

"""LLM 投资事件抽取主流程。

输入:Article(M2/M3 输出)
输出:ExtractedEvent(含 Pydantic 校验过的 EventExtraction)

设计:
    1. 加载 prompts/event_extraction.md,字符串替换填入文章字段;
    2. 调用 LLM JSON mode (response_format={"type":"json_object"});
    3. 解析 JSON -> Pydantic EventExtraction(强校验)+ 重试;
    4. 限制正文长度避免触顶 context window。
"""

from __future__ import annotations

import asyncio
import json
from pathlib import Path
from typing import Any, Protocol

from loguru import logger
from openai import AsyncOpenAI, OpenAI

from extractor import Article

from .client import LLMConfig
from .models import EventExtraction, ExtractedEvent, LLMCallError

# Prompt 模板默认路径
DEFAULT_PROMPT_PATH = Path("prompts/event_extraction.md")

# 文章正文截断长度(防止超出上下文窗口,DeepSeek/Qwen 都支持 32K+,这里保守取 8K 字符)
MAX_CONTENT_CHARS = 8000

# 重试设置
DEFAULT_MAX_ATTEMPTS = 3
RETRY_BASE_WAIT_SEC = 1.0
RETRY_MAX_WAIT_SEC = 8.0


# --------------------------------------------------------------------------- #
# Prompt 渲染
# --------------------------------------------------------------------------- #

class PromptTemplate:
    """Prompt 模板加载器,支持 {placeholder} 字符串替换。"""

    def __init__(self, template_path: str | Path = DEFAULT_PROMPT_PATH) -> None:
        self._path = Path(template_path)
        self._template = self._path.read_text(encoding="utf-8")

    def render(self, article: Article) -> str:
        content = article.content
        if len(content) > MAX_CONTENT_CHARS:
            logger.debug(
                "文章 {} 超长截断: {} -> {}",
                article.url_hash, len(content), MAX_CONTENT_CHARS,
            )
            content = content[:MAX_CONTENT_CHARS] + "\n\n[正文过长已截断]"

        publish_time_str = (
            article.publish_time.strftime("%Y-%m-%d %H:%M")
            if article.publish_time
            else "未知"
        )
        return (
            self._template
            .replace("{title}", article.title)
            .replace("{publish_time}", publish_time_str)
            .replace("{source_name}", article.source_name or article.source_id)
            .replace("{content}", content)
        )


# --------------------------------------------------------------------------- #
# JSON 提取(LLM 偶尔会包 ```json 围栏)
# --------------------------------------------------------------------------- #

def _extract_json_object(text: str) -> str:
    """从 LLM 输出中提取首个 JSON 对象字符串(去围栏 / 取首个 {...})。"""
    s = text.strip()
    if s.startswith("```"):
        # 去除 ```json ... ``` 围栏
        s = s.strip("`")
        # 可能形如 "json\n{...}"
        if s.lower().startswith("json"):
            s = s[4:].lstrip("\n").lstrip()
        # 末尾可能还有 ```
        if s.endswith("```"):
            s = s[:-3]
    # 取首个 { 到对应 }
    start = s.find("{")
    if start < 0:
        return s
    depth = 0
    for i in range(start, len(s)):
        if s[i] == "{":
            depth += 1
        elif s[i] == "}":
            depth -= 1
            if depth == 0:
                return s[start : i + 1]
    return s[start:]


def parse_event_json(raw: str) -> EventExtraction:
    """把 LLM 输出文本解析为 EventExtraction(可能抛 LLMCallError)。"""
    payload = _extract_json_object(raw)
    try:
        obj = json.loads(payload)
    except json.JSONDecodeError as e:
        raise LLMCallError(f"JSON 解析失败: {e}") from e
    if not isinstance(obj, dict):
        raise LLMCallError(f"JSON 顶层非对象: {type(obj).__name__}")
    try:
        return EventExtraction.model_validate(obj)
    except Exception as e:  # noqa: BLE001 - pydantic ValidationError 等多类型
        raise LLMCallError(f"事件 schema 校验失败: {e}") from e


# --------------------------------------------------------------------------- #
# Article -> ExtractedEvent
# --------------------------------------------------------------------------- #

class _SyncChat(Protocol):
    def chat(self, *args: Any, **kwargs: Any) -> Any: ...


def _call_llm_sync(
    client: OpenAI,
    config: LLMConfig,
    prompt: str,
) -> tuple[str, dict[str, int | None]]:
    """同步单次 LLM 调用,返回 (raw_text, usage)。usage 含 prompt_tokens / completion_tokens。"""
    resp = client.chat.completions.create(
        model=config.model,
        messages=[
            {
                "role": "system",
                "content": "你是 A 股投资研究助手,严格按用户指定的 JSON 格式输出。",
            },
            {"role": "user", "content": prompt},
        ],
        temperature=config.temperature,
        response_format={"type": "json_object"},
    )
    text = resp.choices[0].message.content or ""
    usage = {
        "prompt_tokens": getattr(resp.usage, "prompt_tokens", None) if resp.usage else None,
        "completion_tokens": (
            getattr(resp.usage, "completion_tokens", None) if resp.usage else None
        ),
    }
    return text, usage


async def _call_llm_async(
    client: AsyncOpenAI,
    config: LLMConfig,
    prompt: str,
) -> tuple[str, dict[str, int | None]]:
    """异步单次 LLM 调用。"""
    resp = await client.chat.completions.create(
        model=config.model,
        messages=[
            {
                "role": "system",
                "content": "你是 A 股投资研究助手,严格按用户指定的 JSON 格式输出。",
            },
            {"role": "user", "content": prompt},
        ],
        temperature=config.temperature,
        response_format={"type": "json_object"},
    )
    text = resp.choices[0].message.content or ""
    usage = {
        "prompt_tokens": getattr(resp.usage, "prompt_tokens", None) if resp.usage else None,
        "completion_tokens": (
            getattr(resp.usage, "completion_tokens", None) if resp.usage else None
        ),
    }
    return text, usage


def extract_event(
    client: OpenAI,
    config: LLMConfig,
    article: Article,
    *,
    template: PromptTemplate | None = None,
    max_attempts: int | None = None,
    sources: list[str] | None = None,
) -> ExtractedEvent:
    """同步抽取单篇文章的事件(带重试)。

    max_attempts 为 None 时使用 config.max_attempts(来自 YAML/环境变量配置)。
    sources 为该新闻全部来源(来自去重层多源记录);None 时兜底 [article.source_id]。
    """
    tpl = template or PromptTemplate()
    prompt = tpl.render(article)
    max_attempts = max_attempts or config.max_attempts

    last_err: Exception | None = None
    for attempt in range(1, max_attempts + 1):
        try:
            raw, usage = _call_llm_sync(client, config, prompt)
            event = parse_event_json(raw)
            return ExtractedEvent(
                source_id=article.source_id,
                url=article.url,
                url_hash=article.url_hash,
                title=article.title,
                publish_time=article.publish_time,
                sources=sources or [article.source_id],
                event=event,
                provider=config.provider,
                model=config.model,
                attempts=attempt,
                prompt_tokens=usage.get("prompt_tokens"),
                completion_tokens=usage.get("completion_tokens"),
            )
        except LLMCallError as e:
            last_err = e
            logger.warning(
                "LLM 抽取失败 url={} 尝试 {}/{}: {}",
                article.url, attempt, max_attempts, e.reason,
            )
        except Exception as e:  # noqa: BLE001 - 网络/限流等
            last_err = e
            logger.warning(
                "LLM 调用异常 url={} 尝试 {}/{}: {}: {}",
                article.url, attempt, max_attempts, type(e).__name__, e,
            )
        if attempt < max_attempts:
            wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC)
            import time

            time.sleep(wait)

    raise LLMCallError(
        f"LLM 抽取放弃,共 {max_attempts} 次尝试: {last_err}",
        attempts=max_attempts,
    )


async def extract_event_async(
    client: AsyncOpenAI,
    config: LLMConfig,
    article: Article,
    *,
    template: PromptTemplate | None = None,
    max_attempts: int | None = None,
    sources: list[str] | None = None,
    semaphore: asyncio.Semaphore | None = None,
) -> ExtractedEvent:
    """异步抽取(批处理用),与同步版逻辑等价。

    max_attempts 为 None 时使用 config.max_attempts。
    sources 为该新闻全部来源;None 时兜底 [article.source_id]。
    """
    tpl = template or PromptTemplate()
    prompt = tpl.render(article)
    max_attempts = max_attempts or config.max_attempts

    async def _run() -> ExtractedEvent:
        last_err: Exception | None = None
        for attempt in range(1, max_attempts + 1):
            try:
                raw, usage = await _call_llm_async(client, config, prompt)
                event = parse_event_json(raw)
                return ExtractedEvent(
                    source_id=article.source_id,
                    url=article.url,
                    url_hash=article.url_hash,
                    title=article.title,
                    publish_time=article.publish_time,
                    sources=sources or [article.source_id],
                    event=event,
                    provider=config.provider,
                    model=config.model,
                    attempts=attempt,
                    prompt_tokens=usage.get("prompt_tokens"),
                    completion_tokens=usage.get("completion_tokens"),
                )
            except LLMCallError as e:
                last_err = e
                logger.warning(
                    "LLM 抽取失败 url={} 尝试 {}/{}: {}",
                    article.url, attempt, max_attempts, e.reason,
                )
            except Exception as e:  # noqa: BLE001
                last_err = e
                logger.warning(
                    "LLM 调用异常 url={} 尝试 {}/{}: {}: {}",
                    article.url, attempt, max_attempts, type(e).__name__, e,
                )
            if attempt < max_attempts:
                wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC)
                await asyncio.sleep(wait)
        raise LLMCallError(
            f"LLM 抽取放弃,共 {max_attempts} 次尝试: {last_err}",
            attempts=max_attempts,
        )

    if semaphore is None:
        return await _run()
    async with semaphore:
        return await _run()