@vespermcp/mcp-server 1.1.2 → 1.2.0

package/src/python/fusion_engine.py

@@ -0,0 +1,368 @@
+import sys
+import json
+import os
+import hashlib
+import subprocess
+from typing import List, Optional, Union, Dict, Any
+
+try:
+    import polars as pl
+    HAS_POLARS = True
+except ImportError:
+    HAS_POLARS = False
+
+
+def _safe_suffix(source_path: str, idx: int) -> str:
+    base = os.path.basename(source_path)
+    base = os.path.splitext(base)[0].replace(" ", "_")
+    if not base:
+        base = f"source{idx+1}"
+    return base
+
+
+def _load_with_polars(path: str):
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".csv":
+        return pl.read_csv(path, ignore_errors=True)
+    if ext in [".parquet", ".pq"]:
+        return pl.read_parquet(path)
+    if ext in [".feather", ".ftr", ".arrow", ".ipc"]:
+        return pl.read_ipc(path)
+    if ext in [".jsonl", ".ndjson"]:
+        return pl.read_ndjson(path)
+    if ext == ".json":
+        try:
+            return pl.read_json(path)
+        except Exception:
+            return pl.read_ndjson(path)
+    raise ValueError(f"Unsupported source format: {ext} ({path})")
+
+
+def _write_with_polars(df, output_path: str, fmt: str, compression: Optional[str]):
+    fmt = fmt.lower()
+    comp = compression
+    if comp in ["none", "None", "uncompressed"]:
+        comp = None
+
+    if fmt == "csv":
+        df.write_csv(output_path)
+    elif fmt == "parquet":
+        df.write_parquet(output_path, compression=(compression or "snappy"))
+    elif fmt == "feather":
+        ipc_comp = "lz4" if compression is None else compression
+        if ipc_comp == "uncompressed":
+            ipc_comp = None
+        df.write_ipc(output_path, compression=ipc_comp)
+    elif fmt in ["arrow", "ipc"]:
+        df.write_ipc(output_path, compression=comp)
+    elif fmt == "jsonl":
+        df.write_ndjson(output_path)
+    else:
+        raise ValueError(f"Unsupported output format: {fmt}")
+
+
+def _ensure_parent(path: str):
+    parent = os.path.dirname(path)
+    if parent and not os.path.exists(parent):
+        os.makedirs(parent, exist_ok=True)
+
+
+def _compute_null_ratio(df) -> float:
+    if len(df) == 0 or len(df.columns) == 0:
+        return 0.0
+    nulls = 0
+    for col in df.columns:
+        nulls += int(df[col].null_count())
+    total_cells = len(df) * len(df.columns)
+    return (nulls / total_cells) * 100 if total_cells else 0.0
+
+
+def _leakage_report(df, source_col: str = "_vesper_source", id_col: Optional[str] = None) -> Dict[str, Any]:
+    report = {
+        "leakage_detected": False,
+        "leakage_count": 0,
+        "id_column": id_col,
+        "warnings": []
+    }
+
+    if source_col not in df.columns:
+        report["warnings"].append("Source marker column missing; leakage check skipped.")
+        return report
+
+    if id_col and id_col in df.columns:
+        overlap = (
+            df.group_by(id_col)
+            .agg(pl.col(source_col).n_unique().alias("source_count"))
+            .filter(pl.col("source_count") > 1)
+        )
+        overlap_count = len(overlap)
+        if overlap_count > 0:
+            report["leakage_detected"] = True
+            report["leakage_count"] = overlap_count
+            report["warnings"].append(f"Found {overlap_count} IDs appearing across multiple sources")
+    else:
+        # Fallback: hash rows (excluding source marker) and check if same row appears in multiple sources
+        compare_cols = [c for c in df.columns if c != source_col]
+        if not compare_cols:
+            return report
+
+        row_sig = df.select(compare_cols).with_columns(
+            pl.concat_str([pl.col(c).cast(pl.Utf8, strict=False) for c in compare_cols], separator="||").alias("_row_sig")
+        )
+        tmp = row_sig.with_columns(df[source_col]).select(["_row_sig", source_col])
+
+        overlap = (
+            tmp.group_by("_row_sig")
+            .agg(pl.col(source_col).n_unique().alias("source_count"))
+            .filter(pl.col("source_count") > 1)
+        )
+        overlap_count = len(overlap)
+        if overlap_count > 0:
+            report["leakage_detected"] = True
+            report["leakage_count"] = overlap_count
+            report["warnings"].append(f"Found {overlap_count} duplicate rows across multiple sources")
+
+    return report
+
+
+def _run_quality_engine(output_path: str) -> Dict[str, Any]:
+    # Reuse existing quality engine script (same folder)
+    try:
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        quality_script = os.path.join(script_dir, "quality_engine.py")
+        cmd = [sys.executable, quality_script, output_path]
+        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
+        if proc.returncode != 0:
+            return {"error": f"quality_engine failed: {proc.stderr.strip()}"}
+        return json.loads(proc.stdout)
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def _concat_polars(dfs, source_names: List[str]):
+    # Schema alignment by union columns
+    all_cols = []
+    col_set = set()
+    for df in dfs:
+        for col in df.columns:
+            if col not in col_set:
+                all_cols.append(col)
+                col_set.add(col)
+    if "_vesper_source" not in col_set:
+        all_cols.append("_vesper_source")
+
+    aligned = []
+    for i, df in enumerate(dfs):
+        tmp = df
+        if "_vesper_source" not in tmp.columns:
+            tmp = tmp.with_columns(pl.lit(source_names[i]).alias("_vesper_source"))
+        missing = [c for c in all_cols if c not in tmp.columns]
+        for m in missing:
+            tmp = tmp.with_columns(pl.lit(None).alias(m))
+        tmp = tmp.select(all_cols)
+        aligned.append(tmp)
+
+    return pl.concat(aligned, how="vertical_relaxed")
+
+
+def _join_polars(dfs, source_names: List[str], join_on: List[str], how: str):
+    if len(dfs) < 2:
+        raise ValueError("Join strategy requires at least 2 sources")
+
+    conflict_renames = []
+    left = dfs[0]
+    if "_vesper_source" not in left.columns:
+        left = left.with_columns(pl.lit(source_names[0]).alias("_vesper_source"))
+
+    for i in range(1, len(dfs)):
+        right = dfs[i]
+        if "_vesper_source" not in right.columns:
+            right = right.with_columns(pl.lit(source_names[i]).alias("_vesper_source"))
+
+        overlaps = [
+            c for c in right.columns
+            if c in left.columns and c not in join_on
+        ]
+        rename_map = {}
+        suffix = _safe_suffix(source_names[i], i)
+        for c in overlaps:
+            new_name = f"{c}_{suffix}"
+            rename_map[c] = new_name
+            conflict_renames.append({"source": source_names[i], "from": c, "to": new_name})
+
+        if rename_map:
+            right = right.rename(rename_map)
+
+        left = left.join(right, on=join_on, how=how, coalesce=True)
+
+    return left, conflict_renames
+
+
+def fuse_datasets(
+    sources: List[Union[str, Any]],
+    strategy: str = "concat",
+    join_on: Optional[Union[str, List[str]]] = None,
+    how: str = "inner",
+    dedup: bool = True,
+    run_quality_after: bool = True,
+    leakage_check: bool = True,
+    output_path: Optional[str] = None,
+    output_format: str = "feather",
+    compression: Optional[str] = "lz4",
+    preview: bool = True,
+    preview_rows: int = 500,
+    id_column: Optional[str] = None,
+    **kwargs,
+):
+    if not HAS_POLARS:
+        return {"error": "Polars is required for dataset fusion. Install with: pip install polars"}
+
+    if not sources or len(sources) < 2:
+        return {"error": "Need at least 2 sources to fuse"}
+
+    source_paths: List[str] = []
+    source_names: List[str] = []
+
+    for i, src in enumerate(sources):
+        if isinstance(src, str):
+            source_paths.append(src)
+            source_names.append(src)
+        elif isinstance(src, dict):
+            p = src.get("path") or src.get("local_path")
+            if not p:
+                return {"error": f"Source {i} missing path"}
+            source_paths.append(p)
+            source_names.append(src.get("name") or p)
+        else:
+            return {"error": f"Unsupported source type at index {i}"}
+
+    for p in source_paths:
+        if not os.path.exists(p):
+            return {"error": f"Source not found: {p}"}
+
+    if output_path is None:
+        ext_map = {
+            "feather": ".feather",
+            "parquet": ".parquet",
+            "csv": ".csv",
+            "jsonl": ".jsonl",
+            "arrow": ".arrow",
+            "ipc": ".arrow"
+        }
+        ext = ext_map.get(output_format, ".feather")
+        output_path = os.path.abspath(f"fused_dataset{ext}")
+
+    _ensure_parent(output_path)
+
+    try:
+        dfs = [_load_with_polars(p) for p in source_paths]
+    except Exception as e:
+        return {"error": f"Failed to load sources: {str(e)}"}
+
+    rows_before = sum(len(df) for df in dfs)
+    null_before = sum(_compute_null_ratio(df) for df in dfs) / len(dfs)
+
+    strategy = (strategy or "concat").lower()
+    how = (how or "inner").lower()
+    conflict_renames = []
+
+    try:
+        if strategy == "concat":
+            fused = _concat_polars(dfs, source_names)
+        elif strategy == "join":
+            if not join_on:
+                return {"error": "join_on is required when strategy='join'"}
+            join_keys = [join_on] if isinstance(join_on, str) else list(join_on)
+            for key in join_keys:
+                for idx, df in enumerate(dfs):
+                    if key not in df.columns:
+                        return {"error": f"Join key '{key}' missing in source {source_paths[idx]}"}
+            fused, conflict_renames = _join_polars(dfs, source_names, join_keys, how)
+        else:
+            return {"error": f"Unknown strategy: {strategy}. Use concat or join."}
+
+        duplicates_removed = 0
+        if dedup:
+            before = len(fused)
+            fused = fused.unique(maintain_order=True)
+            duplicates_removed = before - len(fused)
+
+        leakage = None
+        if leakage_check:
+            leakage = _leakage_report(fused, source_col="_vesper_source", id_col=id_column)
+
+        _write_with_polars(fused, output_path, output_format, compression)
+
+        preview_path = None
+        if preview:
+            preview_path = os.path.splitext(output_path)[0] + "_preview.csv"
+            fused.head(min(preview_rows, len(fused))).write_csv(preview_path)
+
+        quality_report = None
+        if run_quality_after:
+            quality_report = _run_quality_engine(output_path)
+
+        rows_after = len(fused)
+        null_after = _compute_null_ratio(fused)
+
+        return {
+            "success": True,
+            "output_path": output_path,
+            "preview_path": preview_path,
+            "stats": {
+                "sources_count": len(source_paths),
+                "rows_before": rows_before,
+                "rows_after": rows_after,
+                "columns_after": len(fused.columns),
+                "duplicates_removed": duplicates_removed,
+                "null_ratio_before": round(null_before, 3),
+                "null_ratio_after": round(null_after, 3),
+                "null_delta": round(null_after - null_before, 3),
+                "conflict_renames": conflict_renames,
+            },
+            "quality_report": quality_report,
+            "leakage_report": leakage,
+        }
+
+    except Exception as e:
+        return {"error": f"Fusion failed: {str(e)}"}
+
+
+def main():
+    if len(sys.argv) < 3:
+        print(json.dumps({
+            "error": "Usage: fusion_engine.py <sources_json> <output_path> [config_json]"
+        }))
+        sys.exit(1)
+
+    try:
+        sources = json.loads(sys.argv[1])
+        output_path = sys.argv[2]
+        config = {}
+        if len(sys.argv) > 3:
+            config = json.loads(sys.argv[3])
+
+        result = fuse_datasets(
+            sources=sources,
+            output_path=output_path,
+            strategy=config.get("strategy", "concat"),
+            join_on=config.get("join_on"),
+            how=config.get("how", "inner"),
+            dedup=config.get("dedup", True),
+            run_quality_after=config.get("run_quality_after", True),
+            leakage_check=config.get("leakage_check", True),
+            output_format=config.get("output_format", "feather"),
+            compression=config.get("compression", "lz4"),
+            preview=config.get("preview", True),
+            preview_rows=config.get("preview_rows", 500),
+            id_column=config.get("id_column"),
+        )
+
+        print(json.dumps(result))
+    except Exception as e:
+        print(json.dumps({"error": str(e)}))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

