tools/ai_act_generator.py

#!/usr/bin/env python
# coding=utf-8
import csv
import datetime
import os
import re
from fpdf import FPDF
from langdetect import detect

import gradio as gr
from tools.common import prepend_metadata_questions

# === PDF Export Function with Language Option ===
def export_text_to_pdf(text, answers, output_path=None, language="fr"):
    if output_path is None:
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"ai_act_register_{timestamp}.pdf"

    pdf = FPDF()
    pdf.add_page()
    pdf.set_auto_page_break(auto=True, margin=15)

    # Title
    pdf.set_font("Arial", 'B', 16)
    pdf.set_text_color(0, 51, 102)
    title = "Documentation Record for High-Risk AI Systems" if language == "en" else "Registre de Conformité AI Act"
    pdf.cell(0, 15, title, ln=True, align='C')
    pdf.ln(5)

    # Metadata below title
    pdf.set_font("Arial", 'I', 11)
    pdf.set_text_color(80, 80, 80)
    name = answers.get("user_name", "N/A")
    role = answers.get("user_role", "N/A")
    org = answers.get("organization_name", "N/A")
    timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    pdf.multi_cell(0, 10, f"Completed by {name} ({role}) at {org} on {timestamp}", align="C")
    pdf.ln(5)

    # Content
    pdf.set_font("Arial", '', 12)
    pdf.set_text_color(0, 0, 0)
    for line in text.strip().split('\n'):
        line = line.strip()
        if line.startswith("## "):
            section_title = line.replace("## ", "").strip()
            pdf.set_font("Arial", 'B', 13)
            pdf.set_text_color(30, 30, 120)
            pdf.ln(8)
            pdf.cell(0, 10, section_title, ln=True)
            pdf.set_font("Arial", '', 12)
            pdf.set_text_color(0, 0, 0)
        elif line.startswith("- **"):
            match = re.match(r"- \*\*(.+?)\*\*: (.+)", line)
            if match:
                label, value = match.groups()
                pdf.set_font("Arial", 'B', 12)
                pdf.cell(0, 10, f"{label}:", ln=True)
                pdf.set_font("Arial", '', 12)
                pdf.multi_cell(0, 10, value)
                pdf.ln(2)
        elif line == "---":
            pdf.line(10, pdf.get_y(), 200, pdf.get_y())
            pdf.ln(5)
        else:
            pdf.multi_cell(0, 10, line)
            pdf.ln(2)

    pdf.output(output_path)
    return output_path

# === Sequential Questions ===
QUESTIONS = prepend_metadata_questions([
    ("responsible_person", "Who is responsible for this AI system?"),
    ("deployment_date", "When is the AI system scheduled to be deployed?"),
    ("ai_type", "What type of AI system is it?"),
    ("ai_description", "Please briefly describe what the system does."),
    ("risk_level", "What is the risk level of this system (e.g., high, medium)?"),
    ("risk_justification", "Why do you consider it this risk level?"),
    ("data_evaluation", "How have you evaluated the training data?"),
    ("technical_docs", "What technical documentation is available?"),
    ("human_oversight", "What kind of human oversight is planned?"),
    ("transparency_measures", "What transparency mechanisms are in place?"),
    ("audit_frequency", "How often will the system be audited?"),
    ("compliance_contact", "Who is the contact person for compliance (email or name)?")
])

# === Interactive Collection Flow ===
def step_by_step_agent(user_input, state):
    if state is None:
        state = {"step": 0, "answers": {}}

    step = state["step"]
    answers = state["answers"]

    if step > 0:
        key, _ = QUESTIONS[step - 1]
        answers[key] = user_input

    if step < len(QUESTIONS):
        next_q = QUESTIONS[step][1]
        state["step"] += 1
        return next_q, state, None

    # Build filled template
    filled = f"""
# AI Act Compliance Register

## General Information
- **Responsible Person**: {answers['responsible_person']}
- **Deployment Date**: {answers['deployment_date']}
- **System Description**: {answers['ai_description']}

## Risk Category
- **Type**: {answers['ai_type']}
- **Risk Level**: {answers['risk_level']}
- **Justification**: {answers['risk_justification']}

## Compliance Measures
- **Data Evaluation**: {answers['data_evaluation']}
- **Technical Docs**: {answers['technical_docs']}
- **Human Oversight**: {answers['human_oversight']}
- **Transparency Measures**: {answers['transparency_measures']}

## Audit & Follow-up
- **Audit Frequency**: {answers['audit_frequency']}
- **Compliance Contact**: {answers['compliance_contact']}

---
Generated by AI Act Assistant.
"""

    detected_lang = detect(filled) if filled.strip() else "en"

    csv_file = "ai_act_registers.csv"
    fieldnames = [key for key, _ in QUESTIONS] + ["timestamp"]
    row_data = {**answers, "timestamp": datetime.datetime.now().isoformat()}
    file_exists = os.path.isfile(csv_file)

    with open(csv_file, mode="a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        if not file_exists:
            writer.writeheader()
        writer.writerow(row_data)

    pdf_path = export_text_to_pdf(filled, answers, language=detected_lang)
    return f"✅ Your PDF is ready for download.", {"done": True, "pdf": pdf_path}, pdf_path

# === Gradio Interface ===
def launch_step_by_step_ui():
    with gr.Blocks(title="AI Act Assistant", css="""footer, a[href*="gradio.app"], a[href*="huggingface.co"] { display: none !important; }""") as demo:
        gr.Markdown("### 🔒 GDPR Notice\nThis assistant does not store personal data. Use responsibly.")
        chatbot = gr.Chatbot(type="messages", value=[])
        msg = gr.Textbox(label="Your answer")
        state = gr.State()
        file_output = gr.File(label="Download PDF", visible=True)
        restart = gr.Button("🔁 Restart")

        def chat_logic(user_msg, state):
            reply, updated_state, file_path = step_by_step_agent(user_msg, state)
            messages = [gr.ChatMessage(role="user", content=user_msg)]
            if isinstance(reply, str):
                messages.append(gr.ChatMessage(role="assistant", content=reply))
            return messages, updated_state, file_path if file_path else None

        def reset():
            first_q = QUESTIONS[0][1]
            return [gr.ChatMessage(role="assistant", content=f"👋 Let's get started.\n\n{first_q}")], {"step": 0, "answers": {}}, None

        msg.submit(chat_logic, [msg, state], [chatbot, state, file_output])
        restart.click(reset, outputs=[chatbot, state, file_output])

    demo.launch(show_api=False)

def get_questions():
    return QUESTIONS

def run_tool():
    return launch_step_by_step_ui()