Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,9 +1,6 @@
 import os
 import json
 import logging
-import uuid
-import shutil
-import tempfile
 from typing import Dict, List, Optional, Any, Union
 from datetime import datetime
 from pathlib import Path
@@ -73,77 +70,6 @@ async def startup_event():
     os.makedirs(OUTPUT_DIR, exist_ok=True)
     logger.info(f"Created output directory at {OUTPUT_DIR}")
 
-# Helper function to ensure enhancement report has the expected structure
-def ensure_enhancement_report_structure(enhancement_report):
-    """Ensure the enhancement report has all the expected fields for templates."""
-    if enhancement_report is None:
-        enhancement_report = {}
-    
-    # Ensure ai_enhanced exists
-    if "ai_enhanced" not in enhancement_report:
-        enhancement_report["ai_enhanced"] = False
-    
-    # Ensure ai_model exists
-    if "ai_model" not in enhancement_report:
-        enhancement_report["ai_model"] = "Default AI Model"
-    
-    # Ensure original_score exists and has completeness_score
-    if "original_score" not in enhancement_report:
-        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 exists and has completeness_score
-    if "final_score" not in enhancement_report:
-        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)
-    
-    # Ensure improvement exists
-    if "improvement" not in enhancement_report:
-        enhancement_report["improvement"] = 0
-    
-    return enhancement_report
-
-# Helper function to ensure completeness_score has the expected structure
-def ensure_completeness_score_structure(completeness_score):
-    """Ensure the completeness_score has all the expected fields for templates."""
-    if completeness_score is None:
-        completeness_score = {}
-    
-    # Ensure total_score exists
-    if "total_score" not in completeness_score:
-        completeness_score["total_score"] = 0
-    
-    # Ensure completeness_score exists
-    if "completeness_score" not in completeness_score:
-        completeness_score["completeness_score"] = completeness_score.get("total_score", 0)
-    
-    # Ensure section_scores exists
-    if "section_scores" not in completeness_score:
-        completeness_score["section_scores"] = {}
-    
-    # Ensure max_scores exists
-    if "max_scores" not in completeness_score:
-        completeness_score["max_scores"] = {}
-    
-    # Ensure field_checklist exists
-    if "field_checklist" not in completeness_score:
-        completeness_score["field_checklist"] = {}
-    
-    # Ensure field_tiers exists
-    if "field_tiers" not in completeness_score:
-        completeness_score["field_tiers"] = {}
-    
-    # Ensure completeness_profile exists
-    if "completeness_profile" not in completeness_score:
-        completeness_score["completeness_profile"] = {
-            "name": "Incomplete",
-            "description": "This AI SBOM needs improvement to meet transparency requirements."
-        }
-    
-    return completeness_score
-
 # Root route serves the UI
 @app.get("/", response_class=HTMLResponse)
 async def root(request: Request):
@@ -188,15 +114,14 @@ async def ui(request: Request):
     """Serve the web UI interface (kept for backward compatibility)."""
     return await root(request)
 
-# Status endpoint - both at root and /api path for compatibility
+# Status endpoint
 @app.get("/status", response_model=StatusResponse)
-@app.get("/api/status", response_model=StatusResponse)
 async def get_status():
     """Get the API status and version information."""
     return {
         "status": "operational",
         "version": "1.0.0",
-        "generator_version": "1.0.0",
+        "generator_version": "0.1.0",
     }
 
 # Form-based generate endpoint for web UI
@@ -254,23 +179,13 @@ async def generate_form(
         </script>
         """ % (download_url, filename)
         
-        # Process enhancement report for template
-        enhancement_report = ensure_enhancement_report_structure(enhancement_report)
-        logger.info(f"Enhancement report structure: {enhancement_report}")
-        
         # Get completeness score
         completeness_score = None
         if hasattr(generator, 'get_completeness_score'):
             try:
                 completeness_score = generator.get_completeness_score(model_id)
-                completeness_score = ensure_completeness_score_structure(completeness_score)
-                logger.info(f"Completeness score structure: {completeness_score}")
             except Exception as e:
                 logger.error(f"Error getting completeness score: {str(e)}")
-                completeness_score = ensure_completeness_score_structure(None)
-        else:
-            # Create a default completeness score if the method doesn't exist
-            completeness_score = ensure_completeness_score_structure(None)
         
         # Render result template
         return templates.TemplateResponse(
@@ -352,9 +267,6 @@ async def generate_with_report_api(request: GenerateWithReportRequest):
             use_best_practices=request.use_best_practices
         )
         
-        # Process enhancement report for API response
-        enhancement_report = ensure_enhancement_report_structure(enhancement_report)
-        
         return {
             "aibom": aibom,
             "enhancement_report": enhancement_report
@@ -362,100 +274,6 @@ async def generate_with_report_api(request: GenerateWithReportRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
 
-@app.post("/api/batch-generate")
-async def batch_generate_api(
-    request: BatchGenerateRequest,
-    background_tasks: BackgroundTasks
-):
-    """Start a batch job to generate multiple AI SBOMs."""
-    try:
-        # Generate job ID
-        job_id = str(uuid.uuid4())
-        
-        # Create job directory
-        job_dir = os.path.join(OUTPUT_DIR, job_id)
-        os.makedirs(job_dir, exist_ok=True)
-        
-        # Create job status file
-        status_file = os.path.join(job_dir, "status.json")
-        with open(status_file, "w") as f:
-            json.dump({
-                "job_id": job_id,
-                "status": "pending",
-                "model_ids": request.model_ids,
-                "completed": 0,
-                "total": len(request.model_ids),
-                "results": {}
-            }, f, indent=2)
-        
-        # Start background task
-        background_tasks.add_task(
-            process_batch_job,
-            job_id=job_id,
-            model_ids=request.model_ids,
-            include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices
-        )
-        
-        return {
-            "job_id": job_id,
-            "status": "pending",
-            "model_ids": request.model_ids,
-            "message": f"Batch job started. {len(request.model_ids)} models will be processed."
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error starting batch job: {str(e)}")
-
-@app.get("/api/batch-status/{job_id}")
-async def batch_status_api(job_id: str):
-    """Get the status of a batch job."""
-    try:
-        # Check if job exists
-        job_dir = os.path.join(OUTPUT_DIR, job_id)
-        if not os.path.exists(job_dir):
-            raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
-        
-        # Read job status
-        status_file = os.path.join(job_dir, "status.json")
-        with open(status_file, "r") as f:
-            status = json.load(f)
-        
-        return status
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error getting job status: {str(e)}")
-
-@app.get("/api/batch-results/{job_id}")
-async def batch_results_api(job_id: str):
-    """Get the results of a completed batch job."""
-    try:
-        # Check if job exists
-        job_dir = os.path.join(OUTPUT_DIR, job_id)
-        if not os.path.exists(job_dir):
-            raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
-        
-        # Read job status
-        status_file = os.path.join(job_dir, "status.json")
-        with open(status_file, "r") as f:
-            status = json.load(f)
-        
-        # Check if job is completed
-        if status["status"] != "completed":
-            raise HTTPException(status_code=400, detail=f"Job {job_id} is not completed yet")
-        
-        # Read results
-        results = {}
-        for model_id, result_file in status["results"].items():
-            with open(os.path.join(job_dir, result_file), "r") as f:
-                results[model_id] = json.load(f)
-        
-        return results
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error getting job results: {str(e)}")
-
 @app.get("/api/model-score/{model_id}")
 async def model_score_api(model_id: str):
     """Get the completeness score for a model."""
@@ -478,58 +296,12 @@ async def model_score_api(model_id: str):
         # Get completeness score
         if hasattr(generator, 'get_completeness_score'):
             completeness_score = generator.get_completeness_score(model_id)
-            completeness_score = ensure_completeness_score_structure(completeness_score)
             return completeness_score
         else:
-            # Create a default completeness score if the method doesn't exist
-            completeness_score = ensure_completeness_score_structure(None)
-            return completeness_score
+            raise HTTPException(status_code=501, detail="Completeness score calculation not implemented")
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error getting model score: {str(e)}")
 
-@app.post("/api/upload-model-card")
-async def upload_model_card_api(
-    model_id: str = Form(...),
-    model_card: UploadFile = File(...)
-):
-    """Upload a model card file for a model."""
-    try:
-        # Create temporary file
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".md") as temp_file:
-            # Write uploaded file to temporary file
-            shutil.copyfileobj(model_card.file, temp_file)
-        
-        # Import the generator here to avoid circular imports
-        try:
-            from src.aibom_generator.generator import AIBOMGenerator
-        except ImportError:
-            try:
-                from aibom_generator.generator import AIBOMGenerator
-            except ImportError:
-                try:
-                    from generator import AIBOMGenerator
-                except ImportError:
-                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
-        
-        # Create generator instance
-        generator = AIBOMGenerator()
-        
-        # Process model card
-        if hasattr(generator, 'process_model_card'):
-            result = generator.process_model_card(model_id, temp_file.name)
-            
-            # Clean up temporary file
-            os.unlink(temp_file.name)
-            
-            return result
-        else:
-            # Clean up temporary file
-            os.unlink(temp_file.name)
-            
-            raise HTTPException(status_code=501, detail="Model card processing not implemented")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing model card: {str(e)}")
-
 @app.get("/download/{filename}")
 async def download_file(filename: str):
     """Download a generated AI SBOM file."""
@@ -547,99 +319,3 @@ async def download_file(filename: str):
         raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error downloading file: {str(e)}")
-
-# Background task for batch processing
-async def process_batch_job(
-    job_id: str,
-    model_ids: List[str],
-    include_inference: bool,
-    use_best_practices: bool
-):
-    """Process a batch job in the background."""
-    try:
-        # Import the generator here to avoid circular imports
-        try:
-            from src.aibom_generator.generator import AIBOMGenerator
-        except ImportError:
-            try:
-                from aibom_generator.generator import AIBOMGenerator
-            except ImportError:
-                try:
-                    from generator import AIBOMGenerator
-                except ImportError:
-                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
-        
-        # Create generator instance
-        generator = AIBOMGenerator()
-        
-        # Create job directory
-        job_dir = os.path.join(OUTPUT_DIR, job_id)
-        
-        # Read job status
-        status_file = os.path.join(job_dir, "status.json")
-        with open(status_file, "r") as f:
-            status = json.load(f)
-        
-        # Update status to running
-        status["status"] = "running"
-        with open(status_file, "w") as f:
-            json.dump(status, f, indent=2)
-        
-        # Process each model
-        for i, model_id in enumerate(model_ids):
-            try:
-                # Generate AIBOM
-                aibom, enhancement_report = generator.generate(
-                    model_id=model_id,
-                    include_inference=include_inference,
-                    use_best_practices=use_best_practices
-                )
-                
-                # Process enhancement report for API response
-                enhancement_report = ensure_enhancement_report_structure(enhancement_report)
-                
-                # Save result to file
-                result_file = f"{model_id.replace('/', '_')}_aibom.json"
-                result_path = os.path.join(job_dir, result_file)
-                with open(result_path, "w") as f:
-                    json.dump({
-                        "aibom": aibom,
-                        "enhancement_report": enhancement_report
-                    }, f, indent=2)
-                
-                # Update status
-                status["completed"] += 1
-                status["results"][model_id] = result_file
-                with open(status_file, "w") as f:
-                    json.dump(status, f, indent=2)
-            except Exception as e:
-                # Log error
-                logger.error(f"Error processing model {model_id}: {str(e)}")
-                
-                # Update status
-                status["completed"] += 1
-                status["results"][model_id] = f"Error: {str(e)}"
-                with open(status_file, "w") as f:
-                    json.dump(status, f, indent=2)
-        
-        # Update status to completed
-        status["status"] = "completed"
-        with open(status_file, "w") as f:
-            json.dump(status, f, indent=2)
-    except Exception as e:
-        # Log error
-        logger.error(f"Error processing batch job {job_id}: {str(e)}")
-        
-        try:
-            # Update status to failed
-            status_file = os.path.join(OUTPUT_DIR, job_id, "status.json")
-            with open(status_file, "r") as f:
-                status = json.load(f)
-            
-            status["status"] = "failed"
-            status["error"] = str(e)
-            
-            with open(status_file, "w") as f:
-                json.dump(status, f, indent=2)
-        except Exception as status_e:
-            logger.error(f"Error updating job status: {str(status_e)}")