Update src/aibom_generator/generator.py

src/aibom_generator/generator.py

@@ -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,
@@ -92,8 +93,8 @@ class AIBOMGenerator:
             with open(output_file, 'w') as f:
                 json.dump(aibom, f, indent=2)
 
-        # Create enhancement report for UI display
-        enhancement_report = {
+        # 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,
@@ -101,7 +102,12 @@ class AIBOMGenerator:
             "improvement": round(final_score["total_score"] - original_score["total_score"], 2) if ai_enhanced else 0
         }
 
-        return aibom, enhancement_report
+        # Return only the AIBOM to maintain compatibility with existing code
+        return aibom
+
+    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: