Update app.py

app.py

@@ -55,14 +55,18 @@ def initialize_gemini_model():
         # Running in Mock Mode. Skipping Gemini initialization.
         return True  # Allow proceeding in mock mode
     elif model is not None:
-        # Gemini Model already initialized.
         return True
     return False
 
 # --- Helper Functions ---
 
 def estimate_token_count(text):
-    """Roughly estimate token count (assuming ~3 characters per token)."""
+    """
+    Roughly estimate token count (assuming ~3 characters per token).
+    If an integer is provided (i.e. a character count), simply divide by 3.
+    """
+    if isinstance(text, int):
+        return text // 3
     return len(text) // 3
 
 # --- OPTIMIZATION: Cache ZIP processing ---
@@ -81,7 +85,7 @@ def process_zip_file_cached(file_id, file_size, file_content_bytes):
             members = zip_ref.infolist()
             total_members = len(members)
             for i, member in enumerate(members):
-                # Update progress bar periodically (every 10 files)
+                # Update progress bar every 10 files to reduce overhead
                 if i % 10 == 0:
                     progress_bar.progress(int((i / total_members) * 100))
 
@@ -128,7 +132,7 @@ def process_zip_file_cached(file_id, file_size, file_content_bytes):
     elif file_count == 0 and ignored_files:
         st.warning("No files with recognized code extensions found. Some files were skipped.")
 
-    print(f"Cache miss or new file: Processed ZIP {file_id}")  # Debug print
+    print(f"Cache miss or new file: Processed ZIP {file_id}")
     return code_files, total_chars, file_count, ignored_files
 
 def construct_analysis_prompt(code_files_dict, requested_analyses):
@@ -136,11 +140,8 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
     prompt_parts = ["Analyze the following codebase provided as a collection of file paths and their content.\n\n"]
     current_token_estimate = estimate_token_count(prompt_parts[0])
     included_files = []
-
-    # Use join for potentially faster concatenation
     code_segments = []
 
-    # Provide feedback for large codebases
     prompt_status = st.empty()
     if len(code_files_dict) > 50:
         prompt_status.write("Constructing prompt (processing files)...")
@@ -160,7 +161,7 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
             st.warning(f"⚠️ Codebase may exceed context window estimate (~{MAX_PROMPT_TOKENS_ESTIMATE} tokens). Analysis performed only on the first {len(included_files)} files ({current_token_estimate:,} tokens).")
             break
 
-    prompt_status.empty()  # Clear status message
+    prompt_status.empty()
 
     if not included_files:
         st.error("🚨 No code files could be included within the estimated token limit.")
@@ -169,7 +170,6 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
     concatenated_code = "".join(code_segments)
     prompt_parts.append(concatenated_code)
 
-    # Generate the expected JSON structure description based on selected analyses
     json_structure_description = "{\n"
     structure_parts = []
     if "generate_docs" in requested_analyses:
@@ -209,9 +209,8 @@ def call_gemini_api(prompt):
     # MOCK MODE LOGIC
     if st.session_state.mock_api_call:
         st.info("MOCK MODE: Simulating API call...")
-        st.write("...")  # Minimal feedback in mock mode
-        time.sleep(1)  # Shorter mock delay
-
+        st.write("...")
+        time.sleep(1)
         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"}],
@@ -400,6 +399,7 @@ results_placeholder = st.container()        # Container for results display
 if uploaded_file:
     st.success(f"✅ File '{uploaded_file.name}' uploaded.")
 
+    # Read file bytes once for caching
     uploaded_file_bytes = uploaded_file.getvalue()
     file_id = f"{uploaded_file.name}-{uploaded_file.size}"
 
@@ -448,7 +448,6 @@ if st.session_state.analysis_requested:
             display_results(st.session_state.analysis_results, selected_analyses)
         else:
             st.info("Analysis initiated, but no results or errors were stored. Please try again.")
-
 elif not uploaded_file:
     results_placeholder.info("Upload a ZIP file containing your source code to begin.")