lhl
0a498e4f7a
- template_mapper: 按 design §6.5 仅 Heading level<=1 起章，H2/H3 归入
父章 sub_headings（此前每个 heading 独立成章，无 {{section:id}} 锚点的
6 个子章生成后被静默丢弃）
- models: ChapterSpec.sub_headings 字段；to_vars 暴露 sub_headings 变量
- writer_agent: prompt 新增【小节约束】（有子节时按小节顺序以 level=2
heading 组织，不得遗漏或新增）
- docx_injector: 注入后删除同章内与已生成标题同名且完全无内容的模板裸
H2/H3（保守策略：带内容的模板子节保留）
- 真实试运行验证：7 章 / 无静默丢弃告警 / 14 个 H2 无重复
242 lines
8.9 KiB
Python
242 lines
8.9 KiB
Python
import pytest
|
||
from types import SimpleNamespace
|
||
|
||
from genesis.writer.writer_agent import (
|
||
WriterAgent,
|
||
WRITER_PROMPT_TEMPLATE,
|
||
CHAPTER_OUTPUT_SCHEMA,
|
||
)
|
||
from genesis.writer.models import GenerationContext, ChapterSpec
|
||
from genesis.writer.writer_state import WriterState
|
||
from genesis.writer.exceptions import WriterGenerationError
|
||
|
||
|
||
class FakeEngine:
|
||
def __init__(self):
|
||
self.calls = 0
|
||
|
||
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
self.calls += 1
|
||
return SimpleNamespace(
|
||
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
|
||
status="ok",
|
||
)
|
||
|
||
|
||
class FakePromptRegistry:
|
||
@staticmethod
|
||
def get_or_create(name, template):
|
||
return SimpleNamespace(name=name, version="1", template=template)
|
||
|
||
|
||
def _ctx(cid, title):
|
||
return GenerationContext(
|
||
chapter_id=cid,
|
||
title=title,
|
||
template_marker=ChapterSpec(chapter_id=cid, title=title, section_placeholder=None),
|
||
structured_source="source text",
|
||
write_rules=["W1"],
|
||
design_rules=["D1"],
|
||
template_styles={"Heading1"},
|
||
prior_state=None,
|
||
)
|
||
|
||
|
||
def test_generate_chapter_success():
|
||
agent = WriterAgent(
|
||
session_id="s",
|
||
engine=FakeEngine(),
|
||
prompt_registry=FakePromptRegistry(),
|
||
state=WriterState(["db_design"]),
|
||
)
|
||
content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
||
assert content.chapter_id == "db_design"
|
||
assert content.blocks and content.blocks[0].type == "paragraph"
|
||
|
||
|
||
def test_prompt_template_enforces_title_language():
|
||
# 语言约束：正文需与章节标题语言一致（模板为日文时输出日文）
|
||
assert "语言约束" in WRITER_PROMPT_TEMPLATE
|
||
assert "{{title}}" in WRITER_PROMPT_TEMPLATE
|
||
|
||
|
||
def test_prompt_template_has_impact_context_var():
|
||
# 影响调查结果作为生成主上下文：模板必须包含 impact 变量（无影响书时渲染为空串）
|
||
assert "影响调查上下文" in WRITER_PROMPT_TEMPLATE
|
||
assert "{{impact}}" in WRITER_PROMPT_TEMPLATE
|
||
|
||
|
||
def test_prompt_template_enforces_topic_and_chapter_data():
|
||
# 主题约束：正文必须围绕本章标题主题，依据章节数据，禁止套用系统整体架构
|
||
assert "主题约束" in WRITER_PROMPT_TEMPLATE
|
||
assert "参考资料（本章对应数据）" in WRITER_PROMPT_TEMPLATE
|
||
assert "{{data}}" in WRITER_PROMPT_TEMPLATE
|
||
# 旧的全量 source 变量不再使用
|
||
assert "{{source}}" not in WRITER_PROMPT_TEMPLATE
|
||
|
||
|
||
def test_prompt_template_has_sub_headings_var():
|
||
# 子节结构（design.md §6.5：H2/H3 归入本章）：模板必须提供小节变量并指导按小节组织
|
||
assert "{{sub_headings}}" in WRITER_PROMPT_TEMPLATE
|
||
assert "小节" in WRITER_PROMPT_TEMPLATE
|
||
|
||
|
||
class RecordingEngine(FakeEngine):
|
||
"""记录 variables 以断言 prompt 渲染输入。"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.last_vars = None
|
||
|
||
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
self.last_vars = dict(variables)
|
||
return super().chat_structured(
|
||
session_id=session_id, prompt=prompt,
|
||
variables=variables, schema=schema, retry_count=retry_count,
|
||
)
|
||
|
||
|
||
class RealStylePromptRegistry:
|
||
"""模拟真实 PromptRegistry：无 get_or_create，get() 返回模板字符串（非 Prompt）。"""
|
||
|
||
def __init__(self):
|
||
self._tpl = {}
|
||
|
||
def register(self, name, version, template):
|
||
self._tpl[(name, version)] = template
|
||
|
||
def get(self, name, version=None):
|
||
return self._tpl[(name, version)]
|
||
|
||
|
||
class RenderingEngine:
|
||
"""模拟真实引擎渲染行为：仅当收到 Prompt 对象时才用 jinja2 渲染（engine._render_prompt 契约）。"""
|
||
|
||
def __init__(self):
|
||
self.last_rendered = None
|
||
|
||
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
from jinja2 import Template
|
||
from genesis.inference.types import Prompt as P
|
||
tpl = prompt.template if isinstance(prompt, P) else prompt # str 原样返回 → 不渲染
|
||
self.last_rendered = Template(tpl).render(**variables) if isinstance(prompt, P) else tpl
|
||
return SimpleNamespace(
|
||
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
|
||
status="ok",
|
||
)
|
||
|
||
|
||
def test_prompt_is_rendered_with_real_style_registry():
|
||
"""回归：真实 PromptRegistry.get 返回模板字符串 → 引擎按 str 原样发送，
|
||
{{data}}/{{title}} 等占位符从未被替换（真实试运行暴露）。
|
||
WriterAgent 必须保证交给引擎的是可渲染的 Prompt 对象。"""
|
||
agent = WriterAgent(
|
||
session_id="s",
|
||
engine=RenderingEngine(),
|
||
prompt_registry=RealStylePromptRegistry(),
|
||
state=WriterState(["db_design"]),
|
||
)
|
||
ctx = _ctx("db_design", "DB 設計")
|
||
agent.generate_chapter(ctx)
|
||
assert "{{" not in agent.engine.last_rendered # 无残留占位符
|
||
assert "DB 設計" in agent.engine.last_rendered # title 已渲染
|
||
assert "W1" in agent.engine.last_rendered # write_rules 已渲染
|
||
|
||
|
||
def test_generate_chapter_passes_chapter_scoped_data_var():
|
||
from genesis.data_models import (
|
||
CellValue, ExcelTable, ParsedTemplate, ChapterMarker,
|
||
Provenance, SheetType, StructuredSource,
|
||
)
|
||
|
||
def cell(v):
|
||
return CellValue(value=v, provenance=Provenance("f.xlsx", "s", 1, "A", "列"))
|
||
|
||
src = StructuredSource(
|
||
tables=[
|
||
ExcelTable("DB定義", SheetType.DATABASE, "openpyxl",
|
||
["テーブルID", "テーブル名"],
|
||
[{"テーブルID": cell("T001"), "テーブル名": cell("trade_order")}]),
|
||
ExcelTable("機能一覧", SheetType.FUNCTION, "openpyxl",
|
||
["機能ID", "機能名"],
|
||
[{"機能ID": cell("F001"), "機能名": cell("止损风控")}]),
|
||
],
|
||
template=ParsedTemplate("t.docx",
|
||
[ChapterMarker(type="heading", name="x", level=1)],
|
||
{}, {"used": []}),
|
||
rule_docs=[], image_analyses=[], existing_system=None, comments=[],
|
||
)
|
||
ctx = GenerationContext(
|
||
chapter_id="db_design", title="DB 設計",
|
||
template_marker=ChapterSpec(chapter_id="db_design", title="DB 設計", section_placeholder=None),
|
||
structured_source=src, write_rules=["W1"], design_rules=["D1"],
|
||
template_styles={"Heading1"},
|
||
)
|
||
engine = RecordingEngine()
|
||
agent = WriterAgent(session_id="s", engine=engine,
|
||
prompt_registry=FakePromptRegistry(),
|
||
state=WriterState(["db_design"]))
|
||
agent.generate_chapter(ctx)
|
||
data = engine.last_vars["data"]
|
||
assert "T001" in data and "trade_order" in data # 本章对应数据
|
||
assert "止损风控" not in data # 其他章数据被排除
|
||
|
||
|
||
def test_generate_chapter_retries():
|
||
class Boom(FakeEngine):
|
||
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
raise RuntimeError("boom")
|
||
|
||
agent = WriterAgent(
|
||
session_id="s",
|
||
engine=Boom(),
|
||
prompt_registry=FakePromptRegistry(),
|
||
state=WriterState(["db_design"]),
|
||
max_retries=2,
|
||
)
|
||
with pytest.raises(WriterGenerationError):
|
||
agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
||
|
||
|
||
class AsyncFakeEngine:
|
||
async def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
return SimpleNamespace(
|
||
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
|
||
status="ok",
|
||
)
|
||
|
||
|
||
def test_generate_chapter_with_async_engine():
|
||
agent = WriterAgent(
|
||
session_id="s",
|
||
engine=AsyncFakeEngine(),
|
||
prompt_registry=FakePromptRegistry(),
|
||
state=WriterState(["db_design"]),
|
||
)
|
||
content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
||
assert content.chapter_id == "db_design"
|
||
assert content.blocks and content.blocks[0].type == "paragraph"
|
||
|
||
|
||
def test_generate_chapter_propagates_engine_error_detail():
|
||
class FailingEngine:
|
||
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
return SimpleNamespace(
|
||
data={},
|
||
status="failed",
|
||
error="LLM HTTP 401: invalid key",
|
||
error_code="LLM_NETWORK_ERROR",
|
||
)
|
||
|
||
agent = WriterAgent(
|
||
session_id="s",
|
||
engine=FailingEngine(),
|
||
prompt_registry=FakePromptRegistry(),
|
||
state=WriterState(["db_design"]),
|
||
max_retries=1,
|
||
)
|
||
with pytest.raises(WriterGenerationError) as exc:
|
||
agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
||
assert "LLM HTTP 401" in str(exc.value)
|
||
assert "LLM_NETWORK_ERROR" in str(exc.value)
|