simon
9aa44e610d
- system.yaml: llm.max_retries(死配置) 改为 llm.max_attempts；embedding 段新增 max_attempts - llm/client.py: LLMConfig 新增 max_attempts（默认 3 兜底），load_llm_config 读取配置 - llm/extractor.py: translate_and_extract(_async) max_attempts=None 时取 config.max_attempts - embedding/client.py: load_embedding_config 读取 embedding.max_attempts - scheduler/reporter.py: _call_llm_simple max_retries=None 时取 max_attempts-1 - 新增 2 个配置读取测试；已同步 pi5 验证
629 lines
21 KiB
Python
629 lines
21 KiB
Python
"""M4 LLM 翻译 + 事件抽取模块单元测试。"""
|
||
|
||
import json
|
||
from pathlib import Path
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
from openai import OpenAI
|
||
|
||
from extractor.models import ProcessedArticle
|
||
from llm.client import LLMConfig, load_llm_config
|
||
from llm.extractor import (
|
||
MAX_CONTENT_CHARS,
|
||
PromptTemplate,
|
||
_extract_json_object,
|
||
parse_translation_json,
|
||
translate_and_extract,
|
||
)
|
||
from llm.models import (
|
||
INTERNATIONAL_EVENT_TYPES,
|
||
EnTranslatedArticle,
|
||
EventExtraction,
|
||
LLMCallError,
|
||
LLMTranslationOutput,
|
||
Sentiment,
|
||
)
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 辅助工厂
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
def _make_article(
|
||
*,
|
||
url: str = "https://www.reuters.com/business/1",
|
||
url_hash: str = "abc1234567890000",
|
||
source_id: str = "reuters",
|
||
source_name: str = "Reuters",
|
||
title: str = "Fed Holds Rates Steady as Markets Rally",
|
||
content: str = (
|
||
"The Federal Reserve held interest rates steady on Wednesday, "
|
||
"citing solid economic growth and a strong labor market. "
|
||
"Markets rallied in response, with the S&P 500 gaining 1.2 percent."
|
||
),
|
||
publish_time: str = "2026-06-16T10:00:00",
|
||
word_count: int = 0,
|
||
) -> ProcessedArticle:
|
||
return ProcessedArticle(
|
||
source_id=source_id,
|
||
source_name=source_name,
|
||
url=url,
|
||
url_hash=url_hash,
|
||
title=title,
|
||
content=content,
|
||
publish_time=publish_time,
|
||
word_count=word_count or len(content.split()),
|
||
)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Sentiment / 枚举
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestSentiment:
|
||
"""Sentiment 枚举测试。"""
|
||
|
||
def test_values(self):
|
||
assert Sentiment.POSITIVE == "positive"
|
||
assert Sentiment.NEUTRAL == "neutral"
|
||
assert Sentiment.NEGATIVE == "negative"
|
||
|
||
def test_from_string(self):
|
||
assert Sentiment("positive") == Sentiment.POSITIVE
|
||
assert Sentiment("neutral") == Sentiment.NEUTRAL
|
||
assert Sentiment("negative") == Sentiment.NEGATIVE
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# EventExtraction 模型校验
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestEventExtraction:
|
||
"""EventExtraction 模型测试。"""
|
||
|
||
def test_valid_event(self):
|
||
ev = EventExtraction(
|
||
event_type="财报披露",
|
||
stock_codes=["AAPL", "TSLA"],
|
||
sentiment="positive",
|
||
importance=4,
|
||
summary_zh="苹果第三季度营收超预期",
|
||
)
|
||
assert ev.event_type == "财报披露"
|
||
assert ev.stock_codes == ["AAPL", "TSLA"]
|
||
assert ev.sentiment == Sentiment.POSITIVE
|
||
assert ev.importance == 4
|
||
|
||
def test_stock_codes_filtered_and_uppercased(self):
|
||
ev = EventExtraction(
|
||
event_type="并购收购",
|
||
stock_codes=["aapl", " msft ", "", "INVALID123", "GOOGL"],
|
||
sentiment="neutral",
|
||
importance=3,
|
||
summary_zh="微软收购测试",
|
||
)
|
||
# INVALID123 > 5 chars → 过滤；空 → 过滤；小写 → 大写；去重
|
||
assert ev.stock_codes == ["AAPL", "MSFT", "GOOGL"]
|
||
|
||
def test_empty_stock_codes(self):
|
||
ev = EventExtraction(
|
||
event_type="宏观经济",
|
||
stock_codes=[],
|
||
sentiment="neutral",
|
||
importance=2,
|
||
summary_zh="GDP 数据发布",
|
||
)
|
||
assert ev.stock_codes == []
|
||
|
||
def test_importance_bounds(self):
|
||
# 1 和 5 都能通过
|
||
ev1 = EventExtraction(
|
||
event_type="宏观经济", sentiment="neutral", importance=1, summary_zh="t1"
|
||
)
|
||
assert ev1.importance == 1
|
||
ev5 = EventExtraction(
|
||
event_type="央行决议", sentiment="negative", importance=5, summary_zh="t5"
|
||
)
|
||
assert ev5.importance == 5
|
||
|
||
def test_importance_out_of_range_rejected(self):
|
||
with pytest.raises(Exception):
|
||
EventExtraction(
|
||
event_type="其他", sentiment="neutral", importance=0, summary_zh="t"
|
||
)
|
||
with pytest.raises(Exception):
|
||
EventExtraction(
|
||
event_type="其他", sentiment="neutral", importance=6, summary_zh="t"
|
||
)
|
||
|
||
def test_invalid_sentiment_rejected(self):
|
||
with pytest.raises(Exception):
|
||
EventExtraction(
|
||
event_type="其他", stock_codes=[], sentiment="happy", importance=3,
|
||
summary_zh="t",
|
||
)
|
||
|
||
def test_event_type_normalized(self):
|
||
"""空 event_type 默认"其他"。"""
|
||
ev = EventExtraction(
|
||
event_type="", sentiment="neutral", importance=2, summary_zh="测试"
|
||
)
|
||
assert ev.event_type == "其他"
|
||
|
||
def test_summary_zh_max_length_rejected(self):
|
||
"""超长 summary_zh（>200 字符）直接拒绝。"""
|
||
long_summary = "测试" * 150 # 300 chars > 200
|
||
with pytest.raises(Exception):
|
||
EventExtraction(
|
||
event_type="行业动态",
|
||
sentiment="neutral",
|
||
importance=2,
|
||
summary_zh=long_summary,
|
||
)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# LLMTranslationOutput
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestLLMTranslationOutput:
|
||
"""LLMTranslationOutput 模型测试。"""
|
||
|
||
def test_valid_full_output(self):
|
||
data = {
|
||
"title_zh": "美联储维持利率不变，市场上涨",
|
||
"content_zh": "美联储周三维持利率不变...",
|
||
"events": [
|
||
{
|
||
"event_type": "央行决议",
|
||
"stock_codes": [],
|
||
"sentiment": "positive",
|
||
"importance": 5,
|
||
"summary_zh": "美联储维持利率不变",
|
||
}
|
||
],
|
||
}
|
||
out = LLMTranslationOutput.model_validate(data)
|
||
assert out.title_zh == data["title_zh"]
|
||
assert len(out.events) == 1
|
||
assert out.events[0].event_type == "央行决议"
|
||
|
||
def test_no_events(self):
|
||
data = {
|
||
"title_zh": "每日市场简报",
|
||
"content_zh": "今日市场整体平淡...",
|
||
"events": [],
|
||
}
|
||
out = LLMTranslationOutput.model_validate(data)
|
||
assert out.events == []
|
||
|
||
def test_missing_title_zh_rejected(self):
|
||
data = {
|
||
"content_zh": "正文...",
|
||
"events": [],
|
||
}
|
||
with pytest.raises(Exception):
|
||
LLMTranslationOutput.model_validate(data)
|
||
|
||
def test_missing_content_zh_rejected(self):
|
||
data = {
|
||
"title_zh": "标题",
|
||
"events": [],
|
||
}
|
||
with pytest.raises(Exception):
|
||
LLMTranslationOutput.model_validate(data)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# EnTranslatedArticle
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestEnTranslatedArticle:
|
||
"""EnTranslatedArticle 模型测试。"""
|
||
|
||
def test_minimal_construction(self):
|
||
article = EnTranslatedArticle(
|
||
source_id="reuters",
|
||
source_name="Reuters",
|
||
url="https://example.com/1",
|
||
url_hash="abc123",
|
||
title="Fed Holds Rates",
|
||
title_zh="美联储维持利率",
|
||
content_en="The Fed held rates steady.",
|
||
content_zh="美联储维持利率不变。",
|
||
provider="deepseek",
|
||
model="deepseek-v4-flash",
|
||
)
|
||
assert article.events == []
|
||
assert article.word_count_zh == 0
|
||
|
||
def test_with_events(self):
|
||
article = EnTranslatedArticle(
|
||
source_id="reuters",
|
||
source_name="Reuters",
|
||
url="https://example.com/1",
|
||
url_hash="abc123",
|
||
title="Apple Earnings",
|
||
title_zh="苹果财报",
|
||
content_en="Apple reported record earnings.",
|
||
content_zh="苹果公布了创纪录的财报。",
|
||
events=[
|
||
EventExtraction(
|
||
event_type="财报披露",
|
||
stock_codes=["AAPL"],
|
||
sentiment="positive",
|
||
importance=4,
|
||
summary_zh="苹果财报超预期",
|
||
)
|
||
],
|
||
provider="deepseek",
|
||
model="deepseek-v4-flash",
|
||
)
|
||
assert len(article.events) == 1
|
||
assert "AAPL" in article.short_summary()
|
||
|
||
def test_short_summary_no_events(self):
|
||
article = EnTranslatedArticle(
|
||
source_id="reuters",
|
||
source_name="Reuters",
|
||
url="https://example.com/1",
|
||
url_hash="abc123",
|
||
title="Market Wrap",
|
||
title_zh="市场综述",
|
||
content_en="Markets were flat today.",
|
||
content_zh="今日市场持平。",
|
||
provider="deepseek",
|
||
model="deepseek-v4-flash",
|
||
)
|
||
assert "-" in article.short_summary() or "0events" in article.short_summary()
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# PromptTemplate
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestPromptTemplate:
|
||
"""PromptTemplate 测试。"""
|
||
|
||
def test_parse_and_render(self, tmp_path: Path):
|
||
"""测试模板解析和渲染。"""
|
||
prompt_content = """# 测试标题
|
||
|
||
## System Prompt
|
||
|
||
你是翻译助手。
|
||
|
||
---
|
||
|
||
## User Input
|
||
|
||
标题: {title}
|
||
来源: {source_name}
|
||
正文: {content}
|
||
"""
|
||
prompt_path = tmp_path / "test_prompt.md"
|
||
prompt_path.write_text(prompt_content, encoding="utf-8")
|
||
|
||
tpl = PromptTemplate(template_path=prompt_path)
|
||
article = _make_article()
|
||
system, user = tpl.render(article)
|
||
|
||
assert "翻译助手" in system
|
||
assert article.title in user
|
||
assert article.source_name in user
|
||
assert article.content in user
|
||
|
||
def test_content_truncation(self, tmp_path: Path):
|
||
"""测试超长正文截断。"""
|
||
prompt_content = """## System Prompt
|
||
你是助手。
|
||
---
|
||
|
||
## User Input
|
||
正文: {content}
|
||
"""
|
||
prompt_path = tmp_path / "test_prompt.md"
|
||
prompt_path.write_text(prompt_content, encoding="utf-8")
|
||
|
||
tpl = PromptTemplate(template_path=prompt_path)
|
||
long_content = "X" * (MAX_CONTENT_CHARS + 500)
|
||
article = _make_article(content=long_content)
|
||
|
||
_system, user = tpl.render(article)
|
||
assert "正文过长已截断" in user
|
||
assert len("X" * MAX_CONTENT_CHARS) + len("\n\n[正文过长已截断]") < len(user)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# JSON 提取
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestExtractJsonObject:
|
||
"""_extract_json_object 函数测试。"""
|
||
|
||
def test_plain_json(self):
|
||
raw = '{"key": "value"}'
|
||
assert _extract_json_object(raw) == '{"key": "value"}'
|
||
|
||
def test_json_with_fence(self):
|
||
raw = '```json\n{"key": "value"}\n```'
|
||
assert _extract_json_object(raw) == '{"key": "value"}'
|
||
|
||
def test_json_with_text_before(self):
|
||
raw = 'Here is the result:\n{"key": "value"}'
|
||
assert _extract_json_object(raw) == '{"key": "value"}'
|
||
|
||
def test_nested_braces(self):
|
||
raw = '{"outer": {"inner": [1, 2, 3]}}'
|
||
assert _extract_json_object(raw) == raw.strip()
|
||
|
||
def test_empty_string(self):
|
||
assert _extract_json_object("") == ""
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# parse_translation_json
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestParseTranslationJson:
|
||
"""parse_translation_json 函数测试。"""
|
||
|
||
def test_valid_json(self):
|
||
raw = json.dumps({
|
||
"title_zh": "测试标题",
|
||
"content_zh": "测试正文",
|
||
"events": [],
|
||
})
|
||
result = parse_translation_json(raw)
|
||
assert result.title_zh == "测试标题"
|
||
assert result.content_zh == "测试正文"
|
||
|
||
def test_invalid_json(self):
|
||
with pytest.raises(LLMCallError, match="JSON 解析失败"):
|
||
parse_translation_json("not valid json {{{")
|
||
|
||
def test_non_object(self):
|
||
with pytest.raises(LLMCallError, match="非对象"):
|
||
parse_translation_json("[1, 2, 3]")
|
||
|
||
def test_schema_validation_fails(self):
|
||
"""缺少必填字段时抛出 LLMCallError。"""
|
||
raw = json.dumps({"title_zh": "标题"}) # 缺少 content_zh
|
||
with pytest.raises(LLMCallError, match="schema 校验失败"):
|
||
parse_translation_json(raw)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# LLMConfig / load_llm_config
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestLLMConfig:
|
||
"""LLMConfig 测试。"""
|
||
|
||
def test_valid_config(self):
|
||
cfg = LLMConfig(
|
||
provider="deepseek",
|
||
model="deepseek-chat",
|
||
api_key="sk-test",
|
||
base_url="https://api.deepseek.com",
|
||
)
|
||
assert cfg.provider == "deepseek"
|
||
|
||
def test_empty_api_key_raises(self):
|
||
with pytest.raises(ValueError, match="API key 为空"):
|
||
LLMConfig(
|
||
provider="deepseek",
|
||
model="deepseek-chat",
|
||
api_key="",
|
||
base_url="https://api.deepseek.com",
|
||
)
|
||
|
||
|
||
class TestLoadLLMConfig:
|
||
"""load_llm_config 函数测试。"""
|
||
|
||
def test_deepseek_from_env(self, monkeypatch):
|
||
"""从环境变量构造 DeepSeek 配置。"""
|
||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek-test-key")
|
||
config = load_llm_config(provider="deepseek")
|
||
assert config.provider == "deepseek"
|
||
assert config.api_key == "sk-deepseek-test-key"
|
||
assert "deepseek" in config.base_url
|
||
|
||
def test_qwen_fallback_to_dashscope_key(self, monkeypatch):
|
||
"""Qwen 的 API Key 可回退到 DASHSCOPE_API_KEY。"""
|
||
monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-dashscope-key")
|
||
monkeypatch.delenv("QWEN_API_KEY", raising=False)
|
||
config = load_llm_config(provider="qwen")
|
||
assert config.provider == "qwen"
|
||
assert config.api_key == "sk-dashscope-key"
|
||
|
||
def test_unknown_provider_raises(self):
|
||
with pytest.raises(ValueError, match="未知 LLM provider"):
|
||
load_llm_config(provider="openai")
|
||
|
||
def test_missing_api_key_raises(self, monkeypatch):
|
||
"""未配置 API Key 时应抛出明确错误。"""
|
||
monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False)
|
||
with pytest.raises(ValueError, match="API key 未配置"):
|
||
load_llm_config(provider="deepseek")
|
||
|
||
def test_max_attempts_from_system_yaml(self, monkeypatch):
|
||
"""重试次数取自 system.yaml llm.max_attempts（当前 3）。"""
|
||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek-test-key")
|
||
config = load_llm_config(provider="deepseek")
|
||
assert config.max_attempts == 3
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# translate_and_extract（mock LLM）
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestTranslateAndExtract:
|
||
"""translate_and_extract 测试（mock LLM 响应）。"""
|
||
|
||
def test_successful_translation(self, monkeypatch):
|
||
"""Mock LLM 返回有效 JSON。"""
|
||
config = LLMConfig(
|
||
provider="deepseek",
|
||
model="test-model",
|
||
api_key="sk-test",
|
||
base_url="https://test.api",
|
||
)
|
||
article = _make_article()
|
||
|
||
# Mock OpenAI client
|
||
mock_client = MagicMock(spec=OpenAI)
|
||
mock_response = MagicMock()
|
||
mock_response.choices = [MagicMock()]
|
||
mock_response.choices[0].message.content = json.dumps({
|
||
"title_zh": "美联储维持利率不变",
|
||
"content_zh": "美联储周三维持利率不变，理由是经济增长稳健。",
|
||
"events": [
|
||
{
|
||
"event_type": "央行决议",
|
||
"stock_codes": [],
|
||
"sentiment": "positive",
|
||
"importance": 5,
|
||
"summary_zh": "美联储维持利率不变，市场反弹",
|
||
}
|
||
],
|
||
})
|
||
mock_response.usage = MagicMock()
|
||
mock_response.usage.prompt_tokens = 500
|
||
mock_response.usage.completion_tokens = 200
|
||
mock_client.chat.completions.create.return_value = mock_response
|
||
|
||
result = translate_and_extract(
|
||
client=mock_client,
|
||
config=config,
|
||
article=article,
|
||
)
|
||
|
||
assert result.title_zh == "美联储维持利率不变"
|
||
assert len(result.events) == 1
|
||
assert result.events[0].event_type == "央行决议"
|
||
assert result.provider == "deepseek"
|
||
assert result.prompt_tokens == 500
|
||
assert result.completion_tokens == 200
|
||
assert result.word_count_zh > 0
|
||
|
||
def test_empty_content_zh_retries(self, monkeypatch):
|
||
"""LLM 返回空 content_zh 时触发重试。"""
|
||
config = LLMConfig(
|
||
provider="deepseek",
|
||
model="test-model",
|
||
api_key="sk-test",
|
||
base_url="https://test.api",
|
||
)
|
||
article = _make_article()
|
||
|
||
mock_client = MagicMock(spec=OpenAI)
|
||
# 始终返回空 content_zh
|
||
mock_response = MagicMock()
|
||
mock_response.choices = [MagicMock()]
|
||
mock_response.choices[0].message.content = json.dumps({
|
||
"title_zh": "标题",
|
||
"content_zh": "",
|
||
"events": [],
|
||
})
|
||
mock_response.usage = MagicMock()
|
||
mock_response.usage.prompt_tokens = 100
|
||
mock_response.usage.completion_tokens = 10
|
||
mock_client.chat.completions.create.return_value = mock_response
|
||
|
||
with pytest.raises(LLMCallError, match="放弃"):
|
||
translate_and_extract(
|
||
client=mock_client,
|
||
config=config,
|
||
article=article,
|
||
max_attempts=2, # 减少重试加速测试
|
||
)
|
||
|
||
def test_retry_on_json_parse_failure(self, monkeypatch):
|
||
"""前 N-1 次返回无效 JSON，最后一次成功。"""
|
||
config = LLMConfig(
|
||
provider="deepseek",
|
||
model="test-model",
|
||
api_key="sk-test",
|
||
base_url="https://test.api",
|
||
)
|
||
article = _make_article()
|
||
|
||
mock_client = MagicMock(spec=OpenAI)
|
||
# 第一次失败，第二次成功
|
||
mock_client.chat.completions.create.side_effect = [
|
||
_mock_chat_response("not valid {{{ json"),
|
||
_mock_chat_response(json.dumps({
|
||
"title_zh": "测试标题",
|
||
"content_zh": "测试正文",
|
||
"events": [],
|
||
})),
|
||
]
|
||
|
||
result = translate_and_extract(
|
||
client=mock_client,
|
||
config=config,
|
||
article=article,
|
||
max_attempts=3,
|
||
)
|
||
|
||
assert result.attempts == 2 # 第二次成功
|
||
assert result.title_zh == "测试标题"
|
||
|
||
|
||
def _mock_chat_response(content: str) -> MagicMock:
|
||
"""Helper：构造 mock OpenAI chat completion 响应。"""
|
||
resp = MagicMock()
|
||
resp.choices = [MagicMock()]
|
||
resp.choices[0].message.content = content
|
||
resp.usage = MagicMock()
|
||
resp.usage.prompt_tokens = 100
|
||
resp.usage.completion_tokens = 50
|
||
return resp
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# INTERNATIONAL_EVENT_TYPES
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestEventTypes:
|
||
"""事件类型常量测试。"""
|
||
|
||
def test_has_expected_types(self):
|
||
assert "财报披露" in INTERNATIONAL_EVENT_TYPES
|
||
assert "并购收购" in INTERNATIONAL_EVENT_TYPES
|
||
assert "央行决议" in INTERNATIONAL_EVENT_TYPES
|
||
assert "地缘政治" in INTERNATIONAL_EVENT_TYPES
|
||
assert len(INTERNATIONAL_EVENT_TYPES) >= 10
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# LLMCallError
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
class TestLLMCallError:
|
||
"""LLMCallError 异常测试。"""
|
||
|
||
def test_basic(self):
|
||
err = LLMCallError("测试错误", attempts=3)
|
||
assert err.reason == "测试错误"
|
||
assert err.attempts == 3
|
||
assert str(err) == "测试错误"
|
||
|
||
def test_default_attempts(self):
|
||
err = LLMCallError("错误")
|
||
assert err.attempts == 0
|