Upload api.py

src/aibom_generator/api.py

@@ -0,0 +1,416 @@
+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.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
+
+# Add parent directory to path to import generator module
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# 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:
+            try:
+                from generator import AIBOMGenerator
+            except ImportError:
+                # Last resort: try to import from the aibom_generator module
+                try:
+                    from aibom_generator.generator import AIBOMGenerator
+                except ImportError:
+                    raise ImportError("Could not import AIBOMGenerator from any known location")
+
+# Create FastAPI app
+app = FastAPI(
+    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 all origins in development
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Create output directory for AIBOMs
+output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
+os.makedirs(output_dir, exist_ok=True)
+app.mount("/output", StaticFiles(directory=output_dir), name="output")
+
+# Create a global generator instance
+generator = AIBOMGenerator(use_best_practices=True)
+
+# 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 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 = 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")
+
+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 = {}
+
+# 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",
+    }
+
+@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
+        )
+        
+        # Generate a request ID
+        request_id = str(uuid.uuid4())
+        
+        # Create output file path
+        safe_model_id = request.model_id.replace("/", "_")
+        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)}")
+
+@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:
+        # 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_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
+        )
+        
+        # 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,
+            "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:
+        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("/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)}")
+
+@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(output_dir, 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)
+
+# 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
+    batch_output_dir = os.path.join(output_dir, job_id)
+    os.makedirs(batch_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(batch_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_jobs[job_id]["completed"] += 1
+    
+    # Update job status
+    batch_jobs[job_id]["status"] = "completed"
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)