Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 import os
 import json
 import logging
@@ -7,11 +8,20 @@ from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel
+from datetime import datetime
+from datasets import Dataset, load_dataset, concatenate_datasets
+import os
+import logging
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# --- Add Counter Configuration ---
+HF_REPO = "aetheris-ai/aisbom-usage-log"  # User needs to create this private repo
+HF_TOKEN = os.getenv("HF_TOKEN")  # User must set this environment variable
+# --- End Counter Configuration ---
+
 # Define directories
 templates_dir = "templates"
 OUTPUT_DIR = "/tmp/aibom_output"
@@ -34,6 +44,68 @@ class StatusResponse(BaseModel):
     version: str
     generator_version: str
 
+# --- Add Counter Helper Functions ---
+def log_sbom_generation(model_id: str):
+    """Logs a successful SBOM generation event to the Hugging Face dataset."""
+    if not HF_TOKEN:
+        logger.warning("HF_TOKEN not set. Skipping SBOM generation logging.")
+        return
+
+    try:
+        log_data = {
+            "timestamp": [datetime.utcnow().isoformat()],
+            "event": ["generated"],
+            "model_id": [model_id] # Log the model ID
+        }
+        ds_new_log = Dataset.from_dict(log_data)
+
+        # Try to load existing dataset to append
+        try:
+            # Use trust_remote_code=True if required by the dataset/model on HF
+            existing_ds = load_dataset(HF_REPO, token=HF_TOKEN, split=\'train\', trust_remote_code=True)
+            # Check if dataset is empty or has different columns (handle initial creation)
+            if len(existing_ds) > 0 and set(existing_ds.column_names) == set(log_data.keys()):
+                 ds_to_push = concatenate_datasets([existing_ds, ds_new_log])
+            elif len(existing_ds) == 0:
+                 logger.info(f"Dataset {HF_REPO} is empty. Pushing initial data.")
+                 ds_to_push = ds_new_log
+            else:
+                 logger.warning(f"Dataset {HF_REPO} has unexpected columns {existing_ds.column_names} vs {list(log_data.keys())}. Appending new log anyway, structure might differ.")
+                 # Attempt concatenation even if columns differ slightly, HF might handle it
+                 # Or consider more robust schema migration/handling if needed
+                 ds_to_push = concatenate_datasets([existing_ds, ds_new_log])
+
+        except Exception as load_err:
+             # Handle case where dataset doesn\'t exist yet or other loading errors
+             logger.info(f"Could not load existing dataset {HF_REPO} (may not exist yet): {load_err}. Pushing new dataset.")
+             ds_to_push = ds_new_log # ds is already prepared with the new log entry
+
+        # Push the updated or new dataset
+        ds_to_push.push_to_hub(HF_REPO, token=HF_TOKEN, private=True) # Ensure it\'s private
+        logger.info(f"Successfully logged SBOM generation for {model_id} to {HF_REPO}")
+
+    except Exception as e:
+        logger.error(f"Failed to log SBOM generation to {HF_REPO}: {e}")
+
+def get_sbom_count() -> str:
+    """Retrieves the total count of generated SBOMs from the Hugging Face dataset."""
+    if not HF_TOKEN:
+        logger.warning("HF_TOKEN not set. Cannot retrieve SBOM count.")
+        return "N/A"
+    try:
+        # Load the dataset - assumes \'train\' split exists after first push
+        # Use trust_remote_code=True if required by the dataset/model on HF
+        ds = load_dataset(HF_REPO, token=HF_TOKEN, split=\'train\', trust_remote_code=True)
+        count = len(ds)
+        logger.info(f"Retrieved SBOM count: {count} from {HF_REPO}")
+        # Format count for display (e.g., add commas for large numbers)
+        return f"{count:,}"
+    except Exception as e:
+        logger.error(f"Failed to retrieve SBOM count from {HF_REPO}: {e}")
+        # Return "N/A" or similar indicator on error
+        return "N/A"
+# --- End Counter Helper Functions ---
+
 @app.on_event("startup")
 async def startup_event():
     os.makedirs(OUTPUT_DIR, exist_ok=True)
@@ -42,11 +114,18 @@ async def startup_event():
 
 @app.get("/", response_class=HTMLResponse)
 async def root(request: Request):
+    sbom_count = get_sbom_count() # Get count
     try:
-        return templates.TemplateResponse("index.html", {"request": request})
+        return templates.TemplateResponse("index.html", {"request": request, "sbom_count": sbom_count}) # Pass to template
     except Exception as e:
         logger.error(f"Error rendering template: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Template rendering error: {str(e)}")
+        # Attempt to render error page even if main page fails
+        try:
+            return templates.TemplateResponse("error.html", {"request": request, "error": f"Template rendering error: {str(e)}", "sbom_count": sbom_count})
+        except Exception as template_err:
+             # Fallback if error template also fails
+             logger.error(f"Error rendering error template: {template_err}")
+             raise HTTPException(status_code=500, detail=f"Template rendering error: {str(e)}")
 
 @app.get("/status", response_model=StatusResponse)
 async def get_status():
@@ -58,7 +137,7 @@ def import_utils():
     try:
         # Try different import paths
         sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-        
+
         # Try direct import first
         try:
             from utils import calculate_completeness_score
@@ -66,7 +145,7 @@ def import_utils():
             return calculate_completeness_score
         except ImportError:
             pass
-        
+
         # Try from src
         try:
             from src.aibom_generator.utils import calculate_completeness_score
@@ -74,7 +153,7 @@ def import_utils():
             return calculate_completeness_score
         except ImportError:
             pass
-        
+
         # Try from aibom_generator
         try:
             from aibom_generator.utils import calculate_completeness_score
@@ -82,7 +161,7 @@ def import_utils():
             return calculate_completeness_score
         except ImportError:
             pass
-        
+
         # If all imports fail, use the default implementation
         logger.warning("Could not import calculate_completeness_score, using default implementation")
         return None
@@ -106,7 +185,7 @@ def create_comprehensive_completeness_score(aibom=None):
             return calculate_completeness_score(aibom, validate=True, use_best_practices=True)
         except Exception as e:
             logger.error(f"Error calculating completeness score: {str(e)}")
-    
+
     # Otherwise, return a default comprehensive structure
     return {
         "total_score": 75.5,  # Default score for better UI display
@@ -134,7 +213,7 @@ def create_comprehensive_completeness_score(aibom=None):
             "metadata.tools": "✔ ★★",
             "metadata.authors": "✔ ★★",
             "metadata.component": "✔ ★★",
-            
+
             # Component basic info
             "component.type": "✔ ★★",
             "component.name": "✔ ★★★",
@@ -142,15 +221,15 @@ def create_comprehensive_completeness_score(aibom=None):
             "component.purl": "✔ ★★",
             "component.description": "✔ ★★",
             "component.licenses": "✔ ★★",
-            
+
             # Model card
             "modelCard.modelParameters": "✔ ★★",
             "modelCard.quantitativeAnalysis": "✘ ★★",
             "modelCard.considerations": "✔ ★★",
-            
+
             # External references
             "externalReferences": "✔ ★",
-            
+
             # Additional fields from FIELD_CLASSIFICATION
             "name": "✔ ★★★",
             "downloadLocation": "✔ ★★★",
@@ -184,7 +263,7 @@ def create_comprehensive_completeness_score(aibom=None):
             "metadata.tools": "important",
             "metadata.authors": "important",
             "metadata.component": "important",
-            
+
             # Component basic info
             "component.type": "important",
             "component.name": "critical",
@@ -192,15 +271,15 @@ def create_comprehensive_completeness_score(aibom=None):
             "component.purl": "important",
             "component.description": "important",
             "component.licenses": "important",
-            
+
             # Model card
             "modelCard.modelParameters": "important",
             "modelCard.quantitativeAnalysis": "important",
             "modelCard.considerations": "important",
-            
+
             # External references
             "externalReferences": "supplementary",
-            
+
             # Additional fields from FIELD_CLASSIFICATION
             "name": "critical",
             "downloadLocation": "critical",
@@ -227,8 +306,8 @@ def create_comprehensive_completeness_score(aibom=None):
         "missing_fields": {
             "critical": [],
             "important": ["modelCard.quantitativeAnalysis", "energyConsumption", "safetyRiskAssessment"],
-            "supplementary": ["modelExplainability", "standardCompliance", "energyQuantity", "energyUnit", 
-                             "metric", "metricDecisionThreshold", "modelDataPreprocessing", 
+            "supplementary": ["modelExplainability", "standardCompliance", "energyQuantity", "energyUnit",
+                             "metric", "metricDecisionThreshold", "modelDataPreprocessing",
                              "autonomyType", "useSensitivePersonalInformation"]
         },
         "completeness_profile": {
@@ -267,6 +346,7 @@ async def generate_form(
     include_inference: bool = Form(False),
     use_best_practices: bool = Form(True)
 ):
+    sbom_count = get_sbom_count() # Get count early for context
     try:
         # Try different import paths for AIBOMGenerator
         generator = None
@@ -294,53 +374,58 @@ async def generate_form(
         enhancement_report = generator.get_enhancement_report()
 
         # Save AIBOM to file
-        filename = f"{model_id.replace('/', '_')}_aibom.json"
+        filename = f"{model_id.replace(\'/\', \'_\')}_aibom.json"
         filepath = os.path.join(OUTPUT_DIR, filename)
 
         with open(filepath, "w") as f:
             json.dump(aibom, f, indent=2)
 
+        # --- Log Generation Event ---
+        log_sbom_generation(model_id)
+        sbom_count = get_sbom_count() # Refresh count after logging
+        # --- End Log ---
+
         download_url = f"/output/{filename}"
 
         # Create download and UI interaction scripts
         download_script = f"""
         <script>
             function downloadJSON() {{
-                const a = document.createElement('a');
-                a.href = '{download_url}';
-                a.download = '{filename}';
+                const a = document.createElement(\'a\');
+                a.href = \'{download_url}\';
+                a.download = \'{filename}\';
                 document.body.appendChild(a);
                 a.click();
                 document.body.removeChild(a);
             }}
-            
+
             function switchTab(tabId) {{
                 // Hide all tabs
-                document.querySelectorAll('.tab-content').forEach(tab => {{
-                    tab.classList.remove('active');
+                document.querySelectorAll(\'.tab-content\').forEach(tab => {{
+                    tab.classList.remove(\'active\');
                 }});
-                
+
                 // Deactivate all tab buttons
-                document.querySelectorAll('.aibom-tab').forEach(button => {{
-                    button.classList.remove('active');
+                document.querySelectorAll(\'.aibom-tab\').forEach(button => {{
+                    button.classList.remove(\'active\');
                 }});
-                
+
                 // Show the selected tab
-                document.getElementById(tabId).classList.add('active');
-                
+                document.getElementById(tabId).classList.add(\'active\');
+
                 // Activate the clicked button
-                event.currentTarget.classList.add('active');
+                event.currentTarget.classList.add(\'active\');
             }}
-            
+
             function toggleCollapsible(element) {{
-                element.classList.toggle('active');
+                element.classList.toggle(\'active\');
                 var content = element.nextElementSibling;
                 if (content.style.maxHeight) {{
                     content.style.maxHeight = null;
-                    content.classList.remove('active');
+                    content.classList.remove(\'active\');
                 }} else {{
                     content.style.maxHeight = content.scrollHeight + "px";
-                    content.classList.add('active');
+                    content.classList.add(\'active\');
                 }}
             }}
         </script>
@@ -348,15 +433,15 @@ async def generate_form(
 
         # Get completeness score or create a comprehensive one if not available
         completeness_score = None
-        if hasattr(generator, 'get_completeness_score'):
+        if hasattr(generator, \'get_completeness_score\'):
             try:
                 completeness_score = generator.get_completeness_score(model_id)
                 logger.info("Successfully retrieved completeness_score from generator")
             except Exception as e:
                 logger.error(f"Completeness score error from generator: {str(e)}")
-        
-        # If completeness_score is None or doesn't have field_checklist, use comprehensive one
-        if completeness_score is None or not isinstance(completeness_score, dict) or 'field_checklist' not in completeness_score:
+
+        # If completeness_score is None or doesn\'t have field_checklist, use comprehensive one
+        if completeness_score is None or not isinstance(completeness_score, dict) or \'field_checklist\' not in completeness_score:
             logger.info("Using comprehensive completeness_score with field_checklist")
             completeness_score = create_comprehensive_completeness_score(aibom)
 
@@ -375,7 +460,7 @@ async def generate_form(
                 enhancement_report["original_score"] = {"total_score": 0, "completeness_score": 0}
             elif "completeness_score" not in enhancement_report["original_score"]:
                 enhancement_report["original_score"]["completeness_score"] = enhancement_report["original_score"].get("total_score", 0)
-            
+
             # Ensure final_score has completeness_score
             if "final_score" not in enhancement_report or enhancement_report["final_score"] is None:
                 enhancement_report["final_score"] = {"total_score": 0, "completeness_score": 0}
@@ -390,15 +475,15 @@ async def generate_form(
             "component_model_card": "Model Card",
             "external_references": "External References"
         }
-        
+
         tooltips = {
-            "required_fields": "Basic required fields for a valid AI SBOM",
-            "metadata": "Information about the AI SBOM itself",
+            "required_fields": "Basic required fields for a valid AIBOM",
+            "metadata": "Information about the AIBOM itself",
             "component_basic": "Basic information about the AI model component",
             "component_model_card": "Detailed model card information",
             "external_references": "Links to external resources"
         }
-        
+
         weights = {
             "required_fields": 20,
             "metadata": 20,
@@ -420,28 +505,39 @@ async def generate_form(
                 "download_script": download_script,
                 "display_names": display_names,
                 "tooltips": tooltips,
-                "weights": weights
+                "weights": weights,
+                "sbom_count": sbom_count # Pass count
             }
         )
     except Exception as e:
         logger.error(f"Error generating AI SBOM: {str(e)}")
+        # Ensure count is passed to error template as well
+        sbom_count = get_sbom_count() # Refresh count just in case
         return templates.TemplateResponse(
-            "error.html", {"request": request, "error": str(e)}
+            "error.html", {"request": request, "error": str(e), "sbom_count": sbom_count} # Pass count
         )
 
 @app.get("/download/{filename}")
 async def download_file(filename: str):
     """
-    Download a generated AI SBOM file.
-    
-    This endpoint serves the generated AI SBOM JSON files for download.
+    Download a generated AIBOM file.
+
+    This endpoint serves the generated AIBOM JSON files for download.
     """
     file_path = os.path.join(OUTPUT_DIR, filename)
     if not os.path.exists(file_path):
         raise HTTPException(status_code=404, detail="File not found")
-    
+
     return FileResponse(
         file_path,
         media_type="application/json",
         filename=filename
     )
+
+# If running directly (for local testing)
+if __name__ == "__main__":
+    import uvicorn
+    # Ensure HF_TOKEN is set for local testing if needed
+    if not HF_TOKEN:
+        print("Warning: HF_TOKEN environment variable not set. SBOM count will show N/A and logging will be skipped.")
+    uvicorn.run(app, host="0.0.0.0", port=8000)