Update tools/ai_technical_doc.py

tools/ai_technical_doc.py

@@ -1,13 +1,15 @@
 # tools/ai_technical_doc.py
 import datetime
+import os
 import re
 from fpdf import FPDF
 from langdetect import detect
 import gradio as gr
-import os
+
+from tools.common import prepend_metadata_questions  # Shared metadata question helper
 
 # === PDF Export Function ===
-def export_text_to_pdf(text, output_path=None, language="en"):
+def export_text_to_pdf(text, answers, output_path=None, language="en"):
     if output_path is None:
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         output_path = f"technical_documentation_{timestamp}.pdf"
@@ -20,8 +22,19 @@ def export_text_to_pdf(text, output_path=None, language="en"):
     pdf.set_text_color(0, 51, 102)
     title = "Technical Documentation - AI Act (Art. 11)" if language == "en" else "Documentation Technique - AI Act"
     pdf.cell(0, 15, title, ln=True, align='C')
-    pdf.ln(10)
-
+    pdf.ln(5)
+
+    # Metadata under 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)
+
+    # Main body
     pdf.set_font("Arial", '', 12)
     pdf.set_text_color(0, 0, 0)
     for line in text.strip().split('\n'):
@@ -47,11 +60,12 @@ def export_text_to_pdf(text, output_path=None, language="en"):
             pdf.ln(5)
         else:
             pdf.multi_cell(0, 10, line)
+
     pdf.output(output_path)
     return output_path
 
 # === Questions ===
-QUESTIONS = [
+CORE_QUESTIONS = [
     ("system_name", "What is the name of your AI system?"),
     ("provider", "Who is the provider or developer of the system?"),
     ("intended_purpose", "What is the intended purpose of the system?"),
@@ -66,11 +80,12 @@ QUESTIONS = [
     ("recordkeeping", "How are logs and records maintained?")
 ]
 
+QUESTIONS = prepend_metadata_questions(CORE_QUESTIONS)
+
 def get_questions():
     return QUESTIONS
 
 def run_tool():
-    import gradio as gr
     state = {"step": 0, "answers": {}}
 
     def step_by_step_agent(user_input, state):
@@ -86,13 +101,18 @@ def run_tool():
             state["step"] += 1
             return next_question, state, None
 
-        content = "\n".join([f"- **{label}**: {answers.get(key, '')}" for key, label in QUESTIONS])
+        # Final content for PDF
+        content = "\n".join([f"- **{label}**: {answers.get(key, '')}" for key, label in QUESTIONS if key not in ["user_name", "user_role", "organization_name"]])
         detected_lang = detect(content)
-        pdf_path = export_text_to_pdf(content, language=detected_lang)
+        pdf_path = export_text_to_pdf(content, answers, language=detected_lang)
         return "✅ Completed. Download your documentation below.", {"done": True}, pdf_path
 
     with gr.Blocks(title="AI Technical Documentation Tool") as demo:
-        chatbot = gr.Chatbot(label="🧠 Technical Doc Assistant", value=[{"role": "assistant", "content": QUESTIONS[0][1]}], type="messages")
+        chatbot = gr.Chatbot(
+            label="🧠 Technical Doc Assistant",
+            value=[{"role": "assistant", "content": QUESTIONS[0][1]}],
+            type="messages"
+        )
         msg = gr.Textbox(label="Your answer")
         state_var = gr.State(state)
         file_output = gr.File(label="Download PDF", visible=True)