Update app.py

app.py

@@ -1,126 +1,139 @@
 import streamlit as st
 import pandas as pd
-from pathlib import Path
 from processor import DataProcessor
-from llm_handler import LLMHandler
+from llm_handler import DDoSInference
 
 class DDoSResistanceHelper:
     def __init__(self):
-        # Configure Streamlit app
+        # Set up Streamlit page configuration
         st.set_page_config(
-            page_title="DDoS Resistance Helper URF LLM Network Analyzer",
-            page_icon=":shield:",
-            layout="wide",
-            initial_sidebar_state="expanded"
+            page_title="DDoS Resistance Helper | LLM Network Analyzer",
+            page_icon="🛡️",
+            layout="wide"
         )
 
         # Initialize session state
-        self.initialize_session_state()
+        self.setup_session_state()
 
         # Initialize processor and LLM handler
         self.processor = DataProcessor()
-        self.llm_handler = LLMHandler()
-
-    def initialize_session_state(self):
-        """Set up Streamlit session state variables."""
-        session_keys = [
-            'current_file', 'preprocessed_data', 'analysis_results', 'chat_history'
-        ]
+        self.llm_handler = DDoSInference()
+
+        # Define themes
+        self.themes = {
+            'Light': {
+                'base': 'light',
+                'primaryColor': '#3B82F6',
+                'backgroundColor': '#FFFFFF',
+                'secondaryBackgroundColor': '#F3F4F6',
+                'textColor': '#000000'
+            },
+            'Dark': {
+                'base': 'dark',
+                'primaryColor': '#10B981',
+                'backgroundColor': '#1F2937',
+                'secondaryBackgroundColor': '#374151',
+                'textColor': '#FFFFFF'
+            },
+            'Neon': {
+                'base': 'dark',
+                'primaryColor': '#22D3EE',
+                'backgroundColor': '#000000',
+                'secondaryBackgroundColor': '#1F2937',
+                'textColor': '#22D3EE'
+            }
+        }
+
+    def setup_session_state(self):
+        """Initialize Streamlit session state variables."""
+        st.write("Initializing session state...")
+        session_keys = ['chat_history', 'current_file', 'analysis_complete', 'selected_theme', 'preview_data']
         for key in session_keys:
             if key not in st.session_state:
-                st.session_state[key] = None if key != 'chat_history' else []
-
-    def render_top_bar(self):
-        """Render the top bar with theme and upload options."""
-        col1, col2 = st.columns([8, 2])
-
-        with col1:
-            st.title("🛡️ DDoS Resistance Helper URF LLM Network Analyzer")
-
-        with col2:
-            st.markdown("### Theme Selector")
-            if st.button("Light"):
-                st.markdown("<style>.stApp { background-color: #ffffff; }</style>", unsafe_allow_html=True)
-            elif st.button("Dark"):
-                st.markdown("<style>.stApp { background-color: #1f1f1f; color: white; }</style>", unsafe_allow_html=True)
+                st.session_state[key] = [] if key == 'chat_history' else None
+        st.write("Session state initialized.")
+
+    def apply_theme(self):
+        """Apply selected theme to Streamlit app."""
+        theme_name = st.session_state.get('selected_theme', 'Light')
+        theme = self.themes[theme_name]
+        st.markdown(f"""
+        <style>
+        .stApp {{
+            background-color: {theme['backgroundColor']};
+            color: {theme['textColor']};
+        }}
+        </style>
+        """, unsafe_allow_html=True)
+
+    def render_theme_selector(self):
+        """Render theme selection icons on the top-right corner."""
+        st.sidebar.empty()
+        st.write(f'<div style="position: absolute; top: 15px; right: 15px;">', unsafe_allow_html=True)
+        if st.button("🌞 Light", key="light_theme"):
+            st.session_state.selected_theme = 'Light'
+            self.apply_theme()
+        if st.button("🌙 Dark", key="dark_theme"):
+            st.session_state.selected_theme = 'Dark'
+            self.apply_theme()
+        if st.button("💡 Neon", key="neon_theme"):
+            st.session_state.selected_theme = 'Neon'
+            self.apply_theme()
 
     def render_file_upload(self):
-        """Render the file upload component."""
-        uploaded_file = st.file_uploader("Upload Network Traffic Data (CSV)", type=["csv"],
-                                         label_visibility="collapsed")
+        """Render file upload section with an upload icon in the main header."""
+        st.title("DDoS Resistance Helper | LLM Network Analyzer")
+        uploaded_file = st.file_uploader("Upload Data File 📂", type=['csv'])
         if uploaded_file:
