Create utils.py

src/utils.py

@@ -0,0 +1,183 @@
+import os
+import time
+import pandas as pd
+import numpy as np
+from io import StringIO
+from huggingface_hub import InferenceClient
+import google.generativeai as genai
+
+# ======================================================
+# 🔧 HELPER FUNCTIONS
+# ======================================================
+
+def safe_hf_generate(client, prompt, temperature=0.3, max_tokens=512, retries=2):
+    """Safely call Hugging Face text generation with retry and graceful fallback."""
+    for attempt in range(retries + 1):
+        try:
+            resp = client.text_generation(
+                prompt,
+                temperature=temperature,
+                max_new_tokens=max_tokens,
+                return_full_text=False,
+            )
+            return resp.strip()
+        except Exception as e:
+            err = str(e)
+            if "503" in err or "Service Temporarily Unavailable" in err:
+                time.sleep(2)
+                if attempt < retries:
+                    continue
+                else:
+                    return "⚠️ The Hugging Face model is temporarily unavailable. Please try again later."
+            elif "Supported task: conversational" in err:
+                chat_resp = client.chat_completion(
+                    messages=[{"role": "user", "content": prompt}],
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                )
+                return chat_resp["choices"][0]["message"]["content"].strip()
+            else:
+                raise e
+    return "⚠️ Failed after multiple retries."
+
+# ======================================================
+# 🧼 DATA CLEANING
+# ======================================================
+
+def fallback_clean(df: pd.DataFrame) -> pd.DataFrame:
+    """Perform a basic fallback cleaning if AI-based cleaning fails."""
+    df = df.copy()
+    df.dropna(axis=1, how="all", inplace=True)
+    df.columns = [c.strip().replace(" ", "_").lower() for c in df.columns]
+
+    for col in df.columns:
+        if df[col].dtype == "O":
+            if not df[col].mode().empty:
+                df[col].fillna(df[col].mode()[0], inplace=True)
+            else:
+                df[col].fillna("Unknown", inplace=True)
+        else:
+            df[col].fillna(df[col].median(), inplace=True)
+
+    df.drop_duplicates(inplace=True)
+    return df
+
+
+def ai_clean_dataset(df: pd.DataFrame, cleaner_client: InferenceClient) -> (pd.DataFrame, str):
+    """Clean dataset intelligently using the chosen Hugging Face model."""
+    if len(df) > 50:
+        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.
+Clean and standardize the dataset below dynamically:
+1. Handle missing values
+2. Fix column name inconsistencies
+3. Convert data types (dates, numbers, categories)
+4. Remove irrelevant or duplicate rows
+Return ONLY a valid CSV text (no markdown, no explanations).
+
+Dataset:
+{csv_text}
+"""
+    try:
+        cleaned_str = safe_hf_generate(cleaner_client, prompt, temperature=0.1, max_tokens=4096)
+        cleaned_str = cleaned_str.replace("```csv", "").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."
+    except Exception as e:
+        return fallback_clean(df), f"⚠️ AI cleaning failed, used fallback cleaning instead: {str(e)}"
+
+# ======================================================
+# 📊 DATA SUMMARIZATION
+# ======================================================
+
+def summarize_for_analysis(df: pd.DataFrame, sample_rows: int = 10) -> str:
+    """Generate a concise textual summary of the dataset for AI models."""
+    summary = [f"Rows: {len(df)}, Columns: {len(df.columns)}"]
+
+    for col in df.columns:
+        non_null = int(df[col].notnull().sum())
+        if pd.api.types.is_numeric_dtype(df[col]):
+            desc = df[col].describe().to_dict()
+            summary.append(
+                f"- {col}: mean={desc.get('mean', np.nan):.2f}, median={df[col].median():.2f}, non_null={non_null}"
+            )
+        else:
+            top = df[col].value_counts().head(3).to_dict()
+            summary.append(f"- {col}: top_values={top}, non_null={non_null}")
+
+    sample = df.head(sample_rows).to_csv(index=False)
+    summary.append("--- Sample Data ---")
+    summary.append(sample)
+    return "\n".join(summary)
+
+# ======================================================
+# 🧠 ANALYSIS LOGIC
+# ======================================================
+
+def query_analysis_model(
+    df: pd.DataFrame,
+    user_query: str,
+    dataset_name: str,
+    analyst_model: str,
+    hf_client: InferenceClient = None,
+    temperature: float = 0.3,
+    max_tokens: int = 1024,
+    gemini_api_key: str = None
+) -> str:
+    """Query the selected AI model (Hugging Face or Gemini) to analyze the dataset."""
+    prompt_summary = summarize_for_analysis(df)
+    prompt = f"""
+You are a professional data analyst.
+Analyze the dataset '{dataset_name}' and answer the user's question.
+
+--- DATA SUMMARY ---
+{prompt_summary}
+
+--- USER QUESTION ---
+{user_query}
+
+Respond with:
+1. Key insights and patterns
+2. Quantitative findings
+3. Notable relationships or anomalies
+4. Data-driven recommendations
+"""
+
+    try:
+        if analyst_model == "Gemini 2.5 Flash (Google)":
+            if not gemini_api_key:
+                return "⚠️ Gemini API key missing. Cannot use Gemini."
+            genai.configure(api_key=gemini_api_key)
+            response = genai.GenerativeModel("gemini-2.5-flash").generate_content(
+                prompt,
+                generation_config={
+                    "temperature": temperature,
+                    "max_output_tokens": max_tokens
+                }
+            )
+            return response.text if hasattr(response, "text") else "No valid text response."
+
+        # Otherwise, use Hugging Face model
+        result = safe_hf_generate(hf_client, prompt, temperature=temperature, max_tokens=max_tokens)
+
+        # fallback to Gemini if Hugging Face fails
+        if "temporarily unavailable" in result.lower() and gemini_api_key:
+            genai.configure(api_key=gemini_api_key)
+            alt = genai.GenerativeModel("gemini-2.5-flash").generate_content(prompt)
+            return f"🔄 Fallback to Gemini:\n\n{alt.text}"
+        return result
+
+    except Exception as e:
+        if "503" in str(e) and gemini_api_key:
+            genai.configure(api_key=gemini_api_key)
+            response = genai.GenerativeModel("gemini-2.5-flash").generate_content(prompt)
+            return f"🔄 Fallback to Gemini due to 503 error:\n\n{response.text}"
+        return f"⚠️ Analysis failed: {str(e)}"
+
+# ======================================================
+# 🔍 END OF MODULE
+# ======================================================