tools/high_risk_ai_register.py

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

from tools.common import prepend_metadata_questions  # ✅ Add common metadata helper

# === PDF Export Function ===
def export_text_to_pdf(text, metadata=None, output_path=None, language="en"):
    if output_path is None:
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"high_risk_ai_summary_{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 = "High-Risk AI System Documentation" if language == "en" else "Documentation des Systèmes IA à Haut Risque"
    pdf.cell(0, 15, title, ln=True, align='C')
    pdf.ln(10)

    # Metadata
    if metadata:
        pdf.set_font("Arial", '', 12)
        pdf.set_text_color(90, 90, 90)
        pdf.multi_cell(0, 10, f"Organization: {metadata.get('organization', 'N/A')}")
        pdf.multi_cell(0, 10, f"Completed by: {metadata.get('completed_by', 'N/A')} ({metadata.get('role', 'N/A')})")
        pdf.multi_cell(0, 10, f"Timestamp: {metadata.get('timestamp', 'N/A')}")
        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, answer = 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, f"{answer}")
        elif line == "---":
            pdf.line(10, pdf.get_y(), 200, pdf.get_y())
            pdf.ln(5)
        else:
            pdf.multi_cell(0, 10, line)
    pdf.output(output_path)
    return output_path

# === Questions (core, then prepend metadata) ===
BASE_QUESTIONS = [
    ("system_name", "What is the name of your AI system?"),
    ("system_purpose", "What is its intended purpose?"),
    ("category", "Which Annex III category does it fall under?"),
    ("developer", "Who developed it?"),
    ("users", "Who will use it?"),
    ("input_types", "What types of input data does it use?"),
    ("output", "What actions does it perform?"),
    ("dependencies", "List any critical dependencies (e.g., APIs, models)."),
    ("context", "What is the intended deployment environment?"),
    ("justification", "Why is it high-risk under the AI Act?")
]

QUESTIONS = prepend_metadata_questions(BASE_QUESTIONS)

# === Step-by-step Conversation ===
def step_by_step(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

    # === Format content ===
    content = f"""
# High-Risk AI System Summary

## General Info
- **System Name**: {answers.get('system_name', '')}
- **Purpose**: {answers.get('system_purpose', '')}
- **Annex III Category**: {answers.get('category', '')}
- **Developer**: {answers.get('developer', '')}
- **Intended Users**: {answers.get('users', '')}
- **Deployment Context**: {answers.get('context', '')}

## Technical Info
- **Input Types**: {answers.get('input_types', '')}
- **System Output**: {answers.get('output', '')}
- **Dependencies**: {answers.get('dependencies', '')}

## Risk Classification
- **Justification for High-Risk**: {answers.get('justification', '')}

---
Generated by AI Act Assistant.
"""
    lang = detect(content)
    metadata = {
        "organization": answers.get("organization_name", "N/A"),
        "completed_by": answers.get("user_name", "N/A"),
        "role": answers.get("user_role", "N/A"),
        "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }

    pdf_path = export_text_to_pdf(content, metadata=metadata, language=lang)

    return "✅ All data collected!\n📝 Your summary is ready. Click below to download your PDF.", {"done": True}, pdf_path

# === Gradio UI ===
def launch_ui():
    with gr.Blocks(title="High-Risk AI Summary Tool") as demo:
        chatbot = gr.Chatbot(label="🛡️ High-Risk AI Summary Assistant")
        user_input = gr.Textbox(placeholder="Your answer...", label="Answer")
        state = gr.State()
        file_output = gr.File(label="Download PDF", visible=True)
        reset_btn = gr.Button("🔁 Start Over")

        first_q = QUESTIONS[0][1]
        chatbot.value = [gr.ChatMessage(role="assistant", content=first_q)]

        def run_chat(msg, state):
            reply, state, pdf = step_by_step(msg, state)
            messages = [gr.ChatMessage(role="user", content=msg)]
            if reply:
                messages.append(gr.ChatMessage(role="assistant", content=reply))
            return messages, state, pdf

        def reset_all():
            return [gr.ChatMessage(role="assistant", content=QUESTIONS[0][1])], {"step": 0, "answers": {}}, None

        user_input.submit(run_chat, [user_input, state], [chatbot, state, file_output])
        reset_btn.click(reset_all, outputs=[chatbot, state, file_output])

    demo.launch(show_api=False)

def get_questions():
    return QUESTIONS

def run_tool():
    launch_ui()