Update app.py

app.py

@@ -6,13 +6,14 @@ import json
 import os  # Still needed for API key potentially, but not model names
 from pathlib import Path
 import time
+import plotly.express as px
+import pandas as pd
 
 # --- Configuration ---
-# Model names are now discovered dynamically. Remove hardcoded names.
-MAX_PROMPT_TOKENS_ESTIMATE = 800000  # Keep this estimate
+MAX_PROMPT_TOKENS_ESTIMATE = 800000  # Token estimate limit
 RESULTS_PAGE_SIZE = 25
 
-AVAILABLE_ANALYSES = {  # Keep analyses config
+AVAILABLE_ANALYSES = {
     "generate_docs": "Generate Missing Docstrings/Comments",
     "find_bugs": "Identify Potential Bugs & Anti-patterns",
     "check_style": "Check Style Guide Compliance (General)",
@@ -20,12 +21,11 @@ AVAILABLE_ANALYSES = {  # Keep analyses config
     "suggest_refactoring": "Suggest Refactoring Opportunities",
 }
 CODE_EXTENSIONS = {
-    '.py', '.js', '.java', '.c', '.cpp', '.h', '.cs', '.go', '.rb', 
+    '.py', '.js', '.java', '.c', '.cpp', '.h', '.cs', '.go', '.rb',
     '.php', '.swift', '.kt', '.ts', '.html', '.css', '.scss', '.sql'
-}  # Keep extensions
+}
 
 # --- Session State Initialization ---
-# (Keep most session state, add one for the selected model)
 if 'mock_api_call' not in st.session_state:
     st.session_state.mock_api_call = False
 if 'analysis_results' not in st.session_state:
@@ -35,79 +35,73 @@ if 'error_message' not in st.session_state:
 if 'analysis_requested' not in st.session_state:
     st.session_state.analysis_requested = False
 if 'selected_model_name' not in st.session_state:
-    st.session_state.selected_model_name = None  # Will hold the "models/..." name
+    st.session_state.selected_model_name = None  # Internal model name
 if 'available_models_dict' not in st.session_state:
-    st.session_state.available_models_dict = {}  # Store display_name -> name mapping
+    st.session_state.available_models_dict = {}  # Mapping: display -> internal
+if 'file_count' not in st.session_state:
+    st.session_state.file_count = 0
+if 'total_chars' not in st.session_state:
+    st.session_state.total_chars = 0
 
 # --- Gemini API Setup & Model Discovery ---
 model = None  # Global variable for the initialized model instance
 
-# --- NEW: Function to list available models ---
-@st.cache_data(ttl=3600)  # Cache model list for an hour
+@st.cache_data(ttl=3600)
 def get_available_models():
     """Lists models supporting 'generateContent' using the API key."""
     model_dict = {}
     try:
         if 'GEMINI_API_KEY' not in st.secrets:
-            # Don't stop here, let the main part handle it, but return empty
             print("API key not found in secrets during model listing attempt.")
             return {}
-        # Configure API key temporarily just for listing
         genai.configure(api_key=st.secrets["GEMINI_API_KEY"])
         print("Listing available models via API...")
         for m in genai.list_models():
-            # Check if the model supports the 'generateContent' method
             if 'generateContent' in m.supported_generation_methods:
-                # Store mapping: user-friendly name -> internal name
                 model_dict[m.display_name] = m.name
         print(f"Found {len(model_dict)} compatible models.")
         return model_dict
     except Exception as e:
         st.error(f"🚨 Error listing available models: {e}")
-        return {}  # Return empty on error
+        return {}
 
 def initialize_gemini_model():
     """Initializes the Gemini model based on the selected name."""
     global model
     selected_name = st.session_state.get('selected_model_name')
-
     if selected_name and model is None and not st.session_state.mock_api_call:
         try:
             if 'GEMINI_API_KEY' not in st.secrets:
                 st.error("🚨 Gemini API Key not found. Add it to `.streamlit/secrets.toml`.")
-                st.stop()  # Stop if key missing for initialization
-            # Configure API key (might be redundant if list_models worked, but safe)
+                st.stop()
             genai.configure(api_key=st.secrets["GEMINI_API_KEY"])
             print(f"Initializing Gemini Model: {selected_name}")
-            # Use the selected model name from session state
             model = genai.GenerativeModel(model_name=selected_name)
             print(f"Gemini Model Initialized ({selected_name}).")
             return True
         except Exception as e:
             st.error(f"🚨 Error initializing selected Gemini model '{selected_name}': {e}")
-            st.session_state.selected_model_name = None  # Reset selection on error
+            st.session_state.selected_model_name = None
             st.stop()
             return False
     elif st.session_state.mock_api_call:
-        return True  # No init needed for mock mode
+        return True
     elif model is not None and model.model_name == selected_name:
-        return True  # Already initialized with the correct model
+        return True
     elif model is not None and model.model_name != selected_name:
         print("Model changed. Re-initializing...")
-        model = None  # Reset model instance
-        return initialize_gemini_model()  # Recurse to re-initialize with new name
+        model = None
+        return initialize_gemini_model()
     elif not selected_name and not st.session_state.mock_api_call:
-        # This case happens if no model is selected yet
-        return False  # Cannot initialize without a selection
-    return False  # Default case
+        return False
+    return False
 
 # --- Helper Functions ---
-# Updated estimate_token_count to support integers and strings
 def estimate_token_count(text):
     """
     Estimates the token count.
-    If a string is provided, it calculates based on its length.
-    If an integer is provided (e.g., total character count), it uses that directly.
+    If a string is provided, calculates based on length.
+    If an integer is provided (e.g., total character count), uses that directly.
     """
     if isinstance(text, int):
         return text // 3
@@ -117,7 +111,7 @@ def estimate_token_count(text):
 def process_zip_file_cached(file_id, file_size, file_content_bytes):
     """
     Processes a ZIP file and extracts code files.
-    Returns a tuple of (code_files dict, total_chars, file_count, ignored_files list).
+    Returns (code_files dict, total_chars, file_count, ignored_files list).
     """
     code_files = {}
     total_chars = 0
@@ -166,16 +160,13 @@ def process_zip_file_cached(file_id, file_size, file_content_bytes):
         st.error(f"🚨 ZIP Error: {e}")
         return None, 0, 0, []
     if file_count == 0:
-        if not ignored_files:
-            st.warning("No code files found.")
-        else:
-            st.warning("No code files found; some skipped.")
+        st.warning("No code files found." if not ignored_files else "No code files found; some skipped.")
     return code_files, total_chars, file_count, ignored_files
 
 def construct_analysis_prompt(code_files_dict, requested_analyses):
     """
-    Constructs the prompt for analysis by including code files and JSON structure for expected output.
-    Returns the full prompt and a list of included files.
+    Constructs the prompt for analysis by including code files and JSON structure.
+    Returns the full prompt and list of included files.
     """
     prompt_parts = ["Analyze the following codebase...\n\n"]
     current_token_estimate = estimate_token_count(prompt_parts[0])
@@ -227,12 +218,11 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
 def call_gemini_api(prompt):
     """
     Calls the Gemini API using the provided prompt.
-    Returns the parsed JSON insights or an error message.
+    Returns parsed JSON insights or an error message.
     """
     if not prompt:
         return None, "Prompt generation failed."
     
-    # MOCK MODE
     if st.session_state.mock_api_call:
         st.info(" MOCK MODE: Simulating API call...")
         time.sleep(1)
@@ -245,12 +235,11 @@ def call_gemini_api(prompt):
         })
         st.success("Mock response generated.")
         return json.loads(mock_json_response), None
