Upload test_generator.py

src/aibom_generator/test_generator.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+"""
+Simple test script to verify the AIBOM generator functionality without the API layer.
+This helps isolate whether the issue is with the generator or the API.
+"""
+
+import os
+import json
+import logging
+import sys
+from pathlib import Path
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Define output directory
+OUTPUT_DIR = "/tmp/aibom_output"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+def test_generator(model_id="meta-llama/Llama-2-7b-chat-hf"):
+    """Test the AIBOM generator functionality directly."""
+    logger.info(f"Testing AIBOM generator with model_id: {model_id}")
+    
+    # Try different import paths to find the generator
+    try:
+        logger.info("Trying to import from src.aibom_generator.generator")
+        from src.aibom_generator.generator import AIBOMGenerator
+        logger.info("Successfully imported from src.aibom_generator.generator")
+    except ImportError as e:
+        logger.warning(f"Import from src.aibom_generator.generator failed: {e}")
+        try:
+            logger.info("Trying to import from aibom_generator.generator")
+            from aibom_generator.generator import AIBOMGenerator
+            logger.info("Successfully imported from aibom_generator.generator")
+        except ImportError as e:
+            logger.warning(f"Import from aibom_generator.generator failed: {e}")
+            try:
+                logger.info("Trying to import from generator")
+                from generator import AIBOMGenerator
+                logger.info("Successfully imported from generator")
+            except ImportError as e:
+                logger.error(f"All import attempts failed: {e}")
+                logger.error("Could not import AIBOMGenerator. Please check your installation.")
+                return False
+    
+    try:
+        # Create generator instance
+        logger.info("Creating AIBOMGenerator instance")
+        generator = AIBOMGenerator()
+        
+        # Generate AIBOM
+        logger.info(f"Generating AIBOM for model: {model_id}")
+        aibom, enhancement_report = generator.generate(
+            model_id=model_id,
+            include_inference=True,
+            use_best_practices=True
+        )
+        
+        # Save AIBOM to file
+        filename = f"{model_id.replace('/', '_')}_aibom.json"
+        filepath = os.path.join(OUTPUT_DIR, filename)
+        logger.info(f"Saving AIBOM to: {filepath}")
+        with open(filepath, "w") as f:
+            json.dump(aibom, f, indent=2)
+        
+        # Print summary
+        logger.info("AIBOM generation successful!")
+        logger.info(f"AIBOM saved to: {filepath}")
+        
+        # Check if completeness score is available
+        if hasattr(generator, 'get_completeness_score'):
+            try:
+                logger.info("Getting completeness score")
+                completeness_score = generator.get_completeness_score(model_id)
+                logger.info(f"Completeness score: {completeness_score}")
+            except Exception as e:
+                logger.error(f"Error getting completeness score: {e}")
+        else:
+            logger.warning("get_completeness_score method not available")
+        
+        return True
+    except Exception as e:
+        logger.error(f"Error in test_generator: {e}", exc_info=True)
+        return False
+
+if __name__ == "__main__":
+    # Use command line argument for model_id if provided
+    model_id = sys.argv[1] if len(sys.argv) > 1 else "meta-llama/Llama-2-7b-chat-hf"
+    
+    logger.info("Starting AIBOM generator test")
+    success = test_generator(model_id)
+    
+    if success:
+        logger.info("Test completed successfully!")
+        sys.exit(0)
+    else:
+        logger.error("Test failed!")
+        sys.exit(1)