Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -12,6 +12,7 @@ from datetime import datetime
 from datasets import Dataset, load_dataset, concatenate_datasets
 import os
 import logging
+from urllib.parse import urlparse
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -44,6 +45,22 @@ class StatusResponse(BaseModel):
     version: str
     generator_version: str
 
+# --- Add Model ID Normalization Helper --- 
+def _normalise_model_id(raw_id: str) -> str:
+    """
+    Accept either 'owner/model' or a full URL like
+    'https://huggingface.co/owner/model'. Return 'owner/model'.
+    """
+    if raw_id.startswith(("http://", "https://") ):
+        path = urlparse(raw_id).path.lstrip("/")
+        # path can contain extra segments (e.g. /commit/...), keep first two
+        parts = path.split("/")
+        if len(parts) >= 2:
+            return "/".join(parts[:2])
+        return path # Fallback if path doesn't have owner/model
+    return raw_id
+# --- End Model ID Normalization Helper ---
+
 # --- Add Counter Helper Functions ---
 def log_sbom_generation(model_id: str):
     """Logs a successful SBOM generation event to the Hugging Face dataset."""
@@ -52,10 +69,12 @@ def log_sbom_generation(model_id: str):
         return
 
     try:
+        # Normalize model_id before logging
+        normalized_model_id_for_log = _normalise_model_id(model_id) # added to normalize id
         log_data = {
             "timestamp": [datetime.utcnow().isoformat()],
             "event": ["generated"],
-            "model_id": [model_id] # Log the model ID
+            "model_id": [normalized_model_id_for_log] # use normalized_model_id_for_log
         }
         ds_new_log = Dataset.from_dict(log_data)
 
@@ -85,7 +104,7 @@ def log_sbom_generation(model_id: str):
         # Push the updated or new dataset
         # Corrected: Removed unnecessary backslash in it's
         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}")
+        logger.info(f"Successfully logged SBOM generation for {normalized_model_id_for_log} to {HF_REPO}") # use normalized model id
 
     except Exception as e:
         logger.error(f"Failed to log SBOM generation to {HF_REPO}: {e}")
@@ -351,6 +370,9 @@ async def generate_form(
     use_best_practices: bool = Form(True)
 ):
     sbom_count = get_sbom_count() # Get count early for context
+    # --- Normalize the model ID for display and filename --- 
+    normalized_model_id = _normalise_model_id(model_id) 
+    # --- End Normalization ---
     try:
         # Try different import paths for AIBOMGenerator
         generator = None
@@ -379,7 +401,8 @@ async def generate_form(
 
         # Save AIBOM to file
         # Corrected: Removed unnecessary backslashes around '/' and '_'
-        filename = f"{model_id.replace('/', '_')}_aibom.json"
+        # Save AIBOM to file using normalized ID
+        filename = f"{normalized_model_id.replace('/', '_')}_aibom.json"
         filepath = os.path.join(OUTPUT_DIR, filename)
 
         with open(filepath, "w") as f:
@@ -500,12 +523,12 @@ async def generate_form(
             "external_references": 10
         }
 
-        # Render the template with all necessary data
+        # Render the template with all necessary data, with normalized model ID
         return templates.TemplateResponse(
             "result.html",
             {
                 "request": request,
-                "model_id": model_id,
+                "model_id": normalized_model_id,
                 "aibom": aibom,
                 "enhancement_report": enhancement_report,
                 "completeness_score": completeness_score,
@@ -522,7 +545,7 @@ async def generate_form(
         # 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), "sbom_count": sbom_count} # Pass count
+            "error.html", {"request": request, "error": str(e), "sbom_count": sbom_count, "model_id": normalized_model_id} # Pass count, added normalized model ID
         )
 
 @app.get("/download/{filename}")