Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,321 +1,109 @@
-import os
-import json
 import logging
+import os
 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 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__)
-
-# Define templates directory
-templates_dir = "templates"
-
-# Initialize templates with a simple path
-templates = Jinja2Templates(directory=templates_dir)
-
-# Create app
-app = FastAPI(title="AI SBOM Generator API")
-
-# 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)
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi.middleware.cors import CORSMiddleware
 
-# 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
+from pydantic import BaseModel
 
-class GenerateWithReportRequest(BaseModel):
-    model_id: str
-    include_inference: bool = False
-    use_best_practices: bool = True
+from aibom_generator.generator import AIBOMGenerator
+    allow_headers=["*"],
+)
 
-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
+# Create generator instance
+generator = AIBOMGenerator(
+    hf_token=os.environ.get("HF_TOKEN"),
     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:
-        # 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)
-
-# Status endpoint
-@app.get("/status", response_model=StatusResponse)
-async def get_status():
-    """Get the API status and version information."""
+# Define API endpoints
+@app.get("/", response_model=StatusResponse)
+async def root():
+    """Get API status."""
     return {
-        "status": "operational",
+        "status": "ok",
         "version": "1.0.0",
-        "generator_version": "0.1.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 from form data and render the result template."""
-    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=model_id,
-            include_inference=include_inference,
-            use_best_practices=use_best_practices
-        )
-        
-        # 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)}")
-        
-        # 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 generating AI SBOM: {str(e)}")
-        return templates.TemplateResponse(
-            "error.html",
-            {
-                "request": request,
-                "error_message": f"Error generating AI SBOM: {str(e)}"
-            }
-        )
+@app.post("/generate", response_model=GenerateResponse)
 
-# JSON API endpoints
-@app.post("/api/generate")
-async def generate_api(request: GenerateRequest):
-    """Generate an AI SBOM and return it as JSON."""
+async def generate_aibom(request: GenerateRequest):
+    """Generate an AI SBOM for a Hugging Face 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()
-        
-        # Generate AIBOM
-        aibom, _ = generator.generate(
+        # Generate the AIBOM
+        aibom = generator.generate_aibom(
             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()
+        # Calculate completeness score
+        completeness_score = calculate_completeness_score(aibom)
         
-        # Generate AIBOM
-        aibom, enhancement_report = generator.generate(
-            model_id=request.model_id,
-            include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices
-        )
+        # 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,
-            "enhancement_report": enhancement_report
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
+            "completeness_score": completeness_score,
+        )
+
+
+@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,
+    )
+    
+    return {
+        "status": "accepted",
+        "message": f"AI SBOM generation for {request.model_id} started in the background",
+    }
 
-@app.get("/api/model-score/{model_id}")
-async def model_score_api(model_id: str):
-    """Get the completeness score for a model."""
+
+async def _generate_aibom_background(
+    model_id: str,
+    include_inference: Optional[bool] = None,
+    completeness_threshold: Optional[int] = 0,
+):
+    """Generate an AI SBOM in the background."""
     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.")
+        # Generate the AIBOM
+        aibom = generator.generate_aibom(
+            model_id=model_id,
+            include_inference=include_inference,
+        )
         
-        # Create generator instance
-        generator = AIBOMGenerator()
+        # Calculate completeness score
+        completeness_score = calculate_completeness_score(aibom)
         
-        # 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")
+        # 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}")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error getting model score: {str(e)}")
+        logger.error(f"Error in background AI SBOM generation for {model_id}: {e}")
 
-@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)}")
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {"status": "healthy"}
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 5000)))