Initial commit
This commit is contained in:
@@ -0,0 +1,293 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user