Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -2,16 +2,11 @@ import os
 import json
 import logging
 import sys
-import uuid
-import asyncio
-import shutil
-from datetime import datetime
-from typing import List, Dict, Any, Optional
-from fastapi import FastAPI, HTTPException, Request, Form, File, UploadFile, BackgroundTasks
+from fastapi import FastAPI, HTTPException, Request, Form
 from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -20,8 +15,6 @@ logger = logging.getLogger(__name__)
 # Define directories
 templates_dir = "templates"
 OUTPUT_DIR = "/tmp/aibom_output"
-UPLOAD_DIR = "/tmp/aibom_uploads"
-BATCH_DIR = "/tmp/aibom_batch_jobs"
 
 # Initialize templates
 templates = Jinja2Templates(directory=templates_dir)
@@ -29,10 +22,8 @@ templates = Jinja2Templates(directory=templates_dir)
 # Create app
 app = FastAPI(title="AI SBOM Generator API")
 
-# Ensure directories exist
+# Ensure output directory exists
 os.makedirs(OUTPUT_DIR, exist_ok=True)
-os.makedirs(UPLOAD_DIR, exist_ok=True)
-os.makedirs(BATCH_DIR, exist_ok=True)
 
 # Mount output directory as static files
 app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
@@ -43,41 +34,10 @@ class StatusResponse(BaseModel):
     version: str
     generator_version: str
 
-# Generate request model
-class GenerateRequest(BaseModel):
-    model_id: str
-    include_inference: bool = True
-    use_best_practices: bool = True
-    hf_token: Optional[str] = None
-
-# Batch request model
-class BatchRequest(BaseModel):
-    model_ids: List[str]
-    include_inference: bool = True
-    use_best_practices: bool = True
-    hf_token: Optional[str] = None
-
-# Batch job status model
-class BatchJobStatus(BaseModel):
-    job_id: str
-    status: str
-    model_ids: List[str]
-    created_at: str
-    completed: int = 0
-    total: int = 0
-    results: Dict[str, Any] = {}
-
-# Dictionary to store batch jobs
-batch_jobs = {}
-
 @app.on_event("startup")
 async def startup_event():
     os.makedirs(OUTPUT_DIR, exist_ok=True)
-    os.makedirs(UPLOAD_DIR, exist_ok=True)
-    os.makedirs(BATCH_DIR, exist_ok=True)
     logger.info(f"Output directory ready at {OUTPUT_DIR}")
-    logger.info(f"Upload directory ready at {UPLOAD_DIR}")
-    logger.info(f"Batch directory ready at {BATCH_DIR}")
     logger.info(f"Registered routes: {[route.path for route in app.routes]}")
 
 @app.get("/", response_class=HTMLResponse)
@@ -92,11 +52,6 @@ async def root(request: Request):
 async def get_status():
     return StatusResponse(status="operational", version="1.0.0", generator_version="1.0.0")
 
-# API version of status endpoint
-@app.get("/api/status", response_model=StatusResponse)
-async def api_get_status():
-    return StatusResponse(status="operational", version="1.0.0", generator_version="1.0.0")
-
 # Import utils module for completeness score calculation
 def import_utils():
     """Import utils module with fallback paths."""
