@vespermcp/mcp-server 1.2.4 → 1.2.6

package/build/python/vesper/core/asset_downloader.py

@@ -0,0 +1,415 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import mimetypes
+import os
+import shutil
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Dict, Iterable, List, Optional
+
+import aiohttp
+
+from vesper.core.download_recipe import get_download_recipe
+
+try:
+    import aiofiles
+except Exception:  # pragma: no cover
+    aiofiles = None
+
+try:
+    import webdataset as wds
+except Exception:  # pragma: no cover
+    wds = None
+
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif", ".tiff"}
+
+
+@dataclass
+class DownloadResult:
+    dataset_id: str
+    source: str
+    output_dir: str
+    downloaded_assets: int
+    failed_assets: int
+    errors_file: str
+    metadata_file: str
+    output_format: str
+
+
+class AssetDownloader:
+    def __init__(
+        self,
+        output_root: str,
+        workers: int = 8,
+        recipes_dir: Optional[str] = None,
+        progress_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None] | None]] = None,
+    ) -> None:
+        self.output_root = Path(output_root)
+        self.workers = max(1, min(workers, 32))
+        self.recipes_dir = recipes_dir
+        self.progress_callback = progress_callback
+
+    async def _emit(self, stage: str, payload: Dict[str, Any]) -> None:
+        if not self.progress_callback:
+            return
+        maybe = self.progress_callback(stage, payload)
+        if asyncio.iscoroutine(maybe):
+            await maybe
+
+    @staticmethod
+    def find_image_column(dataset: Any) -> Optional[str]:
+        features = getattr(dataset, "features", None)
+        if features:
+            for name, feature in features.items():
+                feature_name = feature.__class__.__name__.lower()
+                feature_repr = str(feature).lower()
+                if feature_name == "image" or "image(" in feature_repr:
+                    return str(name)
+                lower = str(name).lower()
+                if lower in {"image", "images", "img", "image_path", "image_url", "url"}:
+                    return str(name)
+
+        candidate_columns = ["image", "images", "img", "image_path", "image_url", "url", "file_name", "filepath"]
+        cols = getattr(dataset, "column_names", []) or []
+        for c in candidate_columns:
+            if c in cols:
+                return c
+        return None
+
+    async def download_assets(
+        self,
+        dataset_id: str,
+        source: Optional[str] = None,
+        repo_id: Optional[str] = None,
+        kaggle_ref: Optional[str] = None,
+        urls: Optional[List[str]] = None,
+        output_format: str = "webdataset",
+        max_items: Optional[int] = None,
+        image_column: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        recipe = get_download_recipe(dataset_id, self.recipes_dir)
+        if recipe:
+            source = source or recipe.get("source")
+            repo_id = repo_id or recipe.get("repo_id")
+            image_column = image_column or recipe.get("image_column")
+
+        source = (source or "").lower()
+        if source not in {"huggingface", "kaggle", "url"}:
+            raise ValueError("source must be one of: huggingface, kaggle, url")
+
+        # --- Validate imports and args BEFORE creating any directories ---
+        if source == "huggingface":
+            if not repo_id:
+                raise ValueError("repo_id is required for source=huggingface")
+            try:
+                from datasets import load_dataset as _ld  # noqa: F401
+            except Exception as e:
+                raise RuntimeError(
+                    f"datasets package is required for HuggingFace downloads. "
+                    f"Install with: pip install datasets. Details: {e}"
+                )
+        elif source == "kaggle":
+            ref = kaggle_ref or repo_id
+            if not ref:
+                raise ValueError("kaggle_ref is required for source=kaggle")
+            try:
+                from kaggle.api.kaggle_api_extended import KaggleApi as _Ka  # noqa: F401
+            except Exception as e:
+                raise RuntimeError(
+                    f"kaggle package is required for Kaggle downloads. "
+                    f"Install with: pip install kaggle. Details: {e}"
+                )
+        else:
+            if not urls:
+                raise ValueError("urls are required for source=url")
+
+        if output_format == "webdataset" and wds is None:
+            raise RuntimeError(
+                "webdataset package is required for webdataset output. "
+                "Install with: pip install webdataset"
+            )
+
+        # --- Now safe to create directories ---
+        dataset_dir = self.output_root / dataset_id.replace("/", "_").replace(":", "_")
+        images_dir = dataset_dir / "images"
+        dataset_dir.mkdir(parents=True, exist_ok=True)
+        images_dir.mkdir(parents=True, exist_ok=True)
+
+        errors_file = dataset_dir / "errors.jsonl"
+        metadata_file = dataset_dir / "metadata.jsonl"
+
+        try:
+            if source == "huggingface":
+                summary = await self._download_huggingface(repo_id, dataset_id, images_dir, metadata_file, errors_file, max_items, image_column)
+            elif source == "kaggle":
+                ref = kaggle_ref or repo_id
+                summary = await self._download_kaggle(ref, dataset_id, images_dir, metadata_file, errors_file, max_items)
+            else:
+                summary = await self._download_urls(urls, dataset_id, images_dir, metadata_file, errors_file, max_items)
+        except Exception:
+            # Clean up empty directories on failure so we don't leave ghost artifacts
+            if images_dir.exists() and not any(images_dir.iterdir()):
+                shutil.rmtree(dataset_dir, ignore_errors=True)
+            raise
+
+        if output_format == "webdataset":
+            await self._write_webdataset(dataset_dir, images_dir, metadata_file)
+        elif output_format == "parquet":
+            await self._write_parquet(dataset_dir, metadata_file)
+
+        result = DownloadResult(
+            dataset_id=dataset_id,
+            source=source,
+            output_dir=str(dataset_dir),
+            downloaded_assets=summary["downloaded"],
+            failed_assets=summary["failed"],
+            errors_file=str(errors_file),
+            metadata_file=str(metadata_file),
+            output_format=output_format,
+        )
+        return result.__dict__
+
+    async def _download_huggingface(
+        self,
+        repo_id: str,
+        dataset_id: str,
+        images_dir: Path,
+        metadata_file: Path,
+        errors_file: Path,
+        max_items: Optional[int],
+        image_column: Optional[str],
+    ) -> Dict[str, int]:
+        from datasets import load_dataset  # validated in download_assets()
+
+        await self._emit("start", {"source": "huggingface", "repo_id": repo_id})
+
+        try:
+            ds = load_dataset(repo_id, split="train")
+        except Exception:
+            dd = load_dataset(repo_id)
+            first_split = list(dd.keys())[0]
+            ds = dd[first_split]
+
+        col = image_column or self.find_image_column(ds)
+        if not col:
+            raise RuntimeError(
+                "No image column detected in HuggingFace dataset. Provide image_column or use fallback strategy with URL column."
+            )
+
+        total = len(ds) if hasattr(ds, "__len__") else 0
+        target = min(total, max_items) if max_items and total else (max_items or total or 0)
+
+        downloaded = 0
+        failed = 0
+
+        with metadata_file.open("w", encoding="utf-8") as mf, errors_file.open("w", encoding="utf-8") as ef:
+            for idx, row in enumerate(ds):
+                if max_items and idx >= max_items:
+                    break
+                try:
+                    out_name = f"{idx:08d}.jpg"
+                    out_path = images_dir / out_name
+                    self._save_image_value(row.get(col), out_path)
+
+                    record = {
+                        "dataset_id": dataset_id,
+                        "index": idx,
+                        "image_path": str(out_path),
+                        "source": "huggingface",
+                        "repo_id": repo_id,
+                    }
+                    mf.write(json.dumps(record, ensure_ascii=False) + "\n")
+                    downloaded += 1
+                    if downloaded % 50 == 0:
+                        await self._emit("progress", {"downloaded": downloaded, "failed": failed, "target": target})
+                except Exception as e:
+                    failed += 1
+                    ef.write(json.dumps({"index": idx, "error": str(e)}, ensure_ascii=False) + "\n")
+
+        await self._emit("done", {"downloaded": downloaded, "failed": failed})
+        return {"downloaded": downloaded, "failed": failed}
+
+    async def _download_kaggle(
+        self,
+        kaggle_ref: str,
+        dataset_id: str,
+        images_dir: Path,
+        metadata_file: Path,
+        errors_file: Path,
+        max_items: Optional[int],
+    ) -> Dict[str, int]:
+        from kaggle.api.kaggle_api_extended import KaggleApi  # validated in download_assets()
+
+        await self._emit("start", {"source": "kaggle", "dataset": kaggle_ref})
+
+        api = KaggleApi()
+        api.authenticate()
+
+        tmp_dir = Path(tempfile.mkdtemp(prefix="vesper_kaggle_assets_"))
+        downloaded = 0
+        failed = 0
+
+        try:
+            api.dataset_download_files(kaggle_ref, path=str(tmp_dir), unzip=True, quiet=True)
+            candidates = [p for p in tmp_dir.rglob("*") if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS]
+            if max_items:
+                candidates = candidates[:max_items]
+
+            with metadata_file.open("w", encoding="utf-8") as mf, errors_file.open("w", encoding="utf-8") as ef:
+                for idx, src_path in enumerate(candidates):
+                    try:
+                        out_name = f"{idx:08d}{src_path.suffix.lower()}"
+                        out_path = images_dir / out_name
+                        shutil.copy2(src_path, out_path)
+                        record = {
+                            "dataset_id": dataset_id,
+                            "index": idx,
+                            "image_path": str(out_path),
+                            "source": "kaggle",
+                            "repo_id": kaggle_ref,
+                        }
+                        mf.write(json.dumps(record, ensure_ascii=False) + "\n")
+                        downloaded += 1
+                    except Exception as e:
+                        failed += 1
+                        ef.write(json.dumps({"file": str(src_path), "error": str(e)}, ensure_ascii=False) + "\n")
+        finally:
+            shutil.rmtree(tmp_dir, ignore_errors=True)
+
+        await self._emit("done", {"downloaded": downloaded, "failed": failed})
+        return {"downloaded": downloaded, "failed": failed}
+
+    async def _download_urls(
+        self,
+        urls: List[str],
+        dataset_id: str,
+        images_dir: Path,
+        metadata_file: Path,
+        errors_file: Path,
+        max_items: Optional[int],
+    ) -> Dict[str, int]:
+        if aiofiles is None:
+            raise RuntimeError("aiofiles is required for URL downloads. Install with: pip install aiofiles")
+
+        selected = urls[:max_items] if max_items else urls
+        sem = asyncio.Semaphore(self.workers)
+
+        downloaded = 0
+        failed = 0
+        metadata_lock = asyncio.Lock()
+
+        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=180)) as session:
+            async def worker(idx: int, url: str) -> None:
+                nonlocal downloaded, failed
+                async with sem:
+                    try:
+                        local_path = await self._download_one_url(session, idx, url, images_dir)
+                        async with metadata_lock:
+                            async with aiofiles.open(metadata_file, "a", encoding="utf-8") as mf:
+                                await mf.write(json.dumps({
+                                    "dataset_id": dataset_id,
+                                    "index": idx,
+                                    "image_path": str(local_path),
+                                    "source": "url",
+                                    "url": url,
+                                }, ensure_ascii=False) + "\n")
+                        downloaded += 1
+                    except Exception as e:
+                        failed += 1
+                        async with metadata_lock:
+                            async with aiofiles.open(errors_file, "a", encoding="utf-8") as ef:
+                                await ef.write(json.dumps({"index": idx, "url": url, "error": str(e)}, ensure_ascii=False) + "\n")
+
+            tasks = [asyncio.create_task(worker(i, u)) for i, u in enumerate(selected)]
+            await asyncio.gather(*tasks)
+
+        await self._emit("done", {"downloaded": downloaded, "failed": failed})
+        return {"downloaded": downloaded, "failed": failed}
+
+    async def _download_one_url(self, session: aiohttp.ClientSession, idx: int, url: str, images_dir: Path) -> Path:
+        ext = Path(url.split("?")[0]).suffix.lower()
+        if ext not in IMAGE_EXTENSIONS:
+            ext = ".jpg"
+        out_path = images_dir / f"{idx:08d}{ext}"
+
+        existing_size = out_path.stat().st_size if out_path.exists() else 0
+        headers: Dict[str, str] = {}
+        if existing_size > 0:
+            headers["Range"] = f"bytes={existing_size}-"
+
+        async with session.get(url, headers=headers) as response:
+            if response.status not in (200, 206):
+                raise RuntimeError(f"HTTP {response.status}")
+
+            mode = "ab" if response.status == 206 and existing_size > 0 else "wb"
+            async with aiofiles.open(out_path, mode) as f:
+                async for chunk in response.content.iter_chunked(1024 * 256):
+                    await f.write(chunk)
+
+        return out_path
+
+    @staticmethod
+    def _save_image_value(value: Any, out_path: Path) -> None:
+        if value is None:
+            raise ValueError("empty image value")
+
+        if hasattr(value, "save"):
+            value.save(out_path)
+            return
+
+        if isinstance(value, dict):
+            if value.get("bytes"):
+                out_path.write_bytes(value["bytes"])
+                return
+            if value.get("path") and os.path.exists(value["path"]):
+                shutil.copy2(value["path"], out_path)
+                return
+            if value.get("url"):
+                raise ValueError("image URL requires URL downloader fallback")
+
+        if isinstance(value, str):
+            if os.path.exists(value):
+                shutil.copy2(value, out_path)
+                return
+            raise ValueError("string image value is not a local path")
+
+        raise ValueError(f"Unsupported image value type: {type(value)}")
+
+    async def _write_webdataset(self, dataset_dir: Path, images_dir: Path, metadata_file: Path) -> None:
+        if wds is None:
+            raise RuntimeError("webdataset package is required for webdataset output. Install with: pip install webdataset")
+
+        shard_pattern = str(dataset_dir / "shard-%06d.tar")
+        with metadata_file.open("r", encoding="utf-8") as mf, wds.ShardWriter(shard_pattern, maxcount=5000) as sink:
+            for line in mf:
+                row = json.loads(line)
+                image_path = Path(row["image_path"])
+                if not image_path.exists():
+                    continue
+                key = image_path.stem
+                ext = image_path.suffix.lstrip(".") or "jpg"
+                sample = {
+                    "__key__": key,
+                    ext: image_path.read_bytes(),
+                    "json": json.dumps(row, ensure_ascii=False).encode("utf-8"),
+                }
+                sink.write(sample)
+
+    async def _write_parquet(self, dataset_dir: Path, metadata_file: Path) -> None:
+        try:
+            import pyarrow as pa
+            import pyarrow.parquet as pq
+        except Exception as e:
+            raise RuntimeError(f"pyarrow is required for parquet output: {e}")
+
+        rows: List[Dict[str, Any]] = []
+        with metadata_file.open("r", encoding="utf-8") as mf:
+            for line in mf:
+                rows.append(json.loads(line))
+
+        table = pa.Table.from_pylist(rows)
+        pq.write_table(table, str(dataset_dir / "metadata.parquet"))

