Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,15 +1,12 @@
 import os
 import json
 import logging
-from typing import Dict, List, Optional, Any, Union
-from datetime import datetime
-from pathlib import Path
-
-from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks, Query, File, UploadFile, Form, Request
+from typing import Dict, List, Optional, Any
+from fastapi import FastAPI, HTTPException, Request, Form
 from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -18,11 +15,11 @@ logger = logging.getLogger(__name__)
 # Define templates directory
 templates_dir = "templates"
 
-# Initialize templates with a simple path
+# Initialize templates
 templates = Jinja2Templates(directory=templates_dir)
 
 # Create app
-app = FastAPI(title="AI SBOM Generator API")
+app = FastAPI(title="AI SBOM Generator API - Simplified Version")
 
 # Define output directory for generated AIBOMs
 OUTPUT_DIR = "/tmp/aibom_output"
@@ -34,90 +31,35 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
 
 # Define models
-class GenerateRequest(BaseModel):
-    model_id: str
-    include_inference: bool = False
-    use_best_practices: bool = True
-
-class GenerateWithReportRequest(BaseModel):
-    model_id: str
-    include_inference: bool = False
-    use_best_practices: bool = True
-
-class BatchGenerateRequest(BaseModel):
-    model_ids: List[str]
-    include_inference: bool = False
-    use_best_practices: bool = True
-
-class ModelScoreRequest(BaseModel):
-    model_id: str
-
 class StatusResponse(BaseModel):
     status: str
     version: str
     generator_version: str
 
-class BatchJobResponse(BaseModel):
-    job_id: str
-    status: str
-    model_ids: List[str]
-    message: 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"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:
-        # Check if templates directory exists
-        if os.path.exists(templates_dir):
-            logger.info(f"Templates directory exists")
-            logger.info(f"Templates directory contents: {os.listdir(templates_dir)}")
-        else:
-            logger.warning(f"Templates directory does not exist: {templates_dir}")
-        
-        # Try to render the template
-        try:
-            return templates.TemplateResponse("index.html", {"request": request})
-        except Exception as e:
-            logger.error(f"Error rendering template: {str(e)}")
-            
-            # Try direct file reading as fallback
-            try:
-                with open(os.path.join(templates_dir, "index.html"), "r") as f:
-                    html_content = f.read()
-                return HTMLResponse(content=html_content)
-            except Exception as file_e:
-                logger.error(f"Error reading file directly: {str(file_e)}")
-                
-                # Last resort: try absolute path
-                try:
-                    with open("/app/templates/index.html", "r") as f:
-                        html_content = f.read()
-                    return HTMLResponse(content=html_content)
-                except Exception as abs_e:
-                    logger.error(f"Error reading file from absolute path: {str(abs_e)}")
-                    raise HTTPException(status_code=500, detail=f"Error in UI endpoint: {str(e)}")
-    except Exception as outer_e:
-        logger.error(f"Outer error in UI endpoint: {str(outer_e)}")
-        raise HTTPException(status_code=500, detail=f"Error in UI endpoint: {str(outer_e)}")
-
-# UI route for backward compatibility
-@app.get("/ui", response_class=HTMLResponse)
-async def ui(request: Request):
-    """Serve the web UI interface (kept for backward compatibility)."""
-    return await root(request)
+        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)}")
 
 # 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",
@@ -134,31 +76,43 @@ async def generate_form(
 ):
     """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
-        except ImportError:
+            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
-            except ImportError:
+                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
-                except ImportError:
+                    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")
         generator = AIBOMGenerator()
         
         # Generate AIBOM
+        logger.info(f"Generating AIBOM for model: {model_id}")
         aibom, enhancement_report = generator.generate(
             model_id=model_id,
             include_inference=include_inference,
             use_best_practices=use_best_practices
         )
+        logger.info("AIBOM generation successful")
         
         # 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)
         
@@ -183,11 +137,14 @@ async def generate_form(
         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")
         return templates.TemplateResponse(
             "result.html",
             {
@@ -201,7 +158,7 @@ async def generate_form(
             }
         )
     except Exception as e:
-        logger.error(f"Error generating AI SBOM: {str(e)}")
+        logger.error(f"Error generating AI SBOM: {str(e)}", exc_info=True)
         return templates.TemplateResponse(
             "error.html",
             {
@@ -210,106 +167,18 @@ async def generate_form(
             }
         )
 
-# JSON API endpoints
-@app.post("/api/generate")
-async def generate_api(request: GenerateRequest):
-    """Generate an AI SBOM and return it as JSON."""
-    try:
-        # Import the generator here to avoid circular imports
-        try:
-            from src.aibom_generator.generator import AIBOMGenerator
-        except ImportError:
-            try:
-                from aibom_generator.generator import AIBOMGenerator
-            except ImportError:
-                try:
-                    from generator import AIBOMGenerator
-                except ImportError:
-                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
-        
-        # Create generator instance
-        generator = AIBOMGenerator()
-        
-        # Generate AIBOM
-        aibom, _ = generator.generate(
-            model_id=request.model_id,
-            include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices
-        )
-        
-        return aibom
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
-
-@app.post("/api/generate-with-report")
-async def generate_with_report_api(request: GenerateWithReportRequest):
-    """Generate an AI SBOM with enhancement report and return both as JSON."""
-    try:
-        # Import the generator here to avoid circular imports
-        try:
-            from src.aibom_generator.generator import AIBOMGenerator
-        except ImportError:
-            try:
-                from aibom_generator.generator import AIBOMGenerator
-            except ImportError:
-                try:
-                    from generator import AIBOMGenerator
-                except ImportError:
-                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
-        
-        # Create generator instance
-        generator = AIBOMGenerator()
-        
-        # Generate AIBOM
-        aibom, enhancement_report = generator.generate(
-            model_id=request.model_id,
-            include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices
-        )
-        
-        return {
-            "aibom": aibom,
-            "enhancement_report": enhancement_report
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
-
-@app.get("/api/model-score/{model_id}")
-async def model_score_api(model_id: str):
-    """Get the completeness score for a model."""
-    try:
-        # Import the generator here to avoid circular imports
-        try:
-            from src.aibom_generator.generator import AIBOMGenerator
-        except ImportError:
-            try:
-                from aibom_generator.generator import AIBOMGenerator
-            except ImportError:
-                try:
-                    from generator import AIBOMGenerator
-                except ImportError:
-                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
-        
-        # Create generator instance
-        generator = AIBOMGenerator()
-        
-        # Get completeness score
-        if hasattr(generator, 'get_completeness_score'):
-            completeness_score = generator.get_completeness_score(model_id)
-            return completeness_score
-        else:
-            raise HTTPException(status_code=501, detail="Completeness score calculation not implemented")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error getting model score: {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",
@@ -318,4 +187,5 @@ async def download_file(filename: str):
     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)}")