Update tools/ai_act_generator.py

tools/ai_act_generator.py

@@ -2,25 +2,16 @@
 # coding=utf-8
 import csv
 import datetime
-import mimetypes
 import os
 import re
-import shutil
-from typing import Optional
-
-from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
-from smolagents.agents import ActionStep, MultiStepAgent
-from smolagents.memory import MemoryStep
-from smolagents.utils import _is_package_available
-
-import gradio as gr
 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, output_path=None, language="fr"):
+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"
@@ -34,20 +25,21 @@ def export_text_to_pdf(text, output_path=None, language="fr"):
     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(10)
+    pdf.ln(5)
 
-    # 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')
+    # 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)
 
-    # Reset style for body
+    # Content
     pdf.set_font("Arial", '', 12)
     pdf.set_text_color(0, 0, 0)
-
-    # Parse and format sections
     for line in text.strip().split('\n'):
         line = line.strip()
         if line.startswith("## "):
@@ -56,19 +48,18 @@ def export_text_to_pdf(text, output_path=None, language="fr"):
             pdf.set_text_color(30, 30, 120)
             pdf.ln(8)
             pdf.cell(0, 10, section_title, ln=True)
-            pdf.ln(2)
             pdf.set_font("Arial", '', 12)
             pdf.set_text_color(0, 0, 0)
         elif line.startswith("- **"):
-            key_value = re.match(r"- \*\*(.+?)\*\*: (.+)", line)
-            if key_value:
-                label, answer = key_value.groups()
+            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, f"{answer}")
+                pdf.multi_cell(0, 10, value)
                 pdf.ln(2)
-        elif line.strip() == "---":
+        elif line == "---":
             pdf.line(10, pdf.get_y(), 200, pdf.get_y())
             pdf.ln(5)
         else:
@@ -79,8 +70,7 @@ def export_text_to_pdf(text, output_path=None, language="fr"):
     return output_path
 
 # === Sequential Questions ===
-QUESTIONS = [
-    ("organization_name", "What is the name of your organization?"),
+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?"),
@@ -93,9 +83,7 @@ QUESTIONS = [
     ("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)?")
-]
-
-RESPONSES = {}
+])
 
 # === Interactive Collection Flow ===
 def step_by_step_agent(user_input, state):
@@ -110,15 +98,15 @@ def step_by_step_agent(user_input, state):
         answers[key] = user_input
 
     if step < len(QUESTIONS):
-        next_question = QUESTIONS[step][1]
+        next_q = QUESTIONS[step][1]
         state["step"] += 1
-        return next_question, state, None  # Ensure return matches expected 3 outputs
+        return next_q, state, None
 
-    filled_template = f"""
+    # Build filled template
+    filled = f"""
 # AI Act Compliance Register
 
 ## General Information
-- **Organization**: {answers['organization_name']}
 - **Responsible Person**: {answers['responsible_person']}
 - **Deployment Date**: {answers['deployment_date']}
 - **System Description**: {answers['ai_description']}
@@ -142,97 +130,50 @@ def step_by_step_agent(user_input, state):
 Generated by AI Act Assistant.
 """
 
-    
+    detected_lang = detect(filled) if filled.strip() else "en"
 
-    try:
-        detected_lang = detect(filled_template)
-    except:
-        detected_lang = "en"
-
-    flag = "🇬🇧" if detected_lang == "en" else "🇫🇷"
-    # Save to CSV
     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 csvfile:
-        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+
+    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_template, language=detected_lang)
-    return f"""{flag} Language detected: {detected_lang.upper()}
-✅ All answers received. Your PDF is ready below:
-👉 Click the download button to save your AI Act compliance register.""", {"done": True, "pdf": pdf_path}, pdf_path
+    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(
-        fill_height=True,
-        css="""
-        footer, 
-        a[href*="gradio.app"],
-        a[href*="huggingface.co"] {
-            display: none !important;
-        }
-
-        /* Disable interaction with share buttons but keep visible */
-        button[aria-label="Share"],
-        button[data-testid="share-button"],
-        button[title="Share"] {
-            pointer-events: none !important;
-            opacity: 0.4 !important;
-            cursor: default !important;
-        }
-        """
-    ) as demo:
-
-        # GDPR Notice
-        gr.Markdown("""
-        ### 🔒 GDPR Notice  
-        This assistant does not store any personal data.  
-        All responses are processed temporarily to generate your document.  
-        You are responsible for the accuracy and legality of the information submitted.
-        """)
-
-        initial_question = QUESTIONS[0][1]
-        initial_message = gr.ChatMessage(role="assistant", content=f"""
-👋 Welcome! I will guide you through the AI Act compliance form.
-Let's begin with a few questions to generate your compliance register.
-{initial_question}
-""")
-        stored_messages = [initial_message]
-        chatbot = gr.Chatbot(type="messages", value=stored_messages)
+    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_button = gr.Button("🔁 Restart")
+        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
 
-            if not file_path or not os.path.isfile(file_path):
-                file_path = None
-            return messages, updated_state, file_path
-
-        def restart_conversation():
-            return [initial_message], {"step": 0, "answers": {}}, 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_button.click(restart_conversation, outputs=[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()
+    return launch_step_by_step_ui()