-            try:
-                df = pd.read_csv(uploaded_file)
-                st.session_state.current_file = df
-                st.success("File uploaded successfully!")
-            except Exception as e:
-                st.error(f"Error reading file: {e}")
-
-    def render_analysis(self):
-        """Render the analysis results."""
-        if st.session_state.current_file is None:
-            st.info("Please upload a CSV file to start analysis.")
-            return
-
-        # Preprocess the data
-        st.subheader("Preprocessing Data")
-        with st.spinner("Preprocessing data..."):
-            try:
-                preprocessed_data = self.processor.preprocess_data(st.session_state.current_file)
-                st.session_state.preprocessed_data = preprocessed_data
-                st.success("Data preprocessed successfully!")
-            except Exception as e:
-                st.error(f"Error during preprocessing: {e}")
-
-        # Perform LLM analysis
-        st.subheader("Performing LLM Analysis")
-        with st.spinner("Analyzing data with LLM..."):
-            try:
-                results = self.llm_handler.analyze_data(st.session_state.preprocessed_data)
-                st.session_state.analysis_results = results
-                st.success("Analysis completed successfully!")
-            except Exception as e:
-                st.error(f"Error during LLM analysis: {e}")
-
-        # Show results
-        if st.session_state.analysis_results is not None:
-            st.subheader("Analysis Results")
-            st.dataframe(st.session_state.analysis_results)
-            csv_path = Path("~/.dataset/PROBABILITY_OF_EACH_ROW_DDOS_AND_BENGNIN.csv").expanduser()
-            st.download_button("Download Results as CSV", csv_path.read_bytes(), "analysis_results.csv")
+            st.session_state.current_file = uploaded_file
+            st.session_state.preview_data = pd.read_csv(uploaded_file).head(5)
+            if st.button("🚀 Process Data"):
+                self.process_uploaded_file(uploaded_file)
+
+    def process_uploaded_file(self, file):
+        """Process the uploaded file and perform analysis."""
+        try:
+            st.write("Processing file...")
+            with st.spinner("Analyzing data..."):
+                # Preprocess the uploaded file
+                df = self.processor.preprocess_data(file)
+
+                # Perform LLM analysis
+                results = self.llm_handler.analyze_data(df)
+
+                st.session_state.analysis_complete = True
+                st.success("Analysis completed!")
+        except Exception as e:
+            st.error(f"Error: {e}")
 
     def render_chat_interface(self):
-        """Render a chat interface for interacting with the LLM."""
+        """Render chat interface in the sidebar."""
         st.sidebar.header("💬 Chat Interface")
-        # Display chat history
-        for message in st.session_state.chat_history:
-            with st.chat_message(message['role']):
-                st.write(message['content'])
-
-        # Chat input
-        if prompt := st.sidebar.text_input("Ask about the analysis or mitigation steps..."):
-            # Add user message to chat history
-            st.session_state.chat_history.append({
-                'role': 'user',
-                'content': prompt
-            })
-
-            # Get LLM response
-            response = self.llm_handler.get_chat_response(prompt)
-
-            # Add LLM response to chat history
-            st.session_state.chat_history.append({
-                'role': 'assistant',
-                'content': response
-            })
-
-    def run(self):
-        """Run the Streamlit app."""
-        self.render_top_bar()
+        for msg in st.session_state.chat_history:
+            st.sidebar.write(f"[{msg['role'].capitalize()}]: {msg['content']}")
+        new_chat = st.sidebar.text_input("Ask a question...")
+        if new_chat:
+            st.session_state.chat_history.append({"role": "user", "content": new_chat})
+            response = self.llm_handler.get_chat_response(new_chat)
+            st.session_state.chat_history.append({"role": "assistant", "content": response})
+
+    def render_analysis_interface(self):
+        """Render analysis results interface."""
+        st.subheader("📋 Data Preview")
+        if st.session_state.preview_data is not None:
+            st.dataframe(st.session_state.preview_data)
+        if st.session_state.analysis_complete:
+            st.subheader("🔍 Analysis Results")
+            st.info("Analysis complete. Visualizations and results here.")
+
+    def main(self):
+        """Main application runner."""
+        self.apply_theme()
+        self.render_theme_selector()
         self.render_file_upload()
-        self.render_analysis()
+        self.render_analysis_interface()
         self.render_chat_interface()
 
 if __name__ == "__main__":
     app = DDoSResistanceHelper()
-    app.run()
+    app.main()