Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,9 +1,4 @@
-import os
-import json
-import logging
-from typing import Dict, List, Optional, Any
-from fastapi import FastAPI, HTTPException, Request, Form
-from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
+from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel
@@ -12,61 +7,46 @@ from pydantic import BaseModel
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Define templates directory
+# Define directories
 templates_dir = "templates"
+OUTPUT_DIR = "/tmp/aibom_output"
 
 # Initialize templates
 templates = Jinja2Templates(directory=templates_dir)
 
 # Create app
-app = FastAPI(title="AI SBOM Generator API - Simplified Version")
-
-# Define output directory for generated AIBOMs
-OUTPUT_DIR = "/tmp/aibom_output"
+app = FastAPI(title="AI SBOM Generator API")
 
-# Create output directory if it doesn't exist
+# Ensure output directory exists
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 # Mount output directory as static files
 app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
 
-# Define models
+# Status response model
 class StatusResponse(BaseModel):
     status: str
     version: str
     generator_version: str
 
-# Startup event to ensure directories exist
 @app.on_event("startup")
 async def startup_event():
-    """Create necessary directories on startup."""
     os.makedirs(OUTPUT_DIR, exist_ok=True)
-    logger.info(f"Created output directory at {OUTPUT_DIR}")
+    logger.info(f"Output directory ready at {OUTPUT_DIR}")
     logger.info(f"Registered routes: {[route.path for route in app.routes]}")
 
-# Root route serves the UI
 @app.get("/", response_class=HTMLResponse)
 async def root(request: Request):
-    """Serve the web UI interface as the default view."""
     try:
-        logger.info("Attempting to render index.html template")
         return templates.TemplateResponse("index.html", {"request": request})
     except Exception as e:
         logger.error(f"Error rendering template: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error in UI endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Template rendering error: {str(e)}")
 
-# Status endpoint
 @app.get("/status", response_model=StatusResponse)
 async def get_status():
-    """Get the API status and version information."""
-    logger.info("Status endpoint called")
-    return {
-        "status": "operational",
-        "version": "1.0.0",
-        "generator_version": "1.0.0",
-    }
-
-# Form-based generate endpoint for web UI
+    return StatusResponse(status="operational", version="1.0.0", generator_version="1.0.0")
+
 @app.post("/generate", response_class=HTMLResponse)
 async def generate_form(
     request: Request,
@@ -74,77 +54,46 @@ async def generate_form(
     include_inference: bool = Form(False),
     use_best_practices: bool = Form(True)
 ):
-    """Generate an AI SBOM from form data and render the result template."""
     try:
-        logger.info(f"Generate form endpoint called with model_id: {model_id}")
-        
-        # Import the generator here to avoid circular imports
-        try:
-            from src.aibom_generator.generator import AIBOMGenerator
-            logger.info("Imported AIBOMGenerator from src.aibom_generator.generator")
-        except ImportError as e:
-            logger.warning(f"Import from src.aibom_generator.generator failed: {e}")
-            try:
-                from aibom_generator.generator import AIBOMGenerator
-                logger.info("Imported AIBOMGenerator from aibom_generator.generator")
-            except ImportError as e:
-                logger.warning(f"Import from aibom_generator.generator failed: {e}")
-                try:
-                    from generator import AIBOMGenerator
-                    logger.info("Imported AIBOMGenerator from generator")
-                except ImportError as e:
-                    logger.error(f"All import attempts failed: {e}")
-                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
-        
-        # Create generator instance
-        logger.info("Creating AIBOMGenerator instance")
+        from src.aibom_generator.generator import AIBOMGenerator
+
         generator = AIBOMGenerator()
-        
-        # Generate AIBOM
-        logger.info(f"Generating AIBOM for model: {model_id}")
-        aibom, enhancement_report = generator.generate(
+
+        aibom = generator.generate_aibom(
             model_id=model_id,
             include_inference=include_inference,
             use_best_practices=use_best_practices
         )
-        logger.info("AIBOM generation successful")
-        
-        # Save AIBOM to file
+        enhancement_report = generator.get_enhancement_report()
+
         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)
-        
-        # Create download URL
+
         download_url = f"/output/{filename}"
-        
-        # Create download script
-        download_script = """
+
+        download_script = f"""
         <script>
-            function downloadJSON() {
+            function downloadJSON() {{
                 const a = document.createElement('a');
-                a.href = '%s';
-                a.download = '%s';
+                a.href = '{download_url}';
+                a.download = '{filename}';
                 document.body.appendChild(a);
                 a.click();
                 document.body.removeChild(a);
-            }
+            }}
         </script>
-        """ % (download_url, filename)
-        
-        # Get completeness score
+        """
+
         completeness_score = None
         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: {str(e)}")
-        
-        # Render result template
-        logger.info("Rendering result.html template")
+                logger.error(f"Completeness score error: {str(e)}")
+
         return templates.TemplateResponse(
             "result.html",
             {
@@ -158,34 +107,7 @@ async def generate_form(
             }
         )
     except Exception as e:
-        logger.error(f"Error generating AI SBOM: {str(e)}", exc_info=True)
+        logger.error(f"Error generating AI SBOM: {str(e)}")
         return templates.TemplateResponse(
-            "error.html",
-            {
-                "request": request,
-                "error_message": f"Error generating AI SBOM: {str(e)}"
-            }
+            "error.html", {"request": request, "error": str(e)}
         )
-
-# Download endpoint
-@app.get("/download/{filename}")
-async def download_file(filename: str):
-    """Download a generated AI SBOM file."""
-    try:
-        logger.info(f"Download endpoint called for file: {filename}")
-        filepath = os.path.join(OUTPUT_DIR, filename)
-        if not os.path.exists(filepath):
-            logger.error(f"File not found: {filepath}")
-            raise HTTPException(status_code=404, detail=f"File {filename} not found")
-        
-        logger.info(f"Returning file: {filepath}")
-        return FileResponse(
-            filepath,
-            media_type="application/json",
-            filename=filename
-        )
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Error downloading file: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Error downloading file: {str(e)}")