package/build/python/vesper/core/download_recipe.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+DEFAULT_RECIPES_DIR = Path.home() / ".vesper" / "recipes"
+
+
+@dataclass
+class DownloadRecipe:
+    dataset_id: str
+    source: str
+    repo_id: str
+    image_column: Optional[str]
+    download_method: str
+    requires_auth: bool
+    estimated_asset_size_gb: float
+    total_images: int
+    fallback_strategy: list[str]
+
+
+def _safe_name(value: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9._-]+", "_", value).strip("_") or "dataset"
+
+
+def build_download_recipe(dataset_info: Dict[str, Any]) -> Dict[str, Any]:
+    dataset_id = str(dataset_info.get("dataset_id") or dataset_info.get("id") or "")
+    source = str(dataset_info.get("source") or "unknown").lower()
+    repo_id = str(dataset_info.get("repo_id") or dataset_info.get("id") or dataset_id)
+
+    image_column = dataset_info.get("image_column")
+    if not image_column:
+        features = dataset_info.get("features") or {}
+        if isinstance(features, dict):
+            for key in features.keys():
+                lower = str(key).lower()
+                if lower in {"image", "images", "img", "img_path", "image_url", "url"}:
+                    image_column = key
+                    break
+
+    download_method = "url_list"
+    if source == "huggingface":
+        download_method = "hf_dataset_image_feature"
+    elif source == "kaggle":
+        download_method = "kaggle_archive"
+    elif source in {"dataworld", "openml"}:
+        download_method = "direct_file_scan"
+
+    requires_auth = bool(dataset_info.get("requires_auth", source in {"kaggle", "dataworld"}))
+
+    total_images = int(dataset_info.get("total_images") or dataset_info.get("total_examples") or 0)
+    if total_images <= 0:
+        total_images = 1000
+
+    estimated_asset_size_gb = round(float(dataset_info.get("estimated_asset_size_gb") or (total_images * 0.0004)), 3)
+
+    fallback_strategy = dataset_info.get("fallback_strategy") or [
+        "scan_archive_for_images",
+        "extract_url_column_and_download",
+        "export_metadata_only_with_actionable_error",
+    ]
+
+    recipe = DownloadRecipe(
+        dataset_id=dataset_id or repo_id,
+        source=source,
+        repo_id=repo_id,
+        image_column=image_column,
+        download_method=download_method,
+        requires_auth=requires_auth,
+        estimated_asset_size_gb=estimated_asset_size_gb,
+        total_images=total_images,
+        fallback_strategy=list(fallback_strategy),
+    )
+
+    return asdict(recipe)
+
+
+def save_recipe(recipe: Dict[str, Any], recipes_dir: Optional[str] = None) -> str:
+    root = Path(recipes_dir) if recipes_dir else DEFAULT_RECIPES_DIR
+    root.mkdir(parents=True, exist_ok=True)
+
+    dataset_id = str(recipe.get("dataset_id") or recipe.get("repo_id") or "dataset")
+    recipe_dir = root / _safe_name(dataset_id)
+    recipe_dir.mkdir(parents=True, exist_ok=True)
+
+    out_path = recipe_dir / "download_recipe.json"
+    out_path.write_text(json.dumps(recipe, indent=2, ensure_ascii=False), encoding="utf-8")
+    return str(out_path)
+
+
+def get_download_recipe(dataset_id: str, recipes_dir: Optional[str] = None) -> Optional[Dict[str, Any]]:
+    root = Path(recipes_dir) if recipes_dir else DEFAULT_RECIPES_DIR
+    path = root / _safe_name(dataset_id) / "download_recipe.json"
+    if not path.exists():
+        return None
+
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except Exception:
+        return None

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@vespermcp/mcp-server",
-  "version": "1.2.4",
+  "version": "1.2.6",
   "description": "AI-powered dataset discovery, quality analysis, and preparation MCP server with multimodal support (text, image, audio, video)",
   "type": "module",
   "main": "build/index.js",
