Update src/streamlit_app.py

src/streamlit_app.py

@@ -114,7 +114,7 @@ def fallback_clean(df: pd.DataFrame) -> pd.DataFrame:
 def ai_clean_dataset(df: pd.DataFrame) -> (pd.DataFrame, str):
     """Returns cleaned df and a status message"""
     if len(df) > 50:
-        return df, "AI cleaning skipped: dataset has more than 50 rows."
+        return df, "⚠️ AI cleaning skipped: dataset has more than 50 rows."
     csv_text = df.to_csv(index=False)
     prompt = f"""
 You are a professional data cleaning assistant.
@@ -133,9 +133,9 @@ Dataset:
         cleaned_str = cleaned_str.replace("```csv", "").replace("```", "").replace("###", "").strip()
         cleaned_df = pd.read_csv(StringIO(cleaned_str), on_bad_lines="skip")
         cleaned_df.columns = [c.strip().replace(" ", "_").lower() for c in cleaned_df.columns]
-        return cleaned_df, "AI cleaning completed successfully."
+        return cleaned_df, "✅ AI cleaning completed successfully."
     except Exception as e:
-        return df, f"AI cleaning failed: {str(e)}"
+        return df, f"⚠️ AI cleaning failed: {str(e)}"
 
 # ======================================================
 # 🧩 DATA SUMMARY FOR TOKEN-EFFICIENT ANALYSIS
@@ -150,7 +150,7 @@ def summarize_for_analysis(df: pd.DataFrame, sample_rows=10) -> str:
         else:
             top = df[col].value_counts().head(3).to_dict()
             summary.append(f"- {col}: top_values={top}, non_null={non_null}")
-    # Include a small sample
+    # Include a small sample for context
     sample = df.head(sample_rows).to_csv(index=False)
     summary.append("--- Sample Data ---")
     summary.append(sample)
@@ -181,13 +181,13 @@ Respond with:
         if ANALYST_MODEL == "Gemini 2.5 Flash (Google)":
             if GEMINI_API_KEY is None:
                 return "⚠️ Gemini API key missing."
-            response = genai.generate_text(
+            response = genai.models.generate(
                 model="gemini-2.5-flash",
-                prompt=prompt,
+                messages=[{"author": "user", "content": prompt}],
                 temperature=temperature,
                 max_output_tokens=max_tokens
             )
-            return getattr(response, "candidates", [{"content": "No response from Gemini."}])[0]["content"]
+            return response.candidates[0].content if response.candidates else "No response from Gemini."
         else:
             return safe_hf_generate(hf_analyst_client, prompt, temperature=temperature, max_tokens=max_tokens)
     except Exception as e: