finrep-report/algo/src/finrep_algo_agent/prompts/builders.py

from __future__ import annotations

import json
from decimal import Decimal
from pathlib import Path

from jinja2 import Environment, FileSystemLoader, select_autoescape

from finrep_algo_agent.schemas.outline import OutlineL1Request, OutlineL2Request
from finrep_algo_agent.schemas.section import SectionRequest

_TPL_DIR = Path(__file__).resolve().parent / "templates"


def _env() -> Environment:
    return Environment(
        loader=FileSystemLoader(_TPL_DIR),
        autoescape=select_autoescape(disabled_extensions=("j2",)),
        trim_blocks=True,
        lstrip_blocks=True,
    )


def _fmt_list(items: list[str] | None) -> str:
    if not items:
        return "无"
    return "、".join(str(x) for x in items)


def _fmt_bool(v: bool | None) -> str:
    if v is True:
        return "是"
    if v is False:
        return "否"
    return "未提供"


def _fmt_decimal(v: Decimal | None) -> str:
    if v is None:
        return "未提供"
    return str(v)


def _format_kv_block(obj: dict[str, object], indent: str = "  ") -> str:
    if not obj:
        return f"{indent}（无扩展字段）"
    lines: list[str] = []
    for k in sorted(obj.keys(), key=lambda x: str(x)):
        v = obj[k]
        if isinstance(v, (dict, list)):
            lines.append(f"{indent}{k}: {json.dumps(v, ensure_ascii=False)}")
        else:
            lines.append(f"{indent}{k}: {v}")
    return "\n".join(lines)


def format_chapter_candidates_block(req: OutlineL1Request) -> str:
    """仅用于请求体显式传入 `chapter_candidates` 时的覆盖文本。"""
    cands = req.chapter_candidates or []
    if not cands:
        return ""
    parts: list[str] = []
    for i, c in enumerate(cands, start=1):
        row = c.model_dump(mode="python", exclude_none=True)
        cid = row.pop("chapter_id", "")
        cname = row.pop("chapter_name", "")
        parts.append(f"【候选 {i}】")
        parts.append(f"  chapter_id: {cid}")
        parts.append(f"  chapter_name: {cname}")
        parts.append("  附加属性（来自预置知识体系，原样供判断）:")
        parts.append(_format_kv_block(row) if row else "  （无扩展字段）")
        parts.append("")
    return "\n".join(parts).strip()


def format_leaf_chapter_candidates_block(leaf: list[dict]) -> str:
    """仅用于请求体显式传入 `leaf_chapter_candidates` 时的覆盖文本。"""
    if not leaf:
        return ""
    parts: list[str] = []
    for i, item in enumerate(leaf, start=1):
        parts.append(f"【末级候选 {i}】")
        parts.append(_format_kv_block(dict(item), indent="  "))
        parts.append("")
    parts.append("--- JSON（机器可读备份，与上一致）---")
    parts.append(json.dumps(leaf, ensure_ascii=False, indent=2))
    return "\n".join(parts).strip()


def _task_background_dict(req: OutlineL1Request) -> dict[str, str]:
    return {
        "report_type": (req.report_type or "").strip(),
        "agreement_amount": _fmt_decimal(req.agreement_amount),
        "enterprise_type": req.enterprise_type or "未提供",
        "group_business_segments": _fmt_list(req.group_business_segments),
        "industry_type": req.industry_type or "未提供",
        "has_independent_report": _fmt_bool(req.has_independent_report),
        "independent_report_types": _fmt_list(req.independent_report_types),
        "candidate_financing_tools": _fmt_list(req.candidate_financing_tools),
        "recommended_financing_tools": _fmt_list(req.recommended_financing_tools),
        "other_requirements": req.other_requirements or "无",
    }


def build_outline_l1_user_prompt(req: OutlineL1Request) -> str:
    ctx = _task_background_dict(req)
    ctx["chapter_candidates_override"] = format_chapter_candidates_block(req)
    return _env().get_template("outline_l1.j2").render(**ctx)


def _outline_l2_background(req: OutlineL2Request) -> OutlineL1Request:
    if req.l1_task_snapshot is not None:
        return req.l1_task_snapshot
    return OutlineL1Request(
        report_type=req.report_type or "未分类",
        agreement_amount=req.agreement_amount,
        enterprise_type=req.enterprise_type,
        group_business_segments=list(req.group_business_segments),
        industry_type=req.industry_type,
        has_independent_report=req.has_independent_report,
        independent_report_types=list(req.independent_report_types),
        candidate_financing_tools=list(req.candidate_financing_tools),
        recommended_financing_tools=list(req.recommended_financing_tools),
        other_requirements=req.other_requirements,
        chapter_candidates=[],
    )


def build_outline_l2_user_prompt(req: OutlineL2Request) -> str:
    ctx = _task_background_dict(_outline_l2_background(req))
    ctx.update(
        {
            "chapter_name": req.chapter_name,
            "chapter_no": req.chapter_no,
            "l1_chapter_id": (req.l1_chapter_id or "").strip(),
            "chapter_presentation_enum": (req.chapter_presentation_enum or "未提供").strip(),
            "chapter_paragraph_count_enum": req.chapter_paragraph_count_enum or "未提供",
            "chapter_reason": req.chapter_reason or "无",
            "overall_logic": req.overall_logic or "无",
            "leaf_chapter_candidates_override": format_leaf_chapter_candidates_block(
                req.leaf_chapter_candidates
            ),
        }
    )
    return _env().get_template("outline_l2.j2").render(**ctx)


def _format_json_block(data: dict) -> str:
    if not data:
        return "（无）"
    return json.dumps(data, ensure_ascii=False, indent=2)


def build_section_user_prompt(req: SectionRequest) -> str:
    tt = (req.template_type or "info").lower()
    name_map = {
        "info": "section_info.j2",
        "analysis": "section_analysis.j2",
        "metric": "section_metric.j2",
        "judgment": "section_judgment.j2",
    }
    tpl_name = name_map.get(tt, "section_info.j2")
    ctx = {
        "report_type": req.report_type or "财务顾问",
        "overall_logic": req.overall_logic or "（未提供整篇报告撰写逻辑，请结合输入自行保持语气一致。）",
        "chapter_logic": req.chapter_logic or "（未提供一级章节写作逻辑。）",
        "paragraph_position": req.paragraph_position or "（未提供段落定位说明。）",
        "task_input_block": _format_json_block(req.task_input),
        "data_block": _format_json_block(req.data_package),
        "paragraph_logic": req.paragraph_logic or "（未提供段落撰写逻辑。）",
        "example_block": req.example or "（无示例）",
        "notes_block": req.notes or "（无补充注意事项。）",
    }
    return _env().get_template(tpl_name).render(**ctx)