package/src/python/kaggle_engine.py

@@ -0,0 +1,204 @@
+import sys
+import os
+import json
+import tempfile
+from typing import Dict, Any, List
+from config import get_all
+
+try:
+    from kaggle.api.kaggle_api_extended import KaggleApi
+    HAS_KAGGLE = True
+except Exception:
+    HAS_KAGGLE = False
+
+
+def _ensure_auth() -> Dict[str, Any]:
+    if not HAS_KAGGLE:
+        return {
+            "ok": False,
+            "error": "kaggle package not installed. Install with: pip install kaggle",
+        }
+
+    # Priority:
+    # 1) Existing env vars
+    # 2) secure local store (keyring or ~/.vesper/config.toml)
+    # 3) ~/.kaggle/kaggle.json handled by KaggleApi.authenticate()
+    if not os.getenv("KAGGLE_USERNAME") or not os.getenv("KAGGLE_KEY"):
+        keys = get_all()
+        if keys.get("kaggle_username") and keys.get("kaggle_key"):
+            os.environ["KAGGLE_USERNAME"] = keys["kaggle_username"]
+            os.environ["KAGGLE_KEY"] = keys["kaggle_key"]
+
+    api = KaggleApi()
+    try:
+        api.authenticate()
+    except Exception as e:
+        return {
+            "ok": False,
+            "error": "Kaggle requires API key — run 'vespermcp config keys' (30 seconds) or provide ~/.kaggle/kaggle.json",
+            "details": str(e),
+        }
+
+    return {"ok": True, "api": api}
+
+
+def _dataset_to_dict(ds) -> Dict[str, Any]:
+    # kaggle API object fields differ by version; use getattr defensively
+    ref = getattr(ds, "ref", None) or getattr(ds, "datasetRef", None) or ""
+    title = getattr(ds, "title", None) or ref
+    subtitle = getattr(ds, "subtitle", None) or ""
+    owner = getattr(ds, "creatorName", None) or getattr(ds, "ownerName", None) or ""
+    votes = int(getattr(ds, "voteCount", 0) or 0)
+    downloads = int(getattr(ds, "downloadCount", 0) or 0)
+    size = int(getattr(ds, "totalBytes", 0) or 0)
+    last_updated = str(getattr(ds, "lastUpdated", ""))
+    tags = []
+    raw_tags = getattr(ds, "tags", None)
+    if raw_tags:
+        for t in raw_tags:
+            tags.append(getattr(t, "name", str(t)))
+
+    return {
+        "id": ref,
+        "source": "kaggle",
+        "name": title,
+        "description": subtitle or title,
+        "downloads": downloads,
+        "likes": votes,
+        "stars": 0,
+        "tags": tags,
+        "last_updated": last_updated,
+        "task": "unknown",
+        "domain": "unknown",
+        "languages": [],
+        "splits": [{"name": "data", "num_examples": 0, "size_bytes": size}],
+        "license": {
+            "id": "unknown",
+            "name": "unknown",
+            "category": "unknown",
+            "usage_restrictions": [],
+            "warnings": ["Kaggle license details may vary by dataset"],
+        },
+        "quality_score": 40,
+        "quality_warnings": ["Review dataset card and competition rules before use"],
+        "download_url": f"https://www.kaggle.com/datasets/{ref}",
+        "format": None,
+        "total_examples": 0,
+        "total_size_bytes": size,
+        "total_size_mb": round(size / (1024 * 1024), 2) if size else 0,
+        "columns": [],
+        "is_structured": False,
+        "has_target_column": False,
+        "is_safe_source": True,
+        "has_personal_data": False,
+        "is_paywalled": False,
+        "is_scraped_web_data": False,
+        "uses_https": True,
+        "has_train_split": False,
+        "has_test_split": False,
+        "has_validation_split": False,
+        "description_length": len(subtitle or title),
+        "has_readme": True,
+    }
+
+
+def discover(query: str, limit: int = 20) -> Dict[str, Any]:
+    auth = _ensure_auth()
+    if not auth.get("ok"):
+        return auth
+
+    api: KaggleApi = auth["api"]
+    try:
+        datasets = api.dataset_list(search=query, page_size=max(1, min(limit, 100)))
+        items = [_dataset_to_dict(ds) for ds in datasets[:limit]]
+        return {"ok": True, "results": items, "count": len(items)}
+    except Exception as e:
+        return {"ok": False, "error": f"Kaggle discover failed: {str(e)}"}
+
+
+def _pick_best_file(root: str) -> str:
+    candidates: List[str] = []
+    for base, _, files in os.walk(root):
+        for name in files:
+            full = os.path.join(base, name)
+            lower = name.lower()
+            if lower.endswith((".parquet", ".csv", ".jsonl", ".json", ".feather", ".arrow")):
+                candidates.append(full)
+
+    if not candidates:
+        raise RuntimeError("No suitable data file found after download")
+
+    # prioritize common tabular formats
+    priorities = [".parquet", ".csv", ".jsonl", ".json", ".feather", ".arrow"]
+    for ext in priorities:
+        for c in candidates:
+            if c.lower().endswith(ext):
+                return c
+    return candidates[0]
+
+
+def download(dataset_ref: str, target_dir: str) -> Dict[str, Any]:
+    auth = _ensure_auth()
+    if not auth.get("ok"):
+        return auth
+
+    api: KaggleApi = auth["api"]
+
+    if not target_dir:
+        target_dir = tempfile.mkdtemp(prefix="vesper_kaggle_")
+
+    os.makedirs(target_dir, exist_ok=True)
+
+    try:
+        if "kaggle.com/datasets/" in dataset_ref:
+            dataset_ref = dataset_ref.split("kaggle.com/datasets/")[1].lstrip("/")
+
+        # unzip in place, remove zip for convenience
+        api.dataset_download_files(dataset_ref, path=target_dir, unzip=True, quiet=True)
+        best_file = _pick_best_file(target_dir)
+        return {
+            "ok": True,
+            "dataset_id": dataset_ref,
+            "target_dir": target_dir,
+            "local_path": best_file,
+        }
+    except Exception as e:
+        msg = str(e)
+        if "401" in msg or "Unauthorized" in msg:
+            return {"ok": False, "error": "Invalid Kaggle credentials (401). Run 'vespermcp config kaggle' again."}
+        if "429" in msg or "Too Many Requests" in msg:
+            return {"ok": False, "error": "Kaggle rate limit reached. Please retry later."}
+        return {"ok": False, "error": f"Kaggle download failed: {msg}"}
+
+
+def main():
+    if len(sys.argv) < 2:
+        print(json.dumps({"ok": False, "error": "Usage: kaggle_engine.py <discover|download> ..."}))
+        sys.exit(1)
+
+    command = sys.argv[1]
+
+    if command == "discover":
+        if len(sys.argv) < 3:
+            print(json.dumps({"ok": False, "error": "Usage: kaggle_engine.py discover <query> [limit]"}))
+            sys.exit(1)
+        query = sys.argv[2]
+        limit = int(sys.argv[3]) if len(sys.argv) > 3 else 20
+        print(json.dumps(discover(query, limit)))
+        return
+
+    if command == "download":
+        if len(sys.argv) < 3:
+            print(json.dumps({"ok": False, "error": "Usage: kaggle_engine.py download <dataset_ref> [target_dir]"}))
+            sys.exit(1)
+        dataset_ref = sys.argv[2]
+        target_dir = sys.argv[3] if len(sys.argv) > 3 else ""
+        print(json.dumps(download(dataset_ref, target_dir)))
+        return
+
+    print(json.dumps({"ok": False, "error": f"Unknown command: {command}"}))
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

