Dave67350/First_agent_template

1# tools/ai_technical_doc.py2import datetime3import os4import re5from fpdf import FPDF6from langdetect import detect7import gradio as gr8 9from tools.common import prepend_metadata_questions  # Shared metadata question helper10 11# === PDF Export Function ===12def export_text_to_pdf(text, answers, output_path=None, language="en"):13    if output_path is None:14        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")15        output_path = f"technical_documentation_{timestamp}.pdf"16 17    pdf = FPDF()18    pdf.add_page()19    pdf.set_auto_page_break(auto=True, margin=15)20 21    pdf.set_font("Arial", 'B', 16)22    pdf.set_text_color(0, 51, 102)23    title = "Technical Documentation - AI Act (Art. 11)" if language == "en" else "Documentation Technique - AI Act"24    pdf.cell(0, 15, title, ln=True, align='C')25    pdf.ln(5)26 27    # Metadata under title28    pdf.set_font("Arial", 'I', 11)29    pdf.set_text_color(80, 80, 80)30    name = answers.get("user_name", "N/A")31    role = answers.get("user_role", "N/A")32    org = answers.get("organization_name", "N/A")33    timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')34    pdf.multi_cell(0, 10, f"Completed by {name} ({role}) at {org} on {timestamp}", align="C")35    pdf.ln(5)36 37    # Main body38    pdf.set_font("Arial", '', 12)39    pdf.set_text_color(0, 0, 0)40    for line in text.strip().split('\n'):41        line = line.strip()42        if line.startswith("## "):43            section = line.replace("## ", "").strip()44            pdf.set_font("Arial", 'B', 13)45            pdf.set_text_color(30, 30, 120)46            pdf.ln(8)47            pdf.cell(0, 10, section, ln=True)48            pdf.set_font("Arial", '', 12)49            pdf.set_text_color(0, 0, 0)50        elif line.startswith("- **"):51            match = re.match(r"- \*\*(.+?)\*\*: (.+)", line)52            if match:53                label, value = match.groups()54                pdf.set_font("Arial", 'B', 12)55                pdf.cell(0, 10, f"{label}:", ln=True)56                pdf.set_font("Arial", '', 12)57                pdf.multi_cell(0, 10, value)58        elif line == "---":59            pdf.line(10, pdf.get_y(), 200, pdf.get_y())60            pdf.ln(5)61        else:62            pdf.multi_cell(0, 10, line)63 64    pdf.output(output_path)65    return output_path66 67# === Questions ===68CORE_QUESTIONS = [69    ("system_name", "What is the name of your AI system?"),70    ("provider", "Who is the provider or developer of the system?"),71    ("intended_purpose", "What is the intended purpose of the system?"),72    ("architecture", "Describe the system architecture."),73    ("training_data", "What kind of training data is used?"),74    ("testing_methodology", "How was the system tested and validated?"),75    ("performance_metrics", "What are the system's performance metrics?"),76    ("risk_management", "What risk management measures were taken?"),77    ("cybersecurity", "What cybersecurity measures are in place?"),78    ("human_oversight", "How is human oversight implemented?"),79    ("versioning", "How is version control maintained?"),80    ("recordkeeping", "How are logs and records maintained?")81]82 83QUESTIONS = prepend_metadata_questions(CORE_QUESTIONS)84 85def get_questions():86    return QUESTIONS87 88def run_tool():89    state = {"step": 0, "answers": {}}90 91    def step_by_step_agent(user_input, state):92        step = state["step"]93        answers = state["answers"]94 95        if step > 0:96            key, _ = QUESTIONS[step - 1]97            answers[key] = user_input98 99        if step < len(QUESTIONS):100            next_question = QUESTIONS[step][1]101            state["step"] += 1102            return next_question, state, None103 104        # Final content for PDF105        content = "\n".join([f"- **{label}**: {answers.get(key, '')}" for key, label in QUESTIONS if key not in ["user_name", "user_role", "organization_name"]])106        detected_lang = detect(content)107        pdf_path = export_text_to_pdf(content, answers, language=detected_lang)108        return "✅ Completed. Download your documentation below.", {"done": True}, pdf_path109 110    with gr.Blocks(title="AI Technical Documentation Tool") as demo:111        chatbot = gr.Chatbot(112            label="🧠 Technical Doc Assistant",113            value=[{"role": "assistant", "content": QUESTIONS[0][1]}],114            type="messages"115        )116        msg = gr.Textbox(label="Your answer")117        state_var = gr.State(state)118        file_output = gr.File(label="Download PDF", visible=True)119        reset_btn = gr.Button("🔁 Restart")120 121        def chat_logic(msg_in, state_in):122            reply, updated_state, file = step_by_step_agent(msg_in, state_in)123            messages = [{"role": "user", "content": msg_in}]124            if reply:125                messages.append({"role": "assistant", "content": reply})126            return messages, updated_state, file127 128        def reset():129            return [{"role": "assistant", "content": QUESTIONS[0][1]}], {"step": 0, "answers": {}}, None130 131        msg.submit(chat_logic, [msg, state_var], [chatbot, state_var, file_output])132        reset_btn.click(reset, outputs=[chatbot, state_var, file_output])133 134    demo.launch(show_api=False)135