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, List
+from typing import Dict, Optional, Any
 
 from huggingface_hub import HfApi, ModelCard
 from .utils import calculate_completeness_score
@@ -19,7 +19,6 @@ class AIBOMGenerator:
         self.inference_model_url = inference_model_url
         self.use_inference = use_inference
         self.cache_dir = cache_dir
-        self.enhancement_report = None  # Store enhancement report as instance variable
 
     def generate_aibom(
         self,
@@ -27,124 +26,82 @@ class AIBOMGenerator:
         output_file: Optional[str] = None,
         include_inference: Optional[bool] = None,
     ) -> Dict[str, Any]:
-        try:
-            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)
-            
-            # Store original metadata before any AI enhancement
-            original_metadata = self._extract_structured_metadata(model_id, model_info, model_card)
-            
-            # Create initial AIBOM with original metadata
-            original_aibom = self._create_aibom_structure(model_id, original_metadata)
-            
-            # Calculate initial score
-            original_score = calculate_completeness_score(original_aibom)
-            
-            # Final metadata starts with original metadata
-            final_metadata = original_metadata.copy() if original_metadata else {}
-            
-            # Apply AI enhancement if requested
-            ai_enhanced = False
-            ai_model_name = None
-            
-            if use_inference and self.inference_model_url:
-                try:
-                    # Extract additional metadata using AI
-                    enhanced_metadata = self._extract_unstructured_metadata(model_card, model_id)
+        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)
+        
+        # Store original metadata before any AI enhancement
+        original_metadata = self._extract_structured_metadata(model_id, model_info, model_card)
+        
+        # Create initial AIBOM with original metadata
+        original_aibom = self._create_aibom_structure(model_id, original_metadata)
+        
+        # Calculate initial score
+        original_score = calculate_completeness_score(original_aibom)
+        
+        # Final metadata starts with original metadata
+        final_metadata = original_metadata.copy()
+        
+        # Apply AI enhancement if requested
+        ai_enhanced = False
+        ai_model_name = None
+        
+        if use_inference and self.inference_model_url:
+            try:
+                # Extract additional metadata using AI
+                enhanced_metadata = self._extract_unstructured_metadata(model_card, model_id)
+                
+                # If we got enhanced metadata, merge it with original
+                if enhanced_metadata:
+                    ai_enhanced = True
+                    ai_model_name = "BERT-base-uncased"  # Will be replaced with actual model name
                     
-                    # If we got enhanced metadata, merge it with original
-                    if enhanced_metadata:
-                        ai_enhanced = True
-                        ai_model_name = "BERT-base-uncased"  # Will be replaced with actual model name
-                        
-                        # Merge enhanced metadata with original (enhanced takes precedence)
-                        for key, value in enhanced_metadata.items():
-                            if value is not None and (key not in final_metadata or not final_metadata[key]):
-                                final_metadata[key] = value
-                except Exception as e:
-                    print(f"Error during AI enhancement: {e}")
-                    # Continue with original metadata if enhancement fails
-            
-            # Create final AIBOM with potentially enhanced metadata
-            aibom = self._create_aibom_structure(model_id, final_metadata)
-            
-            # Calculate final score
-            final_score = calculate_completeness_score(aibom)
-            
-            # Add score and enhancement info to metadata properties
-            if "metadata" in aibom and "properties" not in aibom["metadata"]:
-                aibom["metadata"]["properties"] = []
+                    # Merge enhanced metadata with original (enhanced takes precedence)
+                    for key, value in enhanced_metadata.items():
+                        if value is not None and (key not in final_metadata or not final_metadata[key]):
+                            final_metadata[key] = value
+            except Exception as e:
+                print(f"Error during AI enhancement: {e}")
+                # Continue with original metadata if enhancement fails
+        
+        # Create final AIBOM with potentially enhanced metadata
+        aibom = self._create_aibom_structure(model_id, final_metadata)
+        
+        # Calculate final score
+        final_score = calculate_completeness_score(aibom)
+        
+        # Add score and enhancement info to metadata properties
+        if "metadata" in aibom and "properties" not in aibom["metadata"]:
+            aibom["metadata"]["properties"] = []
+        
+        if "metadata" in aibom and "properties" in aibom["metadata"]:
+            # Add score information
+            aibom["metadata"]["properties"].append({"name": "aibom:quality-score", "value": str(final_score["total_score"])})
+            aibom["metadata"]["properties"].append({"name": "aibom:quality-breakdown", "value": json.dumps(final_score["section_scores"])})
+            aibom["metadata"]["properties"].append({"name": "aibom:max-scores", "value": json.dumps(final_score["max_scores"])})
             
