57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
import pytest
|
|
from types import SimpleNamespace
|
|
from docx import Document
|
|
|
|
from genesis.data_models import ParsedTemplate, ChapterMarker
|
|
from genesis.qa.qa_loop import QALoop
|
|
from genesis.writer.models import ChapterContent
|
|
|
|
|
|
class ImprovingEngine:
|
|
def __init__(self):
|
|
self.calls: dict[str, int] = {}
|
|
|
|
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
|
cid = variables["title"]
|
|
self.calls[cid] = self.calls.get(cid, 0) + 1
|
|
n = self.calls[cid]
|
|
# "a" 首次即通过；"b" 首次太短失败，第二次（n>=2）充分
|
|
text = "充分且规范的说明内容，满足写入规则要求。" if (cid == "A" or n >= 2) else "x"
|
|
return SimpleNamespace(data={"title": cid, "blocks": [{"type": "paragraph", "text": text}]}, status="ok")
|
|
|
|
|
|
def _make_template(path):
|
|
doc = Document()
|
|
doc.add_paragraph("A", style="Heading 1")
|
|
doc.add_paragraph("{{section:a}}")
|
|
doc.add_paragraph("B", style="Heading 1")
|
|
doc.add_paragraph("{{section:b}}")
|
|
doc.save(path)
|
|
|
|
|
|
def _ss(template_path):
|
|
parsed = ParsedTemplate(
|
|
file_name=template_path,
|
|
sections=[
|
|
ChapterMarker(type="heading", name="A", level=1),
|
|
ChapterMarker(type="placeholder", name="section:a", level=0),
|
|
ChapterMarker(type="heading", name="B", level=1),
|
|
ChapterMarker(type="placeholder", name="section:b", level=0),
|
|
],
|
|
placeholders={},
|
|
styles={"defined": ["Heading 1"], "used": ["Heading 1"]},
|
|
)
|
|
return SimpleNamespace(template=parsed)
|
|
|
|
|
|
def test_qa_loop_regenerates_only_failed(tmp_path):
|
|
tpl = tmp_path / "tpl.docx"
|
|
out = tmp_path / "out.docx"
|
|
_make_template(str(tpl))
|
|
engine = ImprovingEngine()
|
|
loop = QALoop(max_rounds=3)
|
|
report = loop.run(_ss(str(tpl)), str(out), samples_dir="nonexistent_dir_xyz", engine=engine)
|
|
assert report.passed is True
|
|
assert engine.calls["A"] == 1 # A 首次即通过，未被重生成
|
|
assert engine.calls["B"] == 2 # B 失败后被重生成一次
|