-    # REAL API CALL
     else:
         if not initialize_gemini_model():
             return None, "Gemini Model Initialization Failed."
         if model is None:
-            return None, "Gemini model not selected or available."  # Added check
+            return None, "Gemini model not selected or available."
         try:
             api_status = st.empty()
             api_status.info(f"📡 Sending request to {model.model_name} (Est. prompt tokens: {estimate_token_count(prompt):,})... Please wait.")
@@ -270,7 +259,6 @@ def call_gemini_api(prompt):
             api_status.empty()
             try:
                 json_response_text = response.text.strip()
-                # Remove markdown code fences if present
                 if json_response_text.startswith("```json"):
                     json_response_text = json_response_text[7:]
                 if json_response_text.startswith("```"):
@@ -320,7 +308,7 @@ def call_gemini_api(prompt):
 
 def display_results(results_json, requested_analyses):
     """
-    Displays the analysis results with pagination and allows JSON download.
+    Displays analysis results with pagination and a downloadable JSON report.
     """
     st.header("📊 Analysis Report")
     if not isinstance(results_json, dict):
@@ -407,8 +395,21 @@ def display_results(results_json, requested_analyses):
 st.set_page_config(page_title="Codebase Audit Assistant", layout="wide")
 st.title("🤖 Codebase Audit & Documentation Assistant")
 
