Update app.py

app.py

@@ -10,10 +10,10 @@ import plotly.express as px
 import pandas as pd
 
 # --- Configuration ---
-MAX_PROMPT_TOKENS_ESTIMATE = 800000  # Token estimate limit
+MAX_PROMPT_TOKENS_ESTIMATE = 800000  # Keep this estimate
 RESULTS_PAGE_SIZE = 25
 
-AVAILABLE_ANALYSES = {
+AVAILABLE_ANALYSES = {  # Keep analyses config
     "generate_docs": "Generate Missing Docstrings/Comments",
     "find_bugs": "Identify Potential Bugs & Anti-patterns",
     "check_style": "Check Style Guide Compliance (General)",
@@ -35,13 +35,9 @@ 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  # Internal model name
+    st.session_state.selected_model_name = None  # Will hold the "models/..." name
 if 'available_models_dict' not in st.session_state:
-    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
+    st.session_state.available_models_dict = {}  # Mapping display_name -> name
 
 # --- Gemini API Setup & Model Discovery ---
 model = None  # Global variable for the initialized model instance
@@ -100,8 +96,8 @@ def initialize_gemini_model():
 def estimate_token_count(text):
     """
     Estimates the token count.
-    If a string is provided, calculates based on length.
-    If an integer is provided (e.g., total character count), uses that directly.
+    If a string is provided, calculates based on its length.
+    If an integer (e.g. total char count) is provided, uses that directly.
     """
     if isinstance(text, int):
         return text // 3
@@ -160,13 +156,16 @@ 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:
-        st.warning("No code files found." if not ignored_files else "No code files found; some skipped.")
+        if not ignored_files:
+            st.warning("No code files found.")
+        else:
+            st.warning("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.
-    Returns the full prompt and list of included files.
+    Constructs the prompt for analysis by including code files and a JSON structure for output.
+    Returns the full prompt and a list of included files.
     """
     prompt_parts = ["Analyze the following codebase...\n\n"]
     current_token_estimate = estimate_token_count(prompt_parts[0])
@@ -218,7 +217,7 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
 def call_gemini_api(prompt):
     """
     Calls the Gemini API using the provided prompt.
-    Returns parsed JSON insights or an error message.
+    Returns the parsed JSON insights or an error message.
     """
     if not prompt:
         return None, "Prompt generation failed."
@@ -249,7 +248,7 @@ def call_gemini_api(prompt):
                 generation_config=genai.types.GenerationConfig(temperature=0.2),
                 safety_settings=[
                     {"category": c, "threshold": "BLOCK_MEDIUM_AND_ABOVE"}
-                    for c in ["HARM_CATEGORY_HARASSMENT", "HARM_CATEGORY_HATE_SPEECH", 
+                    for c in ["HARM_CATEGORY_HARASSMENT", "HARM_CATEGORY_HATE_SPEECH",
                               "HARM_CATEGORY_SEXUALLY_EXPLICIT", "HARM_CATEGORY_DANGEROUS_CONTENT"]
                 ]
             )
@@ -308,7 +307,7 @@ def call_gemini_api(prompt):
 
 def display_results(results_json, requested_analyses):
     """
-    Displays analysis results with pagination and a downloadable JSON report.
+    Displays the analysis results with pagination and allows JSON download.
     """
     st.header("📊 Analysis Report")
     if not isinstance(results_json, dict):
@@ -395,7 +394,7 @@ 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 Enhancements ---
+# --- Sidebar with Enhancements ---
 with st.sidebar:
     # Dark Mode Toggle
     dark_mode = st.checkbox("Enable Dark Mode", value=False)
@@ -403,13 +402,14 @@ with st.sidebar:
         st.markdown(
             """
             <style>
-            .reportview-container { background-color: #2E2E2E; color: white; }
-            .sidebar .sidebar-content { background-color: #1A1A1A; }
+            .reportview-container, .main {
+                background-color: #2E2E2E;
+                color: white;
+            }
             </style>
             """,
             unsafe_allow_html=True
         )
-
     st.header("⚙️ Analysis Controls")
     st.session_state.mock_api_call = st.toggle(
         "🧪 Enable Mock API Mode",
@@ -455,7 +455,7 @@ with st.sidebar:
     st.divider()
     st.header("🔎 Select Analyses")
     selected_analyses = [
-        key for key, name in AVAILABLE_ANALYSES.items() 
+        key for key, name in AVAILABLE_ANALYSES.items()
         if st.checkbox(name, value=True, key=f"cb_{key}")
     ]
     st.divider()
@@ -473,7 +473,7 @@ with st.sidebar:
     st.divider()
     st.warning("⚠️ **Privacy:** Code sent to Google API if Mock Mode is OFF.")
 
-# Update title based on selected model
+# Update title dynamically 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:
@@ -503,13 +503,20 @@ 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):,}")
+        # --- Interactive Metrics Visualization ---
+        metrics = {
+            "Metric": ["Files Analyzed", "Total Characters", "Token Estimate", "Ignored Files"],
+            "Value": [file_count, total_chars, estimate_token_count(total_chars), len(ignored_files)]
+        }
+        df_metrics = pd.DataFrame(metrics)
+        fig = px.bar(df_metrics, x="Metric", y="Value", title="Upload Summary Metrics")
+        st.plotly_chart(fig)
+        # --- End Metrics Visualization ---
         if ignored_files:
             with st.expander(f"View {len(ignored_files)} Skipped/Ignored Files"):
                 st.code("\n".join(ignored_files), language='text')
+        
         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"
@@ -518,7 +525,11 @@ 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
@@ -530,7 +541,11 @@ 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)
@@ -544,7 +559,6 @@ if uploaded_file:
                             st.session_state.error_message = "Failed to generate analysis prompt."
                 st.rerun()
 
-# Display analysis results if available
 if st.session_state.analysis_requested:
     with results_placeholder:
         st.divider()
@@ -560,25 +574,5 @@ 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._")