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, List
 
 from huggingface_hub import HfApi, ModelCard
 from .utils import calculate_completeness_score
@@ -19,6 +19,7 @@ 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,
@@ -26,82 +27,125 @@ class AIBOMGenerator:
         output_file: Optional[str] = None,
         include_inference: Optional[bool] = None,
     ) -> 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)
-        
-        # 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
+        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)
                     
-                    # 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 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 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
+            # 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"])})
+                
+                # 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": "aetheris-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": []
         }
 
-        return aibom, enhancement_report
+    def get_enhancement_report(self):
+        """Return the enhancement report from the last generate_aibom call"""
+        return self.enhancement_report
 
     def _fetch_model_info(self, model_id: str) -> Dict[str, Any]:
         try:
@@ -148,31 +192,37 @@ class AIBOMGenerator:
         metadata = {}
 
         if model_info:
-            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,
-            })
+            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}")
 
-        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
+        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}")
 
         metadata["ai:type"] = "Transformer"
         metadata["ai:task"] = metadata.get("pipeline_tag", "Text Generation")
@@ -198,43 +248,46 @@ 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"):
-            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
+        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
                             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
+                
+                # 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}")
         
         return enhanced_metadata
 
@@ -247,7 +300,7 @@ class AIBOMGenerator:
         }]
 
         authors = []
-        if "author" in metadata and metadata["author"]:
+        if metadata and "author" in metadata and metadata["author"]:
             authors.append({
                 "name": metadata["author"],
                 "url": f"https://huggingface.co/{metadata['author']}"
@@ -255,16 +308,20 @@ class AIBOMGenerator:
 
         component = {
             "type": "machine-learning-model",
-            "name": metadata.get("name", model_id.split("/")[-1]),
+            "name": metadata.get("name", model_id.split("/")[-1]) if metadata else model_id.split("/")[-1],
             "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}"
         }
 
         properties = []
-        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)})
+        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}")
 
         metadata_section = {
             "timestamp": timestamp,
@@ -282,30 +339,30 @@ class AIBOMGenerator:
     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:
+        if metadata and "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]),
+            "name": metadata.get("name", model_id.split("/")[-1]) if metadata else model_id.split("/")[-1],
             "purl": purl
         }
 
-        if "description" in metadata:
+        if metadata and "description" in metadata:
             component["description"] = metadata["description"]
 
-        if "commit" in metadata:
+        if metadata and "commit" in metadata:
             component["version"] = metadata["commit"]
 
-        if "license" in metadata:
+        if metadata and "license" in metadata:
             component["licenses"] = [{"license": {"id": metadata["license"]}}]
 
         external_refs = [{
             "type": "website",
             "url": f"https://huggingface.co/{model_id}"
         }]
-        if "commit_url" in metadata:
+        if metadata and "commit_url" in metadata:
             external_refs.append({
                 "type": "vcs",
                 "url": metadata["commit_url"]
@@ -318,18 +375,32 @@ class AIBOMGenerator:
 
     def _create_model_card_section(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
         model_card_section = {}
-        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
-
+        
+        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}")
+            
         return model_card_section