app.py

import streamlit as st
import streamlit_mermaid as stmd
from io import BytesIO
import time
import re
from report_gen import generate_pdf
from sol_inf import generate_inference  # Assuming this function exists in sol_inf.py

# Set page config
st.set_page_config(
    page_title="Critical Infrastructure - LLM Considerations",
    page_icon="🏙️",
    initial_sidebar_state="collapsed"
)

# Global variables
ARCHITECTURE_VIEW_TYPE = 0  # 0 for SVG, 1 for Mermaid

# Password protection
def check_password():
    def password_entered():
        if st.session_state["password"] == st.secrets["app_password"]:
            st.session_state["password_correct"] = True
            del st.session_state["password"]
        else:
            st.session_state["password_correct"] = False

    if "password_correct" not in st.session_state:
        st.markdown("\n\n")
        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
        st.divider()
        st.info("Developed by Milan Mrdenovic © IBM Norway 2024")
        return False
    elif not st.session_state["password_correct"]:
        st.markdown("\n\n")
        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
        st.divider()
        st.info("Developed by Milan Mrdenovic © IBM Norway 2024")
        st.error("😕 Password incorrect")
        return False
    else:
        return True

if not check_password():
    st.stop()

# Initialize session state
if 'current_page' not in st.session_state:
    st.session_state.current_page = 0

if 'answers' not in st.session_state:
    st.session_state.answers = {}

# Define the content for each page
pages = [
    {
        'title': "Beyond Basic T's & C's",
        'content': "This section evaluates the Acceptable Use Policy of AI solutions.",
        'type': 'description',
        'image': "path_to_image.jpg",
        'expander_content': "Additional information about Acceptable Use Policy"
    },
    {
        'title': "Case Study 1",
        'type': 'quiz',
        'image': "path_to_case1_image.jpg",
        'description': "Description of Case Study 1",
        'input_key': 'case1_answer',
        'slider_key': 'case1_risk',
        'expander_image': "path_to_tc_snippet1.jpg"
    },
    {
        'title': "Case Study 2",
        'type': 'quiz',
        'image': "path_to_case2_image.jpg",
        'description': "Description of Case Study 2",
        'input_key': 'case2_answer',
        'slider_key': 'case2_risk',
        'expander_image': "path_to_tc_snippet2.jpg"
    },
    {
        'title': "Case Study 3",
        'type': 'quiz',
        'image': "path_to_case3_image.jpg",
        'description': "Description of Case Study 3",
        'input_key': 'case3_answer',
        'slider_key': 'case3_risk',
        'expander_image': "path_to_tc_snippet3.jpg"
    },
    {
        'title': "Solution Robustness",
        'content': "This section evaluates the robustness of AI solutions.",
        'type': 'description',
        'image': "path_to_robustness_image.jpg",
        'expander_content': "Additional information about Solution Robustness"
    },
    {
        'title': "Robustness Test 1",
        'type': 'inference',
        'prompt_description': "Description of Robustness Test 1",
        'input_key': 'robustness1_input',
        'output_key': 'robustness1_output',
        'placeholder': "Enter your test input here..."
    },
    {
        'title': "Robustness Test 2",
        'type': 'inference',
        'prompt_description': "Description of Robustness Test 2",
        'input_key': 'robustness2_input',
        'output_key': 'robustness2_output',
        'placeholder': "Enter your test input here..."
    },
    {
        'title': "Data Security",
        'content': "This section evaluates the data security aspects of AI solutions.",
        'type': 'description',
        'image': "path_to_security_image.jpg",
        'expander_content': "Additional information about Data Security"
    },
    {
        'title': "Security Analysis 1",
        'type': 'architecture',
        'svg_path': "path_to_architecture_svg1.svg",
        'mermaid_code': """
        graph TD
            A[Client] --> B[Load Balancer]
            B --> C[Web Server]
            C --> D[Application Server]
            D --> E[Database]
        """,
        'expander_content': "Explanation of the solution architecture",
        'input_key': 'security1_analysis'
    },
    {
        'title': "Security Analysis 2",
        'type': 'architecture',
        'svg_path': "path_to_architecture_svg2.svg",
        'mermaid_code': """
        graph TD
            A[User] --> B[API Gateway]
            B --> C[Authentication Service]
            B --> D[Data Processing Service]
            D --> E[Storage Service]
        """,
        'expander_content': "Explanation of the solution architecture",
        'input_key': 'security2_analysis'
    },
    {
        'title': "Generate Evaluation Report",
        'content': "You have completed the Critical Infrastructure - LLM Considerations evaluation. Click the button below to generate and download your PDF report.",
        'type': 'report'
    }
]

# Streamlit app
st.title("Critical Infrastructure - LLM Considerations")

# Navigation buttons
col1, col2, col3 = st.columns([1, 2, 1])
with col1:
    if st.session_state.current_page > 0:
        if st.button("Back"):
            st.session_state.current_page -= 1
            st.rerun()
with col3:
    if st.session_state.current_page < len(pages) - 1:
        if st.button("Next", use_container_width=True):
            st.session_state.current_page += 1
            st.rerun()

# Display current page
current_page = pages[st.session_state.current_page]
st.header(current_page['title'])

if current_page['type'] == 'description':
    st.write(current_page['content'])
    st.image(current_page['image'])
    with st.expander("Learn More"):
        st.write(current_page['expander_content'])

elif current_page['type'] == 'quiz':
    col1, col2 = st.columns(2)
    with col1:
        st.image(current_page['image'])
        st.write(current_page['description'])
    with col2:
        st.session_state.answers[current_page['input_key']] = st.text_area(
            "Your analysis:",
            key=current_page['input_key'],
            height=200
        )
        st.session_state.answers[current_page['slider_key']] = st.slider(
            "Risk Level",
            1, 10,
            key=current_page['slider_key']
        )
    with st.expander("View T&C Snippet"):
        st.image(current_page['expander_image'])

elif current_page['type'] == 'inference':
    col1, col2 = st.columns(2)
    with col1:
        with st.expander("Prompt Description"):
            st.write(current_page['prompt_description'])
        user_input = st.text_area(
            "Enter your test input:",
            key=current_page['input_key'],
            height=200,
            placeholder=current_page['placeholder']
        )
        if st.button("Infer"):
            st.session_state.answers[current_page['input_key']] = user_input
            st.session_state.answers[current_page['output_key']] = generate_inference(user_input)
    with col2:
        if current_page['output_key'] in st.session_state.answers:
            st.write("Inference Result:")
            st.write(st.session_state.answers[current_page['output_key']])

elif current_page['type'] == 'architecture':
    col1, col2 = st.columns(2)
    with col1:
        if ARCHITECTURE_VIEW_TYPE == 0:
            st.image(current_page['svg_path'])
        else:
            stmd.st_mermaid(current_page['mermaid_code'])
        with st.expander("Solution Explanation"):
            st.write(current_page['expander_content'])
    with col2:
        st.session_state.answers[current_page['input_key']] = st.text_area(
            "Your security analysis:",
            key=current_page['input_key'],
            height=300
        )

elif current_page['type'] == 'report':
    st.write(current_page['content'])
    if st.button("Generate and Download PDF", use_container_width=True):
        pdf = generate_pdf(pages, st.session_state.answers)
        st.download_button(
            label="Download PDF",
            data=pdf,
            file_name="Critical_Infrastructure_LLM_Considerations.pdf",
            mime="application/pdf",
            use_container_width=True
        )

# Display progress
st.progress((st.session_state.current_page + 1) / len(pages))

st.divider()
st.info("Developed by Milan Mrdenovic © IBM Norway 2024")