package/src/python/row_count.py

@@ -0,0 +1,54 @@
+import sys
+import json
+import os
+
+try:
+    import polars as pl
+except Exception:
+    print(json.dumps({"ok": False, "error": "polars is required"}))
+    sys.exit(1)
+
+
+def count_rows(path: str) -> int:
+    ext = os.path.splitext(path)[1].lower()
+
+    if ext == ".csv":
+        # Faster than full read for large csv
+        return int(pl.scan_csv(path, ignore_errors=True).select(pl.len()).collect().item())
+    if ext in [".parquet", ".pq"]:
+        return int(pl.scan_parquet(path).select(pl.len()).collect().item())
+    if ext in [".feather", ".ftr", ".arrow", ".ipc"]:
+        return int(pl.scan_ipc(path).select(pl.len()).collect().item())
+    if ext in [".jsonl", ".ndjson"]:
+        return int(pl.scan_ndjson(path).select(pl.len()).collect().item())
+    if ext == ".json":
+        # fallback to eager for plain JSON arrays
+        try:
+            return int(pl.read_json(path).height)
+        except Exception:
+            return int(pl.scan_ndjson(path).select(pl.len()).collect().item())
+
+    # unknown extension fallback
+    return int(pl.read_csv(path, ignore_errors=True).height)
+
+
+def main():
+    if len(sys.argv) < 2:
+        print(json.dumps({"ok": False, "error": "Usage: row_count.py <file_path>"}))
+        sys.exit(1)
+
+    p = sys.argv[1]
+    if not os.path.exists(p):
+        print(json.dumps({"ok": False, "error": f"File not found: {p}"}))
+        sys.exit(1)
+
+    try:
+        rows = count_rows(p)
+        print(json.dumps({"ok": True, "rows": rows}))
+    except Exception as e:
+        print(json.dumps({"ok": False, "error": str(e)}))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()