@@ -145,7 +100,7 @@ def create_comprehensive_completeness_score(aibom=None):
     If aibom is provided and calculate_completeness_score is available, use it to calculate the score.
     Otherwise, return a default score structure.
     """
-    # If we have the calculate_completeness_score function and an AI SBOM, use it
+    # If we have the calculate_completeness_score function and an AIBOM, use it
     if calculate_completeness_score and aibom:
         try:
             return calculate_completeness_score(aibom, validate=True, use_best_practices=True)
@@ -305,132 +260,55 @@ def create_comprehensive_completeness_score(aibom=None):
         ]
     }
 
-# Helper function to get AI SBOM generator
-def get_generator():
-    """Get AI SBOM generator with fallback paths."""
+@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)
+):
     try:
         # Try different import paths for AIBOMGenerator
+        generator = None
         try:
             from src.aibom_generator.generator import AIBOMGenerator
-            return AIBOMGenerator()
+            generator = AIBOMGenerator()
         except ImportError:
             try:
                 from aibom_generator.generator import AIBOMGenerator
-                return AIBOMGenerator()
+                generator = AIBOMGenerator()
             except ImportError:
                 try:
                     from generator import AIBOMGenerator
-                    return AIBOMGenerator()
+                    generator = AIBOMGenerator()
                 except ImportError:
                     logger.error("Could not import AIBOMGenerator from any known location")
                     raise ImportError("Could not import AIBOMGenerator from any known location")
-    except Exception as e:
-        logger.error(f"Error getting generator: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error getting generator: {str(e)}")
 
-# Helper function to generate AI SBOM
-async def generate_aibom(model_id: str, include_inference: bool = True, use_best_practices: bool = True, hf_token: Optional[str] = None):
-    """Generate AI SBOM for a model."""
-    try:
-        generator = get_generator()
-        
-        # Generate AI SBOM
+        # Generate AIBOM
         aibom = generator.generate_aibom(
             model_id=model_id,
             include_inference=include_inference,
-            use_best_practices=use_best_practices,
-            hf_token=hf_token
+            use_best_practices=use_best_practices
         )
         enhancement_report = generator.get_enhancement_report()
-        
-        # Save AI SBOM to file
-        request_id = str(uuid.uuid4())
-        filename = f"{model_id.replace('/', '_')}_{request_id}.json"
+
+        # Save AIBOM to file
+        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}"
-        
-        # Get completeness score or create a comprehensive one if not available
-        completeness_score = None
-        if hasattr(generator, 'get_completeness_score'):
-            try:
-                completeness_score = generator.get_completeness_score(model_id)
-                logger.info("Successfully retrieved completeness_score from generator")
-            except Exception as e:
-                logger.error(f"Completeness score error from generator: {str(e)}")
-        
-        # If completeness_score is None or doesn't have field_checklist, use comprehensive one
-        if completeness_score is None or not isinstance(completeness_score, dict) or 'field_checklist' not in completeness_score:
-            logger.info("Using comprehensive completeness_score with field_checklist")
-            completeness_score = create_comprehensive_completeness_score(aibom)
-        
-        # Ensure enhancement_report has the right structure
-        if enhancement_report is None:
-            enhancement_report = {
-                "ai_enhanced": False,
-                "ai_model": None,
-                "original_score": {"total_score": 0, "completeness_score": 0},
-                "final_score": {"total_score": 0, "completeness_score": 0},
-                "improvement": 0
-            }
-        else:
-            # Ensure original_score has completeness_score
-            if "original_score" not in enhancement_report or enhancement_report["original_score"] is None:
-                enhancement_report["original_score"] = {"total_score": 0, "completeness_score": 0}
-            elif "completeness_score" not in enhancement_report["original_score"]:
-                enhancement_report["original_score"]["completeness_score"] = enhancement_report["original_score"].get("total_score", 0)
-            
-            # Ensure final_score has completeness_score
-            if "final_score" not in enhancement_report or enhancement_report["final_score"] is None:
-                enhancement_report["final_score"] = {"total_score": 0, "completeness_score": 0}
-            elif "completeness_score" not in enhancement_report["final_score"]:
-                enhancement_report["final_score"]["completeness_score"] = enhancement_report["final_score"].get("total_score", 0)
-        
-        return {
-            "aibom": aibom,
-            "model_id": model_id,
-            "generated_at": datetime.now().isoformat(),
-            "request_id": request_id,
-            "download_url": download_url,
-            "enhancement_report": enhancement_report,
-            "completeness_score": completeness_score
-        }
-    except Exception as e:
-        logger.error(f"Error generating AI SBOM: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
 
-@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),
-    hf_token: Optional[str] = Form(None)
-):
-    try:
-        # Generate AI SBOM
-        result = await generate_aibom(
-            model_id=model_id,
-            include_inference=include_inference,
-            use_best_practices=use_best_practices,
-            hf_token=hf_token
-        )
-        
-        aibom = result["aibom"]
-        enhancement_report = result["enhancement_report"]
-        completeness_score = result["completeness_score"]
-        download_url = result["download_url"]
-        
         # Create download and UI interaction scripts
         download_script = f"""
         <script>
             function downloadJSON() {{
                 const a = document.createElement('a');
                 a.href = '{download_url}';
-                a.download = '{os.path.basename(download_url)}';
+                a.download = '{filename}';
                 document.body.appendChild(a);
                 a.click();
                 document.body.removeChild(a);
@@ -467,7 +345,43 @@ async def generate_form(
             }}
         </script>
         """
+
+        # Get completeness score or create a comprehensive one if not available
+        completeness_score = None
+        if hasattr(generator, 'get_completeness_score'):
+            try:
+                completeness_score = generator.get_completeness_score(model_id)
+                logger.info("Successfully retrieved completeness_score from generator")
+            except Exception as e:
+                logger.error(f"Completeness score error from generator: {str(e)}")
         
+        # If completeness_score is None or doesn't have field_checklist, use comprehensive one
+        if completeness_score is None or not isinstance(completeness_score, dict) or 'field_checklist' not in completeness_score:
+            logger.info("Using comprehensive completeness_score with field_checklist")
+            completeness_score = create_comprehensive_completeness_score(aibom)
+
+        # Ensure enhancement_report has the right structure
+        if enhancement_report is None:
+            enhancement_report = {
+                "ai_enhanced": False,
+                "ai_model": None,
+                "original_score": {"total_score": 0, "completeness_score": 0},
+                "final_score": {"total_score": 0, "completeness_score": 0},
+                "improvement": 0
+            }
+        else:
+            # Ensure original_score has completeness_score
+            if "original_score" not in enhancement_report or enhancement_report["original_score"] is None:
+                enhancement_report["original_score"] = {"total_score": 0, "completeness_score": 0}
+            elif "completeness_score" not in enhancement_report["original_score"]:
+                enhancement_report["original_score"]["completeness_score"] = enhancement_report["original_score"].get("total_score", 0)
+            
+            # Ensure final_score has completeness_score
+            if "final_score" not in enhancement_report or enhancement_report["final_score"] is None:
+                enhancement_report["final_score"] = {"total_score": 0, "completeness_score": 0}
+            elif "completeness_score" not in enhancement_report["final_score"]:
+                enhancement_report["final_score"]["completeness_score"] = enhancement_report["final_score"].get("total_score", 0)
+
         # Add display names and tooltips for score sections
         display_names = {
             "required_fields": "Required Fields",
@@ -478,8 +392,8 @@ async def generate_form(
         }
         
         tooltips = {
-            "required_fields": "Basic required fields for a valid AI SBOM",
-            "metadata": "Information about the AI SBOM itself",
+            "required_fields": "Basic required fields for a valid AIBOM",
+            "metadata": "Information about the AIBOM itself",
             "component_basic": "Basic information about the AI model component",
             "component_model_card": "Detailed model card information",
             "external_references": "Links to external resources"
@@ -492,7 +406,7 @@ async def generate_form(
             "component_model_card": 30,
             "external_references": 10
         }
-        
+
         # Render the template with all necessary data
         return templates.TemplateResponse(
             "result.html",
@@ -515,285 +429,12 @@ async def generate_form(
             "error.html", {"request": request, "error": str(e)}
         )
 
-@app.post("/api/generate", response_model=Dict[str, Any])
-async def api_generate(request: GenerateRequest):
-    """
-    Generate an AI SBOM for a specified Hugging Face model.
-    
-    This endpoint generates an AI SBOM for the specified model and returns it as JSON.
-    """
-    try:
-        # Generate AI SBOM
-        result = await generate_aibom(
-            model_id=request.model_id,
-            include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices,
-            hf_token=request.hf_token
-        )
-        
-        # Return result without completeness_score to keep response size manageable
-        return {
-            "aibom": result["aibom"],
-            "model_id": result["model_id"],
-            "generated_at": result["generated_at"],
-            "request_id": result["request_id"],
-            "download_url": result["download_url"]
-        }
-    except Exception as e:
-        logger.error(f"Error generating AI SBOM: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
-
-@app.post("/api/generate-with-report", response_model=Dict[str, Any])
-async def api_generate_with_report(request: GenerateRequest):
-    """
-    Generate an AI SBOM with a detailed enhancement report.
-    
-    This endpoint generates an AI SBOM for the specified model and returns it as JSON
-    along with an enhancement report.
-    """
-    try:
-        # Generate AI SBOM
-        result = await generate_aibom(
-            model_id=request.model_id,
-            include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices,
-            hf_token=request.hf_token
-        )
-        
-        # Return result with enhancement report but without completeness_score
-        return {
-            "aibom": result["aibom"],
-            "model_id": result["model_id"],
-            "generated_at": result["generated_at"],
-            "request_id": result["request_id"],
-            "download_url": result["download_url"],
-            "enhancement_report": result["enhancement_report"]
-        }
-    except Exception as e:
-        logger.error(f"Error generating AI SBOM: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
-
-@app.get("/api/models/{model_id}/score", response_model=Dict[str, Any])
-async def get_model_score(
-    model_id: str,
-    use_best_practices: bool = True,
-    hf_token: Optional[str] = None
-):
-    """
-    Get the completeness score for a model without generating a full AI SBOM.
-    
-    This endpoint calculates the completeness score for the specified model and returns it as JSON.
-    """
-    try:
-        # Get generator
-        generator = get_generator()
-        
-        # Generate minimal AI SBOM to get score
-        aibom = generator.generate_aibom(
-            model_id=model_id,
-            include_inference=False,  # No need for inference just for scoring
-            use_best_practices=use_best_practices,
-            hf_token=hf_token
-        )
-        
-        # Get completeness score or create a comprehensive one if not available
-        completeness_score = None
-        if hasattr(generator, 'get_completeness_score'):
-            try:
-                completeness_score = generator.get_completeness_score(model_id)
-                logger.info("Successfully retrieved completeness_score from generator")
-            except Exception as e:
-                logger.error(f"Completeness score error from generator: {str(e)}")
-        
-        # If completeness_score is None or doesn't have field_checklist, use comprehensive one
-        if completeness_score is None or not isinstance(completeness_score, dict) or 'field_checklist' not in completeness_score:
-            logger.info("Using comprehensive completeness_score with field_checklist")
-            completeness_score = create_comprehensive_completeness_score(aibom)
-        
-        # Return only the score information
-        return {
-            "total_score": completeness_score["total_score"],
-            "section_scores": completeness_score["section_scores"],
-            "max_scores": completeness_score["max_scores"],
-            "model_id": model_id
-        }
-    except Exception as e:
-        logger.error(f"Error getting model score: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error getting model score: {str(e)}")
-
-# Function to process a batch job
-async def process_batch_job(job_id: str, model_ids: List[str], include_inference: bool, use_best_practices: bool, hf_token: Optional[str] = None):
-    """Process a batch job to generate AI SBOMs for multiple models."""
-    try:
-        # Update job status to processing
-        batch_jobs[job_id]["status"] = "processing"
-        
-        # Process each model
-        for model_id in model_ids:
-            try:
-                # Generate AI SBOM
-                result = await generate_aibom(
-                    model_id=model_id,
-                    include_inference=include_inference,
-                    use_best_practices=use_best_practices,
-                    hf_token=hf_token
-                )
-                
-                # Update job results
-                batch_jobs[job_id]["results"][model_id] = {
-                    "status": "completed",
-                    "download_url": result["download_url"],
-                    "enhancement_report": result["enhancement_report"]
-                }
-                
-                # Update completed count
-                batch_jobs[job_id]["completed"] += 1
-            except Exception as e:
-                logger.error(f"Error processing model {model_id} in batch job {job_id}: {str(e)}")
-                batch_jobs[job_id]["results"][model_id] = {
-                    "status": "failed",
-                    "error": str(e)
-                }
-                batch_jobs[job_id]["completed"] += 1
-        
-        # Update job status to completed
-        batch_jobs[job_id]["status"] = "completed"
-    except Exception as e:
-        logger.error(f"Error processing batch job {job_id}: {str(e)}")
-        batch_jobs[job_id]["status"] = "failed"
-        batch_jobs[job_id]["error"] = str(e)
-
-@app.post("/api/batch", response_model=BatchJobStatus)
-async def create_batch_job(request: BatchRequest, background_tasks: BackgroundTasks):
-    """
-    Start a batch job to generate AI SBOMs for multiple models.
-    
-    This endpoint starts a batch job to generate AI SBOMs for the specified models
-    and returns a job ID that can be used to check the status of the job.
-    """
-    try:
-        # Create job ID
-        job_id = str(uuid.uuid4())
-        
-        # Create job directory
-        job_dir = os.path.join(BATCH_DIR, job_id)
-        os.makedirs(job_dir, exist_ok=True)
-        
-        # Create job status
-        job_status = BatchJobStatus(
-            job_id=job_id,
-            status="queued",
-            model_ids=request.model_ids,
-            created_at=datetime.now().isoformat(),
-            total=len(request.model_ids)
-        )
-        
-        # Store job status
-        batch_jobs[job_id] = job_status.dict()
-        
-        # Start batch job in background
-        background_tasks.add_task(
-            process_batch_job,
-            job_id,
-            request.model_ids,
-            request.include_inference,
-            request.use_best_practices,
-            request.hf_token
-        )
-        
-        return job_status
-    except Exception as e:
-        logger.error(f"Error creating batch job: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error creating batch job: {str(e)}")
-
-@app.get("/api/batch/{job_id}", response_model=BatchJobStatus)
-async def get_batch_job_status(job_id: str):
-    """
-    Get the status of a batch job.
-    
-    This endpoint returns the current status of the specified batch job.
-    """
-    try:
-        # Check if job exists
-        if job_id not in batch_jobs:
-            raise HTTPException(status_code=404, detail=f"Batch job {job_id} not found")
-        
-        # Return job status
-        return BatchJobStatus(**batch_jobs[job_id])
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Error getting batch job status: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error getting batch job status: {str(e)}")
-
-@app.post("/api/upload-model-card", response_model=Dict[str, Any])
-async def upload_model_card(
-    model_id: str = Form(...),
-    model_card: UploadFile = File(...),
-    include_inference: bool = Form(True),
-    use_best_practices: bool = Form(True),
-    hf_token: Optional[str] = Form(None)
-):
-    """
-    Upload a custom model card and generate an AI SBOM.
-    
-    This endpoint allows uploading a custom model card file and generates an AI SBOM
-    using the uploaded model card.
-    """
-    try:
-        # Create upload directory for this model
-        model_upload_dir = os.path.join(UPLOAD_DIR, model_id.replace("/", "_"))
-        os.makedirs(model_upload_dir, exist_ok=True)
-        
-        # Save uploaded model card
-        model_card_path = os.path.join(model_upload_dir, "model_card.md")
-        with open(model_card_path, "wb") as f:
-            content = await model_card.read()
-            f.write(content)
-        
-        # Get generator
-        generator = get_generator()
-        
-        # Generate AI SBOM with custom model card
-        aibom = generator.generate_aibom(
-            model_id=model_id,
-            include_inference=include_inference,
-            use_best_practices=use_best_practices,
-            hf_token=hf_token,
-            custom_model_card_path=model_card_path
-        )
-        enhancement_report = generator.get_enhancement_report()
-        
-        # Save AI SBOM to file
-        request_id = str(uuid.uuid4())
-        filename = f"{model_id.replace('/', '_')}_{request_id}.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}"
-        
-        # Return result
-        return {
-            "aibom": aibom,
-            "model_id": model_id,
-            "generated_at": datetime.now().isoformat(),
-            "request_id": request_id,
-            "download_url": download_url,
-            "enhancement_report": enhancement_report,
-            "model_card_uploaded": True
-        }
-    except Exception as e:
-        logger.error(f"Error uploading model card: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error uploading model card: {str(e)}")
-
 @app.get("/download/{filename}")
 async def download_file(filename: str):
     """
-    Download a generated AI SBOM file.
+    Download a generated AIBOM file.
     
-    This endpoint serves the generated AI SBOM JSON files for download.
+    This endpoint serves the generated AIBOM JSON files for download.
     """
     file_path = os.path.join(OUTPUT_DIR, filename)
     if not os.path.exists(file_path):