Update app.py

app.py

@@ -4,16 +4,36 @@ import numpy as np
 import tensorflow as tf
 from sklearn.preprocessing import MinMaxScaler
 import plotly.express as px
-import os
 from groq import Groq
+import io
 
 # Initialize session state
 if 'model' not in st.session_state:
     st.session_state.model = None
 if 'threshold' not in st.session_state:
     st.session_state.threshold = None
+if 'data' not in st.session_state:
+    st.session_state.data = None
 
-# Autoencoder model definition
+# Generate synthetic data
+def generate_sample_data():
+    time_steps = 500
+    base = np.arange(time_steps)
+    data = pd.DataFrame({
+        'timestamp': pd.date_range(start='2024-01-01', periods=time_steps, freq='H'),
+        'device_count': np.random.poisson(50, time_steps),
+        'connection_attempts': np.random.poisson(30, time_steps),
+        'packet_loss': np.random.uniform(0.1, 2.0, time_steps),
+        'latency': np.random.uniform(10, 100, time_steps)
+    })
+    
+    # Add anomalies
+    anomaly_indices = np.random.choice(time_steps, 10, replace=False)
+    data.loc[anomaly_indices, 'connection_attempts'] *= 10
+    data.loc[anomaly_indices, 'latency'] *= 5
+    return data
+
+# Autoencoder model
 def build_autoencoder(input_dim):
     model = tf.keras.Sequential([
         tf.keras.layers.Dense(64, activation='relu', input_shape=(input_dim,)),
@@ -31,74 +51,96 @@ groq_api_key = st.sidebar.text_input("Groq API Key (optional)", type="password")
 
 # Main interface
 st.title("🛰️ AI Network Anomaly Detector")
-st.write("Upload your network data (CSV) to detect anomalies")
+st.write("Upload network data (CSV) or use synthetic data")
 
-# File uploader
+# Data handling
 uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
-
-# Load or generate sample data
-if uploaded_file is not None:
-    data = pd.read_csv(uploaded_file)
+if uploaded_file:
+    try:
+        st.session_state.data = pd.read_csv(uploaded_file)
+    except Exception as e:
+        st.error(f"Error reading file: {str(e)}")
 else:
-    st.info("Using sample data. Upload a file to use your own dataset.")
-    data = pd.read_csv("sample_wifi_data.csv")  # You should provide this sample file
+    st.session_state.data = generate_sample_data()
+    st.info("Using synthetic data. Upload a CSV file to analyze your own data.")
 
 # Preprocessing
-features = ['device_count', 'connection_attempts', 'packet_loss', 'latency']
-scaler = MinMaxScaler()
-data_scaled = scaler.fit_transform(data[features])
+if st.session_state.data is not None:
+    features = ['device_count', 'connection_attempts', 'packet_loss', 'latency']
+    scaler = MinMaxScaler()
+    data_scaled = scaler.fit_transform(st.session_state.data[features])
 
-# Model training/fine-tuning
-if fine_tune or st.session_state.model is None:
+# Model training
+if st.session_state.model is None or fine_tune:
     with st.spinner("Training model..."):
-        autoencoder = build_autoencoder(data_scaled.shape[1])
-        autoencoder.fit(data_scaled, data_scaled, 
-                       epochs=100, 
-                       batch_size=32,
-                       verbose=0,
-                       validation_split=0.1)
-        st.session_state.model = autoencoder
-        
-        # Calculate threshold
-        reconstructions = autoencoder.predict(data_scaled)
-        mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
-        st.session_state.threshold = np.percentile(mse, 95)
+        try:
+            autoencoder = build_autoencoder(data_scaled.shape[1])
+            autoencoder.fit(data_scaled, data_scaled,
+                          epochs=100,
+                          batch_size=32,
+                          verbose=0,
+                          validation_split=0.2)
+            st.session_state.model = autoencoder
+            
+            # Calculate threshold
+            reconstructions = autoencoder.predict(data_scaled)
+            mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
+            st.session_state.threshold = np.percentile(mse, 95)
+            st.success("Model trained successfully!")
+        except Exception as e:
+            st.error(f"Training error: {str(e)}")
 
 # Anomaly detection
 if st.session_state.model and st.button("Detect Anomalies"):
-    reconstructions = st.session_state.model.predict(data_scaled)
-    mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
-    anomalies = mse > st.session_state.threshold
-    
-    # Visualization
-    fig = px.line(data, x=data.index, y='connection_attempts', 
-                 title='Network Traffic with Anomalies')
-    fig.add_scatter(x=data[anomalies].index, y=data[anomalies]['connection_attempts'], 
-                   mode='markers', name='Anomalies')
-    st.plotly_chart(fig)
-    
-    # Generate alert with Groq/Llama3
-    if groq_api_key:
-        try:
-            client = Groq(api_key=groq_api_key)
-            response = client.chat.completions.create(
-                model="llama3-70b-8192",
-                messages=[{
-                    "role": "user",
-                    "content": f"Generate a network security alert for {sum(anomalies)} anomalies detected. Max connection attempts: {data['connection_attempts'].max()}"
-                }]
-            )
-            st.warning(response.choices[0].message.content)
-        except Exception as e:
-            st.error(f"Groq API Error: {str(e)}")
-    else:
-        st.warning(f"Detected {sum(anomalies)} anomalies! Consider adding Groq API key for detailed analysis.")
+    try:
+        reconstructions = st.session_state.model.predict(data_scaled)
+        mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
+        anomalies = mse > st.session_state.threshold
+        
+        # Visualization
+        fig = px.line(st.session_state.data, x='timestamp', y='connection_attempts',
+                     title='Network Traffic with Anomalies')
+        fig.add_scatter(x=st.session_state.data[anomalies]['timestamp'],
+                       y=st.session_state.data[anomalies]['connection_attempts'],
+                       mode='markers', name='Anomalies',
+                       marker=dict(color='red', size=8))
+        st.plotly_chart(fig)
+
+        # Groq API integration
+        if groq_api_key:
+            try:
+                client = Groq(api_key=groq_api_key)
+                response = client.chat.completions.create(
+                    model="llama3-70b-8192",
+                    messages=[{
+                        "role": "user",
+                        "content": f"""Network anomaly report:
+                        - Total data points: {len(st.session_state.data)}
+                        - Anomalies detected: {sum(anomalies)}
+                        - Max connection attempts: {st.session_state.data['connection_attempts'].max()}
+                        - Average latency: {st.session_state.data['latency'].mean():.2f}ms
+                        Provide a technical analysis and recommendations in bullet points."""
+                    }]
+                )
+                st.subheader("AI Analysis")
+                st.write(response.choices[0].message.content)
+            except Exception as e:
+                st.error(f"Groq API Error: {str(e)}")
+        else:
+            st.warning("Groq API key not provided - using basic detection")
+
+        # Download report
+        report = st.session_state.data[anomalies]
+        csv = report.to_csv(index=False).encode('utf-8')
+        st.download_button(
+            label="Download Anomaly Report",
+            data=csv,
+            file_name='anomaly_report.csv',
+            mime='text/csv'
+        )
+    except Exception as e:
+        st.error(f"Detection error: {str(e)}")
 
-# Download button for results
-if st.session_state.threshold:
-    st.download_button(
-        label="Download Anomaly Report",
-        data=data[anomalies].to_csv().encode('utf-8'),
-        file_name='anomalies_report.csv',
-        mime='text/csv'
-    )
+# Display raw data
+if st.checkbox("Show raw data"):
+    st.write(st.session_state.data)