294 lines
9.9 KiB
Python
294 lines
9.9 KiB
Python
"""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 = DEFAULT_MAX_ATTEMPTS,
|
|
) -> ExtractedEvent:
|
|
"""同步抽取单篇文章的事件(带重试)。"""
|
|
tpl = template or PromptTemplate()
|
|
prompt = tpl.render(article)
|
|
|
|
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,
|
|
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 = DEFAULT_MAX_ATTEMPTS,
|
|
semaphore: asyncio.Semaphore | None = None,
|
|
) -> ExtractedEvent:
|
|
"""异步抽取(批处理用),与同步版逻辑等价。"""
|
|
tpl = template or PromptTemplate()
|
|
prompt = tpl.render(article)
|
|
|
|
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,
|
|
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()
|