-# --- Sidebar ---
+# --- Sidebar Enhancements ---
 with st.sidebar:
+    # Dark Mode Toggle
+    dark_mode = st.checkbox("Enable Dark Mode", value=False)
+    if dark_mode:
+        st.markdown(
+            """
+            <style>
+            .reportview-container { background-color: #2E2E2E; color: white; }
+            .sidebar .sidebar-content { background-color: #1A1A1A; }
+            </style>
+            """,
+            unsafe_allow_html=True
+        )
+
     st.header("⚙️ Analysis Controls")
     st.session_state.mock_api_call = st.toggle(
         "🧪 Enable Mock API Mode",
@@ -418,27 +419,20 @@ with st.sidebar:
 
     st.divider()
     st.header("♊ Select Model")
-    # --- NEW: Dynamic Model Selection ---
     if not st.session_state.mock_api_call:
-        # Get available models (uses cache)
         st.session_state.available_models_dict = get_available_models()
         model_display_names = list(st.session_state.available_models_dict.keys())
-
         if model_display_names:
-            # Try to find the index of the previously selected model
             current_model_display_name = None
             if st.session_state.selected_model_name:
-                # Find display name matching the stored internal name
                 for disp_name, internal_name in st.session_state.available_models_dict.items():
                     if internal_name == st.session_state.selected_model_name:
                         current_model_display_name = disp_name
                         break
-
             try:
                 selected_index = model_display_names.index(current_model_display_name) if current_model_display_name in model_display_names else 0
             except ValueError:
-                selected_index = 0  # Default to first if previous selection not found
-
+                selected_index = 0
             selected_display_name = st.selectbox(
                 "Choose Gemini model:",
                 options=model_display_names,
@@ -446,20 +440,18 @@ with st.sidebar:
                 key="model_selector",
                 help="Select the Gemini model to use for analysis."
             )
-            # Update session state with the internal name based on selection
             st.session_state.selected_model_name = st.session_state.available_models_dict.get(selected_display_name)
             st.info(f"Using REAL Gemini API ({st.session_state.selected_model_name})")
         elif 'GEMINI_API_KEY' in st.secrets:
             st.warning("No compatible models found or error listing models. Check API Key permissions.")
-            st.session_state.selected_model_name = None  # Ensure no model selected
+            st.session_state.selected_model_name = None
         else:
             st.warning("Add GEMINI_API_KEY to secrets to list models.")
             st.session_state.selected_model_name = None
-    else:  # Mock mode is active
+    else:
         st.info("Mock API Mode ACTIVE")
-        st.session_state.selected_model_name = "mock_model"  # Use a placeholder name for mock mode
-    # --- End Dynamic Model Selection ---
-
+        st.session_state.selected_model_name = "mock_model"
+    
     st.divider()
     st.header("🔎 Select Analyses")
     selected_analyses = [
@@ -481,7 +473,7 @@ with st.sidebar:
     st.divider()
     st.warning("⚠️ **Privacy:** Code sent to Google API if Mock Mode is OFF.")
 
-# Update title dynamically based on selected model
+# Update title based on selected model
 if st.session_state.selected_model_name and not st.session_state.mock_api_call:
     st.markdown(f"Upload codebase (`.zip`) for analysis via **{st.session_state.selected_model_name}**.")
 elif st.session_state.mock_api_call:
@@ -511,12 +503,13 @@ if uploaded_file:
         file_id, uploaded_file.size, uploaded_file_bytes
     )
     if code_files is not None:
+        # Save these metrics to session state for dashboard use
+        st.session_state.file_count = file_count
+        st.session_state.total_chars = total_chars
         st.info(f"Found **{file_count}** code files ({total_chars:,} chars). Est. tokens: ~{estimate_token_count(total_chars):,}")
         if ignored_files:
             with st.expander(f"View {len(ignored_files)} Skipped/Ignored Files"):
                 st.code("\n".join(ignored_files), language='text')
-
-        # Disable button if no model selected (and not in mock mode)
         model_ready = bool(st.session_state.selected_model_name) or st.session_state.mock_api_call
         analyze_button_disabled = (not selected_analyses or file_count == 0 or not model_ready)
         analyze_button_label = "Analyze Codebase"
@@ -525,11 +518,7 @@ if uploaded_file:
         elif analyze_button_disabled:
             analyze_button_label = "Select Analyses or Upload Valid Code"
 
-        if analysis_button_placeholder.button(
-            analyze_button_label,
-            type="primary",
-            disabled=analyze_button_disabled
-        ):
+        if analysis_button_placeholder.button(analyze_button_label, type="primary", disabled=analyze_button_disabled):
             st.session_state.analysis_requested = True
             st.session_state.analysis_results = None
             st.session_state.error_message = None
@@ -541,11 +530,7 @@ if uploaded_file:
                 st.warning("Please select a Gemini model from the sidebar.")
             else:
                 with results_placeholder:
-                    spinner_model_name = (
-                        st.session_state.selected_model_name 
-                        if not st.session_state.mock_api_call 
-                        else "Mock Mode"
-                    )
+                    spinner_model_name = st.session_state.selected_model_name if not st.session_state.mock_api_call else "Mock Mode"
                     spinner_msg = f"🚀 Preparing prompt & contacting AI ({spinner_model_name})... Please wait."
                     with st.spinner(spinner_msg):
                         analysis_prompt, included_files_in_prompt = construct_analysis_prompt(code_files, selected_analyses)
@@ -559,7 +544,7 @@ if uploaded_file:
                             st.session_state.error_message = "Failed to generate analysis prompt."
                 st.rerun()
 
-# Display results (Keep the same logic)
+# Display analysis results if available
 if st.session_state.analysis_requested:
     with results_placeholder:
         st.divider()
@@ -575,5 +560,25 @@ if st.session_state.analysis_requested:
 elif not uploaded_file:
     results_placeholder.info("Upload a ZIP file to begin.")
 
+# --- Audit Dashboard ---
+# If analysis results exist, display an interactive dashboard summarizing key metrics.
+if st.session_state.analysis_results:
+    st.subheader("Audit Dashboard")
+    # Compute metrics from session state and analysis results
+    metrics = {
+        "Files Analyzed": st.session_state.file_count,
+        "Total Tokens": estimate_token_count(st.session_state.total_chars)
+    }
+    metrics["Documentation Suggestions"] = len(st.session_state.analysis_results.get("documentation_suggestions", []))
+    metrics["Potential Bugs"] = len(st.session_state.analysis_results.get("potential_bugs", []))
+    metrics["Style Issues"] = len(st.session_state.analysis_results.get("style_issues", []))
+    metrics["Module Summaries"] = len(st.session_state.analysis_results.get("module_summaries", []))
+    metrics["Refactoring Suggestions"] = len(st.session_state.analysis_results.get("refactoring_suggestions", []))
+    
+    # Create a DataFrame and a bar chart using Plotly
+    df_metrics = pd.DataFrame(list(metrics.items()), columns=["Metric", "Value"])
+    fig = px.bar(df_metrics, x="Metric", y="Value", title="Audit Summary Metrics")
+    st.plotly_chart(fig)
+
 results_placeholder.divider()
 results_placeholder.markdown("_Assistant powered by Google Gemini._")