Initial commit

2026-07-18 15:51:01 +08:00
commit f2c80c5a9c
799 changed files with 133475 additions and 0 deletions
@@ -0,0 +1,58 @@
+"""LLM 投资事件抽取模块 (M4)。
+
+公共 API:
+    - load_llm_config / make_sync_client / make_async_client
+    - extract_event / extract_event_async
+    - PromptTemplate / parse_event_json
+    - EventExtraction / ExtractedEvent / Sentiment / EVENT_TYPES / LLMCallError
+"""
+
+from .client import (
+    DEFAULT_TEMPERATURE,
+    DEFAULT_TIMEOUT_SEC,
+    LLMConfig,
+    load_llm_config,
+    make_async_client,
+    make_sync_client,
+)
+from .extractor import (
+    DEFAULT_MAX_ATTEMPTS,
+    DEFAULT_PROMPT_PATH,
+    MAX_CONTENT_CHARS,
+    PromptTemplate,
+    extract_event,
+    extract_event_async,
+    parse_event_json,
+)
+from .models import (
+    EVENT_TYPES,
+    MAX_IMPORTANCE,
+    MIN_IMPORTANCE,
+    EventExtraction,
+    ExtractedEvent,
+    LLMCallError,
+    Sentiment,
+)
+
+__all__ = [
+    "DEFAULT_MAX_ATTEMPTS",
+    "DEFAULT_PROMPT_PATH",
+    "DEFAULT_TEMPERATURE",
+    "DEFAULT_TIMEOUT_SEC",
+    "EVENT_TYPES",
+    "MAX_CONTENT_CHARS",
+    "MAX_IMPORTANCE",
+    "MIN_IMPORTANCE",
+    "EventExtraction",
+    "ExtractedEvent",
+    "LLMCallError",
+    "LLMConfig",
+    "PromptTemplate",
+    "Sentiment",
+    "extract_event",
+    "extract_event_async",
+    "load_llm_config",
+    "make_async_client",
+    "make_sync_client",
+    "parse_event_json",
+]
@@ -0,0 +1,119 @@
+"""LLM 客户端抽象与工厂。
+
+支持 DeepSeek 和 Qwen(百炼),两者均为 OpenAI 兼容接口,共用 openai SDK。
+
+环境变量:
+    LLM_PROVIDER       = deepseek | qwen        (默认 deepseek)
+    DeepSeek: DEEPSEEK_API_KEY / DEEPSEEK_BASE_URL / DEEPSEEK_MODEL
+    Qwen:     QWEN_API_KEY / QWEN_BASE_URL / QWEN_MODEL
+              (QWEN_API_KEY -> DASHSCOPE_API_KEY 兜底)
+    LLM_MODEL           (兜底) LLM_TEMPERATURE / LLM_TIMEOUT_SEC
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+from loguru import logger
+from openai import AsyncOpenAI, OpenAI
+
+# 默认基址
+_DEEPSEEK_DEFAULT_BASE = "https://api.deepseek.com"
+_QWEN_DEFAULT_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+
+# 默认模型
+_DEEPSEEK_DEFAULT_MODEL = "deepseek-chat"
+_QWEN_DEFAULT_MODEL = "qwen-plus"
+
+# 抽取任务默认参数
+DEFAULT_TIMEOUT_SEC = 60.0
+DEFAULT_TEMPERATURE = 0.1
+
+
+@dataclass
+class LLMConfig:
+    """LLM 调用配置(provider / model / api_key / base_url / 参数)。"""
+
+    provider: str             # "deepseek" / "qwen"
+    model: str
+    api_key: str
+    base_url: str
+    timeout_sec: float = DEFAULT_TIMEOUT_SEC
+    temperature: float = DEFAULT_TEMPERATURE
+
+    def __post_init__(self) -> None:
+        if not self.api_key:
+            raise ValueError(f"LLM provider={self.provider} 的 API key 为空")
+
+
+def _read_env(key: str, default: str | None = None) -> str | None:
+    val = os.environ.get(key)
+    if val is None or val.strip() == "":
+        return default
+    return val.strip()
+
+
+def load_llm_config(
+    provider: str | None = None,
+    *,
+    model: str | None = None,
+) -> LLMConfig:
+    """根据环境变量构造 LLMConfig。
+
+    provider 为 None 时读 LLM_PROVIDER 环境变量,默认 deepseek。
+    model 为 None 时读 LLM_MODEL 或 provider 默认。
+    """
+    p = (provider or _read_env("LLM_PROVIDER", "deepseek") or "deepseek").lower()
+
+    if p == "deepseek":
+        api_key = _read_env("DEEPSEEK_API_KEY") or ""
+        base = _read_env("DEEPSEEK_BASE_URL", _DEEPSEEK_DEFAULT_BASE) or _DEEPSEEK_DEFAULT_BASE
+        # DEEPSEEK_MODEL → LLM_MODEL(兜底) → 默认
+        m = model or _read_env("DEEPSEEK_MODEL") or _read_env("LLM_MODEL") or _DEEPSEEK_DEFAULT_MODEL
+    elif p in ("qwen", "dashscope"):
+        api_key = _read_env("QWEN_API_KEY") or _read_env("DASHSCOPE_API_KEY") or ""
+        base = _read_env("QWEN_BASE_URL", _QWEN_DEFAULT_BASE) or _QWEN_DEFAULT_BASE
+        # QWEN_MODEL → LLM_MODEL(兜底) → 默认
+        m = model or _read_env("QWEN_MODEL") or _read_env("LLM_MODEL") or _QWEN_DEFAULT_MODEL
+        p = "qwen"  # 内部统一用 qwen
+    else:
+        raise ValueError(f"未知 LLM provider: {p!r},仅支持 deepseek / qwen")
+
+    timeout = float(_read_env("LLM_TIMEOUT_SEC", str(DEFAULT_TIMEOUT_SEC)) or DEFAULT_TIMEOUT_SEC)
+    temperature = float(_read_env("LLM_TEMPERATURE", str(DEFAULT_TEMPERATURE)) or DEFAULT_TEMPERATURE)
+
+    return LLMConfig(
+        provider=p,
+        model=m,
+        api_key=api_key,
+        base_url=base,
+        timeout_sec=timeout,
+        temperature=temperature,
+    )
+
+
+def make_sync_client(config: LLMConfig) -> OpenAI:
+    """构造同步 OpenAI 客户端(指向 DeepSeek/Qwen 兼容端点)。"""
+    logger.debug(
+        "初始化同步 LLM 客户端: provider={} model={} base_url={}",
+        config.provider, config.model, config.base_url,
+    )
+    return OpenAI(
+        api_key=config.api_key,
+        base_url=config.base_url,
+        timeout=config.timeout_sec,
+    )
+
+
+def make_async_client(config: LLMConfig) -> AsyncOpenAI:
+    """构造异步 OpenAI 客户端(用于批处理高并发)。"""
+    logger.debug(
+        "初始化异步 LLM 客户端: provider={} model={} base_url={}",
+        config.provider, config.model, config.base_url,
+    )
+    return AsyncOpenAI(
+        api_key=config.api_key,
+        base_url=config.base_url,
+        timeout=config.timeout_sec,
+    )
@@ -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()
@@ -0,0 +1,169 @@
+"""LLM 投资事件抽取的数据模型 (M4)。
+
+EventExtraction 是 LLM 严格输出 schema(JSON mode 解析后用 Pydantic 校验)。
+ExtractedEvent 把 EventExtraction 与原文章元数据合并,作为 M4 最终落盘格式。
+"""
+
+from __future__ import annotations
+
+import re
+from datetime import datetime
+from enum import StrEnum
+from typing import Self
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+
+class Sentiment(StrEnum):
+    """事件情绪倾向。"""
+
+    POSITIVE = "positive"  # 利好
+    NEUTRAL = "neutral"    # 中性
+    NEGATIVE = "negative"  # 利空
+
+
+# 事件类型枚举(Prompt 中也会展示给 LLM)
+EVENT_TYPES: tuple[str, ...] = (
+    "业绩预告",
+    "业绩快报",
+    "财报披露",
+    "合作签约",
+    "投资并购",
+    "重大合同",
+    "产品发布",
+    "技术突破",
+    "监管处罚",
+    "诉讼仲裁",
+    "股东减持",
+    "股东增持",
+    "回购",
+    "分红",
+    "高管变动",
+    "资产重组",
+    "停牌复牌",
+    "ST警示",
+    "退市风险",
+    "宏观政策",
+    "行业政策",
+    "国际局势",
+    "其他",
+)
+
+# A 股股票代码:6 位数字(000xxx/300xxx/600xxx 等),也可带 .SH/.SZ/.BJ 后缀
+_STOCK_CODE_RE = re.compile(r"^\d{6}(\.(SH|SZ|BJ))?$")
+
+# 重要程度合理区间(LLM 偶尔会给 0/6/10,这里夹紧)
+MIN_IMPORTANCE = 1
+MAX_IMPORTANCE = 5
+
+
+class EventExtraction(BaseModel):
+    """LLM 输出的 JSON 直接映射到此模型。"""
+
+    stock_codes: list[str] = Field(
+        default_factory=list,
+        description="A 股 6 位代码,允许带 .SH/.SZ/.BJ 后缀;无相关股票时为空",
+    )
+    company_names: list[str] = Field(
+        default_factory=list, description="涉及公司中文简称,无关时为空"
+    )
+    industries: list[str] = Field(
+        default_factory=list, description="所属行业(申万二级粒度优先);无关时为空"
+    )
+    sentiment: Sentiment = Field(..., description="positive/neutral/negative")
+    importance: int = Field(
+        ..., ge=MIN_IMPORTANCE, le=MAX_IMPORTANCE, description="1-5 重要程度"
+    )
+    event_type: str = Field(..., description="事件类型,见 EVENT_TYPES")
+    summary: str = Field(
+        default="",
+        max_length=200,
+        description="一句话事件摘要(≤ 100 字),便于人工浏览",
+    )
+
+    @field_validator("stock_codes")
+    @classmethod
+    def _strip_and_validate_stock_codes(cls, v: list[str]) -> list[str]:
+        """剔除空字符串、统一大写、过滤明显非法格式。"""
+        cleaned: list[str] = []
+        for code in v:
+            s = (code or "").strip().upper().replace(" ", "")
+            if not s:
+                continue
+            if _STOCK_CODE_RE.match(s):
+                cleaned.append(s)
+        # 去重保持顺序
+        seen: set[str] = set()
+        out: list[str] = []
+        for c in cleaned:
+            if c not in seen:
+                seen.add(c)
+                out.append(c)
+        return out
+
+    @field_validator("company_names", "industries")
+    @classmethod
+    def _strip_text_lists(cls, v: list[str]) -> list[str]:
+        cleaned = [(s or "").strip() for s in v]
+        cleaned = [s for s in cleaned if s]
+        seen: set[str] = set()
+        out: list[str] = []
+        for c in cleaned:
+            if c not in seen:
+                seen.add(c)
+                out.append(c)
+        return out
+
+    @field_validator("event_type")
+    @classmethod
+    def _normalize_event_type(cls, v: str) -> str:
+        s = (v or "").strip()
+        if not s:
+            return "其他"
+        return s
+
+    @model_validator(mode="after")
+    def _post_check(self) -> Self:
+        """中性情绪时 importance 应较低(1-3),纠正常见误判。"""
+        # 不强制纠正,留作后续校准。占位,便于以后扩展。
+        return self
+
+
+class ExtractedEvent(BaseModel):
+    """落盘格式:文章元数据 + LLM 抽取结果 + 调用元信息。"""
+
+    # ---- 来源标识 ----
+    source_id: str
+    url: str
+    url_hash: str
+    title: str
+    publish_time: datetime | None = None
+
+    # ---- 抽取结果 ----
+    event: EventExtraction
+
+    # ---- 调用元信息 ----
+    provider: str = Field(..., description="deepseek / qwen 等")
+    model: str
+    extracted_at: datetime = Field(default_factory=datetime.now)
+    attempts: int = Field(default=1, ge=1, description="LLM 实际调用次数(含重试)")
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+
+    def short_summary(self) -> str:
+        ev = self.event
+        codes = ",".join(ev.stock_codes) or "-"
+        return (
+            f"[{self.source_id}] {self.title[:30]} "
+            f"-> {ev.sentiment.value}/{ev.importance}/{ev.event_type} "
+            f"({codes})"
+        )
+
+
+class LLMCallError(Exception):
+    """LLM 调用失败(网络 / 解析 / 校验)。"""
+
+    def __init__(self, reason: str, *, attempts: int = 0) -> None:
+        super().__init__(reason)
+        self.reason = reason
+        self.attempts = attempts