Update app.py

app.py

@@ -5,11 +5,11 @@ import io
 import json
 import os
 from pathlib import Path
-import time # Added for simulating mock delay
+import time  # Added for simulating mock delay
 
 # --- Configuration ---
 GEMINI_MODEL_NAME = "gemini-2.5-pro-preview-03-25"
-MAX_PROMPT_TOKENS_ESTIMATE = 800000 # Adjust as needed
+MAX_PROMPT_TOKENS_ESTIMATE = 800000  # Adjust as needed
 
 AVAILABLE_ANALYSES = {
     "generate_docs": "Generate Missing Docstrings/Comments",
@@ -19,17 +19,20 @@ AVAILABLE_ANALYSES = {
     "suggest_refactoring": "Suggest Refactoring Opportunities"
 }
 
-CODE_EXTENSIONS = {'.py', '.js', '.java', '.c', '.cpp', '.h', '.cs', '.go', '.rb', '.php', '.swift', '.kt', '.ts', '.html', '.css', '.scss', '.sql'}
+CODE_EXTENSIONS = {
+    '.py', '.js', '.java', '.c', '.cpp', '.h', '.cs', '.go',
+    '.rb', '.php', '.swift', '.kt', '.ts', '.html', '.css', '.scss', '.sql'
+}
 
 # --- Session State Initialization ---
-# Initialize session state for mock mode toggle if it doesn't exist
 if 'mock_api_call' not in st.session_state:
-    st.session_state.mock_api_call = False # Default to using the real API
+    st.session_state.mock_api_call = False  # Default to using the real API
 
 # --- Gemini API Setup ---
-# Defer full initialization until needed if mock mode might be used first
 model = None
+
 def initialize_gemini_model():
+    """Initializes the Gemini model if not in mock mode."""
     global model
     if model is None and not st.session_state.mock_api_call:
         try:
@@ -46,21 +49,28 @@ def initialize_gemini_model():
             return False
     elif st.session_state.mock_api_call:
         print("Running in Mock Mode. Skipping Gemini initialization.")
-        return True # Allow proceeding in mock mode
+        return True  # Allow proceeding in mock mode
     elif model is not None:
-         print("Gemini Model already initialized.")
-         return True
+        print("Gemini Model already initialized.")
+        return True
     return False
 
-
 # --- Helper Functions ---
 
 def estimate_token_count(text):
-    """Roughly estimate token count (3-4 chars per token)."""
+    """Roughly estimate token count (assumes ~3-4 characters per token)."""
     return len(text) // 3
 
 def process_zip_file(uploaded_file):
-    """Extracts code files and their content from the uploaded zip file."""
+    """
+    Extracts code files and their content from the uploaded ZIP file.
+    
+    Returns:
+        code_files (dict): Mapping of file paths to content.
+        total_chars (int): Total number of characters in included files.
+        file_count (int): Count of processed code files.
+        ignored_files (list): List of files skipped or not processed.
+    """
     code_files = {}
     total_chars = 0
     file_count = 0
@@ -69,6 +79,7 @@ def process_zip_file(uploaded_file):
     try:
         with zipfile.ZipFile(io.BytesIO(uploaded_file.getvalue()), 'r') as zip_ref:
             for member in zip_ref.infolist():
+                # Skip directories, hidden files, and files with '__' in the name
                 if member.is_dir() or any(part.startswith('.') for part in Path(member.filename).parts) or '__' in member.filename:
                     continue
 
@@ -82,17 +93,17 @@ def process_zip_file(uploaded_file):
                                 try:
                                     content = file.read().decode('latin-1')
                                 except Exception as decode_err:
-                                     ignored_files.append(f"{member.filename} (Decode Error: {decode_err})")
-                                     continue
+                                    ignored_files.append(f"{member.filename} (Decode Error: {decode_err})")
+                                    continue
 
                             code_files[member.filename] = content
                             total_chars += len(content)
                             file_count += 1
                     except Exception as read_err:
-                         ignored_files.append(f"{member.filename} (Read Error: {read_err})")
+                        ignored_files.append(f"{member.filename} (Read Error: {read_err})")
                 else:
                     # Only add to ignored if it's not explicitly ignored by path rules above
-                     if not (any(part.startswith('.') for part in Path(member.filename).parts) or '__' in member.filename):
+                    if not (any(part.startswith('.') for part in Path(member.filename).parts) or '__' in member.filename):
                         ignored_files.append(f"{member.filename} (Skipped Extension: {file_path.suffix})")
 
     except zipfile.BadZipFile:
@@ -105,7 +116,13 @@ def process_zip_file(uploaded_file):
     return code_files, total_chars, file_count, ignored_files
 
 def construct_analysis_prompt(code_files_dict, requested_analyses):
-    """Constructs the prompt for Gemini, including code content and JSON structure request."""
+    """
+    Constructs the prompt for Gemini, including code content and a JSON structure request.
+    
+    Returns:
+        full_prompt (str): The complete prompt.
+        included_files (list): List of file names included in the prompt.
+    """
     prompt_content = "Analyze the following codebase provided as a collection of file paths and their content.\n\n"
     current_token_estimate = estimate_token_count(prompt_content)
     included_files = []
@@ -132,8 +149,8 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
 
     prompt_content += concatenated_code
 
+    # Build the expected JSON structure dynamically based on the selected analyses
     json_structure_description = "{\n"
-    # Dynamically build the JSON structure based on selection
     structure_parts = []
     if "generate_docs" in requested_analyses:
         structure_parts.append('    "documentation_suggestions": [{"file": "path/to/file", "line": number, "suggestion": "Suggested docstring/comment"}]')
@@ -145,7 +162,6 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
         structure_parts.append('    "module_summaries": [{"file": "path/to/file", "summary": "One-paragraph summary of the file purpose/functionality"}]')
     if "suggest_refactoring" in requested_analyses:
         structure_parts.append('    "refactoring_suggestions": [{"file": "path/to/file", "line": number, "area": "e.g., function name, class name", "suggestion": "Description of refactoring suggestion"}]')
-
     json_structure_description += ",\n".join(structure_parts)
     json_structure_description += "\n}"
 
@@ -161,57 +177,44 @@ Respond ONLY with a single, valid JSON object adhering strictly to the following
 **JSON Output Only:**
 """
     full_prompt = prompt_content + prompt_footer
-    # print(f"--- PROMPT (First 500 chars): ---\n{full_prompt[:500]}\n--------------------------")
-    # print(f"--- PROMPT (Last 500 chars): ---\n{full_prompt[-500:]}\n--------------------------")
     return full_prompt, included_files
 
-
 def call_gemini_api(prompt):
-    """Calls the Gemini API or returns mock data based on session state."""
+    """
+    Calls the Gemini API (or simulates it in mock mode) with the provided prompt.
+    
+    Returns:
+        insights (dict): The parsed JSON response from the API.
+        error_message (str): An error message if something went wrong.
+    """
     if not prompt:
         return None, "Prompt generation failed."
 
     # --- MOCK MODE LOGIC ---
     if st.session_state.mock_api_call:
         st.info(" MOCK MODE: Simulating API call...")
-        time.sleep(2) # Simulate network/processing delay
+        time.sleep(2)  # Simulate network/processing delay
 
-        # --- CHOOSE YOUR MOCK RESPONSE ---
-        # Option 1: Simulate successful response with some data
+        # Simulated successful response
         mock_json_response = json.dumps({
             "documentation_suggestions": [{"file": "mock/core.py", "line": 15, "suggestion": "def process_data(data):\n    \"\"\"Processes the input data using mock logic.\"\"\""}],
-            "potential_bugs": [{"file":"mock/utils.py", "line": 22, "description":"Potential division by zero if denominator is not checked.", "severity":"Medium"}],
+            "potential_bugs": [{"file": "mock/utils.py", "line": 22, "description": "Potential division by zero if denominator is not checked.", "severity": "Medium"}],
             "style_issues": [{"file": "mock/core.py", "line": 5, "description": "Variable 'varName' does not follow snake_case convention."}],
-            "module_summaries": [{"file": "mock/core.py", "summary": "This file contains the core mock processing logic."}, {"file":"mock/utils.py", "summary": "Utility functions for mocking."}],
-            "refactoring_suggestions": [{"file":"mock/utils.py", "line": 30, "area":"calculate_metrics function", "suggestion": "Function is too long (> 50 lines), consider breaking it down."}]
+            "module_summaries": [
+                {"file": "mock/core.py", "summary": "This file contains the core mock processing logic."},
+                {"file": "mock/utils.py", "summary": "Utility functions for mocking."}
+            ],
+            "refactoring_suggestions": [{"file": "mock/utils.py", "line": 30, "area": "calculate_metrics function", "suggestion": "Function is too long (> 50 lines), consider breaking it down."}]
         })
         st.success("Mock response generated successfully.")
-        return json.loads(mock_json_response), None # Return insights, no error
-
-        # Option 2: Simulate API error
-        # st.error("Simulating API error.")
-        # return None, "MOCK ERROR: Simulated API Quota Exceeded."
-
-        # Option 3: Simulate invalid JSON response
-        # st.warning("Simulating invalid JSON response from AI.")
-        # return {"raw_response": "{malformed json'"}, "AI response was not valid JSON, showing raw text."
-        #
-        # Option 4: Simulate empty results
-        # mock_empty_json = json.dumps({
-        #     "documentation_suggestions": [], "potential_bugs": [], "style_issues": [],
-        #     "module_summaries": [], "refactoring_suggestions": []
-        # })
-        # st.success("Mock response generated (empty results).")
-        # return json.loads(mock_empty_json), None
-    # --- END MOCK MODE LOGIC ---
-
+        return json.loads(mock_json_response), None
 
     # --- REAL API CALL LOGIC ---
     else:
-        if not initialize_gemini_model(): # Ensure model is ready
-             return None, "Gemini Model Initialization Failed."
-        if model is None: # Should not happen if initialize check passed, but safeguard
-             return None, "Gemini model not available."
+        if not initialize_gemini_model():
+            return None, "Gemini Model Initialization Failed."
+        if model is None:
+            return None, "Gemini model not available."
 
         try:
             st.write(f"📡 Sending request to {GEMINI_MODEL_NAME}...")
@@ -227,21 +230,17 @@ def call_gemini_api(prompt):
             )
             st.write("✅ Response received from AI.")
 
-            # Debug: Print raw response text
-            # print(f"--- RAW API RESPONSE ---\n{response.text}\n------------------------")
-
             try:
-                # Try to extract JSON robustly
                 json_response_text = response.text.strip()
-                # Handle potential markdown code block fences
+                # Remove potential markdown code block fences
                 if json_response_text.startswith("```json"):
                     json_response_text = json_response_text[7:]
-                if json_response_text.startswith("```"): # Handle case where ```json wasn't used
-                     json_response_text = json_response_text[3:]
+                if json_response_text.startswith("```"):
+                    json_response_text = json_response_text[3:]
                 if json_response_text.endswith("```"):
                     json_response_text = json_response_text[:-3]
 
-                # Find the first '{' and the last '}'
+                # Extract JSON object boundaries
                 json_start = json_response_text.find('{')
                 json_end = json_response_text.rfind('}') + 1
 
@@ -259,42 +258,40 @@ def call_gemini_api(prompt):
                 st.code(response.text, language='text')
                 return None, f"AI response was not valid JSON: {json_err}"
             except AttributeError:
-                 # Handle cases where response structure might be different (e.g. blocked)
-                 st.error(f"🚨 Unexpected API response structure.")
-                 st.code(f"Response object: {response}", language='text') # Log the problematic response
-                 # Try to get blocked reason if available
-                 try:
-                     block_reason = response.prompt_feedback.block_reason
-                     if block_reason:
-                          return None, f"Content blocked by API. Reason: {block_reason}"
-                 except Exception:
-                      pass # Ignore if feedback structure isn't as expected
-                 return None, "Unexpected response structure from API."
+                st.error("🚨 Unexpected API response structure.")
+                st.code(f"Response object: {response}", language='text')
+                try:
+                    block_reason = response.prompt_feedback.block_reason
+                    if block_reason:
+                        return None, f"Content blocked by API. Reason: {block_reason}"
+                except Exception:
+                    pass
+                return None, "Unexpected response structure from API."
             except Exception as e:
                 st.error(f"🚨 Unexpected issue processing AI response: {e}")
-                try: st.code(f"Response object: {response}", language='text')
-                except: pass
+                try:
+                    st.code(f"Response object: {response}", language='text')
+                except Exception:
+                    pass
                 return None, f"Unexpected response structure: {e}"
 
         except Exception as e:
             st.error(f"🚨 An error occurred during API call: {e}")
             error_msg = f"API call failed: {e}"
-            # Improved error identification
-            if hasattr(e, 'message'): # For google.api_core.exceptions
-                 if "429" in e.message:
-                     error_msg = "API Quota Exceeded or Rate Limit hit. Check your Google Cloud/AI Studio dashboard."
-                 elif "API key not valid" in e.message:
-                     error_msg = "Invalid Gemini API Key. Please check `.streamlit/secrets.toml`."
-                 elif "blocked" in e.message.lower(): # General check for safety blocks
-                     error_msg = "Content blocked due to safety settings. Review input code or adjust safety settings if appropriate."
-            elif "block_reason: SAFETY" in str(e): # Fallback check
-                 error_msg = "Content blocked due to safety settings. Review input code or adjust safety settings if appropriate."
+            if hasattr(e, 'message'):
+                if "429" in e.message:
+                    error_msg = "API Quota Exceeded or Rate Limit hit. Check your Google Cloud/AI Studio dashboard."
+                elif "API key not valid" in e.message:
+                    error_msg = "Invalid Gemini API Key. Please check `.streamlit/secrets.toml`."
+                elif "blocked" in e.message.lower():
+                    error_msg = "Content blocked due to safety settings. Review input code or adjust safety settings if appropriate."
+            elif "block_reason: SAFETY" in str(e):
+                error_msg = "Content blocked due to safety settings. Review input code or adjust safety settings if appropriate."
 
             return None, error_msg
 
-
 def display_results(results_json, requested_analyses):
-    """Renders the analysis results in Streamlit."""
+    """Renders the analysis results in the Streamlit interface."""
     st.header("📊 Analysis Report")
 
     if not isinstance(results_json, dict):
@@ -307,52 +304,54 @@ def display_results(results_json, requested_analyses):
         st.code(results_json["raw_response"], language='text')
         return
 
-    # Define display functions for clarity
     def display_list_items(items, fields):
         if items:
             for item in items:
                 details = []
                 for field_key, field_label in fields.items():
                     value = item.get(field_key, 'N/A')
-                    if value != 'N/A': # Only show if value exists
-                       details.append(f"**{field_label}:** {value}")
+                    if value != 'N/A':
+                        details.append(f"**{field_label}:** {value}")
                 st.markdown("- " + " - ".join(details))
-                # Handle specific multi-line outputs like suggestions/summaries
+                # Display multi-line outputs when applicable
                 if 'suggestion' in item:
-                     st.code(item['suggestion'], language='text')
+                    st.code(item['suggestion'], language='text')
                 elif 'description' in item:
-                     st.markdown(f"  > {item['description']}") # Indent description
+                    st.markdown(f"  > {item['description']}")
                 elif 'summary' in item:
-                     st.markdown(f"  > {item['summary']}") # Indent summary
+                    st.markdown(f"  > {item['summary']}")
         else:
             st.markdown("_No items found for this category._")
         st.divider()
 
-    # Map keys to display configurations
     display_config = {
         "generate_docs": {
-            "key": "documentation_suggestions", "title": AVAILABLE_ANALYSES["generate_docs"],
-            "fields": {"file": "File", "line": "Line"} # Suggestion shown by st.code
+            "key": "documentation_suggestions",
+            "title": AVAILABLE_ANALYSES["generate_docs"],
+            "fields": {"file": "File", "line": "Line"}
         },
         "find_bugs": {
-            "key": "potential_bugs", "title": AVAILABLE_ANALYSES["find_bugs"],
-            "fields": {"file": "File", "line": "Line", "severity": "Severity"} # Description shown separately
+            "key": "potential_bugs",
+            "title": AVAILABLE_ANALYSES["find_bugs"],
+            "fields": {"file": "File", "line": "Line", "severity": "Severity"}
         },
         "check_style": {
-            "key": "style_issues", "title": AVAILABLE_ANALYSES["check_style"],
-            "fields": {"file": "File", "line": "Line"} # Description shown separately
+            "key": "style_issues",
+            "title": AVAILABLE_ANALYSES["check_style"],
+            "fields": {"file": "File", "line": "Line"}
         },
         "summarize_modules": {
-            "key": "module_summaries", "title": AVAILABLE_ANALYSES["summarize_modules"],
-            "fields": {"file": "File"} # Summary shown separately
+            "key": "module_summaries",
+            "title": AVAILABLE_ANALYSES["summarize_modules"],
+            "fields": {"file": "File"}
         },
         "suggest_refactoring": {
-            "key": "refactoring_suggestions", "title": AVAILABLE_ANALYSES["suggest_refactoring"],
-            "fields": {"file": "File", "line": "Line", "area": "Area"} # Suggestion shown separately
+            "key": "refactoring_suggestions",
+            "title": AVAILABLE_ANALYSES["suggest_refactoring"],
+            "fields": {"file": "File", "line": "Line", "area": "Area"}
         }
     }
 
-    # Iterate and display selected sections
     any_results = False
     for analysis_key in requested_analyses:
         if analysis_key in display_config:
@@ -360,18 +359,18 @@ def display_results(results_json, requested_analyses):
             st.subheader(config["title"])
             items = results_json.get(config["key"], [])
             display_list_items(items, config["fields"])
-            if items: any_results = True
+            if items:
+                any_results = True
 
     if not any_results:
-         st.info("No specific findings were identified in the analysis based on your selections.")
+        st.info("No specific findings were identified in the analysis based on your selections.")
 
-    # Download button
     st.download_button(
-         label="Download Full Report (JSON)",
-         data=json.dumps(results_json, indent=4),
-         file_name="code_audit_report.json",
-         mime="application/json"
-     )
+        label="Download Full Report (JSON)",
+        data=json.dumps(results_json, indent=4),
+        file_name="code_audit_report.json",
+        mime="application/json"
+    )
 
 # --- Streamlit App Main Interface ---
 st.set_page_config(page_title="Codebase Audit Assistant", layout="wide")
@@ -382,9 +381,11 @@ st.markdown(f"Upload your codebase (`.zip`) for analysis using **{GEMINI_MODEL_N
 # Sidebar controls
 with st.sidebar:
     st.header("⚙️ Analysis Controls")
-    # Mock Mode Toggle
-    st.session_state.mock_api_call = st.toggle("🧪 Enable Mock API Mode (for Testing)", value=st.session_state.mock_api_call,
-                                               help="If enabled, uses fake data instead of calling the real Gemini API. Saves cost during testing.")
+    st.session_state.mock_api_call = st.toggle(
+        "🧪 Enable Mock API Mode (for Testing)",
+        value=st.session_state.mock_api_call,
+        help="If enabled, uses fake data instead of calling the real Gemini API. Saves cost during testing."
+    )
     if st.session_state.mock_api_call:
         st.info("Mock API Mode ACTIVE")
     else:
@@ -409,77 +410,63 @@ with st.sidebar:
         "7. Review the report."
     )
     st.info(f"**Note:** Only files with common code extensions ({', '.join(CODE_EXTENSIONS)}) are processed. Analysis might be limited (~{MAX_PROMPT_TOKENS_ESTIMATE:,} est. tokens).")
-
     st.divider()
     st.warning("⚠️ **Privacy:** Code content is sent to the Google Gemini API if Mock Mode is OFF. Do not upload sensitive code if uncomfortable.")
 
-
 # Main content area
 uploaded_file = st.file_uploader("📁 Upload Codebase ZIP File", type=['zip'], key="file_uploader")
-
 analysis_triggered = False
-results_cache = None # To store results briefly
+results_cache = None  # To store results briefly
 
 if uploaded_file:
     st.success(f"✅ File '{uploaded_file.name}' uploaded.")
-
     with st.spinner("Inspecting ZIP file..."):
         code_files, total_chars, file_count, ignored_files = process_zip_file(uploaded_file)
 
     if code_files is not None:
         st.info(f"Found **{file_count}** relevant code files ({total_chars:,} characters). Est. tokens: ~{estimate_token_count(total_chars):,}")
         if ignored_files:
-             with st.expander(f"View {len(ignored_files)} Skipped/Ignored Files"):
-                 # Use st.code for better formatting of list
-                 st.code("\n".join(ignored_files), language='text')
+            with st.expander(f"View {len(ignored_files)} Skipped/Ignored Files"):
+                st.code("\n".join(ignored_files), language='text')
 
         analyze_button_disabled = (not selected_analyses or file_count == 0)
         analyze_button_label = "Analyze Codebase" if not analyze_button_disabled else "Select Analyses or Upload Valid Code"
         if st.button(analyze_button_label, type="primary", disabled=analyze_button_disabled):
             analysis_triggered = True
-
             if not selected_analyses:
-                 st.warning("Please select at least one analysis type from the sidebar.")
+                st.warning("Please select at least one analysis type from the sidebar.")
             elif file_count == 0:
-                 st.warning("No relevant code files found in the ZIP archive to analyze.")
+                st.warning("No relevant code files found in the ZIP archive to analyze.")
             else:
                 st.divider()
                 with st.spinner(f"🚀 Preparing prompt & contacting AI ({'Mock Mode' if st.session_state.mock_api_call else GEMINI_MODEL_NAME})... This may take time."):
-                    # 1. Construct Prompt
                     analysis_prompt, included_files_in_prompt = construct_analysis_prompt(code_files, selected_analyses)
-
                     if analysis_prompt and included_files_in_prompt:
                         st.write(f"Analyzing {len(included_files_in_prompt)} files...")
-                        # 2. Call API (Real or Mock)
                         results_json, error_message = call_gemini_api(analysis_prompt)
-                        results_cache = (results_json, error_message) # Store results
+                        results_cache = (results_json, error_message)
                     elif not included_files_in_prompt:
                         results_cache = (None, "Could not proceed: No files included in prompt (check token limits/errors).")
                     else:
-                         results_cache = (None, "Failed to generate analysis prompt.")
-
-    else: # Error during zip processing
-        pass # Error message already shown
+                        results_cache = (None, "Failed to generate analysis prompt.")
+    else:
+        # Error during ZIP processing (error already displayed)
+        pass
 
-# Display results outside the button click block if analysis was triggered
 if analysis_triggered and results_cache:
     results_json, error_message = results_cache
     st.divider()
     if error_message:
         st.error(f"Analysis Failed: {error_message}")
-        # Display partial results if available (e.g., raw response on JSON error)
         if results_json and isinstance(results_json, dict) and "raw_response" in results_json:
-             st.subheader("Raw AI Response")
-             st.code(results_json["raw_response"], language='text')
-
+            st.subheader("Raw AI Response")
+            st.code(results_json["raw_response"], language='text')
     elif results_json:
         display_results(results_json, selected_analyses)
     else:
         st.error("Analysis did not return results or an unknown error occurred.")
-
-
 elif not uploaded_file:
     st.info("Upload a ZIP file containing your source code to begin.")
 
 st.divider()
-st.markdown("_Assistant powered by Google Gemini._")
+st.markdown("_Assistant powered by Google Gemini._")