Improve metadata tool with config detailed information

pyproject.toml

@@ -40,8 +40,7 @@ where = ["src"]
 [tool.pdm.scripts]
 hf-eda-mcp = {cmd="python -m hf_eda_mcp", env_file= ".env"}
 hf_client_playground = "python -m scripts.playground.hf_client_playground"
-metadata_playground = "python -m scripts.playground.metadata_playground"
-test_merged_metadata = "python -m scripts.playground.test_merged_metadata"
+metadata_playground = "python -m scripts.playground.metadata_tool_playground"
 
 [tool.pdm]
 distribution = true

scripts/playground/metadata_tool_playground.py

@@ -0,0 +1,127 @@
+"""
+Test script to verify the merged metadata from Dataset Service.
+
+This script tests that the DatasetService properly merges data from both
+the Hub API and Dataset Viewer API.
+"""
+
+import os
+import logging
+from pprint import pprint
+from dotenv import load_dotenv
+from hf_eda_mcp.services.dataset_service import DatasetService
+
+load_dotenv()
+
+# Setup logging
+logging.basicConfig(
+    filename="scripts.log",
+    encoding='utf-8',
+    level=logging.DEBUG,
+    filemode="w",
+    format='%(asctime)s - %(levelname)s - %(message)s',
+)
+
+logger = logging.getLogger(__name__)
+
+
+def test_merged_metadata(dataset_name = "rajpurkar/squad"):
+    """Test merged metadata retrieval."""
+    print("=" * 80)
+    print("Testing Merged Metadata from DatasetService")
+    print("=" * 80)
+    
+    # Initialize service
+    service = DatasetService(
+        cache_dir="./cache",
+        token=os.environ.get("HF_TOKEN")
+    )
+    
+    # Clear cache to force fresh fetch
+    service.clear_cache(dataset_name)
+    
+    # Test with squad dataset
+    print(f"\n### Testing: {dataset_name} ###\n")
+    try:
+        metadata = service.load_dataset_info(dataset_name)
+        
+        print("Key Information:")
+        print(f"  Dataset ID: {metadata.get('id')}")
+        print(f"  Author: {metadata.get('author')}")
+        print(f"  Size (bytes): {metadata.get('size_bytes', 'N/A')}")
+        print(f"  Size (human): {metadata.get('size_human', 'N/A')}")
+        print(f"  Download Size: {metadata.get('download_size_human', 'N/A')}")
+        print(f"  Total Examples: {metadata.get('total_examples', 'N/A')}")
+        print(f"  Downloads: {metadata.get('downloads', 0):,}")
+        print(f"  Likes: {metadata.get('likes', 0)}")
+        
+        print("\nSplits:")
+        for split_name, split_info in metadata.get('splits', {}).items():
+            if isinstance(split_info, dict):
+                num_examples = split_info.get('num_examples', 'N/A')
+                num_bytes = split_info.get('num_bytes', 'N/A')
+                print(f"  {split_name}: {num_examples:,} examples, {num_bytes:,} bytes")
+            else:
+                print(f"  {split_name}: {split_info}")
+        
+        print("\nFeatures Schema:")
+        features = metadata.get('features', {})
+        if features:
+            for feature_name, feature_info in features.items():
+                print(f"  {feature_name}: {feature_info}")
+        else:
+            print("  No features available")
+        
+        print("\nSummary:")
+        print(f"  {metadata.get('summary', 'N/A')}")
+        
+        print("\n" + "=" * 80)
+        print("Full Metadata:")
+        print("=" * 80)
+        pprint(metadata, indent=2)
+        
+    except Exception as e:
+        print(f"\n✗ Error: {e}")
+        logger.exception("Failed to retrieve merged metadata")
+
+
+def test_multi_config_dataset(dataset_name = "stanfordnlp/imdb"):
+    """Test with a multi-config dataset."""
+    print("\n\n" + "=" * 80)
+    print("Testing Multi-Config Dataset: ")
+    print("=" * 80)
+    
+    service = DatasetService(
+        cache_dir="./cache",
+        token=os.environ.get("HF_TOKEN")
+    )
+    
+    # Clear cache
+    service.clear_cache(dataset_name)
+    
+    print(f"\n### Testing: {dataset_name} ###\n")
+    try:
+        metadata = service.load_dataset_info(dataset_name)
+        
+        print("Key Information:")
+        print(f"  Dataset ID: {metadata.get('id')}")
+        print(f"  Total Examples: {metadata.get('total_examples', 'N/A')}")
+        print(f"  Size (human): {metadata.get('size_human', 'N/A')}")
+        
+        print("\nSplits:")
+        for split_name, split_info in metadata.get('splits', {}).items():
+            if isinstance(split_info, dict):
+                num_examples = split_info.get('num_examples', 'N/A')
+                print(f"  {split_name}: {num_examples:,} examples")
+        
+        print("\nSummary:")
+        print(f"  {metadata.get('summary', 'N/A')}")
+        
+    except Exception as e:
+        print(f"\n✗ Error: {e}")
+        logger.exception("Failed to retrieve imdb metadata")
+
+
+if __name__ == "__main__":
+    test_merged_metadata()
+    test_multi_config_dataset()

