Create pipeline.py

pipeline.py

@@ -0,0 +1,486 @@
+# app/pipeline.py
+import os
+import json
+import logging
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+import pandas as pd
+import numpy as np
+
+logger = logging.getLogger("pipeline")
+logger.setLevel(logging.INFO)
+
+# ---------------------------
+# Safety / helpers
+# ---------------------------
+
+DISALLOWED_PATTERNS = [
+    r"__\w+__", r"\bimport\s+os\b", r"\bimport\s+sys\b", r"\bimport\s+subprocess\b",
+    r"\bimport\s+socket\b", r"\bimport\s+requests\b", r"\bopen\s*\(", r"\beval\s*\(",
+    r"\bexec\s*\(", r"\bcompile\s*\(", r"\bsystem\s*\(", r"\bPopen\b", r"\bsh\b",
+]
+
+def code_is_safe(code: str) -> Tuple[bool, Optional[str]]:
+    for pat in DISALLOWED_PATTERNS:
+        if re.search(pat, code):
+            return False, f"disallowed pattern: {pat}"
+    return True, None
+
+def to_json_serializable(obj):
+    if isinstance(obj, (np.integer,)):
+        return int(obj)
+    if isinstance(obj, (np.floating,)):
+        return float(obj)
+    if isinstance(obj, (np.ndarray,)):
+        return obj.tolist()
+    if pd.isna(obj):
+        return None
+    return obj
+
+def simple_schema(df: pd.DataFrame) -> Dict[str, Any]:
+    return {
+        "columns": [
+            {"name": c, "dtype": str(df[c].dtype), "n_unique": int(df[c].nunique(dropna=True))}
+            for c in df.columns
+        ],
+        "n_rows": int(len(df)),
+    }
+
+# ---------------------------
+# Ingest / preprocess / sampling
+# ---------------------------
+
+def ingest_file(path: str, sheet: Optional[str] = None) -> Tuple[pd.DataFrame, Dict[str, Any]]:
+    ext = os.path.splitext(path)[1].lower()
+    metadata = {"file_type": ext, "sheet_names": None, "selected_sheet": None}
+    if ext in [".csv", ".txt"]:
+        df = pd.read_csv(path)
+        metadata["selected_sheet"] = "csv"
+    elif ext in [".xls", ".xlsx"]:
+        xls = pd.ExcelFile(path)
+        sheets = xls.sheet_names
+        metadata["sheet_names"] = sheets
+        chosen = sheet if sheet and sheet in sheets else sheets[0]
+        metadata["selected_sheet"] = chosen
+        df = pd.read_excel(xls, sheet_name=chosen)
+    else:
+        raise ValueError("Unsupported file type: " + ext)
+    metadata["file_type"] = ext
+    return df, metadata
+
+def preprocess_df(df: pd.DataFrame) -> Tuple[pd.DataFrame, Dict[str, Any]]:
+    actions = []
+    df = df.copy()
+    df.columns = [str(c).strip() for c in df.columns]
+
+    object_cols = df.select_dtypes(include="object").columns.tolist()
+    for c in object_cols:
+        try:
+            df[c] = df[c].where(df[c].isna(), df[c].astype(str).str.strip())
+        except Exception:
+            pass
+
+    for c in df.columns:
+        if df[c].dtype == object:
+            coerced = pd.to_numeric(df[c], errors="coerce")
+            non_na = coerced.notna().sum()
+            if non_na >= max(1, 0.5 * len(df)):
+                df[c] = coerced
+                actions.append(f"coerced {c} -> numeric")
+
+    num_cols = df.select_dtypes(include=[np.number]).columns.tolist()
+    for c in num_cols:
+        median = df[c].median()
+        if pd.isna(median):
+            median = 0
+        df[c] = df[c].fillna(median)
+        actions.append(f"filled numeric {c} nulls with median {median}")
+
+    for c in object_cols:
+        try:
+            mode = df[c].mode().iloc[0] if not df[c].mode().empty else ""
+        except Exception:
+            mode = ""
+        df[c] = df[c].fillna(mode)
+        actions.append(f"filled object {c} nulls with mode '{mode}'")
+
+    schema = simple_schema(df)
+    return df, {"actions": actions, "schema": schema}
+
+def sample_df(df: pd.DataFrame, n: int = 5) -> Dict[str, Any]:
+    head = df.head(n).to_dict(orient="records")
+    tail = df.tail(n).to_dict(orient="records")
+    rnd = df.sample(n=n, random_state=42).to_dict(orient="records") if len(df) > n else df.sample(frac=1.0).to_dict(orient="records")
+    return {"head": head, "tail": tail, "random": rnd}
+
+# ---------------------------
+# Deterministic chart generator (fallback)
+# ---------------------------
+
+def _pick_categorical(df: pd.DataFrame) -> Optional[str]:
+    for c in df.columns:
+        if df[c].dtype == object or df[c].nunique() < max(50, 0.5 * len(df)):
+            return c
+    return None
+
+def _pick_numeric(df: pd.DataFrame, exclude: List[str] = []) -> Optional[str]:
+    for c in df.select_dtypes(include=[np.number]).columns:
+        if c not in exclude:
+            return c
+    for c in df.columns:
+        try:
+            coerced = pd.to_numeric(df[c], errors="coerce")
+            if coerced.notna().sum() > 0:
+                return c
+        except Exception:
+            continue
+    return None
+
+def _pick_datetime_or_index(df: pd.DataFrame) -> Optional[str]:
+    for c in df.columns:
+        if np.issubdtype(df[c].dtype, np.datetime64):
+            return c
+    for c in df.columns:
+        try:
+            parsed = pd.to_datetime(df[c], errors="coerce")
+            if parsed.notna().sum() > 0:
+                return c
+        except Exception:
+            continue
+    return None
+
+def _box_aggregate(df: pd.DataFrame, cat_col: Optional[str], num_col: str) -> List[Dict[str, Any]]:
+    out = []
+    if cat_col is None:
+        series = df[num_col].dropna()
+        if len(series) == 0:
+            return out
+        q1 = float(series.quantile(0.25))
+        median = float(series.quantile(0.5))
+        q3 = float(series.quantile(0.75))
+        out.append({"category": None, "q1": q1, "median": median, "q3": q3})
+        return out
+    for name, group in df.groupby(cat_col):
+        ser = group[num_col].dropna()
+        if len(ser) == 0:
+            continue
+        q1 = float(ser.quantile(0.25))
+        median = float(ser.quantile(0.5))
+        q3 = float(ser.quantile(0.75))
+        out.append({"category": to_json_serializable(name), "q1": q1, "median": median, "q3": q3})
+    return out
+
+def generate_chart_data_from_spec(df: pd.DataFrame, spec: Dict[str, Any]) -> Tuple[str, List[Dict[str, Any]]]:
+    chart_type = spec.get("chart_type")
+    cols = spec.get("target_columns", [])
+    agg = spec.get("aggregation", None)
+    df_local = df.copy()
+
+    if chart_type == "pie":
+        if len(cols) == 0:
+            col = _pick_categorical(df_local)
+            cols = [col] if col else []
+        if len(cols) == 1:
+            col = cols[0]
+            series = df_local[col].astype(str).value_counts(dropna=True)
+            return "pie", [{"name": k, "value": int(v)} for k, v in series.items()]
+        else:
+            cat, val = cols[0], cols[1]
+            grouped = df_local.groupby(cat)[val].sum().reset_index()
+            return "pie", [{"name": r[cat], "value": to_json_serializable(r[val])} for r in grouped.to_dict(orient="records")]
+
+    if chart_type == "bar":
+        if len(cols) == 0:
+            label = _pick_categorical(df_local)
+            cols = [label] if label else []
+        if len(cols) == 1:
+            label = cols[0]
+            series = df_local[label].astype(str).value_counts().reset_index()
+            series.columns = [label, "count"]
+            return "bar", [{"label": r[label], "count": int(r["count"])} for r in series.to_dict(orient="records")]
+        else:
+            label, metric = cols[0], cols[1]
+            if agg in (None, "", "sum"):
+                grouped = df_local.groupby(label)[metric].sum().reset_index()
+            elif agg == "mean":
+                grouped = df_local.groupby(label)[metric].mean().reset_index()
+            else:
+                grouped = df_local.groupby(label)[metric].sum().reset_index()
+            return "bar", [{"label": r[label], metric: to_json_serializable(r[metric])} for r in grouped.to_dict(orient="records")]
+
+    if chart_type == "line":
+        if len(cols) < 2:
+            y = _pick_numeric(df_local)
+            x = _pick_datetime_or_index(df_local)
+            cols = [x, y]
+        xcol, ycol = cols[0], cols[1]
+        s_x = pd.to_datetime(df_local[xcol], errors="coerce") if xcol in df_local.columns else pd.Series(range(len(df_local)))
+        series = pd.DataFrame({"x": s_x, "y": df_local[ycol]})
+        try:
+            series = series.sort_values("x").reset_index(drop=True)
+        except Exception:
+            pass
+        out = []
+        for r in series.to_dict(orient="records"):
+            x_val = r["x"].isoformat() if hasattr(r["x"], "isoformat") else r["x"]
+            out.append({"x": to_json_serializable(x_val), "y": to_json_serializable(r["y"])})
+        return "line", out
+
+    if chart_type == "scatter":
+        if len(cols) < 2:
+            x = _pick_numeric(df_local)
+            y = _pick_numeric(df_local, exclude=[x]) if x else None
+            cols = [x, y]
+        xcol, ycol = cols[0], cols[1]
+        return "scatter", [{"x": to_json_serializable(r[xcol]), "y": to_json_serializable(r[ycol])} for r in df_local[[xcol, ycol]].to_dict(orient="records")]
+
+    if chart_type == "histogram":
+        col = cols[0] if cols else _pick_numeric(df_local)
+        series = df_local[col].dropna()
+        counts, bin_edges = np.histogram(series, bins=10)
+        out = []
+        for i in range(len(counts)):
+            out.append({"bin": f"{float(bin_edges[i]):.6g}-{float(bin_edges[i+1]):.6g}", "count": int(counts[i])})
+        return "histogram", out
+
+    if chart_type == "boxplot":
+        if len(cols) == 0:
+            num = _pick_numeric(df_local)
+            return "boxplot", _box_aggregate(df_local, None, num)
+        if len(cols) == 1:
+            num = cols[0]
+            return "boxplot", _box_aggregate(df_local, None, num)
+        cat, num = cols[0], cols[1]
+        return "boxplot", _box_aggregate(df_local, cat, num)
+
+    return chart_type or "unknown", []
+
+# ---------------------------
+# Gemini wrapper (optional)
+# ---------------------------
+
+def gemini_generate_json(model: str, system_instruction: str, user_content: str, require_json: bool = True) -> Any:
+    """
+    Attempts to call google-genai if GEMINI_API_KEY is provided.
+    Falls back to returning a safe marker that indicates a deterministic fallback should be used.
+    """
+    api_key = "AIzaSyDfy0E-9b2XjoYHrHX2C1nVLHWyrWUFkMs"
+    if not api_key:
+        logger.info("GEMINI_API_KEY not set; using deterministic fallbacks.")
+        return {"__use_fallback__": True}
+
+    try:
+        # import lazily
+        from google import genai
+        from google.genai import types
+    except Exception as e:
+        logger.exception("google-genai not available: %s", e)
+        return {"__use_fallback__": True}
+
+    client = genai.Client(api_key=api_key)
+    contents = [
+        types.Content(
+            role="user",
+            parts=[types.Part.from_text(text=user_content)],
+        )
+    ]
+    config = types.GenerateContentConfig(
+        thinking_config=types.ThinkingConfig(thinking_budget=0),
+        response_mime_type="application/json",
+        system_instruction=[types.Part.from_text(text=system_instruction)],
+    )
+    full_text = ""
+    try:
+        for chunk in client.models.generate_content_stream(model=model, contents=contents, config=config):
+            if hasattr(chunk, "text") and chunk.text:
+                full_text += chunk.text
+            elif (
+                chunk.candidates
+                and chunk.candidates[0].content
+                and chunk.candidates[0].content.parts
+                and chunk.candidates[0].content.parts[0].text
+            ):
+                full_text += chunk.candidates[0].content.parts[0].text
+        full_text = full_text.strip()
+        if require_json:
+            try:
+                return json.loads(full_text)
+            except Exception:
+                return {"__raw_text__": full_text}
+        return full_text
+    except Exception as e:
+        logger.exception("genai call failed: %s", e)
+        return {"__use_fallback__": True}
+
+# ---------------------------
+# Controller
+# ---------------------------
+
+DEFAULT_MODEL = "gemini-2.5-flash-lite"
+
+def ensure_six_tasks(tasks: List[Dict[str, Any]], df: pd.DataFrame) -> List[Dict[str, Any]]:
+    existing_types = [t.get("chart_type") for t in tasks]
+    candidates = ["pie", "bar", "line", "scatter", "histogram", "boxplot"]
+    out = tasks[:]
+    for ct in candidates:
+        if len(out) >= 6:
+            break
+        if ct not in existing_types:
+            if ct == "pie":
+                col = _pick_categorical(df) or df.columns[0]
+                out.append({"chart_type": "pie", "target_columns": [col], "aggregation": "count", "reasoning": "fallback pie"})
+            elif ct == "bar":
+                label = _pick_categorical(df) or df.columns[0]
+                num = _pick_numeric(df) or df.columns[0]
+                out.append({"chart_type": "bar", "target_columns": [label, num], "aggregation": "sum", "reasoning": "fallback bar"})
+            elif ct == "line":
+                y = _pick_numeric(df) or df.columns[0]
+                x = _pick_datetime_or_index(df) or df.index.name or "index"
+                if x == "index":
+                    out.append({"chart_type": "line", "target_columns": [x, y], "aggregation": None, "reasoning": "fallback line on index"})
+                else:
+                    out.append({"chart_type": "line", "target_columns": [x, y], "aggregation": None, "reasoning": "fallback line"})
+            elif ct == "scatter":
+                x = _pick_numeric(df)
+                y = _pick_numeric(df, exclude=[x]) or x
+                out.append({"chart_type": "scatter", "target_columns": [x, y], "aggregation": None, "reasoning": "fallback scatter"})
+            elif ct == "histogram":
+                num = _pick_numeric(df) or df.columns[0]
+                out.append({"chart_type": "histogram", "target_columns": [num], "aggregation": None, "reasoning": "fallback histogram"})
+            elif ct == "boxplot":
+                num = _pick_numeric(df) or df.columns[0]
+                cat = _pick_categorical(df)
+                if cat:
+                    out.append({"chart_type": "boxplot", "target_columns": [cat, num], "aggregation": None, "reasoning": "fallback boxplot by category"})
+                else:
+                    out.append({"chart_type": "boxplot", "target_columns": [num], "aggregation": None, "reasoning": "fallback boxplot global"})
+    return out
+
+def process_file(path: str, sheet: Optional[str] = None, model: str = DEFAULT_MODEL) -> Dict[str, Any]:
+    df, meta = ingest_file(path, sheet)
+    pre_df, preprocess_meta = preprocess_df(df)
+    samples = sample_df(pre_df, n=5)
+
+    classification_input = json.dumps({"samples": samples, "schema": simple_schema(pre_df), "meta": meta})
+    classification_output = gemini_generate_json(
+        model=model,
+        system_instruction="You are a classification agent. Identify domain and chart tasks.",
+        user_content=classification_input,
+        require_json=True,
+    )
+
+    if not isinstance(classification_output, dict) or "tasks" not in classification_output:
+        fallback = {
+            "domain": "unknown",
+            "tasks": [
+                {"chart_type": "pie", "target_columns": [_pick_categorical(pre_df) or pre_df.columns[0]], "aggregation": "count", "reasoning": "fallback"},
+                {"chart_type": "bar", "target_columns": [_pick_categorical(pre_df) or pre_df.columns[0], _pick_numeric(pre_df) or pre_df.columns[0]], "aggregation": "sum", "reasoning": "fallback"},
+            ],
+        }
+        classification_output = fallback
+
+    planning_input = json.dumps({"classification": classification_output, "schema": simple_schema(pre_df), "samples": samples})
+    planning_output = gemini_generate_json(
+        model=model,
+        system_instruction="You are a planning agent. Produce tasks with code assigning `result` variable.",
+        user_content=planning_input,
+        require_json=True,
+    )
+
+    tasks = []
+    if isinstance(planning_output, dict) and "tasks" in planning_output:
+        tasks = planning_output["tasks"]
+    else:
+        tasks = classification_output.get("tasks", [])
+
+    tasks = ensure_six_tasks(tasks, pre_df)
+
+    final = {"pie": [], "bar": [], "line": [], "scatter": [], "histogram": [], "boxplot": []}
+    execution_errors = []
+
+    for idx, task in enumerate(tasks):
+        chart_type = task.get("chart_type")
+        code_snippet = task.get("code")
+        executed = False
+
+        if code_snippet:
+            safe, reason = code_is_safe(code_snippet)
+            if not safe:
+                logger.warning("Rejected unsafe code snippet: %s", reason)
+            else:
+                allowed_globals = {
+                    "__builtins__": {"None": None, "True": True, "False": False, "len": len, "min": min, "max": max, "sum": sum, "sorted": sorted, "round": round},
+                    "pd": pd, "np": np, "df": pre_df.copy(),
+                }
+                local_vars = {}
+                try:
+                    exec(code_snippet, allowed_globals, local_vars)
+                    result = None
+                    if "result" in local_vars:
+                        result = local_vars["result"]
+                    elif "output" in local_vars:
+                        result = local_vars["output"]
+                    else:
+                        single_expr = (("\n" not in code_snippet) and ("=" not in code_snippet) and ("return" not in code_snippet) and (not code_snippet.strip().startswith("def ")))
+                        if single_expr:
+                            try:
+                                result = eval(code_snippet, allowed_globals, {})
+                            except Exception:
+                                result = None
+
+                    result_json = None
+                    if isinstance(result, pd.DataFrame):
+                        result_json = [{k: to_json_serializable(v) for k, v in r.items()} for r in result.to_dict(orient="records")]
+                    elif isinstance(result, list):
+                        norm = []
+                        for r in result:
+                            if isinstance(r, dict):
+                                norm.append({k: to_json_serializable(v) for k, v in r.items()})
+                            else:
+                                norm.append(to_json_serializable(r))
+                        result_json = norm
+                    elif isinstance(result, dict):
+                        result_json = [{k: to_json_serializable(v) for k, v in result.items()}]
+                    else:
+                        result_json = None
+
+                    if result_json is not None:
+                        normalized = []
+                        for item in result_json:
+                            if isinstance(item, dict):
+                                normalized.append(item)
+                            else:
+                                normalized.append({"value": to_json_serializable(item)})
+                        if chart_type in final:
+                            final[chart_type].extend(normalized)
+                        else:
+                            final.setdefault(chart_type, []).extend(normalized)
+                        executed = True
+
+                    if not executed:
+                        execution_errors.append({"task_index": idx, "reason": "result not list-of-dicts or missing", "code": code_snippet})
+                except Exception as e:
+                    logger.exception("Model code execution failed for task %s: %s", idx, str(e))
+                    execution_errors.append({"task_index": idx, "reason": "exception during exec/eval", "exception": str(e), "code": code_snippet})
+
+        if not executed:
+            ct, res = generate_chart_data_from_spec(pre_df, task)
+            final.setdefault(ct, []).extend(res)
+
+    for k in final:
+        if isinstance(final[k], list):
+            final[k] = final[k][:200]
+
+    output_payload = {
+        "metadata": {
+            "source_file": os.path.basename(path),
+            "ingestion_meta": meta,
+            "preprocess_meta": preprocess_meta,
+            "classification": classification_output if isinstance(classification_output, dict) else {"raw": classification_output},
+            "planning_meta": planning_output if isinstance(planning_output, dict) else {"raw": planning_output},
+            "execution_errors": execution_errors,
+        },
+        "charts": final,
+    }
+    return output_payload