finrep-report/algo/tests/test_prompts.py

from decimal import Decimal

from finrep_algo_agent.prompts.builders import (
    build_outline_l1_user_prompt,
    build_outline_l2_user_prompt,
    build_section_user_prompt,
)
from finrep_algo_agent.schemas.outline import ChapterCandidate, OutlineL1Request, OutlineL2Request
from finrep_algo_agent.schemas.section import SectionRequest


def test_outline_l1_template_renders() -> None:
    req = OutlineL1Request(
        report_type="项目融资",
        agreement_amount=Decimal("100"),
        chapter_candidates=[ChapterCandidate(chapter_id="c1", chapter_name="测试章")],
    )
    text = build_outline_l1_user_prompt(req)
    assert "项目融资" in text
    assert "c1" in text
    assert "一级章节候选清单" in text
    assert "presentation_enum" in text


def test_outline_l2_template_renders() -> None:
    l1 = OutlineL1Request(
        report_type="项目融资",
        agreement_amount=Decimal("50"),
        chapter_candidates=[],
    )
    req = OutlineL2Request(
        chapter_name="融资方案",
        chapter_no="3",
        chapter_presentation_enum="S1",
        chapter_paragraph_count_enum="P2",
        chapter_reason="保留",
        overall_logic="总逻辑",
        leaf_chapter_candidates=[{"name": "末级1", "知识单元类型": "单元呈现型"}],
        l1_task_snapshot=l1,
    )
    text = build_outline_l2_user_prompt(req)
    assert "融资方案" in text
    assert "末级1" in text
    assert "判断说明" in text
    assert "呈现方式" in text
    assert "S1" in text


def test_section_template_renders() -> None:
    req = SectionRequest(
        knowledge_unit_id="ku1",
        report_type="财务顾问",
        template_type="info",
        overall_logic="整体",
        chapter_logic="章逻辑",
        paragraph_position="定位",
        task_input={"a": 1},
        data_package={"b": 2},
        paragraph_logic="撰写",
    )
    text = build_section_user_prompt(req)
    assert "财务顾问" in text
    assert "撰写" in text


def test_section_request_accepts_legacy_data_package_1() -> None:
    req = SectionRequest.model_validate(
        {
            "knowledge_unit_id": "ku-x",
            "template_type": "info",
            "data_package_1": {"auto": {"k": "v"}},
        }
    )
    assert req.auto_data_package == {"auto": {"k": "v"}}


def test_builtin_l1_embedded_in_template() -> None:
    req = OutlineL1Request(report_type="项目融资", chapter_candidates=[])
    text = build_outline_l1_user_prompt(req)
    assert "pf-l1-01" in text
    assert "企业基本情况分析" in text
    assert "融资方案设计" in text


def test_builtin_l2_embedded_branch() -> None:
    l1 = OutlineL1Request(report_type="项目融资", agreement_amount=Decimal("1"))
    req = OutlineL2Request(
        chapter_name="企业财务情况分析",
        chapter_no="2",
        l1_chapter_id="pf-l1-02",
        chapter_presentation_enum="S1",
        chapter_paragraph_count_enum="P2",
        chapter_reason="r",
        overall_logic="g",
        leaf_chapter_candidates=[],
        l1_task_snapshot=l1,
    )
    text = build_outline_l2_user_prompt(req)
    assert "pf-02-ku-1" in text or "财务概况" in text


def test_builtin_l2_pf_l1_01_matches_requirement_doc_list() -> None:
    l1 = OutlineL1Request(
        report_type="项目融资",
        agreement_amount=Decimal("40"),
        enterprise_type="集团企业",
        group_business_segments=["系统建设", "AI产品", "硬件设施"],
    )
    req = OutlineL2Request(
        chapter_name="企业基本情况分析",
        chapter_no="1",
        l1_chapter_id="pf-l1-01",
        chapter_presentation_enum="S1",
        chapter_paragraph_count_enum="P2",
        chapter_reason="无独立报告且为集团企业",
        overall_logic="写作蓝图示例",
        leaf_chapter_candidates=[],
        l1_task_snapshot=l1,
    )
    text = build_outline_l2_user_prompt(req)
    assert "企业基本情况概述" in text
    assert "批次处理类" in text
    assert "pf-01-ku-03" in text
    assert "核心决策规则" in text