Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,153 +1,117 @@
-import logging
 import os
-from typing import Dict, List, Optional, Any, Union
-
-from fastapi import FastAPI, HTTPException, BackgroundTasks
-from fastapi.middleware.cors import CORSMiddleware
+import json
+import logging
+from fastapi import FastAPI, HTTPException, Request, Form
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel
 
-from aibom_generator.generator import AIBOMGenerator
-from aibom_generator.utils import setup_logging, calculate_completeness_score
-
-# Set up logging
-setup_logging()
+# Configure logging
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Create FastAPI app
-app = FastAPI(
-    title="AIBOM Generator API",
-    description="API for generating AI Bills of Materials (AIBOMs) in CycloneDX format for Hugging Face models.",
-    version="0.1.0",
-)
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Create generator instance
-generator = AIBOMGenerator(
-    hf_token=os.environ.get("HF_TOKEN"),
-    inference_model_url=os.environ.get("AIBOM_INFERENCE_URL"),
-    use_inference=os.environ.get("AIBOM_USE_INFERENCE", "true").lower() == "true",
-    cache_dir=os.environ.get("AIBOM_CACHE_DIR"),
-)
+# Define directories
+templates_dir = "templates"
+OUTPUT_DIR = "/tmp/aibom_output"
 
+# Initialize templates
+templates = Jinja2Templates(directory=templates_dir)
 
-# Define request and response models
-class GenerateRequest(BaseModel):
-    model_id: str
-    include_inference: Optional[bool] = None
-    completeness_threshold: Optional[int] = 0
+# Create app
+app = FastAPI(title="AI SBOM Generator API")
 
+# Ensure output directory exists
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-class GenerateResponse(BaseModel):
-    aibom: Dict[str, Any]
-    completeness_score: int
-    model_id: str
-
+# Mount output directory as static files
+app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
 
+# Status response model
 class StatusResponse(BaseModel):
     status: str
     version: str
+    generator_version: str
 
+@app.on_event("startup")
+async def startup_event():
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    logger.info(f"Output directory ready at {OUTPUT_DIR}")
+    logger.info(f"Registered routes: {[route.path for route in app.routes]}")
 
-# Define API endpoints
-@app.get("/", response_model=StatusResponse)
-async def root():
-    """Get API status."""
-    return {
-        "status": "ok",
-        "version": "0.1.0",
-    }
-
-
-@app.post("/generate", response_model=GenerateResponse)
-async def generate_aibom(request: GenerateRequest):
-    """Generate an AIBOM for a Hugging Face model."""
+@app.get("/", response_class=HTMLResponse)
+async def root(request: Request):
     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 it meets the threshold
-        if completeness_score < request.completeness_threshold:
-            raise HTTPException(
-                status_code=400,
-                detail=f"AIBOM completeness score ({completeness_score}) is below threshold ({request.completeness_threshold})",
-            )
-        
-        return {
-            "aibom": aibom,
-            "completeness_score": completeness_score,
-            "model_id": request.model_id,
-        }
+        return templates.TemplateResponse("index.html", {"request": request})
     except Exception as e:
-        logger.error(f"Error generating AIBOM: {e}")
-        raise HTTPException(
-            status_code=500,
-            detail=f"Error generating AIBOM: {str(e)}",
-        )
-
-
-@app.post("/generate/async")
-async def generate_aibom_async(
-    request: GenerateRequest,
-    background_tasks: BackgroundTasks,
-):
-    """Generate an AIBOM 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"AIBOM generation for {request.model_id} started in the background",
-    }
-
-
-async def _generate_aibom_background(
-    model_id: str,
-    include_inference: Optional[bool] = None,
-    completeness_threshold: Optional[int] = 0,
+        logger.error(f"Error rendering template: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Template rendering error: {str(e)}")
+
+@app.get("/status", response_model=StatusResponse)
+async def get_status():
+    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,
+    model_id: str = Form(...),
+    include_inference: bool = Form(False),
+    use_best_practices: bool = Form(True)
 ):
-    """Generate an AIBOM in the background."""
     try:
-        # Generate the AIBOM
+        from src.aibom_generator.generator import AIBOMGenerator
+
+        generator = AIBOMGenerator()
+
         aibom = generator.generate_aibom(
             model_id=model_id,
             include_inference=include_inference,
+            use_best_practices=use_best_practices
+        )
+        enhancement_report = generator.get_enhancement_report()
+
+        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)
+
+        download_url = f"/output/{filename}"
+
+        download_script = f"""
+        <script>
+            function downloadJSON() {{
+                const a = document.createElement('a');
+                a.href = '{download_url}';
+                a.download = '{filename}';
+                document.body.appendChild(a);
+                a.click();
+                document.body.removeChild(a);
+            }}
+        </script>
+        """
+
+        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"Completeness score error: {str(e)}")
+
+        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
+            }
         )
-        
-        # Calculate completeness score
-        completeness_score = calculate_completeness_score(aibom)
-        
-        # TODO: Store the result or notify the user
-        logger.info(f"Background AIBOM generation completed for {model_id}")
-        logger.info(f"Completeness score: {completeness_score}")
     except Exception as e:
-        logger.error(f"Error in background AIBOM generation for {model_id}: {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)))
+        logger.error(f"Error generating AI SBOM: {str(e)}")
+        return templates.TemplateResponse(
+            "error.html", {"request": request, "error": str(e)}
+        )