Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,416 +1,423 @@
-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)
+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=["*"],
+)
+
+# Define output directory - use a directory that will have proper permissions
+# Instead of creating it at the module level, we'll create it in the startup event
+output_dir = "/tmp/aibom_output"  # Using /tmp which should be writable
+
+# 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 = {}
+
+# Startup event to create output directory
+@app.on_event("startup")
+async def startup_event():
+    """Create necessary directories on startup."""
+    os.makedirs(output_dir, exist_ok=True)
+    # Mount the static files directory after it's created
+    app.mount("/output", StaticFiles(directory=output_dir), name="output")
+
+# 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("/tmp", "aibom_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)