Update tools/high_risk_ai_register.py

tools/high_risk_ai_register.py

@@ -8,8 +8,10 @@ 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, output_path=None, language="en"):
+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"
@@ -25,17 +27,18 @@ def export_text_to_pdf(text, output_path=None, language="en"):
     pdf.cell(0, 15, title, ln=True, align='C')
     pdf.ln(10)
 
-    # Subtitle
-    pdf.set_font("Arial", 'I', 12)
-    pdf.set_text_color(90, 90, 90)
-    subtitle = f"Generated on {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
-    pdf.cell(0, 10, subtitle, ln=True, align='C')
-    pdf.ln(5)
+    # 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 formatting
+    # 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("## "):
@@ -62,8 +65,8 @@ def export_text_to_pdf(text, output_path=None, language="en"):
     pdf.output(output_path)
     return output_path
 
-# === Questions for AI System Documentation ===
-QUESTIONS = [
+# === 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?"),
@@ -73,10 +76,12 @@ QUESTIONS = [
     ("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?"),
+    ("justification", "Why is it high-risk under the AI Act?")
 ]
 
-# === Conversation Logic ===
+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": {}}
@@ -93,9 +98,10 @@ def step_by_step(user_input, state):
         state["step"] += 1
         return next_q, state, None
 
-    # === Format document ===
+    # === Format content ===
     content = f"""
 # High-Risk AI System Summary
+
 ## General Info
 - **System Name**: {answers.get('system_name', '')}
 - **Purpose**: {answers.get('system_purpose', '')}
@@ -111,14 +117,21 @@ def step_by_step(user_input, state):
 
 ## Risk Classification
 - **Justification for High-Risk**: {answers.get('justification', '')}
+
 ---
 Generated by AI Act Assistant.
 """
-
     lang = detect(content)
-    pdf_path = export_text_to_pdf(content, language=lang)
+    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 f"✅ All data collected!\n📝 Your summary is ready. Click below to download your PDF.", {"done": True}, pdf_path
+    return "✅ All data collected!\n📝 Your summary is ready. Click below to download your PDF.", {"done": True}, pdf_path
 
 # === Gradio UI ===
 def launch_ui():