Delete src/api.py

src/api.py

@@ -1,416 +0,0 @@
-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)