Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,109 +1,321 @@
-import logging
 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, BackgroundTasks
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks, Query, File, UploadFile, Form, Request
+from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+from pydantic import BaseModel, Field
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-from pydantic import BaseModel
+# Define templates directory
+templates_dir = "templates"
 
-from aibom_generator.generator import AIBOMGenerator
-    allow_headers=["*"],
-)
+# Initialize templates with a simple path
+templates = Jinja2Templates(directory=templates_dir)
 
+# Create app
+app = FastAPI(title="AI SBOM Generator API")
 
-# Create generator instance
-generator = AIBOMGenerator(
-    hf_token=os.environ.get("HF_TOKEN"),
-    version: str
+# Define output directory for generated AIBOMs
+OUTPUT_DIR = "/tmp/aibom_output"
 
+# Create output directory if it doesn't exist
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# Define API endpoints
-@app.get("/", response_model=StatusResponse)
-async def root():
-    """Get API status."""
-    return {
-        "status": "ok",
-        "version": "1.0.0",
-    }
+# Mount output directory as static files
+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
 
-@app.post("/generate", response_model=GenerateResponse)
+class BatchGenerateRequest(BaseModel):
+    model_ids: List[str]
+    include_inference: bool = False
+    use_best_practices: bool = True
 
-async def generate_aibom(request: GenerateRequest):
-    """Generate an AI SBOM for a Hugging Face model."""
+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}")
+
+# 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:
-        # Generate the AIBOM
-        aibom = generator.generate_aibom(
-            model_id=request.model_id,
-            include_inference=request.include_inference,
-        )
-        
-        # Calculate completeness score
-        completeness_score = calculate_completeness_score(aibom)
+        # 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}")
         
-        # Check if it meets the threshold
-        if completeness_score < request.completeness_threshold:
-            raise HTTPException(
-                status_code=400,
-                detail=f"AI SBOM completeness score ({completeness_score}) is below threshold ({request.completeness_threshold})",
-            )
-        
-        return {
-            "aibom": aibom,
-            "completeness_score": completeness_score,
-        )
+        # 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)
 
-@app.post("/generate/async")
-async def generate_aibom_async(
-    request: GenerateRequest,
-    background_tasks: BackgroundTasks,
-):
-    """Generate an AI SBOM asynchronously for a Hugging Face model."""
-    # Add to background tasks
-    background_tasks.add_task(
-        _generate_aibom_background,
-        request.model_id,
-        request.include_inference,
-        request.completeness_threshold,
-    )
-    
+# Status endpoint
+@app.get("/status", response_model=StatusResponse)
+async def get_status():
+    """Get the API status and version information."""
     return {
-        "status": "accepted",
-        "message": f"AI SBOM generation for {request.model_id} started in the background",
+        "status": "operational",
+        "version": "1.0.0",
+        "generator_version": "1.0.0",
     }
 
-
-async def _generate_aibom_background(
-    model_id: str,
-    include_inference: Optional[bool] = None,
-    completeness_threshold: Optional[int] = 0,
+# Form-based generate endpoint for web UI
+@app.post("/generate", response_class=HTMLResponse)
+async def generate_form(
+    request: Request,
+    model_id: str = Form(...),
+    include_inference: bool = Form(False),
+    use_best_practices: bool = Form(True)
 ):
-    """Generate an AI SBOM in the background."""
+    """Generate an AI SBOM from form data and render the result template."""
     try:
-        # Generate the AIBOM
-        aibom = generator.generate_aibom(
+        # 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=model_id,
             include_inference=include_inference,
+            use_best_practices=use_best_practices
         )
         
-        # Calculate completeness score
-        completeness_score = calculate_completeness_score(aibom)
+        # Save AIBOM to file
+        filename = f"{model_id.replace('/', '_')}_aibom.json"
+        filepath = os.path.join(OUTPUT_DIR, filename)
+        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 = """
+        <script>
+            function downloadJSON() {
+                const a = document.createElement('a');
+                a.href = '%s';
+                a.download = '%s';
+                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:
+                completeness_score = generator.get_completeness_score(model_id)
+            except Exception as e:
+                logger.error(f"Error getting completeness score: {str(e)}")
         
-        # TODO: Store the result or notify the user
-        logger.info(f"Background AI SBOM generation completed for {model_id}")
-        logger.info(f"Completeness score: {completeness_score}")
+        # Render result template
+        return templates.TemplateResponse(
+            "result.html",
+            {
+                "request": request,
+                "model_id": model_id,
+                "aibom": aibom,
+                "enhancement_report": enhancement_report,
+                "completeness_score": completeness_score,
+                "download_url": download_url,
+                "download_script": download_script
+            }
+        )
     except Exception as e:
-        logger.error(f"Error in background AI SBOM generation for {model_id}: {e}")
+        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)}"
+            }
+        )
 
+# 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.get("/health")
-async def health():
-    """Health check endpoint."""
-    return {"status": "healthy"}
+@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)}")
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 5000)))
+@app.get("/download/{filename}")
+async def download_file(filename: str):
+    """Download a generated AI SBOM file."""
+    try:
+        filepath = os.path.join(OUTPUT_DIR, filename)
+        if not os.path.exists(filepath):
+            raise HTTPException(status_code=404, detail=f"File {filename} not found")
+        
+        return FileResponse(
+            filepath,
+            media_type="application/json",
+            filename=filename
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error downloading file: {str(e)}")