-            if "metadata" in aibom and "properties" in aibom["metadata"]:
-                # Add score information
-                aibom["metadata"]["properties"].append({"name": "aibom:quality-score", "value": str(final_score["total_score"])})
-                aibom["metadata"]["properties"].append({"name": "aibom:quality-breakdown", "value": json.dumps(final_score["section_scores"])})
-                aibom["metadata"]["properties"].append({"name": "aibom:max-scores", "value": json.dumps(final_score["max_scores"])})
-                
-                # Add AI enhancement information
-                if ai_enhanced:
-                    aibom["metadata"]["properties"].append({"name": "aibom:ai-enhanced", "value": "true"})
-                    aibom["metadata"]["properties"].append({"name": "aibom:ai-model", "value": ai_model_name})
-                    aibom["metadata"]["properties"].append({"name": "aibom:original-score", "value": str(original_score["total_score"])})
-                    aibom["metadata"]["properties"].append({"name": "aibom:score-improvement", 
-                                                          "value": str(round(final_score["total_score"] - original_score["total_score"], 2))})
-
-            if output_file:
-                with open(output_file, 'w') as f:
-                    json.dump(aibom, f, indent=2)
-
-            # Create enhancement report for UI display and store as instance variable
-            self.enhancement_report = {
-                "ai_enhanced": ai_enhanced,
-                "ai_model": ai_model_name if ai_enhanced else None,
-                "original_score": original_score,
-                "final_score": final_score,
-                "improvement": round(final_score["total_score"] - original_score["total_score"], 2) if ai_enhanced else 0
-            }
-
-            # Return only the AIBOM to maintain compatibility with existing code
-            return aibom
-        except Exception as e:
-            print(f"Error generating AIBOM: {e}")
-            # Return a minimal valid AIBOM structure in case of error
-            return self._create_minimal_aibom(model_id)
-
-    def _create_minimal_aibom(self, model_id: str) -> Dict[str, Any]:
-        """Create a minimal valid AIBOM structure in case of errors"""
-        return {
-            "bomFormat": "CycloneDX",
-            "specVersion": "1.6",
-            "serialNumber": f"urn:uuid:{str(uuid.uuid4())}",
-            "version": 1,
-            "metadata": {
-                "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
-                "tools": [{
-                    "vendor": "Aetheris AI",
-                    "name": "aibom-generator",
-                    "version": "0.1.0"
-                }],
-                "component": {
-                    "type": "machine-learning-model",
-                    "name": model_id.split("/")[-1],
-                    "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}"
-                },
-                "properties": [
-                    {"name": "aibom:error", "value": "Error generating complete AIBOM"}
-                ]
-            },
-            "components": [{
-                "type": "machine-learning-model",
-                "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}",
-                "name": model_id.split("/")[-1],
-                "purl": f"pkg:huggingface/{model_id.replace('/', '/')}"
-            }],
-            "dependencies": []
+            # Add AI enhancement information
+            if ai_enhanced:
+                aibom["metadata"]["properties"].append({"name": "aibom:ai-enhanced", "value": "true"})
+                aibom["metadata"]["properties"].append({"name": "aibom:ai-model", "value": ai_model_name})
+                aibom["metadata"]["properties"].append({"name": "aibom:original-score", "value": str(original_score["total_score"])})
+                aibom["metadata"]["properties"].append({"name": "aibom:score-improvement", 
+                                                      "value": str(round(final_score["total_score"] - original_score["total_score"], 2))})
+
+        if output_file:
+            with open(output_file, 'w') as f:
+                json.dump(aibom, f, indent=2)
+
+        # Create enhancement report for UI display
+        enhancement_report = {
+            "ai_enhanced": ai_enhanced,
+            "ai_model": ai_model_name if ai_enhanced else None,
+            "original_score": original_score,
+            "final_score": final_score,
+            "improvement": round(final_score["total_score"] - original_score["total_score"], 2) if ai_enhanced else 0
         }
 
-    def get_enhancement_report(self):
-        """Return the enhancement report from the last generate_aibom call"""
-        return self.enhancement_report
+        return aibom, enhancement_report
 
     def _fetch_model_info(self, model_id: str) -> Dict[str, Any]:
         try:
@@ -191,37 +148,31 @@ class AIBOMGenerator:
         metadata = {}
 
         if model_info:
-            try:
-                metadata.update({
-                    "name": model_info.modelId.split("/")[-1] if hasattr(model_info, "modelId") else model_id.split("/")[-1],
-                    "author": model_info.author if hasattr(model_info, "author") else None,
-                    "tags": model_info.tags if hasattr(model_info, "tags") else [],
-                    "pipeline_tag": model_info.pipeline_tag if hasattr(model_info, "pipeline_tag") else None,
-                    "downloads": model_info.downloads if hasattr(model_info, "downloads") else 0,
-                    "last_modified": model_info.lastModified if hasattr(model_info, "lastModified") else None,
-                    "commit": model_info.sha[:7] if hasattr(model_info, "sha") and model_info.sha else None,
-                    "commit_url": f"https://huggingface.co/{model_id}/commit/{model_info.sha}" if hasattr(model_info, "sha") and model_info.sha else None,
-                })
-            except Exception as e:
-                print(f"Error extracting model info metadata: {e}")
+            metadata.update({
+                "name": model_info.modelId.split("/")[-1] if hasattr(model_info, "modelId") else model_id.split("/")[-1],
+                "author": model_info.author if hasattr(model_info, "author") else None,
+                "tags": model_info.tags if hasattr(model_info, "tags") else [],
+                "pipeline_tag": model_info.pipeline_tag if hasattr(model_info, "pipeline_tag") else None,
+                "downloads": model_info.downloads if hasattr(model_info, "downloads") else 0,
+                "last_modified": model_info.lastModified if hasattr(model_info, "lastModified") else None,
+                "commit": model_info.sha[:7] if hasattr(model_info, "sha") and model_info.sha else None,
+                "commit_url": f"https://huggingface.co/{model_id}/commit/{model_info.sha}" if hasattr(model_info, "sha") and model_info.sha else None,
+            })
 
-        if model_card and hasattr(model_card, "data") and model_card.data:
-            try:
-                card_data = model_card.data.to_dict() if hasattr(model_card.data, "to_dict") else {}
-                metadata.update({
-                    "language": card_data.get("language"),
-                    "license": card_data.get("license"),
-                    "library_name": card_data.get("library_name"),
-                    "base_model": card_data.get("base_model"),
-                    "datasets": card_data.get("datasets"),
-                    "model_name": card_data.get("model_name"),
-                    "tags": card_data.get("tags", metadata.get("tags", [])),
-                    "description": card_data.get("model_summary", None)
-                })
-                if hasattr(model_card.data, "eval_results") and model_card.data.eval_results:
-                    metadata["eval_results"] = model_card.data.eval_results
-            except Exception as e:
-                print(f"Error extracting model card metadata: {e}")
+        if model_card and model_card.data:
+            card_data = model_card.data.to_dict() if hasattr(model_card.data, "to_dict") else {}
+            metadata.update({
+                "language": card_data.get("language"),
+                "license": card_data.get("license"),
+                "library_name": card_data.get("library_name"),
+                "base_model": card_data.get("base_model"),
+                "datasets": card_data.get("datasets"),
+                "model_name": card_data.get("model_name"),
+                "tags": card_data.get("tags", metadata.get("tags", [])),
+                "description": card_data.get("model_summary", None)
+            })
+            if hasattr(model_card.data, "eval_results") and model_card.data.eval_results:
+                metadata["eval_results"] = model_card.data.eval_results
 
         metadata["ai:type"] = "Transformer"
         metadata["ai:task"] = metadata.get("pipeline_tag", "Text Generation")
@@ -247,46 +198,43 @@ class AIBOMGenerator:
         # Since we can't install the required libraries due to space constraints,
         # we'll simulate the enhancement with a placeholder implementation
         
-        if model_card and hasattr(model_card, "text") and model_card.text:
-            try:
-                card_text = model_card.text
-                
-                # Simulate BERT extraction with basic text analysis
-                # In reality, this would be done with NLP models
-                
-                # Extract description if missing
-                if card_text and "description" not in enhanced_metadata:
-                    # Take first paragraph that's longer than 20 chars as description
-                    paragraphs = [p.strip() for p in card_text.split('\n\n')]
-                    for p in paragraphs:
-                        if len(p) > 20 and not p.startswith('#'):
-                            enhanced_metadata["description"] = p
+        if model_card and hasattr(model_card, "text"):
+            card_text = model_card.text
+            
+            # Simulate BERT extraction with basic text analysis
+            # In reality, this would be done with NLP models
+            
+            # Extract description if missing
+            if card_text and "description" not in enhanced_metadata:
+                # Take first paragraph that's longer than 20 chars as description
+                paragraphs = [p.strip() for p in card_text.split('\n\n')]
+                for p in paragraphs:
+                    if len(p) > 20 and not p.startswith('#'):
+                        enhanced_metadata["description"] = p
+                        break
+            
+            # Extract limitations if present
+            if "limitations" not in enhanced_metadata:
+                if "## Limitations" in card_text:
+                    limitations_section = card_text.split("## Limitations")[1].split("##")[0].strip()
+                    if limitations_section:
+                        enhanced_metadata["limitations"] = limitations_section
+            
+            # Extract ethical considerations if present
+            if "ethical_considerations" not in enhanced_metadata:
+                for heading in ["## Ethical Considerations", "## Ethics", "## Bias"]:
+                    if heading in card_text:
+                        section = card_text.split(heading)[1].split("##")[0].strip()
+                        if section:
+                            enhanced_metadata["ethical_considerations"] = section
                             break
-                
-                # Extract limitations if present
-                if "limitations" not in enhanced_metadata:
-                    if "## Limitations" in card_text:
-                        limitations_section = card_text.split("## Limitations")[1].split("##")[0].strip()
-                        if limitations_section:
-                            enhanced_metadata["limitations"] = limitations_section
-                
-                # Extract ethical considerations if present
-                if "ethical_considerations" not in enhanced_metadata:
-                    for heading in ["## Ethical Considerations", "## Ethics", "## Bias"]:
-                        if heading in card_text:
-                            section = card_text.split(heading)[1].split("##")[0].strip()
-                            if section:
-                                enhanced_metadata["ethical_considerations"] = section
-                                break
-                
-                # Extract risks if present
-                if "risks" not in enhanced_metadata:
-                    if "## Risks" in card_text:
-                        risks_section = card_text.split("## Risks")[1].split("##")[0].strip()
-                        if risks_section:
-                            enhanced_metadata["risks"] = risks_section
-            except Exception as e:
-                print(f"Error extracting unstructured metadata: {e}")
+            
+            # Extract risks if present
+            if "risks" not in enhanced_metadata:
+                if "## Risks" in card_text:
+                    risks_section = card_text.split("## Risks")[1].split("##")[0].strip()
+                    if risks_section:
+                        enhanced_metadata["risks"] = risks_section
         
         return enhanced_metadata
 
@@ -294,12 +242,12 @@ class AIBOMGenerator:
         timestamp = datetime.datetime.utcnow().isoformat() + "Z"
         tools = [{
             "vendor": "Aetheris AI",
-            "name": "aibom-generator",
+            "name": "aetheris-aibom-generator",
             "version": "0.1.0"
         }]
 
         authors = []
-        if metadata and "author" in metadata and metadata["author"]:
+        if "author" in metadata and metadata["author"]:
             authors.append({
                 "name": metadata["author"],
                 "url": f"https://huggingface.co/{metadata['author']}"
@@ -307,20 +255,16 @@ class AIBOMGenerator:
 
         component = {
             "type": "machine-learning-model",
-            "name": metadata.get("name", model_id.split("/")[-1]) if metadata else model_id.split("/")[-1],
+            "name": metadata.get("name", model_id.split("/")[-1]),
             "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}"
         }
 
         properties = []
-        if metadata:
-            for key, value in metadata.items():
-                if key not in ["name", "author", "license"] and value is not None:
-                    try:
-                        if isinstance(value, (list, dict)):
-                            value = json.dumps(value)
-                        properties.append({"name": key, "value": str(value)})
-                    except Exception as e:
-                        print(f"Error processing metadata property {key}: {e}")
+        for key, value in metadata.items():
+            if key not in ["name", "author", "license"] and value is not None:
+                if isinstance(value, (list, dict)):
+                    value = json.dumps(value)
+                properties.append({"name": key, "value": str(value)})
 
         metadata_section = {
             "timestamp": timestamp,
@@ -336,27 +280,32 @@ class AIBOMGenerator:
         return metadata_section
 
     def _create_component_section(self, model_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
+        # Create PURL with version information if commit is available
+        purl = f"pkg:huggingface/{model_id.replace('/', '/')}"
+        if "commit" in metadata:
+            purl = f"{purl}@{metadata['commit']}"
+            
         component = {
             "type": "machine-learning-model",
             "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}",
-            "name": metadata.get("name", model_id.split("/")[-1]) if metadata else model_id.split("/")[-1],
-            "purl": f"pkg:huggingface/{model_id.replace('/', '/')}"
+            "name": metadata.get("name", model_id.split("/")[-1]),
+            "purl": purl
         }
 
-        if metadata and "description" in metadata:
+        if "description" in metadata:
             component["description"] = metadata["description"]
 
-        if metadata and "commit" in metadata:
+        if "commit" in metadata:
             component["version"] = metadata["commit"]
 
-        if metadata and "license" in metadata:
+        if "license" in metadata:
             component["licenses"] = [{"license": {"id": metadata["license"]}}]
 
         external_refs = [{
             "type": "website",
             "url": f"https://huggingface.co/{model_id}"
         }]
-        if metadata and "commit_url" in metadata:
+        if "commit_url" in metadata:
             external_refs.append({
                 "type": "vcs",
                 "url": metadata["commit_url"]
@@ -369,32 +318,18 @@ class AIBOMGenerator:
 
     def _create_model_card_section(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
         model_card_section = {}
-        
-        if not metadata:
-            return model_card_section
-            
-        try:
-            # Safely extract model parameters
-            model_parameters = {}
-            for k in ["base_model", "library_name", "pipeline_tag"]:
-                if k in metadata:
-                    model_parameters[k] = metadata[k]
-                    
-            if model_parameters:
-                model_card_section["modelParameters"] = model_parameters
-
-            # Safely extract evaluation results
-            if "eval_results" in metadata:
-                model_card_section["quantitativeAnalysis"] = {"performanceMetrics": metadata["eval_results"]}
-
-            # Safely extract considerations
-            considerations = {}
-            for k in ["limitations", "ethical_considerations", "bias", "risks"]:
-                if k in metadata:
-                    considerations[k] = metadata[k]
-            if considerations:
-                model_card_section["considerations"] = considerations
-        except Exception as e:
-            print(f"Error creating model card section: {e}")
-            
+        model_parameters = {k: metadata[k] for k in ["base_model", "library_name", "pipeline_tag"] if k in metadata}
+        if model_parameters:
+            model_card_section["modelParameters"] = model_parameters
+
+        if "eval_results" in metadata:
+            model_card_section["quantitativeAnalysis"] = {"performanceMetrics": metadata["eval_results"]}
+
+        considerations = {}
+        for k in ["limitations", "ethical_considerations", "bias", "risks"]:
+            if k in metadata:
+                considerations[k] = metadata[k]
+        if considerations:
+            model_card_section["considerations"] = considerations
+
         return model_card_section