Update src/aibom_generator/utils.py

src/aibom_generator/utils.py

@@ -1,5 +1,5 @@
 """
-Utility functions for the AIBOM Generator.
+Utility functions for the AIBOM Generator with restored field_checklist support.
 """
 
 import json
@@ -13,7 +13,6 @@ logger = logging.getLogger(__name__)
 
 
 def setup_logging(level=logging.INFO):
-    """Set up logging configuration."""
     logging.basicConfig(
         level=level,
         format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
@@ -22,28 +21,16 @@ def setup_logging(level=logging.INFO):
 
 
 def ensure_directory(directory_path):
-    """Ensure that a directory exists, creating it if necessary."""
     if not os.path.exists(directory_path):
         os.makedirs(directory_path)
     return directory_path
 
 
 def generate_uuid():
-    """Generate a UUID for the AIBOM serialNumber."""
     return str(uuid.uuid4())
 
 
 def normalize_license_id(license_text):
-    """
-    Normalize a license string to a SPDX license identifier if possible.
-    
-    Args:
-        license_text: The license text to normalize
-        
-    Returns:
-        SPDX license identifier or the original text if no match
-    """
-    # Common license mappings
     license_mappings = {
         "mit": "MIT",
         "apache": "Apache-2.0",
@@ -82,169 +69,133 @@ def normalize_license_id(license_text):
         "proprietary": "NONE",
         "commercial": "NONE",
     }
-    
+
     if not license_text:
         return None
-    
-    # Normalize to lowercase and remove punctuation
+
     normalized = re.sub(r'[^\w\s-]', '', license_text.lower())
-    
-    # Check for direct matches
+
     if normalized in license_mappings:
         return license_mappings[normalized]
-    
-    # Check for partial matches
+
     for key, value in license_mappings.items():
         if key in normalized:
             return value
-    
-    # Return original if no match
+
     return license_text
 
 
+def validate_spdx(license_entry):
+    spdx_licenses = [
+        "MIT", "Apache-2.0", "GPL-3.0-only", "GPL-2.0-only", "LGPL-3.0-only",
+        "BSD-3-Clause", "BSD-2-Clause", "CC-BY-4.0", "CC-BY-SA-4.0", "CC0-1.0",
+        "Unlicense", "NONE"
+    ]
+    if isinstance(license_entry, list):
+        return all(lic in spdx_licenses for lic in license_entry)
+    return license_entry in spdx_licenses
+
+
 def calculate_completeness_score(aibom: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Calculate a completeness score for the AIBOM.
-
-    Args:
-        aibom: The AIBOM dictionary
-
-    Returns:
-        Dictionary containing:
-            - total_score: overall completeness score (0-100)
-            - section_scores: points earned per section
-            - field_checklist: dictionary showing presence (✔) or absence (✘) of key fields
-    """
-    score = 0
-    max_score = 100
-    section_scores = {}
     field_checklist = {}
 
-    # Define scoring weights for different sections
-    weights = {
-        "required_fields": 20,
-        "metadata": 20,
-        "component_basic": 20,
-        "component_model_card": 30,
-        "external_references": 10,
-    }
-
-    # Required Fields
     required_fields = ["bomFormat", "specVersion", "serialNumber", "version"]
-    required_present = [field for field in required_fields if field in aibom]
-    required_score = (len(required_present) / len(required_fields)) * weights["required_fields"]
-    section_scores["required_fields"] = round(required_score)
+    required_score = sum([5 if aibom.get(field) else 0 for field in required_fields])
     for field in required_fields:
-        field_checklist[field] = "✔" if field in required_present else "✘"
-
-    # Metadata Fields
-    metadata_score = 0
-    if "metadata" in aibom:
-        metadata_fields = ["timestamp", "tools", "authors", "component"]
-        present = [field for field in metadata_fields if field in aibom["metadata"]]
-        metadata_score = (len(present) / len(metadata_fields)) * weights["metadata"]
-        for field in metadata_fields:
-            field_checklist[f"metadata.{field}"] = "✔" if field in present else "✘"
-    section_scores["metadata"] = round(metadata_score)
-
-    # Component Basic Info
-    component_score = 0
-    component = aibom.get("components", [{}])[0]
-    component_fields = ["type", "name", "bom-ref", "purl", "description", "licenses"]
-    present = [field for field in component_fields if field in component]
-    component_score = (len(present) / len(component_fields)) * weights["component_basic"]
-    section_scores["component_basic"] = round(component_score)
-    for field in component_fields:
-        field_checklist[f"component.{field}"] = "✔" if field in present else "✘"
-
-    # Model Card Section
-    model_card_score = 0
-    model_card_fields = ["modelParameters", "quantitativeAnalysis", "considerations"]
-    if "modelCard" in component:
-        model_card = component["modelCard"]
-        present = [field for field in model_card_fields if field in model_card]
-        model_card_score = (len(present) / len(model_card_fields)) * weights["component_model_card"]
-        for field in model_card_fields:
-            field_checklist[f"modelCard.{field}"] = "✔" if field in present else "✘"
-    else:
-        for field in model_card_fields:
-            field_checklist[f"modelCard.{field}"] = "✘"
-    section_scores["component_model_card"] = round(model_card_score)
-
-    # External References
-    ext_score = weights["external_references"] if aibom.get("externalReferences") else 0
-    section_scores["external_references"] = round(ext_score)
-    field_checklist["externalReferences"] = "✔" if ext_score else "✘"
-
-    # Final total score
-    total_score = round(sum(section_scores.values()))
+        field_checklist[field] = "✔" if aibom.get(field) else "✘"
+
+    metadata = aibom.get("metadata", {})
+    metadata_fields = ["timestamp", "tools", "authors", "component"]
+    metadata_score = sum([5 if metadata.get(field) else 0 for field in metadata_fields])
+    for field in metadata_fields:
+        field_checklist[f"metadata.{field}"] = "✔" if metadata.get(field) else "✘"
+
+    components = aibom.get("components", [])
+    component_scores = []
+    model_card_scores = []
+
+    for comp in components:
+        comp_fields = ["type", "name", "bom-ref", "purl", "description", "licenses"]
+        comp_score = sum([
+            2 if comp.get("type") else 0,
+            4 if comp.get("name") else 0,
+            2 if comp.get("bom-ref") else 0,
+            4 if comp.get("purl") and re.match(r'^pkg:huggingface/.+', comp["purl"]) else 0,
+            4 if comp.get("description") and len(comp["description"]) > 20 else 0,
+            4 if comp.get("licenses") and validate_spdx(comp["licenses"]) else 0
+        ])
+        component_scores.append(comp_score)
+        for field in comp_fields:
+            field_checklist[f"component.{field}"] = "✔" if comp.get(field) else "✘"
+
+        card = comp.get("modelCard", {})
+        card_fields = ["modelParameters", "quantitativeAnalysis", "considerations"]
+        card_score = sum([
+            10 if card.get("modelParameters") else 0,
+            10 if card.get("quantitativeAnalysis") else 0,
+            10 if card.get("considerations") and len(card["considerations"]) > 50 else 0
+        ])
+        model_card_scores.append(card_score)
+        for field in card_fields:
+            field_checklist[f"modelCard.{field}"] = "✔" if field in card else "✘"
+
+    avg_comp_score = (sum(component_scores) / len(component_scores)) if component_scores else 0
+    avg_model_card_score = (sum(model_card_scores) / len(model_card_scores)) if model_card_scores else 0
+
+    ext_refs = aibom.get("externalReferences", [])
+    ext_score = 0
+    for ref in ext_refs:
+        url = ref.get("url", "").lower()
+        if "modelcard" in url:
+            ext_score += 4
+        elif "huggingface.co" in url or "github.com" in url:
+            ext_score += 3
+        elif "dataset" in url:
+            ext_score += 3
+    ext_score = min(ext_score, 10)
+    field_checklist["externalReferences"] = "✔" if ext_refs else "✘"
+
+    total_score = (
+        (required_score * 0.20) +
+        (metadata_score * 0.20) +
+        (avg_comp_score * 0.20) +
+        (avg_model_card_score * 0.30) +
+        (ext_score * 0.10)
+    )
 
     return {
-        "total_score": total_score,
-        "section_scores": section_scores,
+        "total_score": round(total_score, 2),
+        "section_scores": {
+            "required_fields": required_score,
+            "metadata": metadata_score,
+            "component_basic": avg_comp_score,
+            "component_model_card": avg_model_card_score,
+            "external_references": ext_score
+        },
         "field_checklist": field_checklist
     }
 
+
 def merge_metadata(primary: Dict[str, Any], secondary: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Merge two metadata dictionaries, giving priority to the primary dictionary.
-    
-    Args:
-        primary: Primary metadata dictionary
-        secondary: Secondary metadata dictionary
-        
-    Returns:
-        Merged metadata dictionary
-    """
     result = secondary.copy()
-    
     for key, value in primary.items():
         if value is not None:
             if key in result and isinstance(value, dict) and isinstance(result[key], dict):
                 result[key] = merge_metadata(value, result[key])
             else:
                 result[key] = value
-    
     return result
 
 
 def extract_model_id_parts(model_id: str) -> Dict[str, str]:
-    """
-    Extract parts from a Hugging Face model ID.
-    
-    Args:
-        model_id: Hugging Face model ID (e.g., "google/bert-base-uncased")
-        
-    Returns:
-        Dictionary with parts (owner, name)
-    """
     parts = model_id.split("/")
-    
     if len(parts) == 1:
-        return {
-            "owner": None,
-            "name": parts[0],
-        }
-    else:
-        return {
-            "owner": parts[0],
-            "name": "/".join(parts[1:]),
-        }
+        return {"owner": None, "name": parts[0]}
+    return {"owner": parts[0], "name": "/".join(parts[1:])}
 
 
 def create_purl(model_id: str) -> str:
-    """
-    Create a Package URL (purl) for a Hugging Face model.
-    
-    Args:
-        model_id: Hugging Face model ID
-        
-    Returns:
-        Package URL string
-    """
     parts = extract_model_id_parts(model_id)
-    
     if parts["owner"]:
         return f"pkg:huggingface/{parts['owner']}/{parts['name']}"
-    else:
-        return f"pkg:huggingface/{parts['name']}"
+    return f"pkg:huggingface/{parts['name']}"