memo_generator.py

from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.platypus import (
    SimpleDocTemplate, Paragraph, Spacer, Table, 
    TableStyle, PageBreak, Image
)
from reportlab.lib import colors
from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY, TA_LEFT
from datetime import datetime
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import io
import json
import re


def generate_policy_memo(topic, imf_content, wb_content, user_question, full_answer):
    """
    生成专业的政策备忘录PDF
    
    Args:
        topic: 政策主题
        imf_content: IMF相关内容片段列表
        wb_content: World Bank相关内容片段列表
        user_question: 用户的原始问题
        full_answer: AI生成的完整回答
    
    Returns:
        BytesIO: PDF文件buffer
    """
    
    # 1. 使用LLM生成结构化分析
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    
    analysis_prompt = ChatPromptTemplate.from_template("""
You are a Senior Policy Analyst preparing a briefing memo for an Executive Director at an international financial institution.

Based on the following analysis and source materials, create a structured policy brief:

Original Question: {question}

Full Analysis: {full_answer}

IMF Source Materials ({imf_count} excerpts):
{imf_content}

World Bank Source Materials ({wb_count} excerpts):
{wb_content}

Please provide a structured output in the following JSON format:
{{
  "executive_summary": "2-3 sentence high-level summary",
  "key_findings": [
    "Finding 1",
    "Finding 2",
    "Finding 3"
  ],
  "imf_perspective": "1 paragraph summarizing IMF's position",
  "wb_perspective": "1 paragraph summarizing World Bank's position",
  "comparative_analysis": "1 paragraph comparing and contrasting the two perspectives",
  "policy_implications": [
    "Implication 1",
    "Implication 2"
  ],
  "knowledge_gaps": "What information is missing or unclear?"
}}

Ensure the output is valid JSON only, with no additional text.
""")
    
    # 格式化内容
    imf_text = "\n\n".join([
        f"Excerpt {i+1} (Page {doc.metadata.get('page', 'N/A')}):\n{doc.page_content[:400]}..." 
        for i, doc in enumerate(imf_content[:4])
    ]) if imf_content else "No IMF content available."
    
    wb_text = "\n\n".join([
        f"Excerpt {i+1} (Page {doc.metadata.get('page', 'N/A')}):\n{doc.page_content[:400]}..." 
        for i, doc in enumerate(wb_content[:4])
    ]) if wb_content else "No World Bank content available."
    
    try:
        response = llm.invoke(
            analysis_prompt.format(
                question=user_question,
                full_answer=full_answer,
                imf_content=imf_text,
                wb_content=wb_text,
                imf_count=len(imf_content),
                wb_count=len(wb_content)
            )
        )
        
        # 提取JSON（移除可能的markdown标记）
        json_text = response.content
        json_text = re.sub(r'```json\s*', '', json_text)
        json_text = re.sub(r'```\s*', '', json_text)
        json_text = json_text.strip()
        
        analysis = json.loads(json_text)
    except Exception as e:
        # Fallback：如果JSON解析失败，使用默认结构
        print(f"JSON parsing error: {e}")
        analysis = {
            "executive_summary": full_answer[:300] + "...",
            "key_findings": ["See full analysis below"],
            "imf_perspective": "See source materials",
            "wb_perspective": "See source materials",
            "comparative_analysis": full_answer[:500] + "...",
            "policy_implications": ["Refer to detailed analysis"],
            "knowledge_gaps": "Additional research may be needed"
        }
    
    # 2. 创建PDF
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(
        buffer, 
        pagesize=letter,
        topMargin=0.75*inch, 
        bottomMargin=0.75*inch,
        leftMargin=1*inch, 
        rightMargin=1*inch
    )
    
    # 3. 定义样式
    styles = getSampleStyleSheet()
    
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=18,
        textColor=colors.HexColor('#1f4788'),
        spaceAfter=6,
        alignment=TA_CENTER,
        fontName='Helvetica-Bold'
    )
    
    subtitle_style = ParagraphStyle(
        'Subtitle',
        parent=styles['Normal'],
        fontSize=10,
        textColor=colors.HexColor('#666666'),
        alignment=TA_CENTER,
        spaceAfter=20
    )
    
    heading_style = ParagraphStyle(
        'CustomHeading',
        parent=styles['Heading2'],
        fontSize=13,
        textColor=colors.HexColor('#2d5aa6'),
        spaceAfter=8,
        spaceBefore=16,
        fontName='Helvetica-Bold',
        borderPadding=(0, 0, 5, 0),
        borderColor=colors.HexColor('#2d5aa6'),
        borderWidth=0,
        leftIndent=0
    )
    
    body_style = ParagraphStyle(
        'CustomBody',
        parent=styles['BodyText'],
        fontSize=11,
        alignment=TA_JUSTIFY,
        spaceAfter=10,
        leading=14
    )
    
    bullet_style = ParagraphStyle(
        'Bullet',
        parent=styles['Normal'],
        fontSize=11,
        leftIndent=20,
        spaceAfter=8,
        bulletIndent=10,
        leading=14
    )
    
    # 4. 构建文档内容
    story = []
    
    # === 标题页 ===
    story.append(Spacer(1, 0.5*inch))
    story.append(Paragraph("POLICY ANALYSIS BRIEFING", title_style))
    story.append(Paragraph(
        f"{datetime.now().strftime('%B %d, %Y')}", 
        subtitle_style
    ))
    
    # 主题框
    topic_table = Table(
        [[Paragraph(f"<b>TOPIC:</b> {topic[:150]}", styles['Normal'])]],
        colWidths=[6*inch]
    )
    topic_table.setStyle(TableStyle([
        ('BOX', (0, 0), (-1, -1), 2, colors.HexColor('#1f4788')),
        ('BACKGROUND', (0, 0), (-1, -1), colors.HexColor('#f0f4f8')),
        ('TOPPADDING', (0, 0), (-1, -1), 12),
        ('BOTTOMPADDING', (0, 0), (-1, -1), 12),
        ('LEFTPADDING', (0, 0), (-1, -1), 15),
        ('RIGHTPADDING', (0, 0), (-1, -1), 15),
    ]))
    story.append(topic_table)
    story.append(Spacer(1, 0.3*inch))
    
    # 元数据表格
    metadata_data = [
        ['Prepared by:', 'PolicyGraph AI Analyst'],
        ['Analysis Date:', datetime.now().strftime('%B %d, %Y')],
        ['Sources:', f'IMF ({len(imf_content)} excerpts), World Bank ({len(wb_content)} excerpts)']
    ]
    
    metadata_table = Table(metadata_data, colWidths=[1.5*inch, 4.5*inch])
    metadata_table.setStyle(TableStyle([
        ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
        ('FONTSIZE', (0, 0), (-1, -1), 9),
        ('TEXTCOLOR', (0, 0), (0, -1), colors.HexColor('#555555')),
        ('VALIGN', (0, 0), (-1, -1), 'TOP'),
        ('BOTTOMPADDING', (0, 0), (-1, -1), 6),
    ]))
    story.append(metadata_table)
    story.append(Spacer(1, 0.4*inch))
    
    # 分隔线
    story.append(Table([['']], colWidths=[6*inch], 
                      style=[('LINEABOVE', (0,0), (-1,0), 2, colors.HexColor('#1f4788'))]))
    story.append(Spacer(1, 0.3*inch))
    
    # === EXECUTIVE SUMMARY ===
    story.append(Paragraph("EXECUTIVE SUMMARY", heading_style))
    story.append(Paragraph(analysis['executive_summary'], body_style))
    story.append(Spacer(1, 0.2*inch))
    
    # === KEY FINDINGS ===
    story.append(Paragraph("KEY FINDINGS", heading_style))
    for i, finding in enumerate(analysis['key_findings'], 1):
        story.append(Paragraph(f"• {finding}", bullet_style))
    story.append(Spacer(1, 0.2*inch))
    
    # === ORGANIZATIONAL PERSPECTIVES ===
    story.append(Paragraph("ORGANIZATIONAL PERSPECTIVES", heading_style))
    
    # IMF Perspective
    if imf_content:
        story.append(Paragraph("<b>International Monetary Fund (IMF)</b>", body_style))
        story.append(Paragraph(analysis['imf_perspective'], body_style))
        story.append(Spacer(1, 0.15*inch))
    
    # World Bank Perspective
    if wb_content:
        story.append(Paragraph("<b>World Bank Group</b>", body_style))
        story.append(Paragraph(analysis['wb_perspective'], body_style))
        story.append(Spacer(1, 0.2*inch))
    
    # === COMPARATIVE ANALYSIS ===
    story.append(Paragraph("COMPARATIVE ANALYSIS", heading_style))
    story.append(Paragraph(analysis['comparative_analysis'], body_style))
    story.append(Spacer(1, 0.2*inch))
    
    # === 对比表格 ===
    comparison_data = [
        ['Dimension', 'IMF Focus', 'World Bank Focus'],
        [
            'Primary Concern',
            'Labor market risks\nMacroeconomic stability' if imf_content else 'N/A',
            'Infrastructure development\nDigital inclusion' if wb_content else 'N/A'
        ],
        [
            'Geographic Focus',
            'Advanced & emerging economies' if imf_content else 'N/A',
            'Developing economies' if wb_content else 'N/A'
        ],
        [
            'Policy Tools',
            'Fiscal & monetary frameworks' if imf_content else 'N/A',
            'Investment & capacity building' if wb_content else 'N/A'
        ]
    ]
    
    comparison_table = Table(comparison_data, colWidths=[1.5*inch, 2.25*inch, 2.25*inch])
    comparison_table.setStyle(TableStyle([
        ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#1f4788')),
        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
        ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
        ('VALIGN', (0, 0), (-1, -1), 'TOP'),
        ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
        ('FONTSIZE', (0, 0), (-1, 0), 10),
        ('FONTSIZE', (0, 1), (-1, -1), 9),
        ('BOTTOMPADDING', (0, 0), (-1, -1), 8),
        ('TOPPADDING', (0, 0), (-1, -1), 8),
        ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
        ('ROWBACKGROUNDS', (0, 1), (-1, -1), [colors.white, colors.HexColor('#f8f9fa')])
    ]))
    
    story.append(comparison_table)
    story.append(Spacer(1, 0.3*inch))
    
    # === PAGE BREAK ===
    story.append(PageBreak())
    
    # === POLICY IMPLICATIONS ===
    story.append(Paragraph("POLICY IMPLICATIONS & RECOMMENDATIONS", heading_style))
    for i, implication in enumerate(analysis['policy_implications'], 1):
        story.append(Paragraph(f"{i}. {implication}", bullet_style))
    story.append(Spacer(1, 0.2*inch))
    
    # === KNOWLEDGE GAPS ===
    story.append(Paragraph("KNOWLEDGE GAPS & AREAS FOR FURTHER RESEARCH", heading_style))
    story.append(Paragraph(analysis['knowledge_gaps'], body_style))
    story.append(Spacer(1, 0.3*inch))
    
    # === SOURCES ===
    story.append(Paragraph("SOURCE DOCUMENTS", heading_style))
    
    if imf_content:
        story.append(Paragraph("<b>IMF Sources:</b>", body_style))
        for i, doc in enumerate(imf_content[:5], 1):
            source = doc.metadata.get('source', 'Unknown').split('\\')[-1]
            page = doc.metadata.get('page', 'N/A')
            story.append(Paragraph(
                f"{i}. {source}, Page {page}", 
                bullet_style
            ))
        story.append(Spacer(1, 0.1*inch))
    
    if wb_content:
        story.append(Paragraph("<b>World Bank Sources:</b>", body_style))
        start_num = len(imf_content[:5]) + 1 if imf_content else 1
        for i, doc in enumerate(wb_content[:5], start_num):
            source = doc.metadata.get('source', 'Unknown').split('\\')[-1]
            page = doc.metadata.get('page', 'N/A')
            story.append(Paragraph(
                f"{i}. {source}, Page {page}", 
                bullet_style
            ))
    
    # === FOOTER ===
    story.append(Spacer(1, 0.5*inch))
    footer_style = ParagraphStyle(
        'Footer',
        parent=styles['Normal'],
        fontSize=8,
        textColor=colors.grey,
        alignment=TA_CENTER
    )
    story.append(Table([['']], colWidths=[6*inch], 
                      style=[('LINEABOVE', (0,0), (-1,0), 0.5, colors.grey)]))
    story.append(Spacer(1, 0.1*inch))
    story.append(Paragraph(
        "This document was automatically generated by PolicyGraph AI Analyst | For internal use only",
        footer_style
    ))
    story.append(Paragraph(
        f"Generated on {datetime.now().strftime('%B %d, %Y at %I:%M %p')}",
        footer_style
    ))
    
    # 5. 构建PDF
    doc.build(story)
    buffer.seek(0)
    return buffer


# 测试函数
if __name__ == "__main__":
    print("Memo generator module loaded successfully")