src/hf_eda_mcp/integrations/hf_client.py

@@ -133,7 +133,6 @@ class HfClient:
                 "last_modified": dataset_info.last_modified.isoformat()
                 if dataset_info.last_modified
                 else None,
-                "size_bytes": getattr(dataset_info, "size_in_bytes", 0),
                 "configs": [],
                 "splits": {},
                 "features": {},

src/hf_eda_mcp/services/dataset_service.py

@@ -20,6 +20,7 @@ from hf_eda_mcp.integrations.hf_client import (
     AuthenticationError,
     NetworkError
 )
+from hf_eda_mcp.services.dataset_viewer_adapter import DatasetViewerAdapter
 from hf_eda_mcp.error_handling import (
     retry_with_backoff,
     RetryConfig,
@@ -65,6 +66,7 @@ class DatasetService:
             cache_ttl: Cache time-to-live in seconds (default: 1 hour)
         """
         self.hf_client = HfClient(token=token)
+        self.dataset_viewer = DatasetViewerAdapter(token=token)
         self.cache_ttl = cache_ttl
         
         # Set up cache directory
@@ -134,10 +136,219 @@ class DatasetService:
             logger.warning(f"Failed to load cache file {cache_file}: {e}")
             return None
 
+    def _merge_viewer_data(
+        self, 
+        hub_metadata: Dict[str, Any], 
+        viewer_data: Dict[str, Any],
+        config_name: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Merge Dataset Viewer API data into Hub metadata.
+        
+        Enriches the basic Hub metadata with detailed information from the
+        Dataset Viewer API including split sizes, features schema, and byte sizes.
+        
+        When no config is specified, returns detailed information for all configs.
+        
+        Args:
+            hub_metadata: Basic metadata from Hub API
+            viewer_data: Detailed data from Dataset Viewer API
+            config_name: Optional configuration name to extract
+            
+        Returns:
+            Merged metadata dictionary
+        """
+        merged = hub_metadata.copy()
+        
+        # Extract dataset_info from viewer response
+        dataset_info = viewer_data.get('dataset_info', {})
+        
+        if not dataset_info:
+            logger.warning("No dataset_info in viewer data")
+            return merged
+        
+        # Handle two response formats:
+        # 1. When config is specified in API call: dataset_info is the config data directly
+        # 2. When no config specified: dataset_info is a dict with config names as keys
+        
+        if isinstance(dataset_info, dict) and 'config_name' in dataset_info:
+            # Format 1: Single config data (config was specified in API call)
+            config_data = dataset_info
+            self._enrich_with_single_config(merged, config_data)
+        elif config_name:
+            # Format 2: Specific config requested
+            if config_name in dataset_info:
+                config_data = dataset_info[config_name]
+                self._enrich_with_single_config(merged, config_data)
+            else:
+                logger.warning(f"Config '{config_name}' not found in viewer data")
+                return merged
+        else:
+            # No config specified
+            if len(dataset_info) == 1:
+                # Only one config - use single config format for consistency
+                config_data = next(iter(dataset_info.values()))
+                self._enrich_with_single_config(merged, config_data)
+            else:
+                # Multiple configs - return all configs with detailed information
+                self._enrich_with_all_configs(merged, dataset_info)
+        
+        return merged
+
+    def _enrich_with_single_config(self, merged: Dict[str, Any], config_data: Dict[str, Any]) -> None:
+        """
+        Enrich metadata with a single config's data.
+        
+        Args:
+            merged: Metadata dictionary to enrich (modified in place)
+            config_data: Configuration data from Dataset Viewer API
+        """
+        # Enrich features with detailed schema from viewer
+        if 'features' in config_data:
+            merged['features'] = config_data['features']
+        
+        # Enrich splits with actual sizes
+        if 'splits' in config_data:
+            viewer_splits = config_data['splits']
+            enriched_splits = {}
+            
+            for split_name, split_info in viewer_splits.items():
+                enriched_splits[split_name] = {
+                    'num_examples': split_info.get('num_examples', 0),
+                    'num_bytes': split_info.get('num_bytes', 0)
+                }
+            
+            merged['splits'] = enriched_splits
+            merged['total_splits'] = len(enriched_splits)
+        
+        # Add dataset size information
+        if 'dataset_size' in config_data:
+            merged['dataset_size'] = config_data['dataset_size']
+            merged['size_bytes'] = config_data['dataset_size']
+            
+            # Update human-readable size
+            size_bytes = config_data['dataset_size']
+            if size_bytes > 0:
+                merged['size_human'] = self._format_bytes(size_bytes)
+        
+        if 'download_size' in config_data:
+            merged['download_size'] = config_data['download_size']
+            merged['download_size_human'] = self._format_bytes(config_data['download_size'])
+        
+        # Add builder and version info
+        if 'builder_name' in config_data:
+            merged['builder_name'] = config_data['builder_name']
+        
+        if 'version' in config_data:
+            merged['version'] = config_data['version']
+        
+        # Update summary with enriched information
+        if 'splits' in merged and merged['splits']:
+            total_examples = sum(s.get('num_examples', 0) for s in merged['splits'].values())
+            merged['total_examples'] = total_examples
+            
+            # Update summary string
+            split_names = ', '.join(merged['splits'].keys())
+            size_str = merged.get('size_human', 'Unknown')
+            merged['summary'] = (
+                f"Dataset: {merged['id']} | "
+                f"Author: {merged.get('author', 'Unknown')} | "
+                f"Size: {size_str} | "
+                f"Examples: {total_examples:,} | "
+                f"Downloads: {merged.get('downloads', 0):,} | "
+                f"Likes: {merged.get('likes', 0)} | "
+                f"Splits: {split_names}"
+            )
+
+    def _enrich_with_all_configs(self, merged: Dict[str, Any], dataset_info: Dict[str, Any]) -> None:
+        """
+        Enrich metadata with all configs' data.
+        
+        Creates a detailed 'config_details' list with information for each config.
+        
+        Args:
+            merged: Metadata dictionary to enrich (modified in place)
+            dataset_info: Dict mapping config names to their data
+        """
+        config_details = []
+        total_dataset_size = 0
+        total_download_size = 0
+        total_examples_all_configs = 0
+        
+        for cfg_name, cfg_data in dataset_info.items():
+            config_detail = {
+                'config_name': cfg_name,
+                'features': cfg_data.get('features', {}),
+                'splits': {},
+                'dataset_size': cfg_data.get('dataset_size', 0),
+                'download_size': cfg_data.get('download_size', 0),
+                'builder_name': cfg_data.get('builder_name', ''),
+                'version': cfg_data.get('version', {}),
+            }
+            
+            # Process splits for this config
+            if 'splits' in cfg_data:
+                for split_name, split_info in cfg_data['splits'].items():
+                    config_detail['splits'][split_name] = {
+                        'num_examples': split_info.get('num_examples', 0),
+                        'num_bytes': split_info.get('num_bytes', 0)
+                    }
+            
+            # Calculate totals for this config
+            config_total_examples = sum(
+                s.get('num_examples', 0) for s in config_detail['splits'].values()
+            )
+            config_detail['total_examples'] = config_total_examples
+            config_detail['dataset_size_human'] = self._format_bytes(config_detail['dataset_size'])
+            config_detail['download_size_human'] = self._format_bytes(config_detail['download_size'])
+            
+            config_details.append(config_detail)
+            
+            # Accumulate totals across all configs
+            total_dataset_size += config_detail['dataset_size']
+            total_download_size += config_detail['download_size']
+            total_examples_all_configs += config_total_examples
+        
+        # Add detailed config information
+        merged['config_details'] = config_details
+        
+        # Remove redundant top-level fields since they're in config_details
+        merged.pop('splits', None)
+        merged.pop('features', None)
+        
+        # Add aggregate information
+        merged['total_dataset_size'] = total_dataset_size
+        merged['total_dataset_size_human'] = self._format_bytes(total_dataset_size)
+        merged['total_download_size'] = total_download_size
+        merged['total_download_size_human'] = self._format_bytes(total_download_size)
+        merged['total_examples'] = total_examples_all_configs
+        
+        # Update summary for multi-config datasets
+        merged['summary'] = (
+            f"Dataset: {merged['id']} | "
+            f"Author: {merged.get('author', 'Unknown')} | "
+            f"Configs: {len(config_details)} | "
+            f"Total Size: {merged['total_dataset_size_human']} | "
+            f"Total Examples: {total_examples_all_configs:,} | "
+            f"Downloads: {merged.get('downloads', 0):,} | "
+            f"Likes: {merged.get('likes', 0)}"
+        )
+
+    def _format_bytes(self, size_bytes: int) -> str:
+        """Format bytes into human-readable string."""
+        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
+            if size_bytes < 1024.0:
+                return f"{size_bytes:.2f} {unit}"
+            size_bytes /= 1024.0
+        return f"{size_bytes:.2f} PB"
+
     def load_dataset_info(self, dataset_id: str, config_name: Optional[str] = None) -> Dict[str, Any]:
         """
         Load dataset information from HuggingFace Hub with caching.
         
+        Combines data from both the Hub API and Dataset Viewer API to provide
+        comprehensive metadata including split sizes, features schema, and more.
+        
         Includes automatic retry logic for transient failures and comprehensive
         error handling with helpful suggestions.
         
@@ -171,8 +382,21 @@ class DatasetService:
         # Fetch from HuggingFace Hub with retry logic
         try:
             logger.info(f"Fetching metadata for dataset: {dataset_id}")
+            
+            # Get basic metadata from Hub API
             metadata = self.hf_client.get_dataset_info(dataset_id, config_name)
             
+            # Try to enrich with Dataset Viewer API data
+            # Use the full dataset ID from the metadata response
+            try:
+                full_dataset_id = metadata.get('id', dataset_id)
+                viewer_data = self.dataset_viewer.get_dataset_information(full_dataset_id, config_name)
+                metadata = self._merge_viewer_data(metadata, viewer_data, config_name)
+                logger.debug("Successfully enriched metadata with Dataset Viewer API")
+            except Exception as e:
+                # Log but don't fail if viewer API fails - we still have basic metadata
+                logger.warning(f"Failed to fetch Dataset Viewer data, using basic metadata only: {e}")
+            
             # Add cache timestamp
             metadata['_cached_at'] = time.time()
             

src/hf_eda_mcp/services/dataset_viewer_adapter.py

@@ -0,0 +1,156 @@
+
+import os
+import logging
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetViewerError(Exception):
+    """Base exception for Dataset Viewer API errors."""
+    pass
+
+
+class DatasetViewerAdapter():
+    """
+    Uses the dataset Viewer API from HuggingFace. Implements several endpoints
+    Relevant docs: https://huggingface.co/docs/dataset-viewer/info
+    """
+    
+    def __init__(
+        self, 
+        token: Optional[str] = None,
+    ):
+        """
+        Initialize dataset service with optional caching and authentication.
+        
+        Args:
+            token: HuggingFace authentication token
+        """
+        if token:
+            self.token = token
+        else:
+            self.token = os.environ.get("HF_TOKEN")
+        self.base_url = "https://datasets-server.huggingface.co/"
+
+    def _api_get(self, route: str, params: dict, extra_headers: Optional[dict] = None) -> dict:
+        """
+        Make a GET request to the Dataset Viewer API with retry logic.
+        
+        Args:
+            route: API endpoint route
+            params: Query parameters
+            extra_headers: Additional headers to include
+            
+        Returns:
+            JSON response as dictionary
+            
+        Raises:
+            DatasetViewerError: If request fails after retries
+        """
+        headers = {"Authorization": f"Bearer {self.token}"}
+        if extra_headers:
+            headers.update(extra_headers)
+
+        retry_strategy = Retry(
+            total=3,
+            backoff_factor=1,
+            status_forcelist=[429, 500, 502, 503, 504],
+            allowed_methods=["GET"]
+        )
+
+        # Create session with retry adapter
+        session = requests.Session()
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        session.mount("https://", adapter)
+        
+        # Make the request
+        url = f"{self.base_url}{route}"
+        
+        try:
+            logger.debug(f"Making Dataset Viewer API request to {url} with params {params}")
+            response = session.get(url, params=params, headers=headers, timeout=30)
+            response.raise_for_status()
+            
+            result = response.json()
+            logger.debug("Dataset Viewer API request successful")
+            return result
+            
+        except requests.exceptions.HTTPError as e:
+            status_code = e.response.status_code if e.response else None
+            error_msg = f"Dataset Viewer API HTTP error (status {status_code}): {str(e)}"
+            logger.error(error_msg)
+            raise DatasetViewerError(error_msg) from e
+            
+        except requests.exceptions.Timeout as e:
+            error_msg = f"Dataset Viewer API request timed out: {str(e)}"
+            logger.error(error_msg)
+            raise DatasetViewerError(error_msg) from e
+            
+        except requests.exceptions.ConnectionError as e:
+            error_msg = f"Dataset Viewer API connection error: {str(e)}"
+            logger.error(error_msg)
+            raise DatasetViewerError(error_msg) from e
+            
+        except requests.exceptions.RequestException as e:
+            error_msg = f"Dataset Viewer API request failed: {str(e)}"
+            logger.error(error_msg)
+            raise DatasetViewerError(error_msg) from e
+            
+        except ValueError as e:
+            error_msg = f"Failed to parse Dataset Viewer API response: {str(e)}"
+            logger.error(error_msg)
+            raise DatasetViewerError(error_msg) from e
+        
+        finally:
+            session.close()
+
+    def get_dataset_information(self, dataset_name: str, config: Optional[str] = None) -> dict:
+        """
+        Get detailed dataset information from the Dataset Viewer API.
+        
+        Args:
+            dataset_name: HuggingFace dataset identifier
+            config: Optional configuration name
+            
+        Returns:
+            Dictionary containing detailed dataset information including:
+            - dataset_info: Per-config information with features, splits, sizes
+            - failed: List of failed operations
+            - partial: Whether response is partial
+            - pending: List of pending operations
+            
+        Raises:
+            DatasetViewerError: If the API request fails
+        """
+        params = {"dataset": dataset_name}
+        if config is not None:
+            params["config"] = config
+        
+        logger.info(f"Fetching dataset information from Viewer API: {dataset_name}")
+        
+        try:
+            result = self._api_get(
+                route="info",
+                params=params
+            )
+            
+            # Check for errors in response
+            if result.get('failed'):
+                logger.warning(f"Dataset Viewer API returned failures: {result['failed']}")
+            
+            if result.get('partial'):
+                logger.warning("Dataset Viewer API returned partial data")
+            
+            return result
+            
+        except DatasetViewerError:
+            # Re-raise with context
+            raise
+        except Exception as e:
+            error_msg = f"Unexpected error fetching dataset information: {str(e)}"
+            logger.error(error_msg)
+            raise DatasetViewerError(error_msg) from e

src/hf_eda_mcp/tools/metadata.py

@@ -109,20 +109,39 @@ def get_dataset_metadata(dataset_id: str, config_name: Optional[str] = None) ->
         if config_name:
             metadata['config_name'] = config_name
         
-        # Enhance metadata with additional computed fields
-        metadata['total_configs'] = len(metadata.get('configs', []))
-        metadata['total_splits'] = len(metadata.get('splits', {}))
-        metadata['has_multiple_configs'] = metadata['total_configs'] > 1
+        # Enhance metadata with additional computed fields (only if not already set)
+        if 'total_configs' not in metadata:
+            metadata['total_configs'] = len(metadata.get('configs', []))
         
-        # Format size for human readability
-        size_bytes = metadata.get('size_bytes', 0)
-        if size_bytes > 0:
-            metadata['size_human'] = _format_bytes(size_bytes)
-        else:
-            metadata['size_human'] = 'Unknown'
+        if 'total_splits' not in metadata:
+            # For multi-config datasets (with config_details), calculate total unique splits
+            if 'config_details' in metadata:
+                all_splits = set()
+                for config in metadata['config_details']:
+                    all_splits.update(config.get('splits', {}).keys())
+                metadata['total_splits'] = len(all_splits)
+            else:
+                # For single-config datasets, count splits at top level
+                metadata['total_splits'] = len(metadata.get('splits', {}))
+        
+        if 'has_multiple_configs' not in metadata:
+            metadata['has_multiple_configs'] = metadata.get('total_configs', 0) > 1
+        
+        # Format size for human readability (only if not already set by dataset_service)
+        if 'size_human' not in metadata:
+            # For multi-config datasets, use total_dataset_size_human if available
+            if 'config_details' in metadata and 'total_dataset_size_human' in metadata:
+                metadata['size_human'] = metadata['total_dataset_size_human']
+            else:
+                size_bytes = metadata.get('size_bytes', 0)
+                if size_bytes > 0:
+                    metadata['size_human'] = _format_bytes(size_bytes)
+                else:
+                    metadata['size_human'] = 'Unknown'
         
-        # Add summary information
-        metadata['summary'] = _generate_metadata_summary(metadata)
+        # Add summary information (only if not already set by dataset_service)
+        if 'summary' not in metadata:
+            metadata['summary'] = _generate_metadata_summary(metadata)
         
         logger.info(f"Successfully retrieved metadata for {dataset_id}")
         return metadata