Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,13 +1,15 @@
 """
-FastAPI server for the AIBOM Generator.
+FastAPI server for the AIBOM Generator with minimal UI.
 """
 
 import logging
 import os
-from typing import Dict, List, Optional, Any, Union
+from typing import Dict, List, Optional, Any
 
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException, BackgroundTasks, Request, Form
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel
 
 from aibom_generator.generator import AIBOMGenerator
@@ -33,6 +35,9 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Initialize templates
+templates = Jinja2Templates(directory="templates")
+
 # Create generator instance
 generator = AIBOMGenerator(
     hf_token=os.environ.get("HF_TOKEN"),
@@ -60,36 +65,46 @@ class StatusResponse(BaseModel):
     version: str
 
 
-# Define API endpoints
-@app.get("/", response_model=StatusResponse)
-async def root():
-    """Get API status."""
-    return {
-        "status": "ok",
-        "version": "0.1.0",
-    }
+# Web UI endpoint
+@app.get("/", response_class=HTMLResponse)
+async def home(request: Request):
+    return templates.TemplateResponse("index.html", {"request": request})
+
+
+@app.post("/generate", response_class=HTMLResponse)
+async def generate_from_ui(request: Request, model_id: str = Form(...)):
+    try:
+        aibom = generator.generate_aibom(model_id=model_id)
+        completeness_score = calculate_completeness_score(aibom)
+
+        return templates.TemplateResponse(
+            "result.html",
+            {"request": request, "aibom": aibom, "completeness_score": completeness_score, "model_id": model_id},
+        )
+    except Exception as e:
+        logger.error(f"Error generating AIBOM: {e}")
+        return templates.TemplateResponse(
+            "error.html",
+            {"request": request, "error": str(e)},
+        )
 
 
-@app.post("/generate", response_model=GenerateResponse)
+# Original JSON API endpoints (kept unchanged)
+@app.post("/generate/json", response_model=GenerateResponse)
 async def generate_aibom(request: GenerateRequest):
-    """Generate an AIBOM for a Hugging Face model."""
     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,
@@ -103,55 +118,12 @@ async def generate_aibom(request: GenerateRequest):
         )
 
 
-@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,
-):
-    """Generate an AIBOM in the background."""
-    try:
-        # Generate the AIBOM
-        aibom = generator.generate_aibom(
-            model_id=model_id,
-            include_inference=include_inference,
-        )
-        
-        # 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)))
+
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))