lhl
5091125fbc
- writer_agent: prompt【小节约束】明确空结构时不得输出 heading 块 - writer_agent: generate_chapter 程序化强制——template_marker.sub_headings 为空的章，剔除生成结果中全部 type=heading 的内容块（LLM 遵循度不完美， 以模板结构为准做确定性兜底） - 真实试运行验证：自造与章同名 H2 清零、H2 总数 14→6（仅模板定义子节）、 内容引用不受影响
290 lines
11 KiB
Python
290 lines
11 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
|
||
|
||
|
||
def test_prompt_template_forbids_heading_blocks_when_no_sub_headings():
|
||
# 无子节的章：LLM 不得自造 heading 块（真实试运行暴露：帳票一覧/バッチ一覧
|
||
# 章内出现与章同名的自造 H2）
|
||
assert "不得输出任何 type=heading 的内容块" in WRITER_PROMPT_TEMPLATE
|
||
|
||
|
||
class HeadingEngine(FakeEngine):
|
||
"""返回含 heading 块的输出（模拟 LLM 自造标题行为）。"""
|
||
|
||
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
return SimpleNamespace(
|
||
data={
|
||
"title": variables["title"],
|
||
"blocks": [
|
||
{"type": "heading", "level": 2, "text": variables["title"]},
|
||
{"type": "paragraph", "text": "正文内容"},
|
||
],
|
||
},
|
||
status="ok",
|
||
)
|
||
|
||
|
||
def test_generate_chapter_drops_headings_when_no_sub_headings():
|
||
"""无子节结构的章：LLM 自造的 heading 块被程序化剔除（模板结构为准）。"""
|
||
agent = WriterAgent(
|
||
session_id="s", engine=HeadingEngine(),
|
||
prompt_registry=FakePromptRegistry(),
|
||
state=WriterState(["batch_list"]),
|
||
)
|
||
ctx = _ctx("batch_list", "7. バッチ一覧")
|
||
content = agent.generate_chapter(ctx)
|
||
assert all(b.type != "heading" for b in content.blocks)
|
||
assert [b.text for b in content.blocks] == ["正文内容"]
|
||
|
||
|
||
def test_generate_chapter_keeps_headings_when_sub_headings_exist():
|
||
"""有子节结构的章：heading 块保留（按小节组织的内容）。"""
|
||
agent = WriterAgent(
|
||
session_id="s", engine=HeadingEngine(),
|
||
prompt_registry=FakePromptRegistry(),
|
||
state=WriterState(["function_list"]),
|
||
)
|
||
ctx = _ctx("function_list", "2. 機能一覧")
|
||
ctx.template_marker.sub_headings = ["2.1 機能一覧表"]
|
||
content = agent.generate_chapter(ctx)
|
||
assert any(b.type == "heading" for b in content.blocks)
|
||
|
||
|
||
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)
|