Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -2,128 +2,408 @@
 FastAPI server for the AIBOM Generator with minimal UI.
 """
 
-import logging
-import os
-from typing import Dict, List, Optional, Any
-
-from fastapi import FastAPI, HTTPException, BackgroundTasks, Request, Form
+from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks, Query, File, UploadFile, Form
+from fastapi.responses import JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse, JSONResponse
-from fastapi.templating import Jinja2Templates
-from pydantic import BaseModel
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, Field
+from typing import Optional, Dict, Any, List
+import uvicorn
+import json
+import os
+import sys
+import uuid
+import shutil
+from datetime import datetime
 
-from aibom_generator.generator import AIBOMGenerator
-from aibom_generator.utils import setup_logging, calculate_completeness_score
+# Add parent directory to path to import generator module
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-# Set up logging
-setup_logging()
-logger = logging.getLogger(__name__)
+# Import the AIBOM generator
+try:
+    from aibom_fix.final_generator import AIBOMGenerator
+except ImportError:
+    # If not found, try the mapping directory
+    try:
+        from aibom_mapping.final_generator import AIBOMGenerator
+    except ImportError:
+        # If still not found, use the original generator
+        try:
+            from aibom_fix.generator import AIBOMGenerator
+        except ImportError:
+            from generator import AIBOMGenerator
 
 # 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",
+    title="Aetheris AI SBOM Generator API",
+    description="API for generating CycloneDX JSON AI SBOMs for machine learning models",
+    version="1.0.0",
 )
 
 # Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # Allow all origins in development
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Initialize templates
-templates = Jinja2Templates(directory="templates")
-
-# 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"),
-)
+# Create output directory for AIBOMs
+os.makedirs(os.path.join(os.path.dirname(os.path.abspath(__file__)), "output"), exist_ok=True)
+app.mount("/output", StaticFiles(directory=os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")), name="output")
 
+# Create a global generator instance
+generator = AIBOMGenerator(use_best_practices=True)
 
-# Define request and response models
-class GenerateRequest(BaseModel):
-    model_id: str
-    include_inference: Optional[bool] = None
-    completeness_threshold: Optional[int] = 0
+# Define request models
+class GenerateAIBOMRequest(BaseModel):
+    model_id: str = Field(..., description="The Hugging Face model ID (e.g., 'meta-llama/Llama-4-Scout-17B-16E-Instruct')")
+    hf_token: Optional[str] = Field(None, description="Optional Hugging Face API token for accessing private models")
+    include_inference: Optional[bool] = Field(True, description="Whether to use AI inference to enhance the AIBOM")
+    use_best_practices: Optional[bool] = Field(True, description="Whether to use industry best practices for scoring")
 
+class AIBOMResponse(BaseModel):
+    aibom: Dict[str, Any] = Field(..., description="The generated AIBOM in CycloneDX JSON format")
+    model_id: str = Field(..., description="The model ID for which the AIBOM was generated")
+    generated_at: str = Field(..., description="Timestamp when the AIBOM was generated")
+    request_id: str = Field(..., description="Unique ID for this request")
+    download_url: Optional[str] = Field(None, description="URL to download the AIBOM JSON file")
 
-class GenerateResponse(BaseModel):
-    aibom: Dict[str, Any]
-    completeness_score: int
-    model_id: str
+class EnhancementReport(BaseModel):
+    ai_enhanced: bool = Field(..., description="Whether AI enhancement was applied")
+    ai_model: Optional[str] = Field(None, description="The AI model used for enhancement, if any")
+    original_score: Dict[str, Any] = Field(..., description="Original completeness score before enhancement")
+    final_score: Dict[str, Any] = Field(..., description="Final completeness score after enhancement")
+    improvement: float = Field(..., description="Score improvement from enhancement")
 
+class AIBOMWithReportResponse(AIBOMResponse):
+    enhancement_report: Optional[EnhancementReport] = Field(None, description="Report on AI enhancement results")
 
 class StatusResponse(BaseModel):
-    status: str
-    version: str
+    status: str = Field(..., description="API status")
+    version: str = Field(..., description="API version")
+    generator_version: str = Field(..., description="AIBOM generator version")
 
+class BatchGenerateRequest(BaseModel):
+    model_ids: List[str] = Field(..., description="List of Hugging Face model IDs to generate AIBOMs for")
+    hf_token: Optional[str] = Field(None, description="Optional Hugging Face API token for accessing private models")
+    include_inference: Optional[bool] = Field(True, description="Whether to use AI inference to enhance the AIBOM")
+    use_best_practices: Optional[bool] = Field(True, description="Whether to use industry best practices for scoring")
 
-# Web UI endpoint
-@app.get("/", response_class=HTMLResponse)
-async def home(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request})
+class BatchJobResponse(BaseModel):
+    job_id: str = Field(..., description="Unique ID for the batch job")
+    status: str = Field(..., description="Job status (e.g., 'queued', 'processing', 'completed')")
+    model_ids: List[str] = Field(..., description="List of model IDs in the batch")
+    created_at: str = Field(..., description="Timestamp when the job was created")
 
+# In-memory storage for batch jobs
+batch_jobs = {}
 
-@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)
+# Define API endpoints
+@app.get("/", response_model=StatusResponse)
+async def get_status():
+    """Get the API status and version information."""
+    return {
+        "status": "operational",
+        "version": "1.0.0",
+        "generator_version": "0.1.0",
+    }
 
-        return templates.TemplateResponse(
-            "result.html",
-            {"request": request, "aibom": aibom, "completeness_score": completeness_score, "model_id": model_id},
+@app.post("/generate", response_model=AIBOMResponse)
+async def generate_aibom(request: GenerateAIBOMRequest):
+    """
+    Generate a CycloneDX JSON AI SBOM for a Hugging Face model.
+    
+    This endpoint takes a model ID and optional parameters to generate
+    a comprehensive AI SBOM in CycloneDX format.
+    """
+    try:
+        # Create a new generator instance with the provided token if available
+        gen = AIBOMGenerator(
+            hf_token=request.hf_token,
+            use_inference=request.include_inference,
+            use_best_practices=request.use_best_practices
         )
-    except Exception as e:
-        logger.error(f"Error generating AIBOM: {e}")
-        return templates.TemplateResponse(
-            "error.html",
-            {"request": request, "error": str(e)},
+        
+        # Generate a request ID
+        request_id = str(uuid.uuid4())
+        
+        # Create output file path
+        safe_model_id = request.model_id.replace("/", "_")
+        output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
+        output_file = os.path.join(output_dir, f"{safe_model_id}_{request_id}.json")
+        
+        # Generate the AIBOM
+        aibom = gen.generate_aibom(
+            model_id=request.model_id,
+            include_inference=request.include_inference,
+            use_best_practices=request.use_best_practices,
+            output_file=output_file
         )
+        
+        # Create download URL
+        download_url = f"/output/{os.path.basename(output_file)}"
+        
+        # Create response
+        response = {
+            "aibom": aibom,
+            "model_id": request.model_id,
+            "generated_at": datetime.utcnow().isoformat() + "Z",
+            "request_id": request_id,
+            "download_url": download_url
+        }
+        
+        return response
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating AIBOM: {str(e)}")
 
-
-# Original JSON API endpoints (kept unchanged)
-@app.post("/generate/json", response_model=GenerateResponse)
-async def generate_aibom(request: GenerateRequest):
+@app.post("/generate-with-report", response_model=AIBOMWithReportResponse)
+async def generate_aibom_with_report(request: GenerateAIBOMRequest):
+    """
+    Generate a CycloneDX JSON AI SBOM with an enhancement report.
+    
+    This endpoint is similar to /generate but also includes a report
+    on the AI enhancement results, including before/after scores.
+    """
     try:
-        aibom = generator.generate_aibom(
+        # Create a new generator instance with the provided token if available
+        gen = AIBOMGenerator(
+            hf_token=request.hf_token,
+            use_inference=request.include_inference,
+            use_best_practices=request.use_best_practices
+        )
+        
+        # Generate a request ID
+        request_id = str(uuid.uuid4())
+        
+        # Create output file path
+        safe_model_id = request.model_id.replace("/", "_")
+        output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
+        output_file = os.path.join(output_dir, f"{safe_model_id}_{request_id}.json")
+        
+        # Generate the AIBOM
+        aibom = gen.generate_aibom(
             model_id=request.model_id,
             include_inference=request.include_inference,
+            use_best_practices=request.use_best_practices,
+            output_file=output_file
         )
-        completeness_score = calculate_completeness_score(aibom)
-
-        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 {
+        
+        # Get the enhancement report
+        enhancement_report = gen.get_enhancement_report()
+        
+        # Create download URL
+        download_url = f"/output/{os.path.basename(output_file)}"
+        
+        # Create response
+        response = {
             "aibom": aibom,
-            "completeness_score": completeness_score,
             "model_id": request.model_id,
+            "generated_at": datetime.utcnow().isoformat() + "Z",
+            "request_id": request_id,
+            "download_url": download_url,
+            "enhancement_report": enhancement_report
         }
+        
+        return response
     except Exception as e:
-        logger.error(f"Error generating AIBOM: {e}")
-        raise HTTPException(
-            status_code=500,
-            detail=f"Error generating AIBOM: {str(e)}",
+        raise HTTPException(status_code=500, detail=f"Error generating AIBOM: {str(e)}")
+
+@app.get("/models/{model_id}/score", response_model=Dict[str, Any])
+async def get_model_score(
+    model_id: str,
+    hf_token: Optional[str] = Query(None, description="Optional Hugging Face API token for accessing private models"),
+    use_best_practices: bool = Query(True, description="Whether to use industry best practices for scoring")
+):
+    """
+    Get the completeness score for a model without generating a full AIBOM.
+    
+    This is a lightweight endpoint that only returns the scoring information.
+    """
+    try:
+        # Create a new generator instance with the provided token if available
+        gen = AIBOMGenerator(
+            hf_token=hf_token,
+            use_inference=False,  # Don't use inference for scoring only
+            use_best_practices=use_best_practices
+        )
+        
+        # Generate the AIBOM (needed to calculate score)
+        aibom = gen.generate_aibom(
+            model_id=model_id,
+            include_inference=False,  # Don't use inference for scoring only
+            use_best_practices=use_best_practices
         )
+        
+        # Get the enhancement report for the score
+        enhancement_report = gen.get_enhancement_report()
+        
+        if enhancement_report and "final_score" in enhancement_report:
+            return enhancement_report["final_score"]
+        else:
+            raise HTTPException(status_code=500, detail="Failed to calculate score")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error calculating score: {str(e)}")
 
+@app.post("/batch", response_model=BatchJobResponse)
+async def batch_generate(request: BatchGenerateRequest, background_tasks: BackgroundTasks):
+    """
+    Start a batch job to generate AIBOMs for multiple models.
+    
+    This endpoint queues a background task to generate AIBOMs for all the
+    specified model IDs and returns a job ID that can be used to check status.
+    """
+    try:
+        # Generate a job ID
+        job_id = str(uuid.uuid4())
+        
+        # Create job record
+        job = {
+            "job_id": job_id,
+            "status": "queued",
+            "model_ids": request.model_ids,
+            "created_at": datetime.utcnow().isoformat() + "Z",
+            "completed": 0,
+            "total": len(request.model_ids),
+            "results": {}
+        }
+        
+        # Store job in memory
+        batch_jobs[job_id] = job
+        
+        # Add background task to process the batch
+        background_tasks.add_task(
+            process_batch_job,
+            job_id=job_id,
+            model_ids=request.model_ids,
+            hf_token=request.hf_token,
+            include_inference=request.include_inference,
+            use_best_practices=request.use_best_practices
+        )
+        
+        # Return job info
+        return {
+            "job_id": job_id,
+            "status": "queued",
+            "model_ids": request.model_ids,
+            "created_at": job["created_at"]
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error starting batch job: {str(e)}")
 
-@app.get("/health")
-async def health():
-    return {"status": "healthy"}
+@app.get("/batch/{job_id}", response_model=Dict[str, Any])
+async def get_batch_status(job_id: str):
+    """
+    Get the status of a batch job.
+    
+    This endpoint returns the current status of a batch job, including
+    progress information and results for completed models.
+    """
+    if job_id not in batch_jobs:
+        raise HTTPException(status_code=404, detail="Batch job not found")
+    
+    return batch_jobs[job_id]
 
+@app.post("/upload-model-card")
+async def upload_model_card(
+    model_id: str = Form(...),
+    model_card: UploadFile = File(...),
+    include_inference: bool = Form(True),
+    use_best_practices: bool = Form(True)
+):
+    """
+    Generate an AIBOM from an uploaded model card file.
+    
+    This endpoint allows users to upload a model card file directly
+    instead of requiring the model to be on Hugging Face.
+    """
+    try:
+        # Create a temporary directory to store the uploaded file
+        temp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "temp")
+        os.makedirs(temp_dir, exist_ok=True)
+        
+        # Save the uploaded file
+        file_path = os.path.join(temp_dir, model_card.filename)
+        with open(file_path, "wb") as f:
+            shutil.copyfileobj(model_card.file, f)
+        
+        # TODO: Implement custom model card processing
+        # This would require modifying the AIBOMGenerator to accept a file path
+        # instead of a model ID, which is beyond the scope of this example
+        
+        # For now, return a placeholder response
+        return {
+            "status": "not_implemented",
+            "message": "Custom model card processing is not yet implemented"
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing uploaded model card: {str(e)}")
 
-if __name__ == "__main__":
-    import uvicorn
+@app.get("/download/{filename}")
+async def download_aibom(filename: str):
+    """
+    Download a previously generated AIBOM file.
+    
+    This endpoint allows downloading AIBOM files by filename.
+    """
+    file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output", filename)
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="File not found")
+    
+    return FileResponse(file_path, media_type="application/json", filename=filename)
 
-    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))
+# Background task function for batch processing
+async def process_batch_job(job_id: str, model_ids: List[str], hf_token: Optional[str], include_inference: bool, use_best_practices: bool):
+    """Process a batch job in the background."""
+    # Update job status
+    batch_jobs[job_id]["status"] = "processing"
+    
+    # Create output directory
+    output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output", job_id)
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Process each model
+    for model_id in model_ids:
+        try:
+            # Create a new generator instance
+            gen = AIBOMGenerator(
+                hf_token=hf_token,
+                use_inference=include_inference,
+                use_best_practices=use_best_practices
+            )
+            
+            # Create output file path
+            safe_model_id = model_id.replace("/", "_")
+            output_file = os.path.join(output_dir, f"{safe_model_id}.json")
+            
+            # Generate the AIBOM
+            aibom = gen.generate_aibom(
+                model_id=model_id,
+                include_inference=include_inference,
+                use_best_practices=use_best_practices,
+                output_file=output_file
+            )
+            
+            # Get the enhancement report
+            enhancement_report = gen.get_enhancement_report()
+            
+            # Create download URL
+            download_url = f"/output/{job_id}/{safe_model_id}.json"
+            
+            # Store result
+            batch_jobs[job_id]["results"][model_id] = {
+                "status": "completed",
+                "download_url": download_url,
+                "enhancement_report": enhancement_report
+            }
+        except Exception as e:
+            # Store error
+            batch_jobs[job_id]["results"][model_id] = {
+                "status": "error",
+                "error": str(e)
+            }
+        
+        # Update progress
+        batch_
+(Content truncated due to size limit. Use line ranges to read in chunks)