proposa-ai/scripts/tests/test_step4_difficulty.py

"""F043：合成采购需求及真实 Step4/5/6 写出函数，LLM 默认使用桩。"""

import copy
import base64
import io
import importlib.util
import json
import os
import re
import sys
import tempfile
import unittest
from types import SimpleNamespace
from unittest.mock import Mock

from docx import Document
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from models import BidOutline, Chapter, ProjectData, TenderAnalysis
from step3_outlining.scoring_structure import _persist_heading_mappings
from step4_writing import _persist_all_chapter_docs
from step4_writing.difficulty_rewrite import (
    MODE, ROLE_BUDGETS, identify_targets, rewrite_difficulty_sections, word_count,
)
from step5_reviewing.reviewer import _Reviewer
from step6_exporting.assembler import assemble_step5_document
from step6_exporting.docx_builder import build_document


def selection(*ids):
    return {"targets": [{"id": node_id, "roles": list(ROLE_BUDGETS)} for node_id in ids]}


def load_stage(number):
    directory = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    if directory not in sys.path:
        sys.path.insert(0, directory)
    spec = importlib.util.spec_from_file_location(
        f'difficulty_stage{number}', os.path.join(directory, f'test_step{number}.py'))
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


class DifficultyTests(unittest.TestCase):
    def setUp(self):
        self.target = Chapter(id='3.2', title='二、关键挑战与解决策略', level=2,
                              from_template=True, template_original_id='3.2',
                              template_original_title='二、关键挑战与解决策略')
        self.root = Chapter(id='3', title='需求理解', template_original_id='3',
                            from_template=True, template_chapter_id='3',
                            generated_content='需求理解\n保留章正文\n一、服务认知\n保留相邻正文\n'
                                              '二、关键挑战与解决策略\n旧项目重难点\n'
                                              '三、其他事项\n保留末节正文',
                            children=[Chapter(id='3.1', title='一、服务认知', level=2),
                                      self.target,
                                      Chapter(id='3.3', title='三、其他事项', level=2)])
        self.outline = BidOutline(project_name='合成文化中心物业项目', chapters=[self.root])
        _persist_heading_mappings(self.outline)
        self.outline.evaluation_index_entries = [{
            'entry_type': 'scoring', 'source_id': 'SC-1#1', 'criterion_id': 'SC-1',
            'display_name': '项目重难点及对策', 'requirement': '分析项目特点和难点并提出措施',
            'final_heading_id': '3.2', 'final_heading_title': self.target.title,
        }]
        self.analysis = TenderAnalysis(project_name=self.outline.project_name)
        self.data = ProjectData(project_id='synthetic', project_name=self.outline.project_name,
                                procurement_docs=[SimpleNamespace(content=(
            '本项目为合成文化中心物业服务。服务区域包含展厅、公共走廊和设备间。'
            '展览开放期间持续提供保洁及秩序维护，闭馆后安排深度清洁。'
            '活动期间人流集中，须与场馆管理方协调，保持疏散通道畅通。'
            '设备巡检发现异常须记录并报告，按授权开展处置，不得擅自停运设备。'
            '采购需求未规定面积、人数或设备数量。'))])
        self.client = Mock()
        self.client.extract_json.side_effect = [
            selection('3.2'),
            {'new_titles': ['项目业态与概况', '重点难点成因分析', '逐项应对及检查闭环']},
        ]
        self.serial = 0

        def generate(**kwargs):
            payload = json.loads(kwargs['user_prompt'])
            expected = int(re.search(r'正文(\d+)字', payload['length']).group(1))
            self.serial += 1
            # 不同节点不同文字，避免确定性去重把桩文本删除。
            marker = '专项正文' + chr(0x4e10 + self.serial)
            return marker + chr(0x5000 + self.serial) * (expected - len(marker))
        self.client.generate.side_effect = generate

    def test_semantic_candidates_no_match_and_scope(self):
        other = Chapter(id='4', title='服务方案', children=[
            Chapter(id='4.1', title='重点难点分析及应对措施', level=2)])
        self.outline.chapters.append(other)
        self.root.children.append(Chapter(id='3.4', title='附件', level=2, is_attachment=True,
            children=[Chapter(id='3.4.1', title='重点难点及措施', level=3)]))
        client = Mock()
        client.extract_json.return_value = selection()
        before = copy.deepcopy(self.outline)
        rewrite_difficulty_sections(self.outline, self.analysis, ProjectData('synthetic', '合成'), client)
        self.assertEqual(self.outline, before)
        payload = json.loads(client.extract_json.call_args.kwargs['user_prompt'])
        self.assertIn('4.1', [n['id'] for n in payload['nodes']])
        self.assertEqual(next(n['path'] for n in payload['nodes'] if n['id'] == '4.1'),
                         ['服务方案', '重点难点分析及应对措施'])
        self.assertNotIn('3.4.1', [n['id'] for n in payload['nodes']])
        client.generate.assert_not_called()

    def test_invalid_and_overlapping_ids_rejected(self):
        self.target.children = [Chapter(id='3.2.1', title='服务瓶颈', level=3)]
        for ids in [['4.1'], ['missing'], ['3', '3.2'], ['3.2', '3.2'], ['3.2', '3.2.1']]:
            client = Mock()
            client.extract_json.return_value = selection(*ids)
            with self.subTest(ids=ids), self.assertRaises(ValueError):
                identify_targets(self.outline, client)

    def test_missing_procurement_and_short_generation_fail_without_tree_changes(self):
        before = copy.deepcopy(self.outline)
        with self.assertRaisesRegex(ValueError, '采购需求正文为空'):
            rewrite_difficulty_sections(self.outline, self.analysis, ProjectData('synthetic', '合成'), self.client)
        self.assertEqual(self.outline, before)
        self.setUp()
        self.client.generate.side_effect = None
        self.client.generate.return_value = '内容不足'
        with self.assertRaisesRegex(ValueError, '三次生成'):
            rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertEqual(self.target.children, [])
        self.assertNotEqual(self.target.content_generation_mode, MODE)

    def test_word_budget_and_bindings(self):
        bindings = copy.deepcopy(self.outline.evaluation_index_entries)
        rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertEqual(self.outline.evaluation_index_entries, bindings)
        self.assertEqual(sum(word_count(n.supplement_content)
                             for n in [self.target, *self.target.children]), 10000)
        self.assertTrue(all(n.level == 3 for n in self.target.children))
        self.assertNotIn('旧项目重难点', self.root.generated_content)
        self.assertIn('保留相邻正文', self.root.generated_content)
        self.assertIn(self.data.procurement_docs[0].content,
                      self.client.generate.call_args.kwargs['user_prompt'])
        report = load_stage(4)._build_report(self.outline, self.analysis)
        self.assertEqual(report['difficulty_sections'][0]['body_words'], 10000)
        self.target.supplement_content = ''
        report = load_stage(4)._build_report(self.outline, self.analysis)
        self.assertTrue(any('专项字数不达标' in e for e in report['errors']))

    def test_h4_target_uses_batches_without_h5(self):
        self.target.level = 4
        self.client.extract_json.side_effect = [selection('3.2')]
        rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertEqual(self.target.children, [])
        self.assertEqual(word_count(self.target.supplement_content), 10000)
        self.assertEqual(self.client.generate.call_count, 5)

    def test_duplicate_title_rejected(self):
        self.client.extract_json.side_effect = [
            selection('3.2'), {'new_titles': ['关键挑战与解决策略']}]
        with self.assertRaisesRegex(ValueError, '重复'):
            rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)

    def test_existing_children_keep_ids_and_new_ids_do_not_collide(self):
        child = Chapter(id='3.2.4', title='（一）已有难点分析', level=3)
        self.target.children = [child]
        _persist_heading_mappings(self.outline)
        rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertIs(self.target.children[0], child)
        self.assertEqual([c.id for c in self.target.children], ['3.2.4', '3.2.5', '3.2.6', '3.2.7'])
        self.assertEqual(self.target.children[1].title, '（二）项目业态与概况')
        self.assertEqual(sum(word_count(n.supplement_content)
                             for n in [self.target, *self.target.children]), 10000)

    def test_long_procurement_reads_tail_instead_of_truncating(self):
        self.data.procurement_docs[0].content = '采购事实' * 6000 + '末尾特殊设备停运约束'
        self.client.extract_json.side_effect = [
            selection('3.2'), {'facts': '采购事实'}, {'facts': '末尾特殊设备停运约束'},
            {'new_titles': []}]
        rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertIn('末尾特殊设备停运约束', self.client.generate.call_args.kwargs['user_prompt'])

    def test_short_first_response_is_regenerated(self):
        generate = self.client.generate.side_effect
        calls = [0]

        def short_first(**kwargs):
            calls[0] += 1
            return '过短' if calls[0] == 1 else generate(**kwargs)
        self.client.generate.side_effect = short_first
        rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertIn('上次正文2字', self.client.generate.call_args_list[1].kwargs['user_prompt'])
        self.assertEqual(sum(word_count(n.supplement_content)
                             for n in [self.target, *self.target.children]), 10000)

    def test_step4_to_step6_headings_prose_and_native_objects(self):
        self._assert_pipeline()

    def test_distributed_roles_share_budget_and_survive_export(self):
        roles = list(ROLE_BUDGETS)
        nodes = [Chapter(id='3.1', title='项目业态', level=2),
                 Chapter(id='3.2', title='项目概况', level=2),
                 Chapter(id='5.1', title='关键难点成因', level=2),
                 Chapter(id='6.1', title='针对性应对措施', level=2)]
        self.root.children = nodes[:2] + [Chapter(id='3.3', title='无关事项', level=2)]
        self.root.generated_content = '项目业态\n旧专项一\n项目概况\n旧专项二\n无关事项\n必须保留正文'
        self.outline.chapters += [Chapter(id='5', title='难点分析', children=[nodes[2]],
                                         generated_content='关键难点成因\n旧专项三'),
                                  Chapter(id='6', title='实施措施', children=[nodes[3]],
                                         generated_content='针对性应对措施\n旧专项四')]
        _persist_heading_mappings(self.outline)
        original = {'project': self.outline.project_name,
                    'heading_mappings': copy.deepcopy(self.outline.heading_mappings),
                    'evaluation_index_entries': copy.deepcopy(self.outline.evaluation_index_entries)}
        self.client.extract_json.side_effect = [
            {'targets': [{'id': n.id, 'roles': [r]} for n, r in zip(nodes, roles)]},
            {'new_titles': []}, {'new_titles': []},
            {'new_titles': ['开放时段作业冲突分析']}, {'new_titles': []}]
        rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertEqual([sum(word_count(c.supplement_content) for c in [n, *n.children])
                          for n in nodes], list(ROLE_BUDGETS.values()))
        self.assertIn('必须保留正文', self.root.generated_content)
        report = load_stage(4)._build_report(self.outline, self.analysis)
        self.assertEqual(len(report['difficulty_sections']), 1)
        self.assertEqual(report['difficulty_sections'][0]['body_words'], 10000)
        prompts = [json.loads(c.kwargs['user_prompt']) for c in self.client.generate.call_args_list]
        measure = next(p for p in prompts if p['assigned_roles'] == ['measures'])
        self.assertIn('5.1', measure['previous_sections'])
        self.assertEqual(sum(t['word_budget'] for t in measure['coordinated_targets']), 10000)
        stage5, stage6 = load_stage(5), load_stage(6)
        with tempfile.TemporaryDirectory() as tmp:
            d4, d5 = os.path.join(tmp, 'step4'), os.path.join(tmp, 'step5')
            template = os.path.join(tmp, 'template.docx')
            doc = Document()
            doc.sections[0].header.paragraphs[0].text = '合成项目'
            doc.sections[0].footer.paragraphs[0].text = '合成页脚'
            doc.save(template)
            _persist_all_chapter_docs(self.outline, d4)
            restored = stage5._load_step4_outline(d4, original, stage5._load_chapter_texts(d4))
            self.assertEqual(restored.evaluation_index_entries, original['evaluation_index_entries'])
            self.assertEqual([n.id for n in restored.flatten()], [n.id for n in self.outline.flatten()])
            stage5._write_step5_chapters_from_step4(stage5._load_step4_records(d4), restored, self.analysis, d5)
            output = os.path.join(tmp, 'distributed.docx')
            assemble_step5_document(stage6._load_chapter_records(d5), restored, output, template_path=template)
            text = '\n'.join(p.text for p in Document(output).paragraphs)
            self.assertIn('必须保留正文', text)
            self.assertNotIn('旧专项', text)
            for n in nodes + nodes[2].children:
                self.assertTrue(re.sub(r'\s+', '', n.supplement_content) in re.sub(r'\s+', '', text))
            child = nodes[2].children[0]
            headings = [p for p in Document(output).paragraphs if p.text == child.title]
            self.assertEqual(len(headings), 1)
            self.assertEqual(headings[0].style.name, 'Heading 3')

    def test_incomplete_roles_rejected(self):
        self.client.extract_json.side_effect = [{'targets': [{'id': '3.2', 'roles': ['analysis']}]}]
        with self.assertRaisesRegex(ValueError, '职责不完整'):
            identify_targets(self.outline, self.client)

    def test_late_distributed_failure_does_not_commit_partial_rewrite(self):
        self.client.extract_json.side_effect = [
            {'targets': [{'id': '3.1', 'roles': ['business_type', 'overview', 'analysis']},
                         {'id': '3.2', 'roles': ['measures']}]},
            {'new_titles': []}, {'new_titles': []}]
        generate = self.client.generate.side_effect
        def fail_measures(**kwargs):
            return '过短' if json.loads(kwargs['user_prompt'])['assigned_roles'] == ['measures'] else generate(**kwargs)
        self.client.generate.side_effect = fail_measures
        before = copy.deepcopy(self.outline)
        with self.assertRaisesRegex(ValueError, '三次生成'):
            rewrite_difficulty_sections(self.outline, self.analysis, self.data, self.client)
        self.assertEqual(self.outline, before)

    def test_non_third_chapter_with_different_parent_title_through_export(self):
        self._assert_pipeline(chapter_number=5)

    def test_standalone_top_level_difficulty_chapter_through_export(self):
        self._assert_pipeline(chapter_number=6, standalone=True)

    def _assert_pipeline(self, chapter_number=3, standalone=False):
        from step3_outlining.scoring_structure import _int_to_chinese
        if chapter_number != 3:
            for node in self.outline.flatten():
                node.id = str(chapter_number) + node.id[1:]
                if node.template_original_id:
                    node.template_original_id = str(chapter_number) + node.template_original_id[1:]
            self.root.template_chapter_id = str(chapter_number)
            self.root.title = '项目服务实施方案'
            self.root.generated_content = self.root.generated_content.replace('需求理解', self.root.title)
        if standalone:
            self.target = self.root
            self.root.title = '项目关键挑战及应对策略'
            self.root.template_original_title = self.root.title
            self.root.children = []
            self.root.generated_content = self.root.title + '\n旧项目重难点'
            self.root.preserve_template_layout = True
            self.root.template_fill_completed = True
            self.root.content_blocks = [{'block_type': 'template_base'}, {'block_type': 'native_table_plan'}]
        _persist_heading_mappings(self.outline)
        self.outline.evaluation_index_entries[0]['final_heading_id'] = self.target.id
        self.outline.evaluation_index_entries[0]['final_heading_title'] = self.target.title
        self.client.extract_json.side_effect = [
            selection(self.target.id),
            {'new_titles': ['项目业态与概况', '重点难点成因分析', '逐项应对及检查闭环']}]
        # 可选择以真实 LLM 对同一脱敏夹具复验，不读取业务资料。
        client = None if os.environ.get('F043_LIVE_LLM') == '1' else self.client
        original_report = {'project': self.outline.project_name,
                           'heading_mappings': copy.deepcopy(self.outline.heading_mappings),
                           'evaluation_index_entries': copy.deepcopy(self.outline.evaluation_index_entries)}
        rewrite_difficulty_sections(self.outline, self.analysis, self.data, client)
        stage5, stage6 = load_stage(5), load_stage(6)
        with tempfile.TemporaryDirectory() as tmp:
            template = os.path.join(tmp, 'template.docx')
            doc = Document()
            doc.sections[0].header.paragraphs[0].text = self.outline.project_name
            footer = doc.sections[0].footer.paragraphs[0]
            for kind in ['begin', 'end']:
                field = OxmlElement('w:fldChar')
                field.set(qn('w:fldCharType'), kind)
                footer.add_run()._r.append(field)
                if kind == 'begin':
                    instruction = OxmlElement('w:instrText')
                    instruction.text = ' PAGE '
                    footer.add_run()._r.append(instruction)
            doc.add_paragraph(f'第{_int_to_chinese(chapter_number)}章 {self.root.title}', style='Heading 1')
            if not standalone:
                doc.add_paragraph('保留章正文')
                doc.add_paragraph('一、服务认知', style='Heading 2')
                doc.add_paragraph('保留相邻正文')
                doc.add_paragraph(self.target.title, style='Heading 2')
            doc.add_paragraph('旧项目重难点')
            doc.add_paragraph('旧项目重复泛化说明')
            doc.add_paragraph('原生表前说明')
            doc.add_table(rows=1, cols=1).cell(0, 0).text = '受保护表格'
            doc.add_paragraph('原生表后说明')
            doc.add_paragraph('签署人：测试')
            doc.add_picture(io.BytesIO(base64.b64decode(
                'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+aL1sAAAAASUVORK5CYII=')))
            doc.add_paragraph('图1 原生图片说明')
            protected = doc.add_paragraph()
            bookmark = OxmlElement('w:bookmarkStart')
            bookmark.set(qn('w:id'), '10')
            bookmark.set(qn('w:name'), 'protected')
            protected._p.append(bookmark)
            if not standalone:
                doc.add_paragraph('三、其他事项', style='Heading 2')
                doc.add_paragraph('保留末节正文')
            doc.save(template)

            stage4_dir = os.path.join(tmp, 'step4')
            paths = _persist_all_chapter_docs(self.outline, stage4_dir, template_path=template)
            artifact = Document(paths[0])
            actual = '\n'.join(p.text for p in artifact.paragraphs)
            self.assertNotIn('旧项目', actual)
            preserved = ['原生表前说明', '原生表后说明', '签署人：测试']
            if not standalone:
                preserved += ['保留相邻正文', '保留末节正文']
            for text in preserved:
                self.assertIn(text, actual)
            self.assertEqual(len(artifact.tables), 1)
            self.assertEqual(len(artifact.inline_shapes), 1)
            texts = stage5._load_chapter_texts(stage4_dir)
            reviewed = stage5._load_step4_outline(stage4_dir, original_report, texts)
            self.assertEqual([n.id for n in reviewed.flatten()], [n.id for n in self.outline.flatten()])
            reviewer = _Reviewer.__new__(_Reviewer)
            evidence = reviewer._bound_entry_evidence(reviewed, reviewed.evaluation_index_entries[0])
            for n in self.target.children:
                self.assertTrue(re.sub(r'\s+', '', n.supplement_content) in re.sub(r'\s+', '', evidence))
            self.assertGreaterEqual(word_count(evidence), 9000)
            self.assertLess(word_count(evidence), 11200, '一级专项审核不能将整章与各子节重复计入')
            records = stage5._load_step4_records(stage4_dir)
            stage5_dir = os.path.join(tmp, 'step5')
            stage5._write_step5_chapters_from_step4(records, reviewed, self.analysis, stage5_dir)
            output = os.path.join(tmp, 'final.docx')
            assemble_step5_document(stage6._load_chapter_records(stage5_dir), reviewed,
                                    output, template_path=template)
            result = Document(output)
            for n in self.target.children:
                matching = [p for p in result.paragraphs if p.text == n.title]
                self.assertEqual(len(matching), 1)
                self.assertEqual(matching[0].style.name, f'Heading {n.level}')
                self.assertTrue(len(matching[0]._p.findall(qn('w:bookmarkStart'))) > 0)
            self.assertEqual(len(result.tables), 2)  # 原生表 + 评标索引
            self.assertEqual(len(result.inline_shapes), 1)
            self.assertNotIn('旧项目', '\n'.join(p.text for p in result.paragraphs))

            # CLI 的兼容导出入口同样不能恢复模板旧段落或遗漏新子节。
            cli_output = os.path.join(tmp, 'cli-final.docx')
            original_paths = [c.artifact_path for c in self.outline.chapters]
            build_document(self.outline, template, cli_output, analysis=self.analysis)
            cli_doc = Document(cli_output)
            cli_text = '\n'.join(p.text for p in cli_doc.paragraphs)
            self.assertNotIn('旧项目', cli_text)
            for n in self.target.children:
                self.assertIn(n.title, cli_text)
            self.assertEqual([c.artifact_path for c in self.outline.chapters], original_paths)

            # 不允许新 manifest 携带其他轮次的评分映射或改变原标题树。
            manifest_path = stage5._find_manifest(stage4_dir)
            with open(manifest_path, encoding='utf-8') as file:
                manifest = json.load(file)
            manifest['final_outline']['heading_mappings'][0]['final_title'] = '篡改标题'
            with open(manifest_path, 'w', encoding='utf-8') as file:
                json.dump(manifest, file, ensure_ascii=False)
            with self.assertRaisesRegex(ValueError, '既有标题'):
                stage5._load_step4_outline(stage4_dir, original_report, texts)


if __name__ == '__main__':
    unittest.main()