@@ -17,7 +17,7 @@
     "mcp-config-template.json"
   ],
   "scripts": {
-    "build": "tsc && node -e \"const fs=require('fs');const path=require('path');const src='src/python';const dest='build/python';if(!fs.existsSync(dest))fs.mkdirSync(dest,{recursive:true});fs.readdirSync(src).forEach(f=>{if(f.endsWith('.py'))fs.copyFileSync(path.join(src,f),path.join(dest,f));});console.log('Copied Python scripts to build/python');\"",
+    "build": "tsc && node -e \"const fs=require('fs');const path=require('path');const src='src/python';const dest='build/python';const walk=(d)=>fs.readdirSync(d,{withFileTypes:true}).flatMap(e=>e.isDirectory()?walk(path.join(d,e.name)):[path.join(d,e.name)]);if(!fs.existsSync(dest))fs.mkdirSync(dest,{recursive:true});if(fs.existsSync(src)){for(const f of walk(src)){if(!f.endsWith('.py'))continue;const rel=path.relative(src,f);const out=path.join(dest,rel);fs.mkdirSync(path.dirname(out),{recursive:true});fs.copyFileSync(f,out);}}console.log('Copied Python scripts to build/python');\"",
     "dev": "tsx watch src/index.ts",
     "postinstall": "node scripts/postinstall.cjs",
     "scrape": "tsx src/scripts/scrape-metadata.ts",

package/scripts/postinstall.cjs

@@ -23,7 +23,12 @@ const pythonPackages = [
     'pillow',
     'numpy',
     'librosa',
-    'soundfile'
+    'soundfile',
+    'aiohttp',
+    'aiofiles',
+    'datasets',
+    'webdataset',
+    'kaggle'
 ];
 
 try {

package/src/python/asset_downloader_engine.py

@@ -0,0 +1,73 @@
+import argparse
+import asyncio
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any, Dict
+
+CURRENT_DIR = Path(__file__).resolve().parent
+if str(CURRENT_DIR) not in sys.path:
+    sys.path.insert(0, str(CURRENT_DIR))
+
+from vesper.core.asset_downloader import AssetDownloader
+from vesper.core.download_recipe import build_download_recipe, save_recipe, get_download_recipe
+
+
+def _print(payload: Dict[str, Any]) -> None:
+    print(json.dumps(payload, ensure_ascii=False))
+
+
+async def _run_download(args: argparse.Namespace) -> Dict[str, Any]:
+    payload = json.loads(args.payload)
+    output_root = payload.get("output_root") or str(Path.home() / ".vesper" / "data" / "assets")
+    workers = int(payload.get("workers") or 8)
+    recipes_dir = payload.get("recipes_dir")
+
+    downloader = AssetDownloader(output_root=output_root, workers=workers, recipes_dir=recipes_dir)
+
+    result = await downloader.download_assets(
+        dataset_id=str(payload.get("dataset_id")),
+        source=payload.get("source"),
+        repo_id=payload.get("repo_id"),
+        kaggle_ref=payload.get("kaggle_ref"),
+        urls=payload.get("urls"),
+        output_format=payload.get("output_format", "webdataset"),
+        max_items=payload.get("max_items"),
+        image_column=payload.get("image_column"),
+    )
+    return {"ok": True, "result": result}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Vesper Smart Asset Downloader Engine")
+    parser.add_argument("action", choices=["download", "build_recipe", "get_recipe"])
+    parser.add_argument("payload", help="JSON payload")
+    args = parser.parse_args()
+
+    try:
+        if args.action == "download":
+            response = asyncio.run(_run_download(args))
+            _print(response)
+            return
+
+        payload = json.loads(args.payload)
+        if args.action == "build_recipe":
+            recipe = build_download_recipe(payload)
+            saved = save_recipe(recipe, payload.get("recipes_dir"))
+            _print({"ok": True, "recipe": recipe, "saved_to": saved})
+            return
+
+        if args.action == "get_recipe":
+            dataset_id = str(payload.get("dataset_id"))
+            recipe = get_download_recipe(dataset_id, payload.get("recipes_dir"))
+            _print({"ok": True, "recipe": recipe})
+            return
+
+        _print({"ok": False, "error": f"Unknown action: {args.action}"})
+    except Exception as e:
+        _print({"ok": False, "error": str(e)})
+
+
+if __name__ == "__main__":
+    main()

package/src/python/vesper/__init__.py

@@ -0,0 +1 @@
+"""Vesper Python runtime package."""

package/src/python/vesper/core/__init__.py

@@ -0,0 +1 @@
+"""Core data engines for Vesper."""