Update src/aibom_generator/generator.py

src/aibom_generator/generator.py

@@ -1,7 +1,7 @@
 import json
 import uuid
 import datetime
-from typing import Dict, Optional, Any
+from typing import Dict, Optional, Any, Tuple
 
 from huggingface_hub import HfApi, ModelCard
 from .utils import calculate_completeness_score
@@ -25,17 +25,29 @@ class AIBOMGenerator:
         model_id: str,
         output_file: Optional[str] = None,
         include_inference: Optional[bool] = None,
-    ) -> Dict[str, Any]:
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
         use_inference = include_inference if include_inference is not None else self.use_inference
         model_info = self._fetch_model_info(model_id)
         model_card = self._fetch_model_card(model_id)
         aibom = self._create_aibom_structure(model_id, model_info, model_card, use_inference)
+        
+        # Calculate score after AIBOM is complete
+        score_report = calculate_completeness_score(aibom)
+        
+        # Add score to metadata properties
+        if "metadata" in aibom and not "properties" in aibom["metadata"]:
+            aibom["metadata"]["properties"] = []
+        
+        if "metadata" in aibom and "properties" in aibom["metadata"]:
+            aibom["metadata"]["properties"].append({"name": "aibom:quality-score", "value": str(score_report["total_score"])})
+            aibom["metadata"]["properties"].append({"name": "aibom:quality-breakdown", "value": json.dumps(score_report["section_scores"])})
+            aibom["metadata"]["properties"].append({"name": "aibom:max-scores", "value": json.dumps(score_report["max_scores"])})
 
         if output_file:
             with open(output_file, 'w') as f:
                 json.dump(aibom, f, indent=2)
 
-        return aibom
+        return aibom, score_report
 
     def _fetch_model_info(self, model_id: str) -> Dict[str, Any]:
         try:
@@ -125,16 +137,6 @@ class AIBOMGenerator:
         return {}
 
     def _create_metadata_section(self, model_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
-        aibom_stub = {
-            "bomFormat": "CycloneDX",
-            "specVersion": "1.6",
-            "serialNumber": f"urn:uuid:{str(uuid.uuid4())}",
-            "version": 1,
-            "metadata": {},
-            "components": [self._create_component_section(model_id, metadata)]
-        }
-        score_report = calculate_completeness_score(aibom_stub)
-
         timestamp = datetime.datetime.utcnow().isoformat() + "Z"
         tools = [{
             "vendor": "Aetheris AI",
@@ -162,9 +164,6 @@ class AIBOMGenerator:
                     value = json.dumps(value)
                 properties.append({"name": key, "value": str(value)})
 
-        properties.append({"name": "aibom:quality-score", "value": str(score_report["total_score"])})
-        properties.append({"name": "aibom:quality-breakdown", "value": json.dumps(score_report["section_scores"])})
-
         metadata_section = {
             "timestamp": timestamp,
             "tools": tools,
@@ -183,7 +182,7 @@ class AIBOMGenerator:
             "type": "machine-learning-model",
             "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}",
             "name": metadata.get("name", model_id.split("/")[-1]),
-            "purl": f"pkg:generic/{model_id.replace('/', '%2F')}"
+            "purl": f"pkg:huggingface/{model_id.replace('/', '/')}"
         }
 
         if "description" in metadata: