fix checking statistics capabilities

scripts/playground/metadata_tool_playground.py

@@ -25,7 +25,7 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-def test_merged_metadata(dataset_name = "rajpurkar/squad"):
+def test_merged_metadata(dataset_name = "wikitext", config_name = "wikitext-2-raw-v1"):
     """Test merged metadata retrieval."""
     print("=" * 80)
     print("Testing Merged Metadata from DatasetService")
@@ -43,7 +43,7 @@ def test_merged_metadata(dataset_name = "rajpurkar/squad"):
     # Test with squad dataset
     print(f"\n### Testing: {dataset_name} ###\n")
     try:
-        metadata = service.load_dataset_info(dataset_name)
+        metadata = service.load_dataset_info(dataset_name, config_name)
         
         print("Key Information:")
         print(f"  Dataset ID: {metadata.get('id')}")
@@ -124,4 +124,4 @@ def test_multi_config_dataset(dataset_name = "stanfordnlp/imdb"):
 
 if __name__ == "__main__":
     test_merged_metadata()
-    test_multi_config_dataset()
+    #test_multi_config_dataset()

src/hf_eda_mcp/services/dataset_service.py

@@ -808,7 +808,7 @@ class DatasetService:
     def _check_statistics_availability(
         self, 
         dataset_name: str,
-        config: Optional[str] = None
+        config_name: Optional[str] = None
     ) -> dict:
         """
         Check if statistics are available for a dataset.
@@ -818,7 +818,7 @@ class DatasetService:
         
         Args:
             dataset_name: HuggingFace dataset identifier
-            config: Optional configuration name
+            config_name: Optional configuration name
             
         Returns:
             Dictionary with availability information:
@@ -830,34 +830,40 @@ class DatasetService:
             DatasetViewerError: If the API request fails
         """
         try:
-            dataset_info = self.load_dataset_info(dataset_name, config)
+            dataset_info = self.load_dataset_info(dataset_name, config_name)
             full_dataset_id = dataset_info.get('id', dataset_name)
             
             if len(dataset_info["configs"]) == 1:
                 # Single config format
                 builder_name = dataset_info.get('builder_name', '')
                 is_parquet = builder_name == 'parquet'
-                
-                return {
-                    "available": is_parquet,
-                    "full_dataset_id": full_dataset_id,
-                    "configs": [dataset_info.get('config_name')] if is_parquet else [],
-                    "reason": 'Statistics available' if is_parquet else f'Statistics only available for parquet datasets (found: {builder_name})'
-                }
+                configs = [dataset_info["configs"][0]] if is_parquet else [],
+                reason = 'Statistics available' if is_parquet else f'Statistics only available for parquet datasets (found: {builder_name})'
+
             else:
                 # Multiple configs format
-                parquet_configs = []
-                for cfg_data in dataset_info["config_details"]:
-                    if cfg_data.get('builder_name') == 'parquet':
-                        parquet_configs.append(cfg_data.get("config_name"))
-                
-                return {
-                    "full_dataset_id": full_dataset_id,
-                    "available": len(parquet_configs) > 0,
-                    "configs": parquet_configs,
-                    "reason": f'Statistics available for {len(parquet_configs)} config(s)' if parquet_configs else 'No parquet configs found'
-                }
-                
+                if config_name is None:
+                    # Take every configs
+                    configs = []
+                    for cfg_data in dataset_info["config_details"]:
+                        if cfg_data.get('builder_name') == 'parquet':
+                            configs.append(cfg_data.get("config_name"))
+                    is_parquet = len(configs) > 0
+                    reason = f'Statistics available for {len(configs)} config(s)' if configs else 'No parquet configs found'
+
+                else:
+                    configs = [config_name]
+                    builder_name = dataset_info.get('builder_name', '')
+                    is_parquet = builder_name == 'parquet'
+                    reason = f'Statistics available for provided config {config_name}' if is_parquet else f'No parquet found for config {config_name}'
+
+            return {
+                "available": is_parquet,
+                "full_dataset_id": full_dataset_id,
+                "configs": configs,
+                "reason": reason
+            }
+
         except Exception as e:
             error_msg = f"Unexpected error checking statistics availability: {str(e)}"
             logger.error(error_msg)
@@ -867,8 +873,8 @@ class DatasetService:
 def get_dataset_service(hf_api_token: str) -> DatasetService:
     """Get or create the global dataset service instance using current config."""
     config = get_config()
-    #if hf_api_token is None or len(hf_api_token) == 0:
-    hf_api_token = config.hf_token
+    if hf_api_token is None or len(hf_api_token) == 0:
+        hf_api_token = config.hf_token
     dataset_service = DatasetService(
         cache_dir=config.cache_dir,